zip4-zip6 & friends for Streams
[darcs-mirrors/vector.git] / Data / Vector / Fusion / Stream / Monadic.hs
1 {-# LANGUAGE ExistentialQuantification, Rank2Types #-}
2
3 -- |
4 -- Module : Data.Vector.Fusion.Stream.Monadic
5 -- Copyright : (c) Roman Leshchinskiy 2008-2009
6 -- License : BSD-style
7 --
8 -- Maintainer : Roman Leshchinskiy <rl@cse.unsw.edu.au>
9 -- Stability : experimental
10 -- Portability : non-portable
11 --
12 -- Monadic streams
13 --
14
15 module Data.Vector.Fusion.Stream.Monadic (
16 Stream(..), Step(..),
17
18 -- * Size hints
19 size, sized,
20
21 -- * Length
22 length, null,
23
24 -- * Construction
25 empty, singleton, cons, snoc, replicate, (++),
26
27 -- * Accessing elements
28 head, last, (!!),
29
30 -- * Substreams
31 extract, init, tail, take, drop,
32
33 -- * Mapping
34 map, mapM, mapM_, trans, unbox, concatMap,
35
36 -- * Zipping
37 zipWithM, zipWith3M, zipWith4M, zipWith5M, zipWith6M,
38 zipWith, zipWith3, zipWith4, zipWith5, zipWith6,
39 zip, zip3, zip4, zip5, zip6,
40
41 -- * Filtering
42 filter, filterM, takeWhile, takeWhileM, dropWhile, dropWhileM,
43
44 -- * Searching
45 elem, notElem, find, findM, findIndex, findIndexM,
46
47 -- * Folding
48 foldl, foldlM, foldl1, foldl1M, foldM, fold1M,
49 foldl', foldlM', foldl1', foldl1M', foldM', fold1M',
50 foldr, foldrM, foldr1, foldr1M,
51
52 -- * Specialised folds
53 and, or, concatMapM,
54
55 -- * Unfolding
56 unfoldr, unfoldrM,
57
58 -- * Scans
59 prescanl, prescanlM, prescanl', prescanlM',
60 postscanl, postscanlM, postscanl', postscanlM',
61 scanl, scanlM, scanl', scanlM',
62 scanl1, scanl1M, scanl1', scanl1M',
63
64 -- * Enumerations
65 enumFromTo, enumFromThenTo,
66
67 -- * Conversions
68 toList, fromList
69 ) where
70
71 import Data.Vector.Fusion.Stream.Size
72 import Data.Vector.Fusion.Util ( Box(..) )
73
74 import Control.Monad ( liftM )
75 import Prelude hiding ( length, null,
76 replicate, (++),
77 head, last, (!!),
78 init, tail, take, drop,
79 map, mapM, mapM_, concatMap,
80 zipWith, zipWith3, zip, zip3,
81 filter, takeWhile, dropWhile,
82 elem, notElem,
83 foldl, foldl1, foldr, foldr1,
84 and, or,
85 scanl, scanl1,
86 enumFromTo, enumFromThenTo )
87 import qualified Prelude
88
89 import Data.Int ( Int8, Int16, Int32, Int64 )
90 import Data.Word ( Word8, Word16, Word32, Word, Word64 )
91
92 #if __GLASGOW_HASKELL__ >= 613
93 import SpecConstr ( SpecConstrAnnotation(..) )
94 #endif
95
96 #include "vector.h"
97
98 data SPEC = SPEC | SPEC2
99 #if __GLASGOW_HASKELL__ >= 613
100 {-# ANN type SPEC ForceSpecConstr #-}
101 #endif
102
103
104 -- | Result of taking a single step in a stream
105 data Step s a = Yield a s -- ^ a new element and a new seed
106 | Skip s -- ^ just a new seed
107 | Done -- ^ end of stream
108
109 -- | Monadic streams
110 data Stream m a = forall s. Stream (s -> m (Step s a)) s Size
111
112 -- | 'Size' hint of a 'Stream'
113 size :: Stream m a -> Size
114 {-# INLINE size #-}
115 size (Stream _ _ sz) = sz
116
117 -- | Attach a 'Size' hint to a 'Stream'
118 sized :: Stream m a -> Size -> Stream m a
119 {-# INLINE_STREAM sized #-}
120 sized (Stream step s _) sz = Stream step s sz
121
122 -- Length
123 -- ------
124
125 -- | Length of a 'Stream'
126 length :: Monad m => Stream m a -> m Int
127 {-# INLINE_STREAM length #-}
128 length s = foldl' (\n _ -> n+1) 0 s
129
130 -- | Check if a 'Stream' is empty
131 null :: Monad m => Stream m a -> m Bool
132 {-# INLINE_STREAM null #-}
133 null s = foldr (\_ _ -> False) True s
134
135
136 -- Construction
137 -- ------------
138
139 -- | Empty 'Stream'
140 empty :: Monad m => Stream m a
141 {-# INLINE_STREAM empty #-}
142 empty = Stream (const (return Done)) () (Exact 0)
143
144 -- | Singleton 'Stream'
145 singleton :: Monad m => a -> Stream m a
146 {-# INLINE_STREAM singleton #-}
147 singleton x = Stream (return . step) True (Exact 1)
148 where
149 {-# INLINE_INNER step #-}
150 step True = Yield x False
151 step False = Done
152
153 -- | Replicate a value to a given length
154 replicate :: Monad m => Int -> a -> Stream m a
155 {-# INLINE_STREAM replicate #-}
156 replicate n x = Stream (return . step) n (Exact (max n 0))
157 where
158 {-# INLINE_INNER step #-}
159 step i | i > 0 = Yield x (i-1)
160 | otherwise = Done
161
162 -- | Prepend an element
163 cons :: Monad m => a -> Stream m a -> Stream m a
164 {-# INLINE cons #-}
165 cons x s = singleton x ++ s
166
167 -- | Append an element
168 snoc :: Monad m => Stream m a -> a -> Stream m a
169 {-# INLINE snoc #-}
170 snoc s x = s ++ singleton x
171
172 infixr 5 ++
173 -- | Concatenate two 'Stream's
174 (++) :: Monad m => Stream m a -> Stream m a -> Stream m a
175 {-# INLINE_STREAM (++) #-}
176 Stream stepa sa na ++ Stream stepb sb nb = Stream step (Left sa) (na + nb)
177 where
178 {-# INLINE_INNER step #-}
179 step (Left sa) = do
180 r <- stepa sa
181 case r of
182 Yield x sa' -> return $ Yield x (Left sa')
183 Skip sa' -> return $ Skip (Left sa')
184 Done -> return $ Skip (Right sb)
185 step (Right sb) = do
186 r <- stepb sb
187 case r of
188 Yield x sb' -> return $ Yield x (Right sb')
189 Skip sb' -> return $ Skip (Right sb')
190 Done -> return $ Done
191
192 -- Accessing elements
193 -- ------------------
194
195 -- | First element of the 'Stream' or error if empty
196 head :: Monad m => Stream m a -> m a
197 {-# INLINE_STREAM head #-}
198 head (Stream step s _) = head_loop SPEC s
199 where
200 head_loop SPEC s
201 = do
202 r <- step s
203 case r of
204 Yield x _ -> return x
205 Skip s' -> head_loop SPEC s'
206 Done -> BOUNDS_ERROR(emptyStream) "head"
207
208
209
210 -- | Last element of the 'Stream' or error if empty
211 last :: Monad m => Stream m a -> m a
212 {-# INLINE_STREAM last #-}
213 last (Stream step s _) = last_loop0 SPEC s
214 where
215 last_loop0 SPEC s
216 = do
217 r <- step s
218 case r of
219 Yield x s' -> last_loop1 SPEC x s'
220 Skip s' -> last_loop0 SPEC s'
221 Done -> BOUNDS_ERROR(emptyStream) "last"
222
223 last_loop1 SPEC x s
224 = do
225 r <- step s
226 case r of
227 Yield y s' -> last_loop1 SPEC y s'
228 Skip s' -> last_loop1 SPEC x s'
229 Done -> return x
230
231 -- | Element at the given position
232 (!!) :: Monad m => Stream m a -> Int -> m a
233 {-# INLINE (!!) #-}
234 Stream step s _ !! i | i < 0 = BOUNDS_ERROR(error) "!!" "negative index"
235 | otherwise = index_loop SPEC s i
236 where
237 index_loop SPEC s i
238 = i `seq`
239 do
240 r <- step s
241 case r of
242 Yield x s' | i == 0 -> return x
243 | otherwise -> index_loop SPEC s' (i-1)
244 Skip s' -> index_loop SPEC s' i
245 Done -> BOUNDS_ERROR(emptyStream) "!!"
246
247 -- Substreams
248 -- ----------
249
250 -- | Extract a substream of the given length starting at the given position.
251 extract :: Monad m => Stream m a -> Int -- ^ starting index
252 -> Int -- ^ length
253 -> Stream m a
254 {-# INLINE extract #-}
255 extract s i n = take n (drop i s)
256
257 -- | All but the last element
258 init :: Monad m => Stream m a -> Stream m a
259 {-# INLINE_STREAM init #-}
260 init (Stream step s sz) = Stream step' (Nothing, s) (sz - 1)
261 where
262 {-# INLINE_INNER step' #-}
263 step' (Nothing, s) = liftM (\r ->
264 case r of
265 Yield x s' -> Skip (Just x, s')
266 Skip s' -> Skip (Nothing, s')
267 Done -> BOUNDS_ERROR(emptyStream) "init"
268 ) (step s)
269
270 step' (Just x, s) = liftM (\r ->
271 case r of
272 Yield y s' -> Yield x (Just y, s')
273 Skip s' -> Skip (Just x, s')
274 Done -> Done
275 ) (step s)
276
277 -- | All but the first element
278 tail :: Monad m => Stream m a -> Stream m a
279 {-# INLINE_STREAM tail #-}
280 tail (Stream step s sz) = Stream step' (Left s) (sz - 1)
281 where
282 {-# INLINE_INNER step' #-}
283 step' (Left s) = liftM (\r ->
284 case r of
285 Yield x s' -> Skip (Right s')
286 Skip s' -> Skip (Left s')
287 Done -> BOUNDS_ERROR(emptyStream) "tail"
288 ) (step s)
289
290 step' (Right s) = liftM (\r ->
291 case r of
292 Yield x s' -> Yield x (Right s')
293 Skip s' -> Skip (Right s')
294 Done -> Done
295 ) (step s)
296
297 -- | The first @n@ elements
298 take :: Monad m => Int -> Stream m a -> Stream m a
299 {-# INLINE_STREAM take #-}
300 take n (Stream step s sz) = Stream step' (s, 0) (smaller (Exact n) sz)
301 where
302 {-# INLINE_INNER step' #-}
303 step' (s, i) | i < n = liftM (\r ->
304 case r of
305 Yield x s' -> Yield x (s', i+1)
306 Skip s' -> Skip (s', i)
307 Done -> Done
308 ) (step s)
309 step' (s, i) = return Done
310
311 -- | All but the first @n@ elements
312 drop :: Monad m => Int -> Stream m a -> Stream m a
313 {-# INLINE_STREAM drop #-}
314 drop n (Stream step s sz) = Stream step' (s, Just n) (sz - Exact n)
315 where
316 {-# INLINE_INNER step' #-}
317 step' (s, Just i) | i > 0 = liftM (\r ->
318 case r of
319 Yield x s' -> Skip (s', Just (i-1))
320 Skip s' -> Skip (s', Just i)
321 Done -> Done
322 ) (step s)
323 | otherwise = return $ Skip (s, Nothing)
324
325 step' (s, Nothing) = liftM (\r ->
326 case r of
327 Yield x s' -> Yield x (s', Nothing)
328 Skip s' -> Skip (s', Nothing)
329 Done -> Done
330 ) (step s)
331
332
333 -- Mapping
334 -- -------
335
336 instance Monad m => Functor (Stream m) where
337 {-# INLINE fmap #-}
338 fmap = map
339
340 -- | Map a function over a 'Stream'
341 map :: Monad m => (a -> b) -> Stream m a -> Stream m b
342 {-# INLINE map #-}
343 map f = mapM (return . f)
344
345 -- | Map a monadic function over a 'Stream'
346 mapM :: Monad m => (a -> m b) -> Stream m a -> Stream m b
347 {-# INLINE_STREAM mapM #-}
348 mapM f (Stream step s n) = Stream step' s n
349 where
350 {-# INLINE_INNER step' #-}
351 step' s = do
352 r <- step s
353 case r of
354 Yield x s' -> liftM (`Yield` s') (f x)
355 Skip s' -> return (Skip s')
356 Done -> return Done
357
358 -- | Execute a monadic action for each element of the 'Stream'
359 mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()
360 {-# INLINE_STREAM mapM_ #-}
361 mapM_ m (Stream step s _) = mapM_loop SPEC s
362 where
363 mapM_loop SPEC s
364 = do
365 r <- step s
366 case r of
367 Yield x s' -> do { m x; mapM_loop SPEC s' }
368 Skip s' -> mapM_loop SPEC s'
369 Done -> return ()
370
371 -- | Transform a 'Stream' to use a different monad
372 trans :: (Monad m, Monad m') => (forall a. m a -> m' a)
373 -> Stream m a -> Stream m' a
374 {-# INLINE_STREAM trans #-}
375 trans f (Stream step s n) = Stream (f . step) s n
376
377 unbox :: Monad m => Stream m (Box a) -> Stream m a
378 {-# INLINE_STREAM unbox #-}
379 unbox (Stream step s n) = Stream step' s n
380 where
381 {-# INLINE_INNER step' #-}
382 step' s = do
383 r <- step s
384 case r of
385 Yield (Box x) s' -> return $ Yield x s'
386 Skip s' -> return $ Skip s'
387 Done -> return $ Done
388
389 -- Zipping
390 -- -------
391
392 -- | Zip two 'Stream's with the given monadic function
393 zipWithM :: Monad m => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
394 {-# INLINE_STREAM zipWithM #-}
395 zipWithM f (Stream stepa sa na) (Stream stepb sb nb)
396 = Stream step (sa, sb, Nothing) (smaller na nb)
397 where
398 {-# INLINE_INNER step #-}
399 step (sa, sb, Nothing) = liftM (\r ->
400 case r of
401 Yield x sa' -> Skip (sa', sb, Just x)
402 Skip sa' -> Skip (sa', sb, Nothing)
403 Done -> Done
404 ) (stepa sa)
405
406 step (sa, sb, Just x) = do
407 r <- stepb sb
408 case r of
409 Yield y sb' ->
410 do
411 z <- f x y
412 return $ Yield z (sa, sb', Nothing)
413 Skip sb' -> return $ Skip (sa, sb', Just x)
414 Done -> return $ Done
415
416 -- FIXME: This might expose an opportunity for inplace execution.
417 {-# RULES
418
419 "zipWithM xs xs [Vector.Stream]" forall f xs.
420 zipWithM f xs xs = mapM (\x -> f x x) xs
421
422 #-}
423
424 zipWith3M :: Monad m => (a -> b -> c -> m d) -> Stream m a -> Stream m b -> Stream m c -> Stream m d
425 {-# INLINE_STREAM zipWith3M #-}
426 zipWith3M f (Stream stepa sa na) (Stream stepb sb nb) (Stream stepc sc nc)
427 = Stream step (sa, sb, sc, Nothing) (smaller na (smaller nb nc))
428 where
429 {-# INLINE_INNER step #-}
430 step (sa, sb, sc, Nothing) = do
431 r <- stepa sa
432 return $ case r of
433 Yield x sa' -> Skip (sa', sb, sc, Just (x, Nothing))
434 Skip sa' -> Skip (sa', sb, sc, Nothing)
435 Done -> Done
436
437 step (sa, sb, sc, Just (x, Nothing)) = do
438 r <- stepb sb
439 return $ case r of
440 Yield y sb' -> Skip (sa, sb', sc, Just (x, Just y))
441 Skip sb' -> Skip (sa, sb', sc, Just (x, Nothing))
442 Done -> Done
443
444 step (sa, sb, sc, Just (x, Just y)) = do
445 r <- stepc sc
446 case r of
447 Yield z sc' -> f x y z >>= (\res -> return $ Yield res (sa, sb, sc', Nothing))
448 Skip sc' -> return $ Skip (sa, sb, sc', Just (x, Just y))
449 Done -> return $ Done
450
451 zipWith4M :: Monad m => (a -> b -> c -> d -> m e)
452 -> Stream m a -> Stream m b -> Stream m c -> Stream m d
453 -> Stream m e
454 {-# INLINE zipWith4M #-}
455 zipWith4M f sa sb sc sd
456 = zipWithM (\(a,b) (c,d) -> f a b c d) (zip sa sb) (zip sc sd)
457
458 zipWith5M :: Monad m => (a -> b -> c -> d -> e -> m f)
459 -> Stream m a -> Stream m b -> Stream m c -> Stream m d
460 -> Stream m e -> Stream m f
461 {-# INLINE zipWith5M #-}
462 zipWith5M f sa sb sc sd se
463 = zipWithM (\(a,b,c) (d,e) -> f a b c d e) (zip3 sa sb sc) (zip sd se)
464
465 zipWith6M :: Monad m => (a -> b -> c -> d -> e -> f -> m g)
466 -> Stream m a -> Stream m b -> Stream m c -> Stream m d
467 -> Stream m e -> Stream m f -> Stream m g
468 {-# INLINE zipWith6M #-}
469 zipWith6M fn sa sb sc sd se sf
470 = zipWithM (\(a,b,c) (d,e,f) -> fn a b c d e f) (zip3 sa sb sc)
471 (zip3 sd se sf)
472
473 zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
474 {-# INLINE zipWith #-}
475 zipWith f = zipWithM (\a b -> return (f a b))
476
477 zipWith3 :: Monad m => (a -> b -> c -> d)
478 -> Stream m a -> Stream m b -> Stream m c -> Stream m d
479 {-# INLINE zipWith3 #-}
480 zipWith3 f = zipWith3M (\a b c -> return (f a b c))
481
482 zipWith4 :: Monad m => (a -> b -> c -> d -> e)
483 -> Stream m a -> Stream m b -> Stream m c -> Stream m d
484 -> Stream m e
485 {-# INLINE zipWith4 #-}
486 zipWith4 f = zipWith4M (\a b c d -> return (f a b c d))
487
488 zipWith5 :: Monad m => (a -> b -> c -> d -> e -> f)
489 -> Stream m a -> Stream m b -> Stream m c -> Stream m d
490 -> Stream m e -> Stream m f
491 {-# INLINE zipWith5 #-}
492 zipWith5 f = zipWith5M (\a b c d e -> return (f a b c d e))
493
494 zipWith6 :: Monad m => (a -> b -> c -> d -> e -> f -> g)
495 -> Stream m a -> Stream m b -> Stream m c -> Stream m d
496 -> Stream m e -> Stream m f -> Stream m g
497 {-# INLINE zipWith6 #-}
498 zipWith6 fn = zipWith6M (\a b c d e f -> return (fn a b c d e f))
499
500 zip :: Monad m => Stream m a -> Stream m b -> Stream m (a,b)
501 {-# INLINE zip #-}
502 zip = zipWith (,)
503
504 zip3 :: Monad m => Stream m a -> Stream m b -> Stream m c -> Stream m (a,b,c)
505 {-# INLINE zip3 #-}
506 zip3 = zipWith3 (,,)
507
508 zip4 :: Monad m => Stream m a -> Stream m b -> Stream m c -> Stream m d
509 -> Stream m (a,b,c,d)
510 {-# INLINE zip4 #-}
511 zip4 = zipWith4 (,,,)
512
513 zip5 :: Monad m => Stream m a -> Stream m b -> Stream m c -> Stream m d
514 -> Stream m e -> Stream m (a,b,c,d,e)
515 {-# INLINE zip5 #-}
516 zip5 = zipWith5 (,,,,)
517
518 zip6 :: Monad m => Stream m a -> Stream m b -> Stream m c -> Stream m d
519 -> Stream m e -> Stream m f -> Stream m (a,b,c,d,e,f)
520 {-# INLINE zip6 #-}
521 zip6 = zipWith6 (,,,,,)
522
523 -- Filtering
524 -- ---------
525
526 -- | Drop elements which do not satisfy the predicate
527 filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
528 {-# INLINE filter #-}
529 filter f = filterM (return . f)
530
531 -- | Drop elements which do not satisfy the monadic predicate
532 filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
533 {-# INLINE_STREAM filterM #-}
534 filterM f (Stream step s n) = Stream step' s (toMax n)
535 where
536 {-# INLINE_INNER step' #-}
537 step' s = do
538 r <- step s
539 case r of
540 Yield x s' -> do
541 b <- f x
542 return $ if b then Yield x s'
543 else Skip s'
544 Skip s' -> return $ Skip s'
545 Done -> return $ Done
546
547 -- | Longest prefix of elements that satisfy the predicate
548 takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
549 {-# INLINE takeWhile #-}
550 takeWhile f = takeWhileM (return . f)
551
552 -- | Longest prefix of elements that satisfy the monadic predicate
553 takeWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
554 {-# INLINE_STREAM takeWhileM #-}
555 takeWhileM f (Stream step s n) = Stream step' s (toMax n)
556 where
557 {-# INLINE_INNER step' #-}
558 step' s = do
559 r <- step s
560 case r of
561 Yield x s' -> do
562 b <- f x
563 return $ if b then Yield x s' else Done
564 Skip s' -> return $ Skip s'
565 Done -> return $ Done
566
567 -- | Drop the longest prefix of elements that satisfy the predicate
568 dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
569 {-# INLINE dropWhile #-}
570 dropWhile f = dropWhileM (return . f)
571
572 data DropWhile s a = DropWhile_Drop s | DropWhile_Yield a s | DropWhile_Next s
573
574 -- | Drop the longest prefix of elements that satisfy the monadic predicate
575 dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
576 {-# INLINE_STREAM dropWhileM #-}
577 dropWhileM f (Stream step s n) = Stream step' (DropWhile_Drop s) (toMax n)
578 where
579 -- NOTE: we jump through hoops here to have only one Yield; local data
580 -- declarations would be nice!
581
582 {-# INLINE_INNER step' #-}
583 step' (DropWhile_Drop s)
584 = do
585 r <- step s
586 case r of
587 Yield x s' -> do
588 b <- f x
589 return $ if b then Skip (DropWhile_Drop s')
590 else Skip (DropWhile_Yield x s')
591 Skip s' -> return $ Skip (DropWhile_Drop s')
592 Done -> return $ Done
593
594 step' (DropWhile_Yield x s) = return $ Yield x (DropWhile_Next s)
595
596 step' (DropWhile_Next s)
597 = liftM (\r ->
598 case r of
599 Yield x s' -> Skip (DropWhile_Yield x s')
600 Skip s' -> Skip (DropWhile_Next s')
601 Done -> Done
602 ) (step s)
603
604 -- Searching
605 -- ---------
606
607 infix 4 `elem`
608 -- | Check whether the 'Stream' contains an element
609 elem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
610 {-# INLINE_STREAM elem #-}
611 elem x (Stream step s _) = elem_loop SPEC s
612 where
613 elem_loop SPEC s
614 = do
615 r <- step s
616 case r of
617 Yield y s' | x == y -> return True
618 | otherwise -> elem_loop SPEC s'
619 Skip s' -> elem_loop SPEC s'
620 Done -> return False
621
622 infix 4 `notElem`
623 -- | Inverse of `elem`
624 notElem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
625 {-# INLINE notElem #-}
626 notElem x s = liftM not (elem x s)
627
628 -- | Yield 'Just' the first element that satisfies the predicate or 'Nothing'
629 -- if no such element exists.
630 find :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe a)
631 {-# INLINE find #-}
632 find f = findM (return . f)
633
634 -- | Yield 'Just' the first element that satisfies the monadic predicate or
635 -- 'Nothing' if no such element exists.
636 findM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe a)
637 {-# INLINE_STREAM findM #-}
638 findM f (Stream step s _) = find_loop SPEC s
639 where
640 find_loop SPEC s
641 = do
642 r <- step s
643 case r of
644 Yield x s' -> do
645 b <- f x
646 if b then return $ Just x
647 else find_loop SPEC s'
648 Skip s' -> find_loop SPEC s'
649 Done -> return Nothing
650
651 -- | Yield 'Just' the index of the first element that satisfies the predicate
652 -- or 'Nothing' if no such element exists.
653 findIndex :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe Int)
654 {-# INLINE_STREAM findIndex #-}
655 findIndex f = findIndexM (return . f)
656
657 -- | Yield 'Just' the index of the first element that satisfies the monadic
658 -- predicate or 'Nothing' if no such element exists.
659 findIndexM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe Int)
660 {-# INLINE_STREAM findIndexM #-}
661 findIndexM f (Stream step s _) = findIndex_loop SPEC s 0
662 where
663 findIndex_loop SPEC s i
664 = do
665 r <- step s
666 case r of
667 Yield x s' -> do
668 b <- f x
669 if b then return $ Just i
670 else findIndex_loop SPEC s' (i+1)
671 Skip s' -> findIndex_loop SPEC s' i
672 Done -> return Nothing
673
674 -- Folding
675 -- -------
676
677 -- | Left fold
678 foldl :: Monad m => (a -> b -> a) -> a -> Stream m b -> m a
679 {-# INLINE foldl #-}
680 foldl f = foldlM (\a b -> return (f a b))
681
682 -- | Left fold with a monadic operator
683 foldlM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> m a
684 {-# INLINE_STREAM foldlM #-}
685 foldlM m z (Stream step s _) = foldlM_loop SPEC z s
686 where
687 foldlM_loop SPEC z s
688 = do
689 r <- step s
690 case r of
691 Yield x s' -> do { z' <- m z x; foldlM_loop SPEC z' s' }
692 Skip s' -> foldlM_loop SPEC z s'
693 Done -> return z
694
695 -- | Same as 'foldlM'
696 foldM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> m a
697 {-# INLINE foldM #-}
698 foldM = foldlM
699
700 -- | Left fold over a non-empty 'Stream'
701 foldl1 :: Monad m => (a -> a -> a) -> Stream m a -> m a
702 {-# INLINE foldl1 #-}
703 foldl1 f = foldl1M (\a b -> return (f a b))
704
705 -- | Left fold over a non-empty 'Stream' with a monadic operator
706 foldl1M :: Monad m => (a -> a -> m a) -> Stream m a -> m a
707 {-# INLINE_STREAM foldl1M #-}
708 foldl1M f (Stream step s sz) = foldl1M_loop SPEC s
709 where
710 foldl1M_loop SPEC s
711 = do
712 r <- step s
713 case r of
714 Yield x s' -> foldlM f x (Stream step s' (sz - 1))
715 Skip s' -> foldl1M_loop SPEC s'
716 Done -> BOUNDS_ERROR(emptyStream) "foldl1M"
717
718 -- | Same as 'foldl1M'
719 fold1M :: Monad m => (a -> a -> m a) -> Stream m a -> m a
720 {-# INLINE fold1M #-}
721 fold1M = foldl1M
722
723 -- | Left fold with a strict accumulator
724 foldl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> m a
725 {-# INLINE foldl' #-}
726 foldl' f = foldlM' (\a b -> return (f a b))
727
728 -- | Left fold with a strict accumulator and a monadic operator
729 foldlM' :: Monad m => (a -> b -> m a) -> a -> Stream m b -> m a
730 {-# INLINE_STREAM foldlM' #-}
731 foldlM' m z (Stream step s _) = foldlM'_loop SPEC z s
732 where
733 foldlM'_loop SPEC z s
734 = z `seq`
735 do
736 r <- step s
737 case r of
738 Yield x s' -> do { z' <- m z x; foldlM'_loop SPEC z' s' }
739 Skip s' -> foldlM'_loop SPEC z s'
740 Done -> return z
741
742 -- | Same as 'foldlM''
743 foldM' :: Monad m => (a -> b -> m a) -> a -> Stream m b -> m a
744 {-# INLINE foldM' #-}
745 foldM' = foldlM'
746
747 -- | Left fold over a non-empty 'Stream' with a strict accumulator
748 foldl1' :: Monad m => (a -> a -> a) -> Stream m a -> m a
749 {-# INLINE foldl1' #-}
750 foldl1' f = foldl1M' (\a b -> return (f a b))
751
752 -- | Left fold over a non-empty 'Stream' with a strict accumulator and a
753 -- monadic operator
754 foldl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> m a
755 {-# INLINE_STREAM foldl1M' #-}
756 foldl1M' f (Stream step s sz) = foldl1M'_loop SPEC s
757 where
758 foldl1M'_loop SPEC s
759 = do
760 r <- step s
761 case r of
762 Yield x s' -> foldlM' f x (Stream step s' (sz - 1))
763 Skip s' -> foldl1M'_loop SPEC s'
764 Done -> BOUNDS_ERROR(emptyStream) "foldl1M'"
765
766 -- | Same as 'foldl1M''
767 fold1M' :: Monad m => (a -> a -> m a) -> Stream m a -> m a
768 {-# INLINE fold1M' #-}
769 fold1M' = foldl1M'
770
771 -- | Right fold
772 foldr :: Monad m => (a -> b -> b) -> b -> Stream m a -> m b
773 {-# INLINE foldr #-}
774 foldr f = foldrM (\a b -> return (f a b))
775
776 -- | Right fold with a monadic operator
777 foldrM :: Monad m => (a -> b -> m b) -> b -> Stream m a -> m b
778 {-# INLINE_STREAM foldrM #-}
779 foldrM f z (Stream step s _) = foldrM_loop SPEC s
780 where
781 foldrM_loop SPEC s
782 = do
783 r <- step s
784 case r of
785 Yield x s' -> f x =<< foldrM_loop SPEC s'
786 Skip s' -> foldrM_loop SPEC s'
787 Done -> return z
788
789 -- | Right fold over a non-empty stream
790 foldr1 :: Monad m => (a -> a -> a) -> Stream m a -> m a
791 {-# INLINE foldr1 #-}
792 foldr1 f = foldr1M (\a b -> return (f a b))
793
794 -- | Right fold over a non-empty stream with a monadic operator
795 foldr1M :: Monad m => (a -> a -> m a) -> Stream m a -> m a
796 {-# INLINE_STREAM foldr1M #-}
797 foldr1M f (Stream step s _) = foldr1M_loop0 SPEC s
798 where
799 foldr1M_loop0 SPEC s
800 = do
801 r <- step s
802 case r of
803 Yield x s' -> foldr1M_loop1 SPEC x s'
804 Skip s' -> foldr1M_loop0 SPEC s'
805 Done -> BOUNDS_ERROR(emptyStream) "foldr1M"
806
807 foldr1M_loop1 SPEC x s
808 = do
809 r <- step s
810 case r of
811 Yield y s' -> f x =<< foldr1M_loop1 SPEC y s'
812 Skip s' -> foldr1M_loop1 SPEC x s'
813 Done -> return x
814
815 -- Specialised folds
816 -- -----------------
817
818 and :: Monad m => Stream m Bool -> m Bool
819 {-# INLINE_STREAM and #-}
820 and (Stream step s _) = and_loop SPEC s
821 where
822 and_loop SPEC s
823 = do
824 r <- step s
825 case r of
826 Yield False _ -> return False
827 Yield True s' -> and_loop SPEC s'
828 Skip s' -> and_loop SPEC s'
829 Done -> return True
830
831 or :: Monad m => Stream m Bool -> m Bool
832 {-# INLINE_STREAM or #-}
833 or (Stream step s _) = or_loop SPEC s
834 where
835 or_loop SPEC s
836 = do
837 r <- step s
838 case r of
839 Yield False s' -> or_loop SPEC s'
840 Yield True _ -> return True
841 Skip s' -> or_loop SPEC s'
842 Done -> return False
843
844 concatMap :: Monad m => (a -> Stream m b) -> Stream m a -> Stream m b
845 {-# INLINE concatMap #-}
846 concatMap f = concatMapM (return . f)
847
848 concatMapM :: Monad m => (a -> m (Stream m b)) -> Stream m a -> Stream m b
849 {-# INLINE_STREAM concatMapM #-}
850 concatMapM f (Stream step s _) = Stream concatMap_go (Left s) Unknown
851 where
852 concatMap_go (Left s) = do
853 r <- step s
854 case r of
855 Yield a s' -> do
856 b_stream <- f a
857 return $ Skip (Right (b_stream, s'))
858 Skip s' -> return $ Skip (Left s')
859 Done -> return Done
860 concatMap_go (Right (Stream inner_step inner_s sz, s)) = do
861 r <- inner_step inner_s
862 case r of
863 Yield b inner_s' -> return $ Yield b (Right (Stream inner_step inner_s' sz, s))
864 Skip inner_s' -> return $ Skip (Right (Stream inner_step inner_s' sz, s))
865 Done -> return $ Skip (Left s)
866
867 -- Unfolding
868 -- ---------
869
870 -- | Unfold
871 unfoldr :: Monad m => (s -> Maybe (a, s)) -> s -> Stream m a
872 {-# INLINE_STREAM unfoldr #-}
873 unfoldr f = unfoldrM (return . f)
874
875 -- | Unfold with a monadic function
876 unfoldrM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m a
877 {-# INLINE_STREAM unfoldrM #-}
878 unfoldrM f s = Stream step s Unknown
879 where
880 {-# INLINE_INNER step #-}
881 step s = liftM (\r ->
882 case r of
883 Just (x, s') -> Yield x s'
884 Nothing -> Done
885 ) (f s)
886
887 -- Scans
888 -- -----
889
890 -- | Prefix scan
891 prescanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
892 {-# INLINE prescanl #-}
893 prescanl f = prescanlM (\a b -> return (f a b))
894
895 -- | Prefix scan with a monadic operator
896 prescanlM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a
897 {-# INLINE_STREAM prescanlM #-}
898 prescanlM f z (Stream step s sz) = Stream step' (s,z) sz
899 where
900 {-# INLINE_INNER step' #-}
901 step' (s,x) = do
902 r <- step s
903 case r of
904 Yield y s' -> do
905 z <- f x y
906 return $ Yield x (s', z)
907 Skip s' -> return $ Skip (s', x)
908 Done -> return Done
909
910 -- | Prefix scan with strict accumulator
911 prescanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
912 {-# INLINE prescanl' #-}
913 prescanl' f = prescanlM' (\a b -> return (f a b))
914
915 -- | Prefix scan with strict accumulator and a monadic operator
916 prescanlM' :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a
917 {-# INLINE_STREAM prescanlM' #-}
918 prescanlM' f z (Stream step s sz) = Stream step' (s,z) sz
919 where
920 {-# INLINE_INNER step' #-}
921 step' (s,x) = x `seq`
922 do
923 r <- step s
924 case r of
925 Yield y s' -> do
926 z <- f x y
927 return $ Yield x (s', z)
928 Skip s' -> return $ Skip (s', x)
929 Done -> return Done
930
931 -- | Suffix scan
932 postscanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
933 {-# INLINE postscanl #-}
934 postscanl f = postscanlM (\a b -> return (f a b))
935
936 -- | Suffix scan with a monadic operator
937 postscanlM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a
938 {-# INLINE_STREAM postscanlM #-}
939 postscanlM f z (Stream step s sz) = Stream step' (s,z) sz
940 where
941 {-# INLINE_INNER step' #-}
942 step' (s,x) = do
943 r <- step s
944 case r of
945 Yield y s' -> do
946 z <- f x y
947 return $ Yield z (s',z)
948 Skip s' -> return $ Skip (s',x)
949 Done -> return Done
950
951 -- | Suffix scan with strict accumulator
952 postscanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
953 {-# INLINE postscanl' #-}
954 postscanl' f = postscanlM' (\a b -> return (f a b))
955
956 -- | Suffix scan with strict acccumulator and a monadic operator
957 postscanlM' :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a
958 {-# INLINE_STREAM postscanlM' #-}
959 postscanlM' f z (Stream step s sz) = z `seq` Stream step' (s,z) sz
960 where
961 {-# INLINE_INNER step' #-}
962 step' (s,x) = x `seq`
963 do
964 r <- step s
965 case r of
966 Yield y s' -> do
967 z <- f x y
968 z `seq` return (Yield z (s',z))
969 Skip s' -> return $ Skip (s',x)
970 Done -> return Done
971
972 -- | Haskell-style scan
973 scanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
974 {-# INLINE scanl #-}
975 scanl f = scanlM (\a b -> return (f a b))
976
977 -- | Haskell-style scan with a monadic operator
978 scanlM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a
979 {-# INLINE scanlM #-}
980 scanlM f z s = z `cons` postscanlM f z s
981
982 -- | Haskell-style scan with strict accumulator
983 scanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
984 {-# INLINE scanl' #-}
985 scanl' f = scanlM' (\a b -> return (f a b))
986
987 -- | Haskell-style scan with strict accumulator and a monadic operator
988 scanlM' :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a
989 {-# INLINE scanlM' #-}
990 scanlM' f z s = z `seq` (z `cons` postscanlM f z s)
991
992 -- | Scan over a non-empty 'Stream'
993 scanl1 :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
994 {-# INLINE scanl1 #-}
995 scanl1 f = scanl1M (\x y -> return (f x y))
996
997 -- | Scan over a non-empty 'Stream' with a monadic operator
998 scanl1M :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
999 {-# INLINE_STREAM scanl1M #-}
1000 scanl1M f (Stream step s sz) = Stream step' (s, Nothing) sz
1001 where
1002 {-# INLINE_INNER step' #-}
1003 step' (s, Nothing) = do
1004 r <- step s
1005 case r of
1006 Yield x s' -> return $ Yield x (s', Just x)
1007 Skip s' -> return $ Skip (s', Nothing)
1008 Done -> BOUNDS_ERROR(emptyStream) "scanl1M"
1009
1010 step' (s, Just x) = do
1011 r <- step s
1012 case r of
1013 Yield y s' -> do
1014 z <- f x y
1015 return $ Yield z (s', Just z)
1016 Skip s' -> return $ Skip (s', Just x)
1017 Done -> return Done
1018
1019 -- | Scan over a non-empty 'Stream' with a strict accumulator
1020 scanl1' :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
1021 {-# INLINE scanl1' #-}
1022 scanl1' f = scanl1M' (\x y -> return (f x y))
1023
1024 -- | Scan over a non-empty 'Stream' with a strict accumulator and a monadic
1025 -- operator
1026 scanl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
1027 {-# INLINE_STREAM scanl1M' #-}
1028 scanl1M' f (Stream step s sz) = Stream step' (s, Nothing) sz
1029 where
1030 {-# INLINE_INNER step' #-}
1031 step' (s, Nothing) = do
1032 r <- step s
1033 case r of
1034 Yield x s' -> x `seq` return (Yield x (s', Just x))
1035 Skip s' -> return $ Skip (s', Nothing)
1036 Done -> BOUNDS_ERROR(emptyStream) "scanl1M"
1037
1038 step' (s, Just x) = x `seq`
1039 do
1040 r <- step s
1041 case r of
1042 Yield y s' -> do
1043 z <- f x y
1044 z `seq` return (Yield z (s', Just z))
1045 Skip s' -> return $ Skip (s', Just x)
1046 Done -> return Done
1047
1048 -- Enumerations
1049 -- ------------
1050
1051 -- FIXME: The Enum class is broken for this, there just doesn't seem to be a
1052 -- way to implement this generically. We have to either specialise (which
1053 -- doesn't help in polymorphic loops) or define a new Enum-like class with a
1054 -- proper interface.
1055
1056 -- | Enumerate values from @x@ to @y@
1057 enumFromTo :: (Enum a, Monad m) => a -> a -> Stream m a
1058 {-# INLINE_STREAM enumFromTo #-}
1059 enumFromTo x y = fromList [x .. y]
1060
1061 enumFromTo_small :: (Enum a, Ord a, Monad m) => a -> a -> Stream m a
1062 {-# INLINE_STREAM enumFromTo_small #-}
1063 enumFromTo_small x y = Stream step x (Exact n)
1064 where
1065 n = max (fromEnum y - fromEnum x + 1) 0
1066
1067 {-# INLINE_INNER step #-}
1068 step x | x <= y = return $ Yield x (succ x)
1069 | otherwise = return $ Done
1070
1071 {-# RULES
1072
1073 "enumFromTo<Int> [Stream]"
1074 enumFromTo = enumFromTo_small :: Monad m => Int -> Int -> Stream m Int
1075
1076 "enumFromTo<Char> [Stream]"
1077 enumFromTo = enumFromTo_small :: Monad m => Char -> Char -> Stream m Char
1078
1079 "enumFromTo<Int8> [Stream]"
1080 enumFromTo = enumFromTo_small :: Monad m => Int8 -> Int8 -> Stream m Int8
1081
1082 "enumFromTo<Int16> [Stream]"
1083 enumFromTo = enumFromTo_small :: Monad m => Int16 -> Int16 -> Stream m Int16
1084
1085 "enumFromTo<Int32> [Stream]"
1086 enumFromTo = enumFromTo_small :: Monad m => Int32 -> Int32 -> Stream m Int32
1087
1088 "enumFromTo<Word8> [Stream]"
1089 enumFromTo = enumFromTo_small :: Monad m => Word8 -> Word8 -> Stream m Word8
1090
1091 "enumFromTo<Word16> [Stream]"
1092 enumFromTo = enumFromTo_small :: Monad m => Word16 -> Word16 -> Stream m Word16
1093
1094 #-}
1095
1096 enumFromTo_big :: (Enum a, Integral a, Monad m) => a -> a -> Stream m a
1097 {-# INLINE_STREAM enumFromTo_big #-}
1098 enumFromTo_big x y = Stream step x (Exact n)
1099 where
1100 n | x > y = 0
1101 | y - x < fromIntegral (maxBound :: Int) = fromIntegral (y-x+1)
1102 | otherwise = error $ "vector.enumFromTo_big: Array too large"
1103
1104 {-# INLINE_INNER step #-}
1105 step x | x <= y = return $ Yield x (succ x)
1106 | otherwise = return $ Done
1107
1108 {-# RULES
1109
1110
1111 "enumFromTo<Int64> [Stream]"
1112 enumFromTo = enumFromTo_big :: Monad m => Int64 -> Int64 -> Stream m Int64
1113
1114 "enumFromTo<Word32> [Stream]"
1115 enumFromTo = enumFromTo_big :: Monad m => Word32 -> Word32 -> Stream m Word32
1116
1117 "enumFromTo<Word64> [Stream]"
1118 enumFromTo = enumFromTo_big :: Monad m => Word64 -> Word64 -> Stream m Word64
1119
1120 "enumFromTo<Integer> [Stream]"
1121 enumFromTo = enumFromTo_big :: Monad m => Integer -> Integer -> Stream m Integer
1122
1123 #-}
1124
1125
1126 -- | Enumerate values from @x@ to @y@
1127 enumFromThenTo :: (Enum a, Monad m) => a -> a -> a -> Stream m a
1128 {-# INLINE_STREAM enumFromThenTo #-}
1129 enumFromThenTo x y z = fromList [x, y .. z]
1130
1131 -- Conversions
1132 -- -----------
1133
1134 -- | Convert a 'Stream' to a list
1135 toList :: Monad m => Stream m a -> m [a]
1136 {-# INLINE toList #-}
1137 toList = foldr (:) []
1138
1139 -- | Convert a list to a 'Stream'
1140 fromList :: Monad m => [a] -> Stream m a
1141 {-# INLINE_STREAM fromList #-}
1142 fromList xs = Stream step xs Unknown
1143 where
1144 step (x:xs) = return (Yield x xs)
1145 step [] = return Done
1146