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