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