Implement and test zipWith3, zip3, unzip, unzip3
[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
36 map, mapM, mapM_, 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 -- Zipping
340 -- -------
341
342 -- | Zip two 'Stream's with the given function
343 zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
344 {-# INLINE zipWith #-}
345 zipWith f = zipWithM (\a b -> return (f a b))
346
347 -- | Zip two 'Stream's with the given monadic function
348 zipWithM :: Monad m => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
349 {-# INLINE_STREAM zipWithM #-}
350 zipWithM f (Stream stepa sa na) (Stream stepb sb nb)
351 = Stream step (sa, sb, Nothing) (smaller na nb)
352 where
353 {-# INLINE step #-}
354 step (sa, sb, Nothing) = liftM (\r ->
355 case r of
356 Yield x sa' -> Skip (sa', sb, Just x)
357 Skip sa' -> Skip (sa', sb, Nothing)
358 Done -> Done
359 ) (stepa sa)
360
361 step (sa, sb, Just x) = do
362 r <- stepb sb
363 case r of
364 Yield y sb' ->
365 do
366 z <- f x y
367 return $ Yield z (sa, sb', Nothing)
368 Skip sb' -> return $ Skip (sa, sb', Just x)
369 Done -> return $ Done
370
371 -- | Zip three 'Stream's with the given function
372 zipWith3 :: Monad m => (a -> b -> c -> d) -> Stream m a -> Stream m b -> Stream m c -> Stream m d
373 {-# INLINE zipWith3 #-}
374 zipWith3 f = zipWith3M (\a b c -> return (f a b c))
375
376 -- | Zip three 'Stream's with the given monadic function
377 zipWith3M :: Monad m => (a -> b -> c -> m d) -> Stream m a -> Stream m b -> Stream m c -> Stream m d
378 {-# INLINE_STREAM zipWith3M #-}
379 zipWith3M f (Stream stepa sa na) (Stream stepb sb nb) (Stream stepc sc nc)
380 = Stream step (sa, sb, sc, Nothing) (smaller na (smaller nb nc))
381 where
382 {-# INLINE step #-}
383 step (sa, sb, sc, Nothing) = do
384 r <- stepa sa
385 return $ case r of
386 Yield x sa' -> Skip (sa', sb, sc, Just (x, Nothing))
387 Skip sa' -> Skip (sa', sb, sc, Nothing)
388 Done -> Done
389
390 step (sa, sb, sc, Just (x, Nothing)) = do
391 r <- stepb sb
392 return $ case r of
393 Yield y sb' -> Skip (sa, sb', sc, Just (x, Just y))
394 Skip sb' -> Skip (sa, sb', sc, Just (x, Nothing))
395 Done -> Done
396
397 step (sa, sb, sc, Just (x, Just y)) = do
398 r <- stepc sc
399 case r of
400 Yield z sc' -> f x y z >>= (\res -> return $ Yield res (sa, sb, sc', Nothing))
401 Skip sc' -> return $ Skip (sa, sb, sc', Just (x, Just y))
402 Done -> return $ Done
403
404 -- Filtering
405 -- ---------
406
407 -- | Drop elements which do not satisfy the predicate
408 filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
409 {-# INLINE filter #-}
410 filter f = filterM (return . f)
411
412 -- | Drop elements which do not satisfy the monadic predicate
413 filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
414 {-# INLINE_STREAM filterM #-}
415 filterM f (Stream step s n) = Stream step' s (toMax n)
416 where
417 {-# INLINE step' #-}
418 step' s = do
419 r <- step s
420 case r of
421 Yield x s' -> do
422 b <- f x
423 return $ if b then Yield x s'
424 else Skip s'
425 Skip s' -> return $ Skip s'
426 Done -> return $ Done
427
428 -- | Longest prefix of elements that satisfy the predicate
429 takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
430 {-# INLINE takeWhile #-}
431 takeWhile f = takeWhileM (return . f)
432
433 -- | Longest prefix of elements that satisfy the monadic predicate
434 takeWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
435 {-# INLINE_STREAM takeWhileM #-}
436 takeWhileM f (Stream step s n) = Stream step' s (toMax n)
437 where
438 {-# INLINE step' #-}
439 step' s = do
440 r <- step s
441 case r of
442 Yield x s' -> do
443 b <- f x
444 return $ if b then Yield x s' else Done
445 Skip s' -> return $ Skip s'
446 Done -> return $ Done
447
448 -- | Drop the longest prefix of elements that satisfy the predicate
449 dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
450 {-# INLINE dropWhile #-}
451 dropWhile f = dropWhileM (return . f)
452
453 data DropWhile s a = DropWhile_Drop s | DropWhile_Yield a s | DropWhile_Next s
454
455 -- | Drop the longest prefix of elements that satisfy the monadic predicate
456 dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
457 {-# INLINE_STREAM dropWhileM #-}
458 dropWhileM f (Stream step s n) = Stream step' (DropWhile_Drop s) (toMax n)
459 where
460 -- NOTE: we jump through hoops here to have only one Yield; local data
461 -- declarations would be nice!
462
463 {-# INLINE step' #-}
464 step' (DropWhile_Drop s)
465 = do
466 r <- step s
467 case r of
468 Yield x s' -> do
469 b <- f x
470 return $ if b then Skip (DropWhile_Drop s')
471 else Skip (DropWhile_Yield x s')
472 Skip s' -> return $ Skip (DropWhile_Drop s')
473 Done -> return $ Done
474
475 step' (DropWhile_Yield x s) = return $ Yield x (DropWhile_Next s)
476
477 step' (DropWhile_Next s)
478 = liftM (\r ->
479 case r of
480 Yield x s' -> Skip (DropWhile_Yield x s')
481 Skip s' -> Skip (DropWhile_Next s')
482 Done -> Done
483 ) (step s)
484
485 -- Searching
486 -- ---------
487
488 infix 4 `elem`
489 -- | Check whether the 'Stream' contains an element
490 elem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
491 {-# INLINE_STREAM elem #-}
492 elem x (Stream step s _) = elem_loop s
493 where
494 elem_loop s = do
495 r <- step s
496 case r of
497 Yield y s' | x == y -> return True
498 | otherwise -> elem_loop s'
499 Skip s' -> elem_loop s'
500 Done -> return False
501
502 infix 4 `notElem`
503 -- | Inverse of `elem`
504 notElem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
505 {-# INLINE notElem #-}
506 notElem x s = liftM not (elem x s)
507
508 -- | Yield 'Just' the first element that satisfies the predicate or 'Nothing'
509 -- if no such element exists.
510 find :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe a)
511 {-# INLINE find #-}
512 find f = findM (return . f)
513
514 -- | Yield 'Just' the first element that satisfies the monadic predicate or
515 -- 'Nothing' if no such element exists.
516 findM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe a)
517 {-# INLINE_STREAM findM #-}
518 findM f (Stream step s _) = find_loop s
519 where
520 find_loop s = do
521 r <- step s
522 case r of
523 Yield x s' -> do
524 b <- f x
525 if b then return $ Just x
526 else find_loop s'
527 Skip s' -> find_loop s'
528 Done -> return Nothing
529
530 -- | Yield 'Just' the index of the first element that satisfies the predicate
531 -- or 'Nothing' if no such element exists.
532 findIndex :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe Int)
533 {-# INLINE_STREAM findIndex #-}
534 findIndex f = findIndexM (return . f)
535
536 -- | Yield 'Just' the index of the first element that satisfies the monadic
537 -- predicate or 'Nothing' if no such element exists.
538 findIndexM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe Int)
539 {-# INLINE_STREAM findIndexM #-}
540 findIndexM f (Stream step s _) = findIndex_loop s 0
541 where
542 findIndex_loop s i = do
543 r <- step s
544 case r of
545 Yield x s' -> do
546 b <- f x
547 if b then return $ Just i
548 else findIndex_loop s' (i+1)
549 Skip s' -> findIndex_loop s' i
550 Done -> return Nothing
551
552 -- Folding
553 -- -------
554
555 -- | Left fold
556 foldl :: Monad m => (a -> b -> a) -> a -> Stream m b -> m a
557 {-# INLINE foldl #-}
558 foldl f = foldlM (\a b -> return (f a b))
559
560 -- | Left fold with a monadic operator
561 foldlM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> m a
562 {-# INLINE_STREAM foldlM #-}
563 foldlM m z (Stream step s _) = foldlM_go z s
564 where
565 foldlM_go z s = 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 -- | Same as 'foldlM'
573 foldM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> m a
574 {-# INLINE foldM #-}
575 foldM = foldlM
576
577 -- | Left fold over a non-empty 'Stream'
578 foldl1 :: Monad m => (a -> a -> a) -> Stream m a -> m a
579 {-# INLINE foldl1 #-}
580 foldl1 f = foldl1M (\a b -> return (f a b))
581
582 -- | Left fold over a non-empty 'Stream' with a monadic operator
583 foldl1M :: Monad m => (a -> a -> m a) -> Stream m a -> m a
584 {-# INLINE_STREAM foldl1M #-}
585 foldl1M f (Stream step s sz) = foldl1M_go s
586 where
587 foldl1M_go s = do
588 r <- step s
589 case r of
590 Yield x s' -> foldlM f x (Stream step s' (sz - 1))
591 Skip s' -> foldl1M_go s'
592 Done -> errorEmptyStream "foldl1M"
593
594 -- | Left fold with a strict accumulator
595 foldl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> m a
596 {-# INLINE foldl' #-}
597 foldl' f = foldlM' (\a b -> return (f a b))
598
599 -- | Left fold with a strict accumulator and a monadic operator
600 foldlM' :: Monad m => (a -> b -> m a) -> a -> Stream m b -> m a
601 {-# INLINE_STREAM foldlM' #-}
602 foldlM' m z (Stream step s _) = foldlM'_go z s
603 where
604 foldlM'_go z s = z `seq`
605 do
606 r <- step s
607 case r of
608 Yield x s' -> do { z' <- m z x; foldlM'_go z' s' }
609 Skip s' -> foldlM'_go z s'
610 Done -> return z
611
612 -- | Left fold over a non-empty 'Stream' with a strict accumulator
613 foldl1' :: Monad m => (a -> a -> a) -> Stream m a -> m a
614 {-# INLINE foldl1' #-}
615 foldl1' f = foldl1M' (\a b -> return (f a b))
616
617 -- | Left fold over a non-empty 'Stream' with a strict accumulator and a
618 -- monadic operator
619 foldl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> m a
620 {-# INLINE_STREAM foldl1M' #-}
621 foldl1M' f (Stream step s sz) = foldl1M'_go s
622 where
623 foldl1M'_go s = do
624 r <- step s
625 case r of
626 Yield x s' -> foldlM' f x (Stream step s' (sz - 1))
627 Skip s' -> foldl1M'_go s'
628 Done -> errorEmptyStream "foldl1M'"
629
630 -- | Right fold
631 foldr :: Monad m => (a -> b -> b) -> b -> Stream m a -> m b
632 {-# INLINE foldr #-}
633 foldr f = foldrM (\a b -> return (f a b))
634
635 -- | Right fold with a monadic operator
636 foldrM :: Monad m => (a -> b -> m b) -> b -> Stream m a -> m b
637 {-# INLINE_STREAM foldrM #-}
638 foldrM f z (Stream step s _) = foldrM_go s
639 where
640 foldrM_go s = do
641 r <- step s
642 case r of
643 Yield x s' -> f x =<< foldrM_go s'
644 Skip s' -> foldrM_go s'
645 Done -> return z
646
647 -- | Right fold over a non-empty stream
648 foldr1 :: Monad m => (a -> a -> a) -> Stream m a -> m a
649 {-# INLINE foldr1 #-}
650 foldr1 f = foldr1M (\a b -> return (f a b))
651
652 -- | Right fold over a non-empty stream with a monadic operator
653 foldr1M :: Monad m => (a -> a -> m a) -> Stream m a -> m a
654 {-# INLINE_STREAM foldr1M #-}
655 foldr1M f (Stream step s _) = foldr1M_go0 s
656 where
657 foldr1M_go0 s = do
658 r <- step s
659 case r of
660 Yield x s' -> foldr1M_go1 x s'
661 Skip s' -> foldr1M_go0 s'
662 Done -> errorEmptyStream "foldr1M"
663
664 foldr1M_go1 x s = do
665 r <- step s
666 case r of
667 Yield y s' -> f x =<< foldr1M_go1 y s'
668 Skip s' -> foldr1M_go1 x s'
669 Done -> return x
670
671 -- Specialised folds
672 -- -----------------
673
674 and :: Monad m => Stream m Bool -> m Bool
675 {-# INLINE_STREAM and #-}
676 and (Stream step s _) = and_go s
677 where
678 and_go s = do
679 r <- step s
680 case r of
681 Yield False _ -> return False
682 Yield True s' -> and_go s'
683 Skip s' -> and_go s'
684 Done -> return True
685
686 or :: Monad m => Stream m Bool -> m Bool
687 {-# INLINE_STREAM or #-}
688 or (Stream step s _) = or_go s
689 where
690 or_go s = do
691 r <- step s
692 case r of
693 Yield False s' -> or_go s'
694 Yield True _ -> return True
695 Skip s' -> or_go s'
696 Done -> return False
697
698 concatMap :: Monad m => (a -> Stream m b) -> Stream m a -> Stream m b
699 {-# INLINE concatMap #-}
700 concatMap f = concatMapM (return . f)
701
702 concatMapM :: Monad m => (a -> m (Stream m b)) -> Stream m a -> Stream m b
703 {-# INLINE_STREAM concatMapM #-}
704 concatMapM f (Stream step s _) = Stream concatMap_go (Left s) Unknown
705 where
706 concatMap_go (Left s) = do
707 r <- step s
708 case r of
709 Yield a s' -> do
710 b_stream <- f a
711 return $ Skip (Right (b_stream, s'))
712 Skip s' -> return $ Skip (Left s')
713 Done -> return Done
714 concatMap_go (Right (Stream inner_step inner_s sz, s)) = do
715 r <- inner_step inner_s
716 case r of
717 Yield b inner_s' -> return $ Yield b (Right (Stream inner_step inner_s' sz, s))
718 Skip inner_s' -> return $ Skip (Right (Stream inner_step inner_s' sz, s))
719 Done -> return $ Skip (Left s)
720
721 -- Unfolding
722 -- ---------
723
724 -- | Unfold
725 unfoldr :: Monad m => (s -> Maybe (a, s)) -> s -> Stream m a
726 {-# INLINE_STREAM unfoldr #-}
727 unfoldr f = unfoldrM (return . f)
728
729 -- | Unfold with a monadic function
730 unfoldrM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m a
731 {-# INLINE_STREAM unfoldrM #-}
732 unfoldrM f s = Stream step s Unknown
733 where
734 {-# INLINE step #-}
735 step s = liftM (\r ->
736 case r of
737 Just (x, s') -> Yield x s'
738 Nothing -> Done
739 ) (f s)
740
741 -- Scans
742 -- -----
743
744 -- | Prefix scan
745 prescanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
746 {-# INLINE prescanl #-}
747 prescanl f = prescanlM (\a b -> return (f a b))
748
749 -- | Prefix scan with a monadic operator
750 prescanlM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a
751 {-# INLINE_STREAM prescanlM #-}
752 prescanlM f z (Stream step s sz) = Stream step' (s,z) sz
753 where
754 {-# INLINE step' #-}
755 step' (s,x) = do
756 r <- step s
757 case r of
758 Yield y s' -> do
759 z <- f x y
760 return $ Yield x (s', z)
761 Skip s' -> return $ Skip (s', x)
762 Done -> return Done
763
764 -- | Prefix scan with strict accumulator
765 prescanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
766 {-# INLINE prescanl' #-}
767 prescanl' f = prescanlM' (\a b -> return (f a b))
768
769 -- | Prefix scan with strict accumulator and a monadic operator
770 prescanlM' :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a
771 {-# INLINE_STREAM prescanlM' #-}
772 prescanlM' f z (Stream step s sz) = Stream step' (s,z) sz
773 where
774 {-# INLINE step' #-}
775 step' (s,x) = x `seq`
776 do
777 r <- step s
778 case r of
779 Yield y s' -> do
780 z <- f x y
781 return $ Yield x (s', z)
782 Skip s' -> return $ Skip (s', x)
783 Done -> return Done
784
785 -- Conversions
786 -- -----------
787
788 -- | Convert a 'Stream' to a list
789 toList :: Monad m => Stream m a -> m [a]
790 {-# INLINE toList #-}
791 toList = foldr (:) []
792
793 -- | Convert a list to a 'Stream'
794 fromList :: Monad m => [a] -> Stream m a
795 {-# INLINE_STREAM fromList #-}
796 fromList xs = Stream step xs Unknown
797 where
798 step (x:xs) = return (Yield x xs)
799 step [] = return Done
800
801
802 streamError :: String -> String -> a
803 streamError fn msg = error $ "Data.Vector.Fusion.Stream.Monadic."
804 Prelude.++ fn Prelude.++ ": " Prelude.++ msg
805
806 errorEmptyStream :: String -> a
807 errorEmptyStream fn = streamError fn "empty stream"
808
809 errorNegativeIndex :: String -> a
810 errorNegativeIndex fn = streamError fn "negative index"
811
812 errorIndexOutOfRange :: String -> a
813 errorIndexOutOfRange fn = streamError fn "index out of range"
814