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