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