Function crossMapP for fixing desugaring of comprehensions
[haskell-directory.git] / GHC / PArr.hs
1 {-# OPTIONS_GHC -fparr -funbox-strict-fields #-}
2
3 -----------------------------------------------------------------------------
4 -- |
5 -- Module      :  GHC.PArr
6 -- Copyright   :  (c) 2001-2002 Manuel M T Chakravarty & Gabriele Keller
7 -- License     :  see libraries/base/LICENSE
8 -- 
9 -- Maintainer  :  Manuel M. T. Chakravarty <chak@cse.unsw.edu.au>
10 -- Stability   :  internal
11 -- Portability :  non-portable (GHC Extensions)
12 --
13 --  Basic implementation of Parallel Arrays.
14 --
15 --  This module has two functions: (1) It defines the interface to the
16 --  parallel array extension of the Prelude and (2) it provides a vanilla
17 --  implementation of parallel arrays that does not require to flatten the
18 --  array code.  The implementation is not very optimised.
19 --
20 --- DOCU ----------------------------------------------------------------------
21 --
22 --  Language: Haskell 98 plus unboxed values and parallel arrays
23 --
24 --  The semantic difference between standard Haskell arrays (aka "lazy
25 --  arrays") and parallel arrays (aka "strict arrays") is that the evaluation
26 --  of two different elements of a lazy array is independent, whereas in a
27 --  strict array either non or all elements are evaluated.  In other words,
28 --  when a parallel array is evaluated to WHNF, all its elements will be
29 --  evaluated to WHNF.  The name parallel array indicates that all array
30 --  elements may, in general, be evaluated to WHNF in parallel without any
31 --  need to resort to speculative evaluation.  This parallel evaluation
32 --  semantics is also beneficial in the sequential case, as it facilitates
33 --  loop-based array processing as known from classic array-based languages,
34 --  such as Fortran.
35 --
36 --  The interface of this module is essentially a variant of the list
37 --  component of the Prelude, but also includes some functions (such as
38 --  permutations) that are not provided for lists.  The following list
39 --  operations are not supported on parallel arrays, as they would require the
40 --  availability of infinite parallel arrays: `iterate', `repeat', and `cycle'.
41 --
42 --  The current implementation is quite simple and entirely based on boxed
43 --  arrays.  One disadvantage of boxed arrays is that they require to
44 --  immediately initialise all newly allocated arrays with an error thunk to
45 --  keep the garbage collector happy, even if it is guaranteed that the array
46 --  is fully initialised with different values before passing over the
47 --  user-visible interface boundary.  Currently, no effort is made to use
48 --  raw memory copy operations to speed things up.
49 --
50 --- TODO ----------------------------------------------------------------------
51 --
52 --  * We probably want a standard library `PArray' in addition to the prelude
53 --    extension in the same way as the standard library `List' complements the
54 --    list functions from the prelude.
55 --
56 --  * Currently, functions that emphasis the constructor-based definition of
57 --    lists (such as, head, last, tail, and init) are not supported.  
58 --
59 --    Is it worthwhile to support the string processing functions lines,
60 --    words, unlines, and unwords?  (Currently, they are not implemented.)
61 --
62 --    It can, however, be argued that it would be worthwhile to include them
63 --    for completeness' sake; maybe only in the standard library `PArray'.
64 --
65 --  * Prescans are often more useful for array programming than scans.  Shall
66 --    we include them into the Prelude or the library?
67 --
68 --  * Due to the use of the iterator `loop', we could define some fusion rules
69 --    in this module.
70 --
71 --  * We might want to add bounds checks that can be deactivated.
72 --
73
74 module GHC.PArr (
75   -- [::],              -- Built-in syntax
76
77   mapP,                 -- :: (a -> b) -> [:a:] -> [:b:]
78   (+:+),                -- :: [:a:] -> [:a:] -> [:a:]
79   filterP,              -- :: (a -> Bool) -> [:a:] -> [:a:]
80   concatP,              -- :: [:[:a:]:] -> [:a:]
81   concatMapP,           -- :: (a -> [:b:]) -> [:a:] -> [:b:]
82 --  head, last, tail, init,   -- it's not wise to use them on arrays
83   nullP,                -- :: [:a:] -> Bool
84   lengthP,              -- :: [:a:] -> Int
85   (!:),                 -- :: [:a:] -> Int -> a
86   foldlP,               -- :: (a -> b -> a) -> a -> [:b:] -> a
87   foldl1P,              -- :: (a -> a -> a) ->      [:a:] -> a
88   scanlP,               -- :: (a -> b -> a) -> a -> [:b:] -> [:a:]
89   scanl1P,              -- :: (a -> a -> a) ->      [:a:] -> [:a:]
90   foldrP,               -- :: (a -> b -> b) -> b -> [:a:] -> b
91   foldr1P,              -- :: (a -> a -> a) ->      [:a:] -> a
92   scanrP,               -- :: (a -> b -> b) -> b -> [:a:] -> [:b:]
93   scanr1P,              -- :: (a -> a -> a) ->      [:a:] -> [:a:]
94 --  iterate, repeat,          -- parallel arrays must be finite
95   replicateP,           -- :: Int -> a -> [:a:]
96 --  cycle,                    -- parallel arrays must be finite
97   takeP,                -- :: Int -> [:a:] -> [:a:]
98   dropP,                -- :: Int -> [:a:] -> [:a:]
99   splitAtP,             -- :: Int -> [:a:] -> ([:a:],[:a:])
100   takeWhileP,           -- :: (a -> Bool) -> [:a:] -> [:a:]
101   dropWhileP,           -- :: (a -> Bool) -> [:a:] -> [:a:]
102   spanP,                -- :: (a -> Bool) -> [:a:] -> ([:a:], [:a:])
103   breakP,               -- :: (a -> Bool) -> [:a:] -> ([:a:], [:a:])
104 --  lines, words, unlines, unwords,  -- is string processing really needed
105   reverseP,             -- :: [:a:] -> [:a:]
106   andP,                 -- :: [:Bool:] -> Bool
107   orP,                  -- :: [:Bool:] -> Bool
108   anyP,                 -- :: (a -> Bool) -> [:a:] -> Bool
109   allP,                 -- :: (a -> Bool) -> [:a:] -> Bool
110   elemP,                -- :: (Eq a) => a -> [:a:] -> Bool
111   notElemP,             -- :: (Eq a) => a -> [:a:] -> Bool
112   lookupP,              -- :: (Eq a) => a -> [:(a, b):] -> Maybe b
113   sumP,                 -- :: (Num a) => [:a:] -> a
114   productP,             -- :: (Num a) => [:a:] -> a
115   maximumP,             -- :: (Ord a) => [:a:] -> a
116   minimumP,             -- :: (Ord a) => [:a:] -> a
117   zipP,                 -- :: [:a:] -> [:b:]          -> [:(a, b)   :]
118   zip3P,                -- :: [:a:] -> [:b:] -> [:c:] -> [:(a, b, c):]
119   zipWithP,             -- :: (a -> b -> c)      -> [:a:] -> [:b:] -> [:c:]
120   zipWith3P,            -- :: (a -> b -> c -> d) -> [:a:]->[:b:]->[:c:]->[:d:]
121   unzipP,               -- :: [:(a, b)   :] -> ([:a:], [:b:])
122   unzip3P,              -- :: [:(a, b, c):] -> ([:a:], [:b:], [:c:])
123
124   -- overloaded functions
125   --
126   enumFromToP,          -- :: Enum a => a -> a      -> [:a:]
127   enumFromThenToP,      -- :: Enum a => a -> a -> a -> [:a:]
128
129   -- the following functions are not available on lists
130   --
131   toP,                  -- :: [a] -> [:a:]
132   fromP,                -- :: [:a:] -> [a]
133   sliceP,               -- :: Int -> Int -> [:e:] -> [:e:]
134   foldP,                -- :: (e -> e -> e) -> e -> [:e:] -> e
135   fold1P,               -- :: (e -> e -> e) ->      [:e:] -> e
136   permuteP,             -- :: [:Int:] -> [:e:] ->          [:e:]
137   bpermuteP,            -- :: [:Int:] -> [:e:] ->          [:e:]
138   dpermuteP,            -- :: [:Int:] -> [:e:] -> [:e:] -> [:e:]
139   crossP,               -- :: [:a:] -> [:b:] -> [:(a, b):]
140   crossMapP,            -- :: [:a:] -> (a -> [:b:]) -> [:(a, b):]
141   indexOfP              -- :: (a -> Bool) -> [:a:] -> [:Int:]
142 ) where
143
144 import Prelude
145
146 import GHC.ST   ( ST(..), STRep, runST )
147 import GHC.Exts ( Int#, Array#, Int(I#), MutableArray#, newArray#,
148                   unsafeFreezeArray#, indexArray#, writeArray#, (<#), (>=#) )
149
150 infixl 9  !:
151 infixr 5  +:+
152 infix  4  `elemP`, `notElemP`
153
154
155 -- representation of parallel arrays
156 -- ---------------------------------
157
158 -- this rather straight forward implementation maps parallel arrays to the
159 -- internal representation used for standard Haskell arrays in GHC's Prelude
160 -- (EXPORTED ABSTRACTLY)
161 --
162 -- * This definition *must* be kept in sync with `TysWiredIn.parrTyCon'!
163 --
164 data [::] e = PArr Int# (Array# e)
165
166
167 -- exported operations on parallel arrays
168 -- --------------------------------------
169
170 -- operations corresponding to list operations
171 --
172
173 mapP   :: (a -> b) -> [:a:] -> [:b:]
174 mapP f  = fst . loop (mapEFL f) noAL
175
176 (+:+)     :: [:a:] -> [:a:] -> [:a:]
177 a1 +:+ a2  = fst $ loop (mapEFL sel) noAL (enumFromToP 0 (len1 + len2 - 1))
178                        -- we can't use the [:x..y:] form here for tedious
179                        -- reasons to do with the typechecker and the fact that
180                        -- `enumFromToP' is defined in the same module
181              where
182                len1 = lengthP a1
183                len2 = lengthP a2
184                --
185                sel i | i < len1  = a1!:i
186                      | otherwise = a2!:(i - len1)
187
188 filterP   :: (a -> Bool) -> [:a:] -> [:a:]
189 filterP p  = fst . loop (filterEFL p) noAL
190
191 concatP     :: [:[:a:]:] -> [:a:]
192 concatP xss  = foldlP (+:+) [::] xss
193
194 concatMapP   :: (a -> [:b:]) -> [:a:] -> [:b:]
195 concatMapP f  = concatP . mapP f
196
197 --  head, last, tail, init,   -- it's not wise to use them on arrays
198
199 nullP      :: [:a:] -> Bool
200 nullP [::]  = True
201 nullP _     = False
202
203 lengthP             :: [:a:] -> Int
204 lengthP (PArr n# _)  = I# n#
205
206 (!:) :: [:a:] -> Int -> a
207 (!:)  = indexPArr
208
209 foldlP     :: (a -> b -> a) -> a -> [:b:] -> a
210 foldlP f z  = snd . loop (foldEFL (flip f)) z
211
212 foldl1P        :: (a -> a -> a) -> [:a:] -> a
213 foldl1P f [::]  = error "Prelude.foldl1P: empty array"
214 foldl1P f a     = snd $ loopFromTo 1 (lengthP a - 1) (foldEFL f) (a!:0) a
215
216 scanlP     :: (a -> b -> a) -> a -> [:b:] -> [:a:]
217 scanlP f z  = fst . loop (scanEFL (flip f)) z
218
219 scanl1P        :: (a -> a -> a) -> [:a:] -> [:a:]
220 scanl1P f [::]  = error "Prelude.scanl1P: empty array"
221 scanl1P f a     = fst $ loopFromTo 1 (lengthP a - 1) (scanEFL f) (a!:0) a
222
223 foldrP :: (a -> b -> b) -> b -> [:a:] -> b
224 foldrP  = error "Prelude.foldrP: not implemented yet" -- FIXME
225
226 foldr1P :: (a -> a -> a) -> [:a:] -> a
227 foldr1P  = error "Prelude.foldr1P: not implemented yet" -- FIXME
228
229 scanrP :: (a -> b -> b) -> b -> [:a:] -> [:b:]
230 scanrP  = error "Prelude.scanrP: not implemented yet" -- FIXME
231
232 scanr1P :: (a -> a -> a) -> [:a:] -> [:a:]
233 scanr1P  = error "Prelude.scanr1P: not implemented yet" -- FIXME
234
235 --  iterate, repeat           -- parallel arrays must be finite
236
237 replicateP             :: Int -> a -> [:a:]
238 {-# INLINE replicateP #-}
239 replicateP n e  = runST (do
240   marr# <- newArray n e
241   mkPArr n marr#)
242
243 --  cycle                     -- parallel arrays must be finite
244
245 takeP   :: Int -> [:a:] -> [:a:]
246 takeP n  = sliceP 0 (n - 1)
247
248 dropP     :: Int -> [:a:] -> [:a:]
249 dropP n a  = sliceP n (lengthP a - 1) a
250
251 splitAtP      :: Int -> [:a:] -> ([:a:],[:a:])
252 splitAtP n xs  = (takeP n xs, dropP n xs)
253
254 takeWhileP :: (a -> Bool) -> [:a:] -> [:a:]
255 takeWhileP  = error "Prelude.takeWhileP: not implemented yet" -- FIXME
256
257 dropWhileP :: (a -> Bool) -> [:a:] -> [:a:]
258 dropWhileP  = error "Prelude.dropWhileP: not implemented yet" -- FIXME
259
260 spanP :: (a -> Bool) -> [:a:] -> ([:a:], [:a:])
261 spanP  = error "Prelude.spanP: not implemented yet" -- FIXME
262
263 breakP   :: (a -> Bool) -> [:a:] -> ([:a:], [:a:])
264 breakP p  = spanP (not . p)
265
266 --  lines, words, unlines, unwords,  -- is string processing really needed
267
268 reverseP   :: [:a:] -> [:a:]
269 reverseP a  = permuteP (enumFromThenToP (len - 1) (len - 2) 0) a
270                        -- we can't use the [:x, y..z:] form here for tedious
271                        -- reasons to do with the typechecker and the fact that
272                        -- `enumFromThenToP' is defined in the same module
273               where
274                 len = lengthP a
275
276 andP :: [:Bool:] -> Bool
277 andP  = foldP (&&) True
278
279 orP :: [:Bool:] -> Bool
280 orP  = foldP (||) True
281
282 anyP   :: (a -> Bool) -> [:a:] -> Bool
283 anyP p  = orP . mapP p
284
285 allP :: (a -> Bool) -> [:a:] -> Bool
286 allP p  = andP . mapP p
287
288 elemP   :: (Eq a) => a -> [:a:] -> Bool
289 elemP x  = anyP (== x)
290
291 notElemP   :: (Eq a) => a -> [:a:] -> Bool
292 notElemP x  = allP (/= x)
293
294 lookupP :: (Eq a) => a -> [:(a, b):] -> Maybe b
295 lookupP  = error "Prelude.lookupP: not implemented yet" -- FIXME
296
297 sumP :: (Num a) => [:a:] -> a
298 sumP  = foldP (+) 0
299
300 productP :: (Num a) => [:a:] -> a
301 productP  = foldP (*) 1
302
303 maximumP      :: (Ord a) => [:a:] -> a
304 maximumP [::]  = error "Prelude.maximumP: empty parallel array"
305 maximumP xs    = fold1P max xs
306
307 minimumP :: (Ord a) => [:a:] -> a
308 minimumP [::]  = error "Prelude.minimumP: empty parallel array"
309 minimumP xs    = fold1P min xs
310
311 zipP :: [:a:] -> [:b:] -> [:(a, b):]
312 zipP  = zipWithP (,)
313
314 zip3P :: [:a:] -> [:b:] -> [:c:] -> [:(a, b, c):]
315 zip3P  = zipWith3P (,,)
316
317 zipWithP         :: (a -> b -> c) -> [:a:] -> [:b:] -> [:c:]
318 zipWithP f a1 a2  = let 
319                       len1 = lengthP a1
320                       len2 = lengthP a2
321                       len  = len1 `min` len2
322                     in
323                     fst $ loopFromTo 0 (len - 1) combine 0 a1
324                     where
325                       combine e1 i = (Just $ f e1 (a2!:i), i + 1)
326
327 zipWith3P :: (a -> b -> c -> d) -> [:a:]->[:b:]->[:c:]->[:d:]
328 zipWith3P f a1 a2 a3 = let 
329                         len1 = lengthP a1
330                         len2 = lengthP a2
331                         len3 = lengthP a3
332                         len  = len1 `min` len2 `min` len3
333                       in
334                       fst $ loopFromTo 0 (len - 1) combine 0 a1
335                       where
336                         combine e1 i = (Just $ f e1 (a2!:i) (a3!:i), i + 1)
337
338 unzipP   :: [:(a, b):] -> ([:a:], [:b:])
339 unzipP a  = (fst $ loop (mapEFL fst) noAL a, fst $ loop (mapEFL snd) noAL a)
340 -- FIXME: these two functions should be optimised using a tupled custom loop
341 unzip3P   :: [:(a, b, c):] -> ([:a:], [:b:], [:c:])
342 unzip3P a  = (fst $ loop (mapEFL fst3) noAL a, 
343               fst $ loop (mapEFL snd3) noAL a,
344               fst $ loop (mapEFL trd3) noAL a)
345              where
346                fst3 (a, _, _) = a
347                snd3 (_, b, _) = b
348                trd3 (_, _, c) = c
349
350 -- instances
351 --
352
353 instance Eq a => Eq [:a:] where
354   a1 == a2 | lengthP a1 == lengthP a2 = andP (zipWithP (==) a1 a2)
355            | otherwise                = False
356
357 instance Ord a => Ord [:a:] where
358   compare a1 a2 = case foldlP combineOrdering EQ (zipWithP compare a1 a2) of
359                     EQ | lengthP a1 == lengthP a2 -> EQ
360                        | lengthP a1 <  lengthP a2 -> LT
361                        | otherwise                -> GT
362                   where
363                     combineOrdering EQ    EQ    = EQ
364                     combineOrdering EQ    other = other
365                     combineOrdering other _     = other
366
367 instance Functor [::] where
368   fmap = mapP
369
370 instance Monad [::] where
371   m >>= k  = foldrP ((+:+) . k      ) [::] m
372   m >>  k  = foldrP ((+:+) . const k) [::] m
373   return x = [:x:]
374   fail _   = [::]
375
376 instance Show a => Show [:a:]  where
377   showsPrec _  = showPArr . fromP
378     where
379       showPArr []     s = "[::]" ++ s
380       showPArr (x:xs) s = "[:" ++ shows x (showPArr' xs s)
381
382       showPArr' []     s = ":]" ++ s
383       showPArr' (y:ys) s = ',' : shows y (showPArr' ys s)
384
385 instance Read a => Read [:a:]  where
386   readsPrec _ a = [(toP v, rest) | (v, rest) <- readPArr a]
387     where
388       readPArr = readParen False (\r -> do
389                                           ("[:",s) <- lex r
390                                           readPArr1 s)
391       readPArr1 s = 
392         (do { (":]", t) <- lex s; return ([], t) }) ++
393         (do { (x, t) <- reads s; (xs, u) <- readPArr2 t; return (x:xs, u) })
394
395       readPArr2 s = 
396         (do { (":]", t) <- lex s; return ([], t) }) ++
397         (do { (",", t) <- lex s; (x, u) <- reads t; (xs, v) <- readPArr2 u; 
398               return (x:xs, v) })
399
400 -- overloaded functions
401 -- 
402
403 -- Ideally, we would like `enumFromToP' and `enumFromThenToP' to be members of
404 -- `Enum'.  On the other hand, we really do not want to change `Enum'.  Thus,
405 -- for the moment, we hope that the compiler is sufficiently clever to
406 -- properly fuse the following definitions.
407
408 enumFromToP     :: Enum a => a -> a -> [:a:]
409 enumFromToP x y  = mapP toEnum (eftInt (fromEnum x) (fromEnum y))
410   where
411     eftInt x y = scanlP (+) x $ replicateP (y - x + 1) 1
412
413 enumFromThenToP       :: Enum a => a -> a -> a -> [:a:]
414 enumFromThenToP x y z  = 
415   mapP toEnum (efttInt (fromEnum x) (fromEnum y) (fromEnum z))
416   where
417     efttInt x y z = scanlP (+) x $ 
418                       replicateP (abs (z - x) `div` abs delta + 1) delta
419       where
420        delta = y - x
421
422 -- the following functions are not available on lists
423 --
424
425 -- create an array from a list (EXPORTED)
426 --
427 toP   :: [a] -> [:a:]
428 toP l  = fst $ loop store l (replicateP (length l) ())
429          where
430            store _ (x:xs) = (Just x, xs)
431
432 -- convert an array to a list (EXPORTED)
433 --
434 fromP   :: [:a:] -> [a]
435 fromP a  = [a!:i | i <- [0..lengthP a - 1]]
436
437 -- cut a subarray out of an array (EXPORTED)
438 --
439 sliceP :: Int -> Int -> [:e:] -> [:e:]
440 sliceP from to a = 
441   fst $ loopFromTo (0 `max` from) (to `min` (lengthP a - 1)) (mapEFL id) noAL a
442
443 -- parallel folding (EXPORTED)
444 --
445 -- * the first argument must be associative; otherwise, the result is undefined
446 --
447 foldP :: (e -> e -> e) -> e -> [:e:] -> e
448 foldP  = foldlP
449
450 -- parallel folding without explicit neutral (EXPORTED)
451 --
452 -- * the first argument must be associative; otherwise, the result is undefined
453 --
454 fold1P :: (e -> e -> e) -> [:e:] -> e
455 fold1P  = foldl1P
456
457 -- permute an array according to the permutation vector in the first argument
458 -- (EXPORTED)
459 --
460 permuteP       :: [:Int:] -> [:e:] -> [:e:]
461 permuteP is es 
462   | isLen /= esLen = error "GHC.PArr: arguments must be of the same length"
463   | otherwise      = runST (do
464                        marr <- newArray isLen noElem
465                        permute marr is es
466                        mkPArr isLen marr)
467   where
468     noElem = error "GHC.PArr.permuteP: I do not exist!"
469              -- unlike standard Haskell arrays, this value represents an
470              -- internal error
471     isLen = lengthP is
472     esLen = lengthP es
473
474 -- permute an array according to the back-permutation vector in the first
475 -- argument (EXPORTED)
476 --
477 -- * the permutation vector must represent a surjective function; otherwise,
478 --   the result is undefined
479 --
480 bpermuteP       :: [:Int:] -> [:e:] -> [:e:]
481 bpermuteP is es  = fst $ loop (mapEFL (es!:)) noAL is
482
483 -- permute an array according to the permutation vector in the first
484 -- argument, which need not be surjective (EXPORTED)
485 --
486 -- * any elements in the result that are not covered by the permutation
487 --   vector assume the value of the corresponding position of the third
488 --   argument 
489 --
490 dpermuteP :: [:Int:] -> [:e:] -> [:e:] -> [:e:]
491 dpermuteP is es dft
492   | isLen /= esLen = error "GHC.PArr: arguments must be of the same length"
493   | otherwise      = runST (do
494                        marr <- newArray dftLen noElem
495                        trans 0 (isLen - 1) marr dft copyOne noAL
496                        permute marr is es
497                        mkPArr dftLen marr)
498   where
499     noElem = error "GHC.PArr.permuteP: I do not exist!"
500              -- unlike standard Haskell arrays, this value represents an
501              -- internal error
502     isLen  = lengthP is
503     esLen  = lengthP es
504     dftLen = lengthP dft
505
506     copyOne e _ = (Just e, noAL)
507
508 -- computes the cross combination of two arrays (EXPORTED)
509 --
510 crossP       :: [:a:] -> [:b:] -> [:(a, b):]
511 crossP a1 a2  = fst $ loop combine (0, 0) $ replicateP len ()
512                 where
513                   len1 = lengthP a1
514                   len2 = lengthP a2
515                   len  = len1 * len2
516                   --
517                   combine _ (i, j) = (Just $ (a1!:i, a2!:j), next)
518                                      where
519                                        next | (i + 1) == len1 = (0    , j + 1)
520                                             | otherwise       = (i + 1, j)
521
522 {- An alternative implementation
523    * The one above is certainly better for flattened code, but here where we
524      are handling boxed arrays, the trade off is less clear.  However, I
525      think, the above one is still better.
526
527 crossP a1 a2  = let
528                   len1 = lengthP a1
529                   len2 = lengthP a2
530                   x1   = concatP $ mapP (replicateP len2) a1
531                   x2   = concatP $ replicateP len1 a2
532                 in
533                 zipP x1 x2
534  -}
535
536 -- |Compute a cross of an array and the arrays produced by the given function
537 -- for the elements of the first array.
538 --
539 crossMapP :: [:a:] -> (a -> [:b:]) -> [:(a, b):]
540 crossMapP a f = let
541                   bs   = mapP f a
542                   segd = mapP lengthP bs
543                   as   = zipWithP replicateP segd a
544                 in
545                 zipP (concatP as) (concatP bs)
546
547 {- The following may seem more straight forward, but the above is very cheap
548    with segmented arrays, as `mapP lengthP', `zipP', and `concatP' are
549    constant time, and `map f' uses the lifted version of `f'.
550
551 crossMapP a f = concatP $ mapP (\x -> mapP ((,) x) (f x)) a
552
553  -}
554
555 -- computes an index array for all elements of the second argument for which
556 -- the predicate yields `True' (EXPORTED)
557 --
558 indexOfP     :: (a -> Bool) -> [:a:] -> [:Int:]
559 indexOfP p a  = fst $ loop calcIdx 0 a
560                 where
561                   calcIdx e idx | p e       = (Just idx, idx + 1)
562                                 | otherwise = (Nothing , idx    )
563
564
565 -- auxiliary functions
566 -- -------------------
567
568 -- internally used mutable boxed arrays
569 --
570 data MPArr s e = MPArr Int# (MutableArray# s e)
571
572 -- allocate a new mutable array that is pre-initialised with a given value
573 --
574 newArray             :: Int -> e -> ST s (MPArr s e)
575 {-# INLINE newArray #-}
576 newArray (I# n#) e  = ST $ \s1# ->
577   case newArray# n# e s1# of { (# s2#, marr# #) ->
578   (# s2#, MPArr n# marr# #)}
579
580 -- convert a mutable array into the external parallel array representation
581 --
582 mkPArr                           :: Int -> MPArr s e -> ST s [:e:]
583 {-# INLINE mkPArr #-}
584 mkPArr (I# n#) (MPArr _ marr#)  = ST $ \s1# ->
585   case unsafeFreezeArray# marr# s1#   of { (# s2#, arr# #) ->
586   (# s2#, PArr n# arr# #) }
587
588 -- general array iterator
589 --
590 -- * corresponds to `loopA' from ``Functional Array Fusion'', Chakravarty &
591 --   Keller, ICFP 2001
592 --
593 loop :: (e -> acc -> (Maybe e', acc))    -- mapping & folding, once per element
594      -> acc                              -- initial acc value
595      -> [:e:]                            -- input array
596      -> ([:e':], acc)
597 {-# INLINE loop #-}
598 loop mf acc arr = loopFromTo 0 (lengthP arr - 1) mf acc arr
599
600 -- general array iterator with bounds
601 --
602 loopFromTo :: Int                        -- from index
603            -> Int                        -- to index
604            -> (e -> acc -> (Maybe e', acc))
605            -> acc
606            -> [:e:]
607            -> ([:e':], acc)
608 {-# INLINE loopFromTo #-}
609 loopFromTo from to mf start arr = runST (do
610   marr      <- newArray (to - from + 1) noElem
611   (n', acc) <- trans from to marr arr mf start
612   arr       <- mkPArr n' marr
613   return (arr, acc))
614   where
615     noElem = error "GHC.PArr.loopFromTo: I do not exist!"
616              -- unlike standard Haskell arrays, this value represents an
617              -- internal error
618
619 -- actual loop body of `loop'
620 --
621 -- * for this to be really efficient, it has to be translated with the
622 --   constructor specialisation phase "SpecConstr" switched on; as of GHC 5.03
623 --   this requires an optimisation level of at least -O2
624 --
625 trans :: Int                            -- index of first elem to process
626       -> Int                            -- index of last elem to process
627       -> MPArr s e'                     -- destination array
628       -> [:e:]                          -- source array
629       -> (e -> acc -> (Maybe e', acc))  -- mutator
630       -> acc                            -- initial accumulator
631       -> ST s (Int, acc)                -- final destination length/final acc
632 {-# INLINE trans #-}
633 trans from to marr arr mf start = trans' from 0 start
634   where
635     trans' arrOff marrOff acc 
636       | arrOff > to = return (marrOff, acc)
637       | otherwise   = do
638                         let (oe', acc') = mf (arr `indexPArr` arrOff) acc
639                         marrOff' <- case oe' of
640                                       Nothing -> return marrOff 
641                                       Just e' -> do
642                                         writeMPArr marr marrOff e'
643                                         return $ marrOff + 1
644                         trans' (arrOff + 1) marrOff' acc'
645
646 -- Permute the given elements into the mutable array.
647 --
648 permute :: MPArr s e -> [:Int:] -> [:e:] -> ST s ()
649 permute marr is es = perm 0
650   where
651     perm i
652       | i == n = return ()
653       | otherwise  = writeMPArr marr (is!:i) (es!:i) >> perm (i + 1)
654       where
655         n = lengthP is
656
657
658 -- common patterns for using `loop'
659 --
660
661 -- initial value for the accumulator when the accumulator is not needed
662 --
663 noAL :: ()
664 noAL  = ()
665
666 -- `loop' mutator maps a function over array elements
667 --
668 mapEFL   :: (e -> e') -> (e -> () -> (Maybe e', ()))
669 {-# INLINE mapEFL #-}
670 mapEFL f  = \e a -> (Just $ f e, ())
671
672 -- `loop' mutator that filter elements according to a predicate
673 --
674 filterEFL   :: (e -> Bool) -> (e -> () -> (Maybe e, ()))
675 {-# INLINE filterEFL #-}
676 filterEFL p  = \e a -> if p e then (Just e, ()) else (Nothing, ())
677
678 -- `loop' mutator for array folding
679 --
680 foldEFL   :: (e -> acc -> acc) -> (e -> acc -> (Maybe (), acc))
681 {-# INLINE foldEFL #-}
682 foldEFL f  = \e a -> (Nothing, f e a)
683
684 -- `loop' mutator for array scanning
685 --
686 scanEFL   :: (e -> acc -> acc) -> (e -> acc -> (Maybe acc, acc))
687 {-# INLINE scanEFL #-}
688 scanEFL f  = \e a -> (Just a, f e a)
689
690 -- elementary array operations
691 --
692
693 -- unlifted array indexing 
694 --
695 indexPArr                       :: [:e:] -> Int -> e
696 {-# INLINE indexPArr #-}
697 indexPArr (PArr n# arr#) (I# i#) 
698   | i# >=# 0# && i# <# n# =
699     case indexArray# arr# i# of (# e #) -> e
700   | otherwise = error $ "indexPArr: out of bounds parallel array index; " ++
701                         "idx = " ++ show (I# i#) ++ ", arr len = "
702                         ++ show (I# n#)
703
704 -- encapsulate writing into a mutable array into the `ST' monad
705 --
706 writeMPArr                           :: MPArr s e -> Int -> e -> ST s ()
707 {-# INLINE writeMPArr #-}
708 writeMPArr (MPArr n# marr#) (I# i#) e 
709   | i# >=# 0# && i# <# n# =
710     ST $ \s# ->
711     case writeArray# marr# i# e s# of s'# -> (# s'#, () #)
712   | otherwise = error $ "writeMPArr: out of bounds parallel array index; " ++
713                         "idx = " ++ show (I# i#) ++ ", arr len = "
714                         ++ show (I# n#)