Fuse the fairly common case of zipWith f xs xs
[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 -- FIXME: This might expose an opportunity for inplace execution.
403 {-# RULES
404
405 "zipWithM xs xs [Vector.Stream]" forall f xs.
406 zipWithM f xs xs = mapM (\x -> f x x) xs
407
408 #-}
409
410 -- | Zip three 'Stream's with the given function
411 zipWith3 :: Monad m => (a -> b -> c -> d) -> Stream m a -> Stream m b -> Stream m c -> Stream m d
412 {-# INLINE zipWith3 #-}
413 zipWith3 f = zipWith3M (\a b c -> return (f a b c))
414
415 -- | Zip three 'Stream's with the given monadic function
416 zipWith3M :: Monad m => (a -> b -> c -> m d) -> Stream m a -> Stream m b -> Stream m c -> Stream m d
417 {-# INLINE_STREAM zipWith3M #-}
418 zipWith3M f (Stream stepa sa na) (Stream stepb sb nb) (Stream stepc sc nc)
419 = Stream step (sa, sb, sc, Nothing) (smaller na (smaller nb nc))
420 where
421 {-# INLINE_INNER step #-}
422 step (sa, sb, sc, Nothing) = do
423 r <- stepa sa
424 return $ case r of
425 Yield x sa' -> Skip (sa', sb, sc, Just (x, Nothing))
426 Skip sa' -> Skip (sa', sb, sc, Nothing)
427 Done -> Done
428
429 step (sa, sb, sc, Just (x, Nothing)) = do
430 r <- stepb sb
431 return $ case r of
432 Yield y sb' -> Skip (sa, sb', sc, Just (x, Just y))
433 Skip sb' -> Skip (sa, sb', sc, Just (x, Nothing))
434 Done -> Done
435
436 step (sa, sb, sc, Just (x, Just y)) = do
437 r <- stepc sc
438 case r of
439 Yield z sc' -> f x y z >>= (\res -> return $ Yield res (sa, sb, sc', Nothing))
440 Skip sc' -> return $ Skip (sa, sb, sc', Just (x, Just y))
441 Done -> return $ Done
442
443 -- Filtering
444 -- ---------
445
446 -- | Drop elements which do not satisfy the predicate
447 filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
448 {-# INLINE filter #-}
449 filter f = filterM (return . f)
450
451 -- | Drop elements which do not satisfy the monadic predicate
452 filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
453 {-# INLINE_STREAM filterM #-}
454 filterM f (Stream step s n) = Stream step' s (toMax n)
455 where
456 {-# INLINE_INNER step' #-}
457 step' s = do
458 r <- step s
459 case r of
460 Yield x s' -> do
461 b <- f x
462 return $ if b then Yield x s'
463 else Skip s'
464 Skip s' -> return $ Skip s'
465 Done -> return $ Done
466
467 -- | Longest prefix of elements that satisfy the predicate
468 takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
469 {-# INLINE takeWhile #-}
470 takeWhile f = takeWhileM (return . f)
471
472 -- | Longest prefix of elements that satisfy the monadic predicate
473 takeWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
474 {-# INLINE_STREAM takeWhileM #-}
475 takeWhileM f (Stream step s n) = Stream step' s (toMax n)
476 where
477 {-# INLINE_INNER step' #-}
478 step' s = do
479 r <- step s
480 case r of
481 Yield x s' -> do
482 b <- f x
483 return $ if b then Yield x s' else Done
484 Skip s' -> return $ Skip s'
485 Done -> return $ Done
486
487 -- | Drop the longest prefix of elements that satisfy the predicate
488 dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
489 {-# INLINE dropWhile #-}
490 dropWhile f = dropWhileM (return . f)
491
492 data DropWhile s a = DropWhile_Drop s | DropWhile_Yield a s | DropWhile_Next s
493
494 -- | Drop the longest prefix of elements that satisfy the monadic predicate
495 dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
496 {-# INLINE_STREAM dropWhileM #-}
497 dropWhileM f (Stream step s n) = Stream step' (DropWhile_Drop s) (toMax n)
498 where
499 -- NOTE: we jump through hoops here to have only one Yield; local data
500 -- declarations would be nice!
501
502 {-# INLINE_INNER step' #-}
503 step' (DropWhile_Drop s)
504 = do
505 r <- step s
506 case r of
507 Yield x s' -> do
508 b <- f x
509 return $ if b then Skip (DropWhile_Drop s')
510 else Skip (DropWhile_Yield x s')
511 Skip s' -> return $ Skip (DropWhile_Drop s')
512 Done -> return $ Done
513
514 step' (DropWhile_Yield x s) = return $ Yield x (DropWhile_Next s)
515
516 step' (DropWhile_Next s)
517 = liftM (\r ->
518 case r of
519 Yield x s' -> Skip (DropWhile_Yield x s')
520 Skip s' -> Skip (DropWhile_Next s')
521 Done -> Done
522 ) (step s)
523
524 -- Searching
525 -- ---------
526
527 infix 4 `elem`
528 -- | Check whether the 'Stream' contains an element
529 elem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
530 {-# INLINE_STREAM elem #-}
531 elem x (Stream step s _) = elem_loop s
532 where
533 elem_loop s = do
534 r <- step s
535 case r of
536 Yield y s' | x == y -> return True
537 | otherwise -> elem_loop s'
538 Skip s' -> elem_loop s'
539 Done -> return False
540
541 infix 4 `notElem`
542 -- | Inverse of `elem`
543 notElem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
544 {-# INLINE notElem #-}
545 notElem x s = liftM not (elem x s)
546
547 -- | Yield 'Just' the first element that satisfies the predicate or 'Nothing'
548 -- if no such element exists.
549 find :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe a)
550 {-# INLINE find #-}
551 find f = findM (return . f)
552
553 -- | Yield 'Just' the first element that satisfies the monadic predicate or
554 -- 'Nothing' if no such element exists.
555 findM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe a)
556 {-# INLINE_STREAM findM #-}
557 findM f (Stream step s _) = find_loop s
558 where
559 find_loop s = do
560 r <- step s
561 case r of
562 Yield x s' -> do
563 b <- f x
564 if b then return $ Just x
565 else find_loop s'
566 Skip s' -> find_loop s'
567 Done -> return Nothing
568
569 -- | Yield 'Just' the index of the first element that satisfies the predicate
570 -- or 'Nothing' if no such element exists.
571 findIndex :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe Int)
572 {-# INLINE_STREAM findIndex #-}
573 findIndex f = findIndexM (return . f)
574
575 -- | Yield 'Just' the index of the first element that satisfies the monadic
576 -- predicate or 'Nothing' if no such element exists.
577 findIndexM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe Int)
578 {-# INLINE_STREAM findIndexM #-}
579 findIndexM f (Stream step s _) = findIndex_loop s 0
580 where
581 findIndex_loop s i = do
582 r <- step s
583 case r of
584 Yield x s' -> do
585 b <- f x
586 if b then return $ Just i
587 else findIndex_loop s' (i+1)
588 Skip s' -> findIndex_loop s' i
589 Done -> return Nothing
590
591 -- Folding
592 -- -------
593
594 -- | Left fold
595 foldl :: Monad m => (a -> b -> a) -> a -> Stream m b -> m a
596 {-# INLINE foldl #-}
597 foldl f = foldlM (\a b -> return (f a b))
598
599 -- | Left fold with a monadic operator
600 foldlM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> m a
601 {-# INLINE_STREAM foldlM #-}
602 foldlM m z (Stream step s _) = foldlM_go z s
603 where
604 foldlM_go z s = do
605 r <- step s
606 case r of
607 Yield x s' -> do { z' <- m z x; foldlM_go z' s' }
608 Skip s' -> foldlM_go z s'
609 Done -> return z
610
611 -- | Same as 'foldlM'
612 foldM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> m a
613 {-# INLINE foldM #-}
614 foldM = foldlM
615
616 -- | Left fold over a non-empty 'Stream'
617 foldl1 :: Monad m => (a -> a -> a) -> Stream m a -> m a
618 {-# INLINE foldl1 #-}
619 foldl1 f = foldl1M (\a b -> return (f a b))
620
621 -- | Left fold over a non-empty 'Stream' with a monadic operator
622 foldl1M :: Monad m => (a -> a -> m a) -> Stream m a -> m a
623 {-# INLINE_STREAM foldl1M #-}
624 foldl1M f (Stream step s sz) = foldl1M_go s
625 where
626 foldl1M_go s = do
627 r <- step s
628 case r of
629 Yield x s' -> foldlM f x (Stream step s' (sz - 1))
630 Skip s' -> foldl1M_go s'
631 Done -> BOUNDS_ERROR(emptyStream) "foldl1M"
632
633 -- | Same as 'foldl1M'
634 fold1M :: Monad m => (a -> a -> m a) -> Stream m a -> m a
635 {-# INLINE fold1M #-}
636 fold1M = foldl1M
637
638 -- | Left fold with a strict accumulator
639 foldl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> m a
640 {-# INLINE foldl' #-}
641 foldl' f = foldlM' (\a b -> return (f a b))
642
643 -- | Left fold with a strict accumulator and a monadic operator
644 foldlM' :: Monad m => (a -> b -> m a) -> a -> Stream m b -> m a
645 {-# INLINE_STREAM foldlM' #-}
646 foldlM' m z (Stream step s _) = foldlM'_go z s
647 where
648 foldlM'_go z s = z `seq`
649 do
650 r <- step s
651 case r of
652 Yield x s' -> do { z' <- m z x; foldlM'_go z' s' }
653 Skip s' -> foldlM'_go z s'
654 Done -> return z
655
656 -- | Same as 'foldlM''
657 foldM' :: Monad m => (a -> b -> m a) -> a -> Stream m b -> m a
658 {-# INLINE foldM' #-}
659 foldM' = foldlM'
660
661 -- | Left fold over a non-empty 'Stream' with a strict accumulator
662 foldl1' :: Monad m => (a -> a -> a) -> Stream m a -> m a
663 {-# INLINE foldl1' #-}
664 foldl1' f = foldl1M' (\a b -> return (f a b))
665
666 -- | Left fold over a non-empty 'Stream' with a strict accumulator and a
667 -- monadic operator
668 foldl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> m a
669 {-# INLINE_STREAM foldl1M' #-}
670 foldl1M' f (Stream step s sz) = foldl1M'_go s
671 where
672 foldl1M'_go s = do
673 r <- step s
674 case r of
675 Yield x s' -> foldlM' f x (Stream step s' (sz - 1))
676 Skip s' -> foldl1M'_go s'
677 Done -> BOUNDS_ERROR(emptyStream) "foldl1M'"
678
679 -- | Same as 'foldl1M''
680 fold1M' :: Monad m => (a -> a -> m a) -> Stream m a -> m a
681 {-# INLINE fold1M' #-}
682 fold1M' = foldl1M'
683
684 -- | Right fold
685 foldr :: Monad m => (a -> b -> b) -> b -> Stream m a -> m b
686 {-# INLINE foldr #-}
687 foldr f = foldrM (\a b -> return (f a b))
688
689 -- | Right fold with a monadic operator
690 foldrM :: Monad m => (a -> b -> m b) -> b -> Stream m a -> m b
691 {-# INLINE_STREAM foldrM #-}
692 foldrM f z (Stream step s _) = foldrM_go s
693 where
694 foldrM_go s = do
695 r <- step s
696 case r of
697 Yield x s' -> f x =<< foldrM_go s'
698 Skip s' -> foldrM_go s'
699 Done -> return z
700
701 -- | Right fold over a non-empty stream
702 foldr1 :: Monad m => (a -> a -> a) -> Stream m a -> m a
703 {-# INLINE foldr1 #-}
704 foldr1 f = foldr1M (\a b -> return (f a b))
705
706 -- | Right fold over a non-empty stream with a monadic operator
707 foldr1M :: Monad m => (a -> a -> m a) -> Stream m a -> m a
708 {-# INLINE_STREAM foldr1M #-}
709 foldr1M f (Stream step s _) = foldr1M_go0 s
710 where
711 foldr1M_go0 s = do
712 r <- step s
713 case r of
714 Yield x s' -> foldr1M_go1 x s'
715 Skip s' -> foldr1M_go0 s'
716 Done -> BOUNDS_ERROR(emptyStream) "foldr1M"
717
718 foldr1M_go1 x s = do
719 r <- step s
720 case r of
721 Yield y s' -> f x =<< foldr1M_go1 y s'
722 Skip s' -> foldr1M_go1 x s'
723 Done -> return x
724
725 -- Specialised folds
726 -- -----------------
727
728 and :: Monad m => Stream m Bool -> m Bool
729 {-# INLINE_STREAM and #-}
730 and (Stream step s _) = and_go s
731 where
732 and_go s = do
733 r <- step s
734 case r of
735 Yield False _ -> return False
736 Yield True s' -> and_go s'
737 Skip s' -> and_go s'
738 Done -> return True
739
740 or :: Monad m => Stream m Bool -> m Bool
741 {-# INLINE_STREAM or #-}
742 or (Stream step s _) = or_go s
743 where
744 or_go s = do
745 r <- step s
746 case r of
747 Yield False s' -> or_go s'
748 Yield True _ -> return True
749 Skip s' -> or_go s'
750 Done -> return False
751
752 concatMap :: Monad m => (a -> Stream m b) -> Stream m a -> Stream m b
753 {-# INLINE concatMap #-}
754 concatMap f = concatMapM (return . f)
755
756 concatMapM :: Monad m => (a -> m (Stream m b)) -> Stream m a -> Stream m b
757 {-# INLINE_STREAM concatMapM #-}
758 concatMapM f (Stream step s _) = Stream concatMap_go (Left s) Unknown
759 where
760 concatMap_go (Left s) = do
761 r <- step s
762 case r of
763 Yield a s' -> do
764 b_stream <- f a
765 return $ Skip (Right (b_stream, s'))
766 Skip s' -> return $ Skip (Left s')
767 Done -> return Done
768 concatMap_go (Right (Stream inner_step inner_s sz, s)) = do
769 r <- inner_step inner_s
770 case r of
771 Yield b inner_s' -> return $ Yield b (Right (Stream inner_step inner_s' sz, s))
772 Skip inner_s' -> return $ Skip (Right (Stream inner_step inner_s' sz, s))
773 Done -> return $ Skip (Left s)
774
775 -- Unfolding
776 -- ---------
777
778 -- | Unfold
779 unfoldr :: Monad m => (s -> Maybe (a, s)) -> s -> Stream m a
780 {-# INLINE_STREAM unfoldr #-}
781 unfoldr f = unfoldrM (return . f)
782
783 -- | Unfold with a monadic function
784 unfoldrM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m a
785 {-# INLINE_STREAM unfoldrM #-}
786 unfoldrM f s = Stream step s Unknown
787 where
788 {-# INLINE_INNER step #-}
789 step s = liftM (\r ->
790 case r of
791 Just (x, s') -> Yield x s'
792 Nothing -> Done
793 ) (f s)
794
795 -- Scans
796 -- -----
797
798 -- | Prefix scan
799 prescanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
800 {-# INLINE prescanl #-}
801 prescanl f = prescanlM (\a b -> return (f a b))
802
803 -- | Prefix scan with a monadic operator
804 prescanlM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a
805 {-# INLINE_STREAM prescanlM #-}
806 prescanlM f z (Stream step s sz) = Stream step' (s,z) sz
807 where
808 {-# INLINE_INNER step' #-}
809 step' (s,x) = do
810 r <- step s
811 case r of
812 Yield y s' -> do
813 z <- f x y
814 return $ Yield x (s', z)
815 Skip s' -> return $ Skip (s', x)
816 Done -> return Done
817
818 -- | Prefix scan with strict accumulator
819 prescanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
820 {-# INLINE prescanl' #-}
821 prescanl' f = prescanlM' (\a b -> return (f a b))
822
823 -- | Prefix scan with strict accumulator and a monadic operator
824 prescanlM' :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a
825 {-# INLINE_STREAM prescanlM' #-}
826 prescanlM' f z (Stream step s sz) = Stream step' (s,z) sz
827 where
828 {-# INLINE_INNER step' #-}
829 step' (s,x) = x `seq`
830 do
831 r <- step s
832 case r of
833 Yield y s' -> do
834 z <- f x y
835 return $ Yield x (s', z)
836 Skip s' -> return $ Skip (s', x)
837 Done -> return Done
838
839 -- | Suffix scan
840 postscanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
841 {-# INLINE postscanl #-}
842 postscanl f = postscanlM (\a b -> return (f a b))
843
844 -- | Suffix scan with a monadic operator
845 postscanlM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a
846 {-# INLINE_STREAM postscanlM #-}
847 postscanlM f z (Stream step s sz) = Stream step' (s,z) sz
848 where
849 {-# INLINE_INNER step' #-}
850 step' (s,x) = do
851 r <- step s
852 case r of
853 Yield y s' -> do
854 z <- f x y
855 return $ Yield z (s',z)
856 Skip s' -> return $ Skip (s',x)
857 Done -> return Done
858
859 -- | Suffix scan with strict accumulator
860 postscanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
861 {-# INLINE postscanl' #-}
862 postscanl' f = postscanlM' (\a b -> return (f a b))
863
864 -- | Suffix scan with strict acccumulator and a monadic operator
865 postscanlM' :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a
866 {-# INLINE_STREAM postscanlM' #-}
867 postscanlM' f z (Stream step s sz) = z `seq` Stream step' (s,z) sz
868 where
869 {-# INLINE_INNER step' #-}
870 step' (s,x) = x `seq`
871 do
872 r <- step s
873 case r of
874 Yield y s' -> do
875 z <- f x y
876 z `seq` return (Yield z (s',z))
877 Skip s' -> return $ Skip (s',x)
878 Done -> return Done
879
880 -- | Haskell-style scan
881 scanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
882 {-# INLINE scanl #-}
883 scanl f = scanlM (\a b -> return (f a b))
884
885 -- | Haskell-style scan with a monadic operator
886 scanlM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a
887 {-# INLINE scanlM #-}
888 scanlM f z s = z `cons` postscanlM f z s
889
890 -- | Haskell-style scan with strict accumulator
891 scanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
892 {-# INLINE scanl' #-}
893 scanl' f = scanlM' (\a b -> return (f a b))
894
895 -- | Haskell-style scan with strict accumulator and a monadic operator
896 scanlM' :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a
897 {-# INLINE scanlM' #-}
898 scanlM' f z s = z `seq` (z `cons` postscanlM f z s)
899
900 -- | Scan over a non-empty 'Stream'
901 scanl1 :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
902 {-# INLINE scanl1 #-}
903 scanl1 f = scanl1M (\x y -> return (f x y))
904
905 -- | Scan over a non-empty 'Stream' with a monadic operator
906 scanl1M :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
907 {-# INLINE_STREAM scanl1M #-}
908 scanl1M f (Stream step s sz) = Stream step' (s, Nothing) sz
909 where
910 {-# INLINE_INNER step' #-}
911 step' (s, Nothing) = do
912 r <- step s
913 case r of
914 Yield x s' -> return $ Yield x (s', Just x)
915 Skip s' -> return $ Skip (s', Nothing)
916 Done -> BOUNDS_ERROR(emptyStream) "scanl1M"
917
918 step' (s, Just x) = do
919 r <- step s
920 case r of
921 Yield y s' -> do
922 z <- f x y
923 return $ Yield z (s', Just z)
924 Skip s' -> return $ Skip (s', Just x)
925 Done -> return Done
926
927 -- | Scan over a non-empty 'Stream' with a strict accumulator
928 scanl1' :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
929 {-# INLINE scanl1' #-}
930 scanl1' f = scanl1M' (\x y -> return (f x y))
931
932 -- | Scan over a non-empty 'Stream' with a strict accumulator and a monadic
933 -- operator
934 scanl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
935 {-# INLINE_STREAM scanl1M' #-}
936 scanl1M' f (Stream step s sz) = Stream step' (s, Nothing) sz
937 where
938 {-# INLINE_INNER step' #-}
939 step' (s, Nothing) = do
940 r <- step s
941 case r of
942 Yield x s' -> x `seq` return (Yield x (s', Just x))
943 Skip s' -> return $ Skip (s', Nothing)
944 Done -> BOUNDS_ERROR(emptyStream) "scanl1M"
945
946 step' (s, Just x) = x `seq`
947 do
948 r <- step s
949 case r of
950 Yield y s' -> do
951 z <- f x y
952 z `seq` return (Yield z (s', Just z))
953 Skip s' -> return $ Skip (s', Just x)
954 Done -> return Done
955
956 -- Enumerations
957 -- ------------
958
959 -- FIXME: The Enum class is broken for this, there just doesn't seem to be a
960 -- way to implement this generically. We have to either specialise (which
961 -- doesn't help in polymorphic loops) or define a new Enum-like class with a
962 -- proper interface.
963
964 -- | Enumerate values from @x@ to @y@
965 enumFromTo :: (Enum a, Monad m) => a -> a -> Stream m a
966 {-# INLINE_STREAM enumFromTo #-}
967 enumFromTo x y = fromList [x .. y]
968
969 enumFromTo_small :: (Enum a, Ord a, Monad m) => a -> a -> Stream m a
970 {-# INLINE_STREAM enumFromTo_small #-}
971 enumFromTo_small x y = Stream step x (Exact n)
972 where
973 n = max (fromEnum y - fromEnum x + 1) 0
974
975 {-# INLINE_INNER step #-}
976 step x | x <= y = return $ Yield x (succ x)
977 | otherwise = return $ Done
978
979 {-# RULES
980
981 "enumFromTo<Int> [Stream]"
982 enumFromTo = enumFromTo_small :: Monad m => Int -> Int -> Stream m Int
983
984 "enumFromTo<Char> [Stream]"
985 enumFromTo = enumFromTo_small :: Monad m => Char -> Char -> Stream m Char
986
987 "enumFromTo<Int8> [Stream]"
988 enumFromTo = enumFromTo_small :: Monad m => Int8 -> Int8 -> Stream m Int8
989
990 "enumFromTo<Int16> [Stream]"
991 enumFromTo = enumFromTo_small :: Monad m => Int16 -> Int16 -> Stream m Int16
992
993 "enumFromTo<Int32> [Stream]"
994 enumFromTo = enumFromTo_small :: Monad m => Int32 -> Int32 -> Stream m Int32
995
996 "enumFromTo<Word8> [Stream]"
997 enumFromTo = enumFromTo_small :: Monad m => Word8 -> Word8 -> Stream m Word8
998
999 "enumFromTo<Word16> [Stream]"
1000 enumFromTo = enumFromTo_small :: Monad m => Word16 -> Word16 -> Stream m Word16
1001
1002 #-}
1003
1004 enumFromTo_big :: (Enum a, Integral a, Monad m) => a -> a -> Stream m a
1005 {-# INLINE_STREAM enumFromTo_big #-}
1006 enumFromTo_big x y = Stream step x (Exact n)
1007 where
1008 n | x > y = 0
1009 | y - x < fromIntegral (maxBound :: Int) = fromIntegral (y-x+1)
1010 | otherwise = error $ "vector.enumFromTo_big: Array too large"
1011
1012 {-# INLINE_INNER step #-}
1013 step x | x <= y = return $ Yield x (succ x)
1014 | otherwise = return $ Done
1015
1016 {-# RULES
1017
1018
1019 "enumFromTo<Int64> [Stream]"
1020 enumFromTo = enumFromTo_big :: Monad m => Int64 -> Int64 -> Stream m Int64
1021
1022 "enumFromTo<Word32> [Stream]"
1023 enumFromTo = enumFromTo_big :: Monad m => Word32 -> Word32 -> Stream m Word32
1024
1025 "enumFromTo<Word64> [Stream]"
1026 enumFromTo = enumFromTo_big :: Monad m => Word64 -> Word64 -> Stream m Word64
1027
1028 "enumFromTo<Integer> [Stream]"
1029 enumFromTo = enumFromTo_big :: Monad m => Integer -> Integer -> Stream m Integer
1030
1031 #-}
1032
1033
1034 -- | Enumerate values from @x@ to @y@
1035 enumFromThenTo :: (Enum a, Monad m) => a -> a -> a -> Stream m a
1036 {-# INLINE_STREAM enumFromThenTo #-}
1037 enumFromThenTo x y z = fromList [x, y .. y]
1038
1039 -- Conversions
1040 -- -----------
1041
1042 -- | Convert a 'Stream' to a list
1043 toList :: Monad m => Stream m a -> m [a]
1044 {-# INLINE toList #-}
1045 toList = foldr (:) []
1046
1047 -- | Convert a list to a 'Stream'
1048 fromList :: Monad m => [a] -> Stream m a
1049 {-# INLINE_STREAM fromList #-}
1050 fromList xs = Stream step xs Unknown
1051 where
1052 step (x:xs) = return (Yield x xs)
1053 step [] = return Done
1054