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