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