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