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