fromVectorStream -> concatVectors
[darcs-mirrors/vector.git] / Data / Vector / Fusion / Stream.hs
1 {-# LANGUAGE FlexibleInstances, Rank2Types, BangPatterns #-}
2
3 -- |
4 -- Module : Data.Vector.Fusion.Stream
5 -- Copyright : (c) Roman Leshchinskiy 2008-2010
6 -- License : BSD-style
7 --
8 -- Maintainer : Roman Leshchinskiy <rl@cse.unsw.edu.au>
9 -- Stability : experimental
10 -- Portability : non-portable
11 --
12 -- Streams for stream fusion
13 --
14
15 module Data.Vector.Fusion.Stream (
16 -- * Types
17 Step(..), Chunk(..), Facets, MFacets,
18
19 -- * In-place markers
20 inplace,
21
22 -- * Size hints
23 size, sized,
24
25 -- * Length information
26 length, null,
27
28 -- * Construction
29 empty, singleton, cons, snoc, replicate, generate, (++),
30
31 -- * Accessing individual elements
32 head, last, (!!), (!?),
33
34 -- * Substreams
35 slice, init, tail, take, drop,
36
37 -- * Mapping
38 map, concatMap, flatten, unbox,
39
40 -- * Zipping
41 indexed, indexedR,
42 zipWith, zipWith3, zipWith4, zipWith5, zipWith6,
43 zip, zip3, zip4, zip5, zip6,
44
45 -- * Filtering
46 filter, takeWhile, dropWhile,
47
48 -- * Searching
49 elem, notElem, find, findIndex,
50
51 -- * Folding
52 foldl, foldl1, foldl', foldl1', foldr, foldr1,
53
54 -- * Specialised folds
55 and, or,
56
57 -- * Unfolding
58 unfoldr, unfoldrN, iterateN,
59
60 -- * Scans
61 prescanl, prescanl',
62 postscanl, postscanl',
63 scanl, scanl',
64 scanl1, scanl1',
65
66 -- * Enumerations
67 enumFromStepN, enumFromTo, enumFromThenTo,
68
69 -- * Conversions
70 toList, fromList, fromListN, unsafeFromList, liftStream,
71 fromVector, reVector, fromVectors, concatVectors,
72
73 -- * Monadic combinators
74 mapM, mapM_, zipWithM, zipWithM_, filterM, foldM, fold1M, foldM', fold1M',
75
76 eq, cmp
77 ) where
78
79 import Data.Vector.Generic.Base ( Vector )
80 import Data.Vector.Fusion.Stream.Size
81 import Data.Vector.Fusion.Util
82 import Data.Vector.Fusion.Stream.Monadic ( Step(..), Chunk(..), SPEC(..) )
83 import qualified Data.Vector.Fusion.Stream.Monadic as M
84
85 import Prelude hiding ( length, null,
86 replicate, (++),
87 head, last, (!!),
88 init, tail, take, drop,
89 map, concatMap,
90 zipWith, zipWith3, zip, zip3,
91 filter, takeWhile, dropWhile,
92 elem, notElem,
93 foldl, foldl1, foldr, foldr1,
94 and, or,
95 scanl, scanl1,
96 enumFromTo, enumFromThenTo,
97 mapM, mapM_ )
98
99 import GHC.Base ( build )
100
101 #include "vector.h"
102
103 -- | The type of pure streams
104 type Facets = M.Facets Id
105
106 -- | Alternative name for monadic streams
107 type MFacets = M.Facets
108
109 inplace :: (forall m. Monad m => M.Facets m v a -> M.Facets m v b)
110 -> Facets v a -> Facets v b
111 {-# INLINE_STREAM inplace #-}
112 inplace f s = s `seq` f s
113
114 {-# RULES
115
116 "inplace/inplace [Vector]"
117 forall (f :: forall m. Monad m => MFacets m v a -> MFacets m v a)
118 (g :: forall m. Monad m => MFacets m v a -> MFacets m v a)
119 s.
120 inplace f (inplace g s) = inplace (f . g) s
121
122 #-}
123
124 -- | Convert a pure stream to a monadic stream
125 liftStream :: Monad m => Facets v a -> M.Facets m v a
126 {-# INLINE_STREAM liftStream #-}
127 liftStream (M.Facets (M.Unf step s) (M.Unf vstep t) v sz)
128 = M.Facets (M.Unf (return . unId . step) s)
129 (M.Unf (return . unId . vstep) t) v sz
130
131 -- | 'Size' hint of a 'Facets'
132 size :: Facets v a -> Size
133 {-# INLINE size #-}
134 size = M.size
135
136 -- | Attach a 'Size' hint to a 'Facets'
137 sized :: Facets v a -> Size -> Facets v a
138 {-# INLINE sized #-}
139 sized = M.sized
140
141 -- Length
142 -- ------
143
144 -- | Length of a 'Facets'
145 length :: Facets v a -> Int
146 {-# INLINE length #-}
147 length = unId . M.length
148
149 -- | Check if a 'Facets' is empty
150 null :: Facets v a -> Bool
151 {-# INLINE null #-}
152 null = unId . M.null
153
154 -- Construction
155 -- ------------
156
157 -- | Empty 'Facets'
158 empty :: Facets v a
159 {-# INLINE empty #-}
160 empty = M.empty
161
162 -- | Singleton 'Facets'
163 singleton :: a -> Facets v a
164 {-# INLINE singleton #-}
165 singleton = M.singleton
166
167 -- | Replicate a value to a given length
168 replicate :: Int -> a -> Facets v a
169 {-# INLINE replicate #-}
170 replicate = M.replicate
171
172 -- | Generate a stream from its indices
173 generate :: Int -> (Int -> a) -> Facets v a
174 {-# INLINE generate #-}
175 generate = M.generate
176
177 -- | Prepend an element
178 cons :: a -> Facets v a -> Facets v a
179 {-# INLINE cons #-}
180 cons = M.cons
181
182 -- | Append an element
183 snoc :: Facets v a -> a -> Facets v a
184 {-# INLINE snoc #-}
185 snoc = M.snoc
186
187 infixr 5 ++
188 -- | Concatenate two 'Facets's
189 (++) :: Facets v a -> Facets v a -> Facets v a
190 {-# INLINE (++) #-}
191 (++) = (M.++)
192
193 -- Accessing elements
194 -- ------------------
195
196 -- | First element of the 'Facets' or error if empty
197 head :: Facets v a -> a
198 {-# INLINE head #-}
199 head = unId . M.head
200
201 -- | Last element of the 'Facets' or error if empty
202 last :: Facets v a -> a
203 {-# INLINE last #-}
204 last = unId . M.last
205
206 infixl 9 !!
207 -- | Element at the given position
208 (!!) :: Facets v a -> Int -> a
209 {-# INLINE (!!) #-}
210 s !! i = unId (s M.!! i)
211
212 infixl 9 !?
213 -- | Element at the given position or 'Nothing' if out of bounds
214 (!?) :: Facets v a -> Int -> Maybe a
215 {-# INLINE (!?) #-}
216 s !? i = unId (s M.!? i)
217
218 -- Substreams
219 -- ----------
220
221 -- | Extract a substream of the given length starting at the given position.
222 slice :: Int -- ^ starting index
223 -> Int -- ^ length
224 -> Facets v a
225 -> Facets v a
226 {-# INLINE slice #-}
227 slice = M.slice
228
229 -- | All but the last element
230 init :: Facets v a -> Facets v a
231 {-# INLINE init #-}
232 init = M.init
233
234 -- | All but the first element
235 tail :: Facets v a -> Facets v a
236 {-# INLINE tail #-}
237 tail = M.tail
238
239 -- | The first @n@ elements
240 take :: Int -> Facets v a -> Facets v a
241 {-# INLINE take #-}
242 take = M.take
243
244 -- | All but the first @n@ elements
245 drop :: Int -> Facets v a -> Facets v a
246 {-# INLINE drop #-}
247 drop = M.drop
248
249 -- Mapping
250 -- ---------------
251
252 -- | Map a function over a 'Facets'
253 map :: (a -> b) -> Facets v a -> Facets v b
254 {-# INLINE map #-}
255 map = M.map
256
257 unbox :: Facets v (Box a) -> Facets v a
258 {-# INLINE unbox #-}
259 unbox = M.unbox
260
261 concatMap :: (a -> Facets v b) -> Facets v a -> Facets v b
262 {-# INLINE concatMap #-}
263 concatMap = M.concatMap
264
265 -- Zipping
266 -- -------
267
268 -- | Pair each element in a 'Facets' with its index
269 indexed :: Facets v a -> Facets v (Int,a)
270 {-# INLINE indexed #-}
271 indexed = M.indexed
272
273 -- | Pair each element in a 'Facets' with its index, starting from the right
274 -- and counting down
275 indexedR :: Int -> Facets v a -> Facets v (Int,a)
276 {-# INLINE_STREAM indexedR #-}
277 indexedR = M.indexedR
278
279 -- | Zip two 'Facets's with the given function
280 zipWith :: (a -> b -> c) -> Facets v a -> Facets v b -> Facets v c
281 {-# INLINE zipWith #-}
282 zipWith = M.zipWith
283
284 -- | Zip three 'Facets's with the given function
285 zipWith3 :: (a -> b -> c -> d) -> Facets v a -> Facets v b -> Facets v c -> Facets v d
286 {-# INLINE zipWith3 #-}
287 zipWith3 = M.zipWith3
288
289 zipWith4 :: (a -> b -> c -> d -> e)
290 -> Facets v a -> Facets v b -> Facets v c -> Facets v d
291 -> Facets v e
292 {-# INLINE zipWith4 #-}
293 zipWith4 = M.zipWith4
294
295 zipWith5 :: (a -> b -> c -> d -> e -> f)
296 -> Facets v a -> Facets v b -> Facets v c -> Facets v d
297 -> Facets v e -> Facets v f
298 {-# INLINE zipWith5 #-}
299 zipWith5 = M.zipWith5
300
301 zipWith6 :: (a -> b -> c -> d -> e -> f -> g)
302 -> Facets v a -> Facets v b -> Facets v c -> Facets v d
303 -> Facets v e -> Facets v f -> Facets v g
304 {-# INLINE zipWith6 #-}
305 zipWith6 = M.zipWith6
306
307 zip :: Facets v a -> Facets v b -> Facets v (a,b)
308 {-# INLINE zip #-}
309 zip = M.zip
310
311 zip3 :: Facets v a -> Facets v b -> Facets v c -> Facets v (a,b,c)
312 {-# INLINE zip3 #-}
313 zip3 = M.zip3
314
315 zip4 :: Facets v a -> Facets v b -> Facets v c -> Facets v d
316 -> Facets v (a,b,c,d)
317 {-# INLINE zip4 #-}
318 zip4 = M.zip4
319
320 zip5 :: Facets v a -> Facets v b -> Facets v c -> Facets v d
321 -> Facets v e -> Facets v (a,b,c,d,e)
322 {-# INLINE zip5 #-}
323 zip5 = M.zip5
324
325 zip6 :: Facets v a -> Facets v b -> Facets v c -> Facets v d
326 -> Facets v e -> Facets v f -> Facets v (a,b,c,d,e,f)
327 {-# INLINE zip6 #-}
328 zip6 = M.zip6
329
330 -- Filtering
331 -- ---------
332
333 -- | Drop elements which do not satisfy the predicate
334 filter :: (a -> Bool) -> Facets v a -> Facets v a
335 {-# INLINE filter #-}
336 filter = M.filter
337
338 -- | Longest prefix of elements that satisfy the predicate
339 takeWhile :: (a -> Bool) -> Facets v a -> Facets v a
340 {-# INLINE takeWhile #-}
341 takeWhile = M.takeWhile
342
343 -- | Drop the longest prefix of elements that satisfy the predicate
344 dropWhile :: (a -> Bool) -> Facets v a -> Facets v a
345 {-# INLINE dropWhile #-}
346 dropWhile = M.dropWhile
347
348 -- Searching
349 -- ---------
350
351 infix 4 `elem`
352 -- | Check whether the 'Facets' contains an element
353 elem :: Eq a => a -> Facets v a -> Bool
354 {-# INLINE elem #-}
355 elem x = unId . M.elem x
356
357 infix 4 `notElem`
358 -- | Inverse of `elem`
359 notElem :: Eq a => a -> Facets v a -> Bool
360 {-# INLINE notElem #-}
361 notElem x = unId . M.notElem x
362
363 -- | Yield 'Just' the first element matching the predicate or 'Nothing' if no
364 -- such element exists.
365 find :: (a -> Bool) -> Facets v a -> Maybe a
366 {-# INLINE find #-}
367 find f = unId . M.find f
368
369 -- | Yield 'Just' the index of the first element matching the predicate or
370 -- 'Nothing' if no such element exists.
371 findIndex :: (a -> Bool) -> Facets v a -> Maybe Int
372 {-# INLINE findIndex #-}
373 findIndex f = unId . M.findIndex f
374
375 -- Folding
376 -- -------
377
378 -- | Left fold
379 foldl :: (a -> b -> a) -> a -> Facets v b -> a
380 {-# INLINE foldl #-}
381 foldl f z = unId . M.foldl f z
382
383 -- | Left fold on non-empty 'Facets's
384 foldl1 :: (a -> a -> a) -> Facets v a -> a
385 {-# INLINE foldl1 #-}
386 foldl1 f = unId . M.foldl1 f
387
388 -- | Left fold with strict accumulator
389 foldl' :: (a -> b -> a) -> a -> Facets v b -> a
390 {-# INLINE foldl' #-}
391 foldl' f z = unId . M.foldl' f z
392
393 -- | Left fold on non-empty 'Facets's with strict accumulator
394 foldl1' :: (a -> a -> a) -> Facets v a -> a
395 {-# INLINE foldl1' #-}
396 foldl1' f = unId . M.foldl1' f
397
398 -- | Right fold
399 foldr :: (a -> b -> b) -> b -> Facets v a -> b
400 {-# INLINE foldr #-}
401 foldr f z = unId . M.foldr f z
402
403 -- | Right fold on non-empty 'Facets's
404 foldr1 :: (a -> a -> a) -> Facets v a -> a
405 {-# INLINE foldr1 #-}
406 foldr1 f = unId . M.foldr1 f
407
408 -- Specialised folds
409 -- -----------------
410
411 and :: Facets v Bool -> Bool
412 {-# INLINE and #-}
413 and = unId . M.and
414
415 or :: Facets v Bool -> Bool
416 {-# INLINE or #-}
417 or = unId . M.or
418
419 -- Unfolding
420 -- ---------
421
422 -- | Unfold
423 unfoldr :: (s -> Maybe (a, s)) -> s -> Facets v a
424 {-# INLINE unfoldr #-}
425 unfoldr = M.unfoldr
426
427 -- | Unfold at most @n@ elements
428 unfoldrN :: Int -> (s -> Maybe (a, s)) -> s -> Facets v a
429 {-# INLINE unfoldrN #-}
430 unfoldrN = M.unfoldrN
431
432 -- | Apply function n-1 times to value. Zeroth element is original value.
433 iterateN :: Int -> (a -> a) -> a -> Facets v a
434 {-# INLINE iterateN #-}
435 iterateN = M.iterateN
436
437 -- Scans
438 -- -----
439
440 -- | Prefix scan
441 prescanl :: (a -> b -> a) -> a -> Facets v b -> Facets v a
442 {-# INLINE prescanl #-}
443 prescanl = M.prescanl
444
445 -- | Prefix scan with strict accumulator
446 prescanl' :: (a -> b -> a) -> a -> Facets v b -> Facets v a
447 {-# INLINE prescanl' #-}
448 prescanl' = M.prescanl'
449
450 -- | Suffix scan
451 postscanl :: (a -> b -> a) -> a -> Facets v b -> Facets v a
452 {-# INLINE postscanl #-}
453 postscanl = M.postscanl
454
455 -- | Suffix scan with strict accumulator
456 postscanl' :: (a -> b -> a) -> a -> Facets v b -> Facets v a
457 {-# INLINE postscanl' #-}
458 postscanl' = M.postscanl'
459
460 -- | Haskell-style scan
461 scanl :: (a -> b -> a) -> a -> Facets v b -> Facets v a
462 {-# INLINE scanl #-}
463 scanl = M.scanl
464
465 -- | Haskell-style scan with strict accumulator
466 scanl' :: (a -> b -> a) -> a -> Facets v b -> Facets v a
467 {-# INLINE scanl' #-}
468 scanl' = M.scanl'
469
470 -- | Scan over a non-empty 'Facets'
471 scanl1 :: (a -> a -> a) -> Facets v a -> Facets v a
472 {-# INLINE scanl1 #-}
473 scanl1 = M.scanl1
474
475 -- | Scan over a non-empty 'Facets' with a strict accumulator
476 scanl1' :: (a -> a -> a) -> Facets v a -> Facets v a
477 {-# INLINE scanl1' #-}
478 scanl1' = M.scanl1'
479
480
481 -- Comparisons
482 -- -----------
483
484 -- | Check if two 'Facets's are equal
485 eq :: Eq a => Facets v a -> Facets v a -> Bool
486 {-# INLINE eq #-}
487 eq x y = unId (M.eq x y)
488
489 -- | Lexicographically compare two 'Facets's
490 cmp :: Ord a => Facets v a -> Facets v a -> Ordering
491 {-# INLINE cmp #-}
492 cmp x y = unId (M.cmp x y)
493
494 instance Eq a => Eq (M.Facets Id v a) where
495 {-# INLINE (==) #-}
496 (==) = eq
497
498 instance Ord a => Ord (M.Facets Id v a) where
499 {-# INLINE compare #-}
500 compare = cmp
501
502 -- Monadic combinators
503 -- -------------------
504
505 -- | Apply a monadic action to each element of the stream, producing a monadic
506 -- stream of results
507 mapM :: Monad m => (a -> m b) -> Facets v a -> M.Facets m v b
508 {-# INLINE mapM #-}
509 mapM f = M.mapM f . liftStream
510
511 -- | Apply a monadic action to each element of the stream
512 mapM_ :: Monad m => (a -> m b) -> Facets v a -> m ()
513 {-# INLINE mapM_ #-}
514 mapM_ f = M.mapM_ f . liftStream
515
516 zipWithM :: Monad m => (a -> b -> m c) -> Facets v a -> Facets v b -> M.Facets m v c
517 {-# INLINE zipWithM #-}
518 zipWithM f as bs = M.zipWithM f (liftStream as) (liftStream bs)
519
520 zipWithM_ :: Monad m => (a -> b -> m c) -> Facets v a -> Facets v b -> m ()
521 {-# INLINE zipWithM_ #-}
522 zipWithM_ f as bs = M.zipWithM_ f (liftStream as) (liftStream bs)
523
524 -- | Yield a monadic stream of elements that satisfy the monadic predicate
525 filterM :: Monad m => (a -> m Bool) -> Facets v a -> M.Facets m v a
526 {-# INLINE filterM #-}
527 filterM f = M.filterM f . liftStream
528
529 -- | Monadic fold
530 foldM :: Monad m => (a -> b -> m a) -> a -> Facets v b -> m a
531 {-# INLINE foldM #-}
532 foldM m z = M.foldM m z . liftStream
533
534 -- | Monadic fold over non-empty stream
535 fold1M :: Monad m => (a -> a -> m a) -> Facets v a -> m a
536 {-# INLINE fold1M #-}
537 fold1M m = M.fold1M m . liftStream
538
539 -- | Monadic fold with strict accumulator
540 foldM' :: Monad m => (a -> b -> m a) -> a -> Facets v b -> m a
541 {-# INLINE foldM' #-}
542 foldM' m z = M.foldM' m z . liftStream
543
544 -- | Monad fold over non-empty stream with strict accumulator
545 fold1M' :: Monad m => (a -> a -> m a) -> Facets v a -> m a
546 {-# INLINE fold1M' #-}
547 fold1M' m = M.fold1M' m . liftStream
548
549 -- Enumerations
550 -- ------------
551
552 -- | Yield a 'Facets' of the given length containing the values @x@, @x+y@,
553 -- @x+y+y@ etc.
554 enumFromStepN :: Num a => a -> a -> Int -> Facets v a
555 {-# INLINE enumFromStepN #-}
556 enumFromStepN = M.enumFromStepN
557
558 -- | Enumerate values
559 --
560 -- /WARNING:/ This operations can be very inefficient. If at all possible, use
561 -- 'enumFromStepN' instead.
562 enumFromTo :: Enum a => a -> a -> Facets v a
563 {-# INLINE enumFromTo #-}
564 enumFromTo = M.enumFromTo
565
566 -- | Enumerate values with a given step.
567 --
568 -- /WARNING:/ This operations is very inefficient. If at all possible, use
569 -- 'enumFromStepN' instead.
570 enumFromThenTo :: Enum a => a -> a -> a -> Facets v a
571 {-# INLINE enumFromThenTo #-}
572 enumFromThenTo = M.enumFromThenTo
573
574 -- Conversions
575 -- -----------
576
577 -- | Convert a 'Facets' to a list
578 toList :: Facets v a -> [a]
579 {-# INLINE toList #-}
580 -- toList s = unId (M.toList s)
581 toList s = build (\c n -> toListFB c n s)
582
583 -- This supports foldr/build list fusion that GHC implements
584 toListFB :: (a -> b -> b) -> b -> Facets v a -> b
585 {-# INLINE [0] toListFB #-}
586 toListFB c n M.Facets{M.sElems = M.Unf step s} = go s
587 where
588 go s = case unId (step s) of
589 Yield x s' -> x `c` go s'
590 Skip s' -> go s'
591 Done -> n
592
593 -- | Create a 'Facets' from a list
594 fromList :: [a] -> Facets v a
595 {-# INLINE fromList #-}
596 fromList = M.fromList
597
598 -- | Create a 'Facets' from the first @n@ elements of a list
599 --
600 -- > fromListN n xs = fromList (take n xs)
601 fromListN :: Int -> [a] -> Facets v a
602 {-# INLINE fromListN #-}
603 fromListN = M.fromListN
604
605 unsafeFromList :: Size -> [a] -> Facets v a
606 {-# INLINE unsafeFromList #-}
607 unsafeFromList = M.unsafeFromList
608
609 fromVector :: Vector v a => v a -> Facets v a
610 {-# INLINE fromVector #-}
611 fromVector = M.fromVector
612
613 reVector :: Facets u a -> Facets v a
614 {-# INLINE reVector #-}
615 reVector = M.reVector
616
617 fromVectors :: Vector v a => [v a] -> Facets v a
618 {-# INLINE fromVectors #-}
619 fromVectors = M.fromVectors
620
621 concatVectors :: Vector v a => Facets u (v a) -> Facets v a
622 {-# INLINE concatVectors #-}
623 concatVectors = M.concatVectors
624
625 -- | Create a 'Facets' of values from a 'Facets' of streamable things
626 flatten :: (a -> s) -> (s -> Step s b) -> Size -> Facets v a -> Facets v b
627 {-# INLINE_STREAM flatten #-}
628 flatten mk istep sz = M.flatten (return . mk) (return . istep) sz . liftStream
629