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