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