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