Add various monadic folds over 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
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 #include "phases.h"
16
17 module Data.Vector.Fusion.Stream.Monadic (
18 Stream(..), Step(..),
19
20 -- * Size hints
21 size, sized,
22
23 -- * Length
24 length, null,
25
26 -- * Construction
27 empty, singleton, cons, snoc, replicate, (++),
28
29 -- * Accessing elements
30 head, last, (!!),
31
32 -- * Substreams
33 extract, init, tail, take, drop,
34
35 -- * Mapping
36 map, mapM, mapM_, trans, concatMap,
37
38 -- * Zipping
39 zipWith, zipWithM, zipWith3, zipWith3M,
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 -- * Conversions
65 toList, fromList
66 ) where
67
68 import Data.Vector.Fusion.Stream.Size
69
70 import Control.Monad ( liftM )
71 import Prelude hiding ( length, null,
72 replicate, (++),
73 head, last, (!!),
74 init, tail, take, drop,
75 map, mapM, mapM_, concatMap,
76 zipWith, zipWith3,
77 filter, takeWhile, dropWhile,
78 elem, notElem,
79 foldl, foldl1, foldr, foldr1,
80 and, or,
81 scanl, scanl1 )
82 import qualified Prelude
83
84 -- | Result of taking a single step in a stream
85 data Step s a = Yield a s -- ^ a new element and a new seed
86 | Skip s -- ^ just a new seed
87 | Done -- ^ end of stream
88
89 -- | Monadic streams
90 data Stream m a = forall s. Stream (s -> m (Step s a)) s Size
91
92 -- | 'Size' hint of a 'Stream'
93 size :: Stream m a -> Size
94 {-# INLINE size #-}
95 size (Stream _ _ sz) = sz
96
97 -- | Attach a 'Size' hint to a 'Stream'
98 sized :: Stream m a -> Size -> Stream m a
99 {-# INLINE_STREAM sized #-}
100 sized (Stream step s _) sz = Stream step s sz
101
102 -- Length
103 -- ------
104
105 -- | Length of a 'Stream'
106 length :: Monad m => Stream m a -> m Int
107 {-# INLINE_STREAM length #-}
108 length s = foldl' (\n _ -> n+1) 0 s
109
110 -- | Check if a 'Stream' is empty
111 null :: Monad m => Stream m a -> m Bool
112 {-# INLINE_STREAM null #-}
113 null s = foldr (\_ _ -> False) True s
114
115
116 -- Construction
117 -- ------------
118
119 -- | Empty 'Stream'
120 empty :: Monad m => Stream m a
121 {-# INLINE_STREAM empty #-}
122 empty = Stream (const (return Done)) () (Exact 0)
123
124 -- | Singleton 'Stream'
125 singleton :: Monad m => a -> Stream m a
126 {-# INLINE_STREAM singleton #-}
127 singleton x = Stream (return . step) True (Exact 1)
128 where
129 {-# INLINE step #-}
130 step True = Yield x False
131 step False = Done
132
133 -- | Replicate a value to a given length
134 replicate :: Monad m => Int -> a -> Stream m a
135 {-# INLINE_STREAM replicate #-}
136 replicate n x = Stream (return . step) n (Exact (max n 0))
137 where
138 {-# INLINE step #-}
139 step i | i > 0 = Yield x (i-1)
140 | otherwise = Done
141
142 -- | Prepend an element
143 cons :: Monad m => a -> Stream m a -> Stream m a
144 {-# INLINE cons #-}
145 cons x s = singleton x ++ s
146
147 -- | Append an element
148 snoc :: Monad m => Stream m a -> a -> Stream m a
149 {-# INLINE snoc #-}
150 snoc s x = s ++ singleton x
151
152 infixr 5 ++
153 -- | Concatenate two 'Stream's
154 (++) :: Monad m => Stream m a -> Stream m a -> Stream m a
155 {-# INLINE_STREAM (++) #-}
156 Stream stepa sa na ++ Stream stepb sb nb = Stream step (Left sa) (na + nb)
157 where
158 step (Left sa) = do
159 r <- stepa sa
160 case r of
161 Yield x sa' -> return $ Yield x (Left sa')
162 Skip sa' -> return $ Skip (Left sa')
163 Done -> return $ Skip (Right sb)
164 step (Right sb) = do
165 r <- stepb sb
166 case r of
167 Yield x sb' -> return $ Yield x (Right sb')
168 Skip sb' -> return $ Skip (Right sb')
169 Done -> return $ Done
170
171 -- Accessing elements
172 -- ------------------
173
174 -- | First element of the 'Stream' or error if empty
175 head :: Monad m => Stream m a -> m a
176 {-# INLINE_STREAM head #-}
177 head (Stream step s _) = head_loop s
178 where
179 head_loop s = do
180 r <- step s
181 case r of
182 Yield x _ -> return x
183 Skip s' -> head_loop s'
184 Done -> errorEmptyStream "head"
185
186 -- | Last element of the 'Stream' or error if empty
187 last :: Monad m => Stream m a -> m a
188 {-# INLINE_STREAM last #-}
189 last (Stream step s _) = last_loop0 s
190 where
191 last_loop0 s = do
192 r <- step s
193 case r of
194 Yield x s' -> last_loop1 x s'
195 Skip s' -> last_loop0 s'
196 Done -> errorEmptyStream "last"
197
198 last_loop1 x s = do
199 r <- step s
200 case r of
201 Yield y s' -> last_loop1 y s'
202 Skip s' -> last_loop1 x s'
203 Done -> return x
204
205 -- | Element at the given position
206 (!!) :: Monad m => Stream m a -> Int -> m a
207 {-# INLINE (!!) #-}
208 Stream step s _ !! i | i < 0 = errorNegativeIndex "!!"
209 | otherwise = loop s i
210 where
211 loop s i = i `seq`
212 do
213 r <- step s
214 case r of
215 Yield x s' | i == 0 -> return x
216 | otherwise -> loop s' (i-1)
217 Skip s' -> loop s' i
218 Done -> errorIndexOutOfRange "!!"
219
220 -- Substreams
221 -- ----------
222
223 -- | Extract a substream of the given length starting at the given position.
224 extract :: Monad m => Stream m a -> Int -- ^ starting index
225 -> Int -- ^ length
226 -> Stream m a
227 {-# INLINE extract #-}
228 extract s i n = take n (drop i s)
229
230 -- | All but the last element
231 init :: Monad m => Stream m a -> Stream m a
232 {-# INLINE_STREAM init #-}
233 init (Stream step s sz) = Stream step' (Nothing, s) (sz - 1)
234 where
235 {-# INLINE step' #-}
236 step' (Nothing, s) = liftM (\r ->
237 case r of
238 Yield x s' -> Skip (Just x, s')
239 Skip s' -> Skip (Nothing, s')
240 Done -> errorEmptyStream "init"
241 ) (step s)
242
243 step' (Just x, s) = liftM (\r ->
244 case r of
245 Yield y s' -> Yield x (Just y, s')
246 Skip s' -> Skip (Just x, s')
247 Done -> Done
248 ) (step s)
249
250 -- | All but the first element
251 tail :: Monad m => Stream m a -> Stream m a
252 {-# INLINE_STREAM tail #-}
253 tail (Stream step s sz) = Stream step' (Left s) (sz - 1)
254 where
255 {-# INLINE step' #-}
256 step' (Left s) = liftM (\r ->
257 case r of
258 Yield x s' -> Skip (Right s')
259 Skip s' -> Skip (Left s')
260 Done -> errorEmptyStream "tail"
261 ) (step s)
262
263 step' (Right s) = liftM (\r ->
264 case r of
265 Yield x s' -> Yield x (Right s')
266 Skip s' -> Skip (Right s')
267 Done -> Done
268 ) (step s)
269
270 -- | The first @n@ elements
271 take :: Monad m => Int -> Stream m a -> Stream m a
272 {-# INLINE_STREAM take #-}
273 take n (Stream step s sz) = Stream step' (s, 0) (smaller (Exact n) sz)
274 where
275 {-# INLINE step' #-}
276 step' (s, i) | i < n = liftM (\r ->
277 case r of
278 Yield x s' -> Yield x (s', i+1)
279 Skip s' -> Skip (s', i)
280 Done -> Done
281 ) (step s)
282 step' (s, i) = return Done
283
284 -- | All but the first @n@ elements
285 drop :: Monad m => Int -> Stream m a -> Stream m a
286 {-# INLINE_STREAM drop #-}
287 drop n (Stream step s sz) = Stream step' (s, Just n) (sz - Exact n)
288 where
289 {-# INLINE step' #-}
290 step' (s, Just i) | i > 0 = liftM (\r ->
291 case r of
292 Yield x s' -> Skip (s', Just (i-1))
293 Skip s' -> Skip (s', Just i)
294 Done -> Done
295 ) (step s)
296 | otherwise = return $ Skip (s, Nothing)
297
298 step' (s, Nothing) = liftM (\r ->
299 case r of
300 Yield x s' -> Yield x (s', Nothing)
301 Skip s' -> Skip (s', Nothing)
302 Done -> Done
303 ) (step s)
304
305
306 -- Mapping
307 -- -------
308
309 instance Monad m => Functor (Stream m) where
310 {-# INLINE fmap #-}
311 fmap = map
312
313 -- | Map a function over a 'Stream'
314 map :: Monad m => (a -> b) -> Stream m a -> Stream m b
315 {-# INLINE map #-}
316 map f = mapM (return . f)
317
318 -- | Map a monadic function over a 'Stream'
319 mapM :: Monad m => (a -> m b) -> Stream m a -> Stream m b
320 {-# INLINE_STREAM mapM #-}
321 mapM f (Stream step s n) = Stream step' s n
322 where
323 {-# INLINE step' #-}
324 step' s = do
325 r <- step s
326 case r of
327 Yield x s' -> liftM (`Yield` s') (f x)
328 Skip s' -> return (Skip s')
329 Done -> return Done
330
331 -- | Execute a monadic action for each element of the 'Stream'
332 mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()
333 {-# INLINE_STREAM mapM_ #-}
334 mapM_ m (Stream step s _) = mapM_go s
335 where
336 mapM_go s = do
337 r <- step s
338 case r of
339 Yield x s' -> do { m x; mapM_go s' }
340 Skip s' -> mapM_go s'
341 Done -> return ()
342
343 -- | Transform a 'Stream' to use a different monad
344 trans :: (Monad m, Monad m') => (forall a. m a -> m' a)
345 -> Stream m a -> Stream m' a
346 {-# INLINE_STREAM trans #-}
347 trans f (Stream step s n) = Stream (f . step) s n
348
349 -- Zipping
350 -- -------
351
352 -- | Zip two 'Stream's with the given function
353 zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
354 {-# INLINE zipWith #-}
355 zipWith f = zipWithM (\a b -> return (f a b))
356
357 -- | Zip two 'Stream's with the given monadic function
358 zipWithM :: Monad m => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
359 {-# INLINE_STREAM zipWithM #-}
360 zipWithM f (Stream stepa sa na) (Stream stepb sb nb)
361 = Stream step (sa, sb, Nothing) (smaller na nb)
362 where
363 {-# INLINE step #-}
364 step (sa, sb, Nothing) = liftM (\r ->
365 case r of
366 Yield x sa' -> Skip (sa', sb, Just x)
367 Skip sa' -> Skip (sa', sb, Nothing)
368 Done -> Done
369 ) (stepa sa)
370
371 step (sa, sb, Just x) = do
372 r <- stepb sb
373 case r of
374 Yield y sb' ->
375 do
376 z <- f x y
377 return $ Yield z (sa, sb', Nothing)
378 Skip sb' -> return $ Skip (sa, sb', Just x)
379 Done -> return $ Done
380
381 -- | Zip three 'Stream's with the given function
382 zipWith3 :: Monad m => (a -> b -> c -> d) -> Stream m a -> Stream m b -> Stream m c -> Stream m d
383 {-# INLINE zipWith3 #-}
384 zipWith3 f = zipWith3M (\a b c -> return (f a b c))
385
386 -- | Zip three 'Stream's with the given monadic function
387 zipWith3M :: Monad m => (a -> b -> c -> m d) -> Stream m a -> Stream m b -> Stream m c -> Stream m d
388 {-# INLINE_STREAM zipWith3M #-}
389 zipWith3M f (Stream stepa sa na) (Stream stepb sb nb) (Stream stepc sc nc)
390 = Stream step (sa, sb, sc, Nothing) (smaller na (smaller nb nc))
391 where
392 {-# INLINE step #-}
393 step (sa, sb, sc, Nothing) = do
394 r <- stepa sa
395 return $ case r of
396 Yield x sa' -> Skip (sa', sb, sc, Just (x, Nothing))
397 Skip sa' -> Skip (sa', sb, sc, Nothing)
398 Done -> Done
399
400 step (sa, sb, sc, Just (x, Nothing)) = do
401 r <- stepb sb
402 return $ case r of
403 Yield y sb' -> Skip (sa, sb', sc, Just (x, Just y))
404 Skip sb' -> Skip (sa, sb', sc, Just (x, Nothing))
405 Done -> Done
406
407 step (sa, sb, sc, Just (x, Just y)) = do
408 r <- stepc sc
409 case r of
410 Yield z sc' -> f x y z >>= (\res -> return $ Yield res (sa, sb, sc', Nothing))
411 Skip sc' -> return $ Skip (sa, sb, sc', Just (x, Just y))
412 Done -> return $ Done
413
414 -- Filtering
415 -- ---------
416
417 -- | Drop elements which do not satisfy the predicate
418 filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
419 {-# INLINE filter #-}
420 filter f = filterM (return . f)
421
422 -- | Drop elements which do not satisfy the monadic predicate
423 filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
424 {-# INLINE_STREAM filterM #-}
425 filterM f (Stream step s n) = Stream step' s (toMax n)
426 where
427 {-# INLINE step' #-}
428 step' s = do
429 r <- step s
430 case r of
431 Yield x s' -> do
432 b <- f x
433 return $ if b then Yield x s'
434 else Skip s'
435 Skip s' -> return $ Skip s'
436 Done -> return $ Done
437
438 -- | Longest prefix of elements that satisfy the predicate
439 takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
440 {-# INLINE takeWhile #-}
441 takeWhile f = takeWhileM (return . f)
442
443 -- | Longest prefix of elements that satisfy the monadic predicate
444 takeWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
445 {-# INLINE_STREAM takeWhileM #-}
446 takeWhileM f (Stream step s n) = Stream step' s (toMax n)
447 where
448 {-# INLINE step' #-}
449 step' s = do
450 r <- step s
451 case r of
452 Yield x s' -> do
453 b <- f x
454 return $ if b then Yield x s' else Done
455 Skip s' -> return $ Skip s'
456 Done -> return $ Done
457
458 -- | Drop the longest prefix of elements that satisfy the predicate
459 dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
460 {-# INLINE dropWhile #-}
461 dropWhile f = dropWhileM (return . f)
462
463 data DropWhile s a = DropWhile_Drop s | DropWhile_Yield a s | DropWhile_Next s
464
465 -- | Drop the longest prefix of elements that satisfy the monadic predicate
466 dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
467 {-# INLINE_STREAM dropWhileM #-}
468 dropWhileM f (Stream step s n) = Stream step' (DropWhile_Drop s) (toMax n)
469 where
470 -- NOTE: we jump through hoops here to have only one Yield; local data
471 -- declarations would be nice!
472
473 {-# INLINE step' #-}
474 step' (DropWhile_Drop s)
475 = do
476 r <- step s
477 case r of
478 Yield x s' -> do
479 b <- f x
480 return $ if b then Skip (DropWhile_Drop s')
481 else Skip (DropWhile_Yield x s')
482 Skip s' -> return $ Skip (DropWhile_Drop s')
483 Done -> return $ Done
484
485 step' (DropWhile_Yield x s) = return $ Yield x (DropWhile_Next s)
486
487 step' (DropWhile_Next s)
488 = liftM (\r ->
489 case r of
490 Yield x s' -> Skip (DropWhile_Yield x s')
491 Skip s' -> Skip (DropWhile_Next s')
492 Done -> Done
493 ) (step s)
494
495 -- Searching
496 -- ---------
497
498 infix 4 `elem`
499 -- | Check whether the 'Stream' contains an element
500 elem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
501 {-# INLINE_STREAM elem #-}
502 elem x (Stream step s _) = elem_loop s
503 where
504 elem_loop s = do
505 r <- step s
506 case r of
507 Yield y s' | x == y -> return True
508 | otherwise -> elem_loop s'
509 Skip s' -> elem_loop s'
510 Done -> return False
511
512 infix 4 `notElem`
513 -- | Inverse of `elem`
514 notElem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
515 {-# INLINE notElem #-}
516 notElem x s = liftM not (elem x s)
517
518 -- | Yield 'Just' the first element that satisfies the predicate or 'Nothing'
519 -- if no such element exists.
520 find :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe a)
521 {-# INLINE find #-}
522 find f = findM (return . f)
523
524 -- | Yield 'Just' the first element that satisfies the monadic predicate or
525 -- 'Nothing' if no such element exists.
526 findM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe a)
527 {-# INLINE_STREAM findM #-}
528 findM f (Stream step s _) = find_loop s
529 where
530 find_loop s = do
531 r <- step s
532 case r of
533 Yield x s' -> do
534 b <- f x
535 if b then return $ Just x
536 else find_loop s'
537 Skip s' -> find_loop s'
538 Done -> return Nothing
539
540 -- | Yield 'Just' the index of the first element that satisfies the predicate
541 -- or 'Nothing' if no such element exists.
542 findIndex :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe Int)
543 {-# INLINE_STREAM findIndex #-}
544 findIndex f = findIndexM (return . f)
545
546 -- | Yield 'Just' the index of the first element that satisfies the monadic
547 -- predicate or 'Nothing' if no such element exists.
548 findIndexM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe Int)
549 {-# INLINE_STREAM findIndexM #-}
550 findIndexM f (Stream step s _) = findIndex_loop s 0
551 where
552 findIndex_loop s i = do
553 r <- step s
554 case r of
555 Yield x s' -> do
556 b <- f x
557 if b then return $ Just i
558 else findIndex_loop s' (i+1)
559 Skip s' -> findIndex_loop s' i
560 Done -> return Nothing
561
562 -- Folding
563 -- -------
564
565 -- | Left fold
566 foldl :: Monad m => (a -> b -> a) -> a -> Stream m b -> m a
567 {-# INLINE foldl #-}
568 foldl f = foldlM (\a b -> return (f a b))
569
570 -- | Left fold with a monadic operator
571 foldlM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> m a
572 {-# INLINE_STREAM foldlM #-}
573 foldlM m z (Stream step s _) = foldlM_go z s
574 where
575 foldlM_go z s = do
576 r <- step s
577 case r of
578 Yield x s' -> do { z' <- m z x; foldlM_go z' s' }
579 Skip s' -> foldlM_go z s'
580 Done -> return z
581
582 -- | Same as 'foldlM'
583 foldM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> m a
584 {-# INLINE foldM #-}
585 foldM = foldlM
586
587 -- | Left fold over a non-empty 'Stream'
588 foldl1 :: Monad m => (a -> a -> a) -> Stream m a -> m a
589 {-# INLINE foldl1 #-}
590 foldl1 f = foldl1M (\a b -> return (f a b))
591
592 -- | Left fold over a non-empty 'Stream' with a monadic operator
593 foldl1M :: Monad m => (a -> a -> m a) -> Stream m a -> m a
594 {-# INLINE_STREAM foldl1M #-}
595 foldl1M f (Stream step s sz) = foldl1M_go s
596 where
597 foldl1M_go s = do
598 r <- step s
599 case r of
600 Yield x s' -> foldlM f x (Stream step s' (sz - 1))
601 Skip s' -> foldl1M_go s'
602 Done -> errorEmptyStream "foldl1M"
603
604 -- | Same as 'foldl1M'
605 fold1M :: Monad m => (a -> a -> m a) -> Stream m a -> m a
606 {-# INLINE fold1M #-}
607 fold1M = foldl1M
608
609 -- | Left fold with a strict accumulator
610 foldl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> m a
611 {-# INLINE foldl' #-}
612 foldl' f = foldlM' (\a b -> return (f a b))
613
614 -- | Left fold with a strict accumulator and a monadic operator
615 foldlM' :: Monad m => (a -> b -> m a) -> a -> Stream m b -> m a
616 {-# INLINE_STREAM foldlM' #-}
617 foldlM' m z (Stream step s _) = foldlM'_go z s
618 where
619 foldlM'_go z s = z `seq`
620 do
621 r <- step s
622 case r of
623 Yield x s' -> do { z' <- m z x; foldlM'_go z' s' }
624 Skip s' -> foldlM'_go z s'
625 Done -> return z
626
627 -- | Same as 'foldlM''
628 foldM' :: Monad m => (a -> b -> m a) -> a -> Stream m b -> m a
629 {-# INLINE foldM' #-}
630 foldM' = foldlM'
631
632 -- | Left fold over a non-empty 'Stream' with a strict accumulator
633 foldl1' :: Monad m => (a -> a -> a) -> Stream m a -> m a
634 {-# INLINE foldl1' #-}
635 foldl1' f = foldl1M' (\a b -> return (f a b))
636
637 -- | Left fold over a non-empty 'Stream' with a strict accumulator and a
638 -- monadic operator
639 foldl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> m a
640 {-# INLINE_STREAM foldl1M' #-}
641 foldl1M' f (Stream step s sz) = foldl1M'_go s
642 where
643 foldl1M'_go s = do
644 r <- step s
645 case r of
646 Yield x s' -> foldlM' f x (Stream step s' (sz - 1))
647 Skip s' -> foldl1M'_go s'
648 Done -> errorEmptyStream "foldl1M'"
649
650 -- | Same as 'foldl1M''
651 fold1M' :: Monad m => (a -> a -> m a) -> Stream m a -> m a
652 {-# INLINE fold1M' #-}
653 fold1M' = foldl1M'
654
655 -- | Right fold
656 foldr :: Monad m => (a -> b -> b) -> b -> Stream m a -> m b
657 {-# INLINE foldr #-}
658 foldr f = foldrM (\a b -> return (f a b))
659
660 -- | Right fold with a monadic operator
661 foldrM :: Monad m => (a -> b -> m b) -> b -> Stream m a -> m b
662 {-# INLINE_STREAM foldrM #-}
663 foldrM f z (Stream step s _) = foldrM_go s
664 where
665 foldrM_go s = do
666 r <- step s
667 case r of
668 Yield x s' -> f x =<< foldrM_go s'
669 Skip s' -> foldrM_go s'
670 Done -> return z
671
672 -- | Right fold over a non-empty stream
673 foldr1 :: Monad m => (a -> a -> a) -> Stream m a -> m a
674 {-# INLINE foldr1 #-}
675 foldr1 f = foldr1M (\a b -> return (f a b))
676
677 -- | Right fold over a non-empty stream with a monadic operator
678 foldr1M :: Monad m => (a -> a -> m a) -> Stream m a -> m a
679 {-# INLINE_STREAM foldr1M #-}
680 foldr1M f (Stream step s _) = foldr1M_go0 s
681 where
682 foldr1M_go0 s = do
683 r <- step s
684 case r of
685 Yield x s' -> foldr1M_go1 x s'
686 Skip s' -> foldr1M_go0 s'
687 Done -> errorEmptyStream "foldr1M"
688
689 foldr1M_go1 x s = do
690 r <- step s
691 case r of
692 Yield y s' -> f x =<< foldr1M_go1 y s'
693 Skip s' -> foldr1M_go1 x s'
694 Done -> return x
695
696 -- Specialised folds
697 -- -----------------
698
699 and :: Monad m => Stream m Bool -> m Bool
700 {-# INLINE_STREAM and #-}
701 and (Stream step s _) = and_go s
702 where
703 and_go s = do
704 r <- step s
705 case r of
706 Yield False _ -> return False
707 Yield True s' -> and_go s'
708 Skip s' -> and_go s'
709 Done -> return True
710
711 or :: Monad m => Stream m Bool -> m Bool
712 {-# INLINE_STREAM or #-}
713 or (Stream step s _) = or_go s
714 where
715 or_go s = do
716 r <- step s
717 case r of
718 Yield False s' -> or_go s'
719 Yield True _ -> return True
720 Skip s' -> or_go s'
721 Done -> return False
722
723 concatMap :: Monad m => (a -> Stream m b) -> Stream m a -> Stream m b
724 {-# INLINE concatMap #-}
725 concatMap f = concatMapM (return . f)
726
727 concatMapM :: Monad m => (a -> m (Stream m b)) -> Stream m a -> Stream m b
728 {-# INLINE_STREAM concatMapM #-}
729 concatMapM f (Stream step s _) = Stream concatMap_go (Left s) Unknown
730 where
731 concatMap_go (Left s) = do
732 r <- step s
733 case r of
734 Yield a s' -> do
735 b_stream <- f a
736 return $ Skip (Right (b_stream, s'))
737 Skip s' -> return $ Skip (Left s')
738 Done -> return Done
739 concatMap_go (Right (Stream inner_step inner_s sz, s)) = do
740 r <- inner_step inner_s
741 case r of
742 Yield b inner_s' -> return $ Yield b (Right (Stream inner_step inner_s' sz, s))
743 Skip inner_s' -> return $ Skip (Right (Stream inner_step inner_s' sz, s))
744 Done -> return $ Skip (Left s)
745
746 -- Unfolding
747 -- ---------
748
749 -- | Unfold
750 unfoldr :: Monad m => (s -> Maybe (a, s)) -> s -> Stream m a
751 {-# INLINE_STREAM unfoldr #-}
752 unfoldr f = unfoldrM (return . f)
753
754 -- | Unfold with a monadic function
755 unfoldrM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m a
756 {-# INLINE_STREAM unfoldrM #-}
757 unfoldrM f s = Stream step s Unknown
758 where
759 {-# INLINE step #-}
760 step s = liftM (\r ->
761 case r of
762 Just (x, s') -> Yield x s'
763 Nothing -> Done
764 ) (f s)
765
766 -- Scans
767 -- -----
768
769 -- | Prefix scan
770 prescanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
771 {-# INLINE prescanl #-}
772 prescanl f = prescanlM (\a b -> return (f a b))
773
774 -- | Prefix scan with a monadic operator
775 prescanlM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a
776 {-# INLINE_STREAM prescanlM #-}
777 prescanlM f z (Stream step s sz) = Stream step' (s,z) sz
778 where
779 {-# INLINE step' #-}
780 step' (s,x) = do
781 r <- step s
782 case r of
783 Yield y s' -> do
784 z <- f x y
785 return $ Yield x (s', z)
786 Skip s' -> return $ Skip (s', x)
787 Done -> return Done
788
789 -- | Prefix scan with strict accumulator
790 prescanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
791 {-# INLINE prescanl' #-}
792 prescanl' f = prescanlM' (\a b -> return (f a b))
793
794 -- | Prefix scan with strict accumulator and a monadic operator
795 prescanlM' :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a
796 {-# INLINE_STREAM prescanlM' #-}
797 prescanlM' f z (Stream step s sz) = Stream step' (s,z) sz
798 where
799 {-# INLINE step' #-}
800 step' (s,x) = x `seq`
801 do
802 r <- step s
803 case r of
804 Yield y s' -> do
805 z <- f x y
806 return $ Yield x (s', z)
807 Skip s' -> return $ Skip (s', x)
808 Done -> return Done
809
810 -- | Suffix scan
811 postscanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
812 {-# INLINE postscanl #-}
813 postscanl f = postscanlM (\a b -> return (f a b))
814
815 -- | Suffix scan with a monadic operator
816 postscanlM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a
817 {-# INLINE_STREAM postscanlM #-}
818 postscanlM f z (Stream step s sz) = Stream step' (s,z) sz
819 where
820 {-# INLINE step' #-}
821 step' (s,x) = do
822 r <- step s
823 case r of
824 Yield y s' -> do
825 z <- f x y
826 return $ Yield z (s',z)
827 Skip s' -> return $ Skip (s',x)
828 Done -> return Done
829
830 -- | Suffix scan with strict accumulator
831 postscanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
832 {-# INLINE postscanl' #-}
833 postscanl' f = postscanlM' (\a b -> return (f a b))
834
835 -- | Suffix scan with strict acccumulator and a monadic operator
836 postscanlM' :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a
837 {-# INLINE_STREAM postscanlM' #-}
838 postscanlM' f z (Stream step s sz) = z `seq` Stream step' (s,z) sz
839 where
840 {-# INLINE step' #-}
841 step' (s,x) = x `seq`
842 do
843 r <- step s
844 case r of
845 Yield y s' -> do
846 z <- f x y
847 z `seq` return (Yield z (s',z))
848 Skip s' -> return $ Skip (s',x)
849 Done -> return Done
850
851 -- | Haskell-style scan
852 scanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
853 {-# INLINE scanl #-}
854 scanl f = scanlM (\a b -> return (f a b))
855
856 -- | Haskell-style scan with a monadic operator
857 scanlM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a
858 {-# INLINE scanlM #-}
859 scanlM f z s = z `cons` postscanlM f z s
860
861 -- | Haskell-style scan with strict accumulator
862 scanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
863 {-# INLINE scanl' #-}
864 scanl' f = scanlM' (\a b -> return (f a b))
865
866 -- | Haskell-style scan with strict accumulator and a monadic operator
867 scanlM' :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a
868 {-# INLINE scanlM' #-}
869 scanlM' f z s = z `seq` (z `cons` postscanlM f z s)
870
871 -- | Scan over a non-empty 'Stream'
872 scanl1 :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
873 {-# INLINE scanl1 #-}
874 scanl1 f = scanl1M (\x y -> return (f x y))
875
876 -- | Scan over a non-empty 'Stream' with a monadic operator
877 scanl1M :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
878 {-# INLINE_STREAM scanl1M #-}
879 scanl1M f (Stream step s sz) = Stream step' (s, Nothing) sz
880 where
881 {-# INLINE step' #-}
882 step' (s, Nothing) = do
883 r <- step s
884 case r of
885 Yield x s' -> return $ Yield x (s', Just x)
886 Skip s' -> return $ Skip (s', Nothing)
887 Done -> errorEmptyStream "scanl1M"
888
889 step' (s, Just x) = do
890 r <- step s
891 case r of
892 Yield y s' -> do
893 z <- f x y
894 return $ Yield z (s', Just z)
895 Skip s' -> return $ Skip (s', Just x)
896 Done -> return Done
897
898 -- | Scan over a non-empty 'Stream' with a strict accumulator
899 scanl1' :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
900 {-# INLINE scanl1' #-}
901 scanl1' f = scanl1M' (\x y -> return (f x y))
902
903 -- | Scan over a non-empty 'Stream' with a strict accumulator and a monadic
904 -- operator
905 scanl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
906 {-# INLINE_STREAM scanl1M' #-}
907 scanl1M' f (Stream step s sz) = Stream step' (s, Nothing) sz
908 where
909 {-# INLINE step' #-}
910 step' (s, Nothing) = do
911 r <- step s
912 case r of
913 Yield x s' -> x `seq` return (Yield x (s', Just x))
914 Skip s' -> return $ Skip (s', Nothing)
915 Done -> errorEmptyStream "scanl1M"
916
917 step' (s, Just x) = x `seq`
918 do
919 r <- step s
920 case r of
921 Yield y s' -> do
922 z <- f x y
923 z `seq` return (Yield z (s', Just z))
924 Skip s' -> return $ Skip (s', Just x)
925 Done -> return Done
926
927 -- Conversions
928 -- -----------
929
930 -- | Convert a 'Stream' to a list
931 toList :: Monad m => Stream m a -> m [a]
932 {-# INLINE toList #-}
933 toList = foldr (:) []
934
935 -- | Convert a list to a 'Stream'
936 fromList :: Monad m => [a] -> Stream m a
937 {-# INLINE_STREAM fromList #-}
938 fromList xs = Stream step xs Unknown
939 where
940 step (x:xs) = return (Yield x xs)
941 step [] = return Done
942
943
944 streamError :: String -> String -> a
945 streamError fn msg = error $ "Data.Vector.Fusion.Stream.Monadic."
946 Prelude.++ fn Prelude.++ ": " Prelude.++ msg
947
948 errorEmptyStream :: String -> a
949 errorEmptyStream fn = streamError fn "empty stream"
950
951 errorNegativeIndex :: String -> a
952 errorNegativeIndex fn = streamError fn "negative index"
953
954 errorIndexOutOfRange :: String -> a
955 errorIndexOutOfRange fn = streamError fn "index out of range"
956