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