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