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