Allow streams to produce entire vectors as well as individual elements
[darcs-mirrors/vector.git] / Data / Vector / Fusion / Stream / Monadic.hs
1 {-# LANGUAGE ExistentialQuantification, MultiParamTypeClasses, FlexibleInstances, Rank2Types, BangPatterns, KindSignatures, GADTs, ScopedTypeVariables #-}
2
3 -- |
4 -- Module : Data.Vector.Fusion.Stream.Monadic
5 -- Copyright : (c) Roman Leshchinskiy 2008-2010
6 -- License : BSD-style
7 --
8 -- Maintainer : Roman Leshchinskiy <rl@cse.unsw.edu.au>
9 -- Stability : experimental
10 -- Portability : non-portable
11 --
12 -- Monadic stream combinators.
13 --
14
15 module Data.Vector.Fusion.Stream.Monadic (
16 Stream(..), Unf(..), Step(..), SPEC(..),
17
18 simple,
19
20 -- * Size hints
21 size, sized,
22
23 -- * Length
24 length, null,
25
26 -- * Construction
27 empty, singleton, cons, snoc, replicate, replicateM, generate, generateM, (++),
28
29 -- * Accessing elements
30 head, last, (!!), (!?),
31
32 -- * Substreams
33 slice, init, tail, take, drop,
34
35 -- * Mapping
36 map, mapM, mapM_, trans, unbox, concatMap, flatten,
37
38 -- * Zipping
39 indexed, indexedR, zipWithM_,
40 zipWithM, zipWith3M, zipWith4M, zipWith5M, zipWith6M,
41 zipWith, zipWith3, zipWith4, zipWith5, zipWith6,
42 zip, zip3, zip4, zip5, zip6,
43
44 -- * Filtering
45 filter, filterM, takeWhile, takeWhileM, dropWhile, dropWhileM,
46
47 -- * Searching
48 elem, notElem, find, findM, findIndex, findIndexM,
49
50 -- * Folding
51 foldl, foldlM, foldl1, foldl1M, foldM, fold1M,
52 foldl', foldlM', foldl1', foldl1M', foldM', fold1M',
53 foldr, foldrM, foldr1, foldr1M,
54
55 vfoldlM, vfoldlM',
56
57 -- * Specialised folds
58 and, or, concatMapM,
59
60 -- * Unfolding
61 unfoldr, unfoldrM,
62 unfoldrN, unfoldrNM,
63 iterateN, iterateNM,
64
65 -- * Scans
66 prescanl, prescanlM, prescanl', prescanlM',
67 postscanl, postscanlM, postscanl', postscanlM',
68 scanl, scanlM, scanl', scanlM',
69 scanl1, scanl1M, scanl1', scanl1M',
70
71 -- * Enumerations
72 enumFromStepN, enumFromTo, enumFromThenTo,
73
74 -- * Conversions
75 toList, fromList, fromListN, unsafeFromList,
76 fromVector, reVector
77 ) where
78
79 import Data.Vector.Generic.Base
80 import Data.Vector.Fusion.Stream.Size
81 import Data.Vector.Fusion.Util ( Box(..), delay_inline )
82
83 import Data.Char ( ord )
84 import GHC.Base ( unsafeChr )
85 import Control.Monad ( liftM )
86 import Prelude hiding ( length, null,
87 replicate, (++),
88 head, last, (!!),
89 init, tail, take, drop,
90 map, mapM, mapM_, concatMap,
91 zipWith, zipWith3, zip, zip3,
92 filter, takeWhile, dropWhile,
93 elem, notElem,
94 foldl, foldl1, foldr, foldr1,
95 and, or,
96 scanl, scanl1,
97 enumFromTo, enumFromThenTo )
98
99 import Data.Int ( Int8, Int16, Int32, Int64 )
100 import Data.Word ( Word8, Word16, Word32, Word, Word64 )
101
102 #if __GLASGOW_HASKELL__ >= 700
103 import GHC.Exts ( SpecConstrAnnotation(..) )
104 #endif
105
106 #include "vector.h"
107
108 data SPEC = SPEC | SPEC2
109 #if __GLASGOW_HASKELL__ >= 700
110 {-# ANN type SPEC ForceSpecConstr #-}
111 #endif
112
113 emptyStream :: String
114 {-# NOINLINE emptyStream #-}
115 emptyStream = "empty stream"
116
117 #define EMPTY_STREAM (\s -> ERROR s emptyStream)
118
119 -- | Result of taking a single step in a stream
120 data Step s a where
121 Yield :: a -> s -> Step s a
122 Skip :: s -> Step s a
123 Done :: Step s a
124
125 instance Functor (Step s) where
126 {-# INLINE fmap #-}
127 fmap f (Yield x s) = Yield (f x) s
128 fmap f (Skip s) = Skip s
129 fmap f Done = Done
130
131 data Unf m a = forall s. Unf (s -> m (Step s a)) s
132
133 instance Monad m => Functor (Unf m) where
134 {-# INLINE fmap #-}
135 fmap f (Unf step s) = Unf step' s
136 where
137 step' s = do r <- step s ; return (fmap f r)
138
139 -- | Monadic streams
140 data Stream m v a = Stream (Unf m a) (Unf m (Either a (v a))) Size
141
142 simple :: Monad m => (s -> m (Step s a)) -> s -> Size -> Stream m v a
143 {-# INLINE simple #-}
144 simple step s sz = Stream (Unf step s) (Unf step' s) sz
145 where
146 step' s = do r <- step s ; return (fmap Left r)
147
148 -- | 'Size' hint of a 'Stream'
149 size :: Stream m v a -> Size
150 {-# INLINE size #-}
151 size (Stream _ _ sz) = sz
152
153 -- | Attach a 'Size' hint to a 'Stream'
154 sized :: Stream m v a -> Size -> Stream m v a
155 {-# INLINE_STREAM sized #-}
156 sized (Stream p q _) sz = Stream p q sz
157
158 -- Length
159 -- ------
160
161 -- | Length of a 'Stream'
162 length :: Monad m => Stream m v a -> m Int
163 {-# INLINE_STREAM length #-}
164 length s = foldl' (\n _ -> n+1) 0 s
165
166 -- | Check if a 'Stream' is empty
167 null :: Monad m => Stream m v a -> m Bool
168 {-# INLINE_STREAM null #-}
169 null s = foldr (\_ _ -> False) True s
170
171
172 -- Construction
173 -- ------------
174
175 -- | Empty 'Stream'
176 empty :: Monad m => Stream m v a
177 {-# INLINE_STREAM empty #-}
178 empty = simple (const (return Done)) () (Exact 0)
179
180 -- | Singleton 'Stream'
181 singleton :: Monad m => a -> Stream m v a
182 {-# INLINE_STREAM singleton #-}
183 singleton x = simple (return . step) True (Exact 1)
184 where
185 {-# INLINE_INNER step #-}
186 step True = Yield x False
187 step False = Done
188
189 -- | Replicate a value to a given length
190 replicate :: Monad m => Int -> a -> Stream m v a
191 {-# INLINE replicate #-}
192 replicate n x = replicateM n (return x)
193
194 -- | Yield a 'Stream' of values obtained by performing the monadic action the
195 -- given number of times
196 replicateM :: Monad m => Int -> m a -> Stream m v a
197 {-# INLINE_STREAM replicateM #-}
198 -- NOTE: We delay inlining max here because GHC will create a join point for
199 -- the call to newArray# otherwise which is not really nice.
200 replicateM n p = simple step n (Exact (delay_inline max n 0))
201 where
202 {-# INLINE_INNER step #-}
203 step i | i <= 0 = return Done
204 | otherwise = do { x <- p; return $ Yield x (i-1) }
205
206 generate :: Monad m => Int -> (Int -> a) -> Stream m v a
207 {-# INLINE generate #-}
208 generate n f = generateM n (return . f)
209
210 -- | Generate a stream from its indices
211 generateM :: Monad m => Int -> (Int -> m a) -> Stream m v a
212 {-# INLINE_STREAM generateM #-}
213 generateM n f = n `seq` simple step 0 (Exact (delay_inline max n 0))
214 where
215 {-# INLINE_INNER step #-}
216 step i | i < n = do
217 x <- f i
218 return $ Yield x (i+1)
219 | otherwise = return Done
220
221 -- | Prepend an element
222 cons :: Monad m => a -> Stream m v a -> Stream m v a
223 {-# INLINE cons #-}
224 cons x s = singleton x ++ s
225
226 -- | Append an element
227 snoc :: Monad m => Stream m v a -> a -> Stream m v a
228 {-# INLINE snoc #-}
229 snoc s x = s ++ singleton x
230
231 infixr 5 ++
232 -- | Concatenate two 'Stream's
233 (++) :: Monad m => Stream m v a -> Stream m v a -> Stream m v a
234 {-# INLINE_STREAM (++) #-}
235 Stream (Unf stepa sa) (Unf vstepa vsa) na
236 ++ Stream (Unf stepb sb) (Unf vstepb vsb) nb
237 = Stream (Unf step (Left sa)) (Unf vstep (Left vsa)) (na + nb)
238 where
239 {-# INLINE_INNER step #-}
240 step (Left sa) = do
241 r <- stepa sa
242 case r of
243 Yield x sa' -> return $ Yield x (Left sa')
244 Skip sa' -> return $ Skip (Left sa')
245 Done -> return $ Skip (Right sb)
246 step (Right sb) = do
247 r <- stepb sb
248 case r of
249 Yield x sb' -> return $ Yield x (Right sb')
250 Skip sb' -> return $ Skip (Right sb')
251 Done -> return $ Done
252
253 {-# INLINE_INNER vstep #-}
254 vstep (Left vsa) = do
255 r <- vstepa vsa
256 case r of
257 Yield x vsa' -> return $ Yield x (Left vsa')
258 Skip vsa' -> return $ Skip (Left vsa')
259 Done -> return $ Skip (Right vsb)
260 vstep (Right vsb) = do
261 r <- vstepb vsb
262 case r of
263 Yield x vsb' -> return $ Yield x (Right vsb')
264 Skip vsb' -> return $ Skip (Right vsb')
265 Done -> return $ Done
266
267 -- Accessing elements
268 -- ------------------
269
270 -- | First element of the 'Stream' or error if empty
271 head :: Monad m => Stream m v a -> m a
272 {-# INLINE_STREAM head #-}
273 head (Stream (Unf step s) _ _) = head_loop SPEC s
274 where
275 head_loop !sPEC s
276 = do
277 r <- step s
278 case r of
279 Yield x _ -> return x
280 Skip s' -> head_loop SPEC s'
281 Done -> EMPTY_STREAM "head"
282
283
284
285 -- | Last element of the 'Stream' or error if empty
286 last :: Monad m => Stream m v a -> m a
287 {-# INLINE_STREAM last #-}
288 last (Stream (Unf step s) _ _) = last_loop0 SPEC s
289 where
290 last_loop0 !sPEC s
291 = do
292 r <- step s
293 case r of
294 Yield x s' -> last_loop1 SPEC x s'
295 Skip s' -> last_loop0 SPEC s'
296 Done -> EMPTY_STREAM "last"
297
298 last_loop1 !sPEC x s
299 = do
300 r <- step s
301 case r of
302 Yield y s' -> last_loop1 SPEC y s'
303 Skip s' -> last_loop1 SPEC x s'
304 Done -> return x
305
306 infixl 9 !!
307 -- | Element at the given position
308 (!!) :: Monad m => Stream m v a -> Int -> m a
309 {-# INLINE (!!) #-}
310 Stream (Unf step s) _ _ !! i | i < 0 = ERROR "!!" "negative index"
311 | otherwise = index_loop SPEC s i
312 where
313 index_loop !sPEC s i
314 = i `seq`
315 do
316 r <- step s
317 case r of
318 Yield x s' | i == 0 -> return x
319 | otherwise -> index_loop SPEC s' (i-1)
320 Skip s' -> index_loop SPEC s' i
321 Done -> EMPTY_STREAM "!!"
322
323 infixl 9 !?
324 -- | Element at the given position or 'Nothing' if out of bounds
325 (!?) :: Monad m => Stream m v a -> Int -> m (Maybe a)
326 {-# INLINE (!?) #-}
327 Stream (Unf step s) _ _ !? i = index_loop SPEC s i
328 where
329 index_loop !sPEC s i
330 = i `seq`
331 do
332 r <- step s
333 case r of
334 Yield x s' | i == 0 -> return (Just x)
335 | otherwise -> index_loop SPEC s' (i-1)
336 Skip s' -> index_loop SPEC s' i
337 Done -> return Nothing
338
339 -- Substreams
340 -- ----------
341
342 -- | Extract a substream of the given length starting at the given position.
343 slice :: Monad m => Int -- ^ starting index
344 -> Int -- ^ length
345 -> Stream m v a
346 -> Stream m v a
347 {-# INLINE slice #-}
348 slice i n s = take n (drop i s)
349
350 -- | All but the last element
351 init :: Monad m => Stream m v a -> Stream m v a
352 {-# INLINE_STREAM init #-}
353 init (Stream (Unf step s) _ sz) = simple step' (Nothing, s) (sz - 1)
354 where
355 {-# INLINE_INNER step' #-}
356 step' (Nothing, s) = liftM (\r ->
357 case r of
358 Yield x s' -> Skip (Just x, s')
359 Skip s' -> Skip (Nothing, s')
360 Done -> EMPTY_STREAM "init"
361 ) (step s)
362
363 step' (Just x, s) = liftM (\r ->
364 case r of
365 Yield y s' -> Yield x (Just y, s')
366 Skip s' -> Skip (Just x, s')
367 Done -> Done
368 ) (step s)
369
370 -- | All but the first element
371 tail :: Monad m => Stream m v a -> Stream m v a
372 {-# INLINE_STREAM tail #-}
373 tail (Stream (Unf step s) _ sz) = simple step' (Left s) (sz - 1)
374 where
375 {-# INLINE_INNER step' #-}
376 step' (Left s) = liftM (\r ->
377 case r of
378 Yield x s' -> Skip (Right s')
379 Skip s' -> Skip (Left s')
380 Done -> EMPTY_STREAM "tail"
381 ) (step s)
382
383 step' (Right s) = liftM (\r ->
384 case r of
385 Yield x s' -> Yield x (Right s')
386 Skip s' -> Skip (Right s')
387 Done -> Done
388 ) (step s)
389
390 -- | The first @n@ elements
391 take :: Monad m => Int -> Stream m v a -> Stream m v a
392 {-# INLINE_STREAM take #-}
393 take n (Stream (Unf step s) _ sz) = simple step' (s, 0) (smaller (Exact n) sz)
394 where
395 {-# INLINE_INNER step' #-}
396 step' (s, i) | i < n = liftM (\r ->
397 case r of
398 Yield x s' -> Yield x (s', i+1)
399 Skip s' -> Skip (s', i)
400 Done -> Done
401 ) (step s)
402 step' (s, i) = return Done
403
404 -- | All but the first @n@ elements
405 drop :: Monad m => Int -> Stream m v a -> Stream m v a
406 {-# INLINE_STREAM drop #-}
407 drop n (Stream (Unf step s) _ sz) = simple step' (s, Just n) (sz - Exact n)
408 where
409 {-# INLINE_INNER step' #-}
410 step' (s, Just i) | i > 0 = liftM (\r ->
411 case r of
412 Yield x s' -> Skip (s', Just (i-1))
413 Skip s' -> Skip (s', Just i)
414 Done -> Done
415 ) (step s)
416 | otherwise = return $ Skip (s, Nothing)
417
418 step' (s, Nothing) = liftM (\r ->
419 case r of
420 Yield x s' -> Yield x (s', Nothing)
421 Skip s' -> Skip (s', Nothing)
422 Done -> Done
423 ) (step s)
424
425 -- Mapping
426 -- -------
427
428 instance Monad m => Functor (Stream m v) where
429 {-# INLINE fmap #-}
430 fmap = map
431
432 -- | Map a function over a 'Stream'
433 map :: Monad m => (a -> b) -> Stream m v a -> Stream m v b
434 {-# INLINE map #-}
435 map f = mapM (return . f)
436
437
438 -- | Map a monadic function over a 'Stream'
439 mapM :: Monad m => (a -> m b) -> Stream m v a -> Stream m v b
440 {-# INLINE_STREAM mapM #-}
441 mapM f (Stream (Unf step s) _ n) = simple step' s n
442 where
443 {-# INLINE_INNER step' #-}
444 step' s = do
445 r <- step s
446 case r of
447 Yield x s' -> liftM (`Yield` s') (f x)
448 Skip s' -> return (Skip s')
449 Done -> return Done
450
451 consume :: Monad m => Stream m v a -> m ()
452 {-# INLINE_STREAM consume #-}
453 consume (Stream _ (Unf step s) _) = consume_loop SPEC s
454 where
455 consume_loop !sPEC s
456 = do
457 r <- step s
458 case r of
459 Yield _ s' -> consume_loop SPEC s'
460 Skip s' -> consume_loop SPEC s'
461 Done -> return ()
462
463 -- | Execute a monadic action for each element of the 'Stream'
464 mapM_ :: Monad m => (a -> m b) -> Stream m v a -> m ()
465 {-# INLINE_STREAM mapM_ #-}
466 mapM_ m = consume . mapM m
467
468 -- | Transform a 'Stream' to use a different monad
469 trans :: (Monad m, Monad m') => (forall a. m a -> m' a)
470 -> Stream m v a -> Stream m' v a
471 {-# INLINE_STREAM trans #-}
472 trans f (Stream (Unf step s) _ n) = simple (f . step) s n
473
474 unbox :: Monad m => Stream m v (Box a) -> Stream m v a
475 {-# INLINE_STREAM unbox #-}
476 unbox (Stream (Unf step s) _ n) = simple step' s n
477 where
478 {-# INLINE_INNER step' #-}
479 step' s = do
480 r <- step s
481 case r of
482 Yield (Box x) s' -> return $ Yield x s'
483 Skip s' -> return $ Skip s'
484 Done -> return $ Done
485
486 -- Zipping
487 -- -------
488
489 -- | Pair each element in a 'Stream' with its index
490 indexed :: Monad m => Stream m v a -> Stream m v (Int,a)
491 {-# INLINE_STREAM indexed #-}
492 indexed (Stream (Unf step s) _ n) = simple step' (s,0) n
493 where
494 {-# INLINE_INNER step' #-}
495 step' (s,i) = i `seq`
496 do
497 r <- step s
498 case r of
499 Yield x s' -> return $ Yield (i,x) (s', i+1)
500 Skip s' -> return $ Skip (s', i)
501 Done -> return Done
502
503 -- | Pair each element in a 'Stream' with its index, starting from the right
504 -- and counting down
505 indexedR :: Monad m => Int -> Stream m v a -> Stream m v (Int,a)
506 {-# INLINE_STREAM indexedR #-}
507 indexedR m (Stream (Unf step s) _ n) = simple step' (s,m) n
508 where
509 {-# INLINE_INNER step' #-}
510 step' (s,i) = i `seq`
511 do
512 r <- step s
513 case r of
514 Yield x s' -> let i' = i-1
515 in
516 return $ Yield (i',x) (s', i')
517 Skip s' -> return $ Skip (s', i)
518 Done -> return Done
519
520 -- | Zip two 'Stream's with the given monadic function
521 zipWithM :: Monad m => (a -> b -> m c) -> Stream m v a -> Stream m v b -> Stream m v c
522 {-# INLINE_STREAM zipWithM #-}
523 zipWithM f (Stream (Unf stepa sa) _ na) (Stream (Unf stepb sb) _ nb)
524 = simple step (sa, sb, Nothing) (smaller na nb)
525 where
526 {-# INLINE_INNER step #-}
527 step (sa, sb, Nothing) = liftM (\r ->
528 case r of
529 Yield x sa' -> Skip (sa', sb, Just x)
530 Skip sa' -> Skip (sa', sb, Nothing)
531 Done -> Done
532 ) (stepa sa)
533
534 step (sa, sb, Just x) = do
535 r <- stepb sb
536 case r of
537 Yield y sb' ->
538 do
539 z <- f x y
540 return $ Yield z (sa, sb', Nothing)
541 Skip sb' -> return $ Skip (sa, sb', Just x)
542 Done -> return $ Done
543
544 -- FIXME: This might expose an opportunity for inplace execution.
545 {-# RULES
546
547 "zipWithM xs xs [Vector.Stream]" forall f xs.
548 zipWithM f xs xs = mapM (\x -> f x x) xs
549
550 #-}
551
552 zipWithM_ :: Monad m => (a -> b -> m c) -> Stream m v a -> Stream m v b -> m ()
553 {-# INLINE zipWithM_ #-}
554 zipWithM_ f sa sb = consume (zipWithM f sa sb)
555
556 zipWith3M :: Monad m => (a -> b -> c -> m d) -> Stream m v a -> Stream m v b -> Stream m v c -> Stream m v d
557 {-# INLINE_STREAM zipWith3M #-}
558 zipWith3M f (Stream (Unf stepa sa) _ na) (Stream (Unf stepb sb) _ nb) (Stream (Unf stepc sc) _ nc)
559 = simple step (sa, sb, sc, Nothing) (smaller na (smaller nb nc))
560 where
561 {-# INLINE_INNER step #-}
562 step (sa, sb, sc, Nothing) = do
563 r <- stepa sa
564 return $ case r of
565 Yield x sa' -> Skip (sa', sb, sc, Just (x, Nothing))
566 Skip sa' -> Skip (sa', sb, sc, Nothing)
567 Done -> Done
568
569 step (sa, sb, sc, Just (x, Nothing)) = do
570 r <- stepb sb
571 return $ case r of
572 Yield y sb' -> Skip (sa, sb', sc, Just (x, Just y))
573 Skip sb' -> Skip (sa, sb', sc, Just (x, Nothing))
574 Done -> Done
575
576 step (sa, sb, sc, Just (x, Just y)) = do
577 r <- stepc sc
578 case r of
579 Yield z sc' -> f x y z >>= (\res -> return $ Yield res (sa, sb, sc', Nothing))
580 Skip sc' -> return $ Skip (sa, sb, sc', Just (x, Just y))
581 Done -> return $ Done
582
583 zipWith4M :: Monad m => (a -> b -> c -> d -> m e)
584 -> Stream m v a -> Stream m v b -> Stream m v c -> Stream m v d
585 -> Stream m v e
586 {-# INLINE zipWith4M #-}
587 zipWith4M f sa sb sc sd
588 = zipWithM (\(a,b) (c,d) -> f a b c d) (zip sa sb) (zip sc sd)
589
590 zipWith5M :: Monad m => (a -> b -> c -> d -> e -> m f)
591 -> Stream m v a -> Stream m v b -> Stream m v c -> Stream m v d
592 -> Stream m v e -> Stream m v f
593 {-# INLINE zipWith5M #-}
594 zipWith5M f sa sb sc sd se
595 = zipWithM (\(a,b,c) (d,e) -> f a b c d e) (zip3 sa sb sc) (zip sd se)
596
597 zipWith6M :: Monad m => (a -> b -> c -> d -> e -> f -> m g)
598 -> Stream m v a -> Stream m v b -> Stream m v c -> Stream m v d
599 -> Stream m v e -> Stream m v f -> Stream m v g
600 {-# INLINE zipWith6M #-}
601 zipWith6M fn sa sb sc sd se sf
602 = zipWithM (\(a,b,c) (d,e,f) -> fn a b c d e f) (zip3 sa sb sc)
603 (zip3 sd se sf)
604
605 zipWith :: Monad m => (a -> b -> c) -> Stream m v a -> Stream m v b -> Stream m v c
606 {-# INLINE zipWith #-}
607 zipWith f = zipWithM (\a b -> return (f a b))
608
609 zipWith3 :: Monad m => (a -> b -> c -> d)
610 -> Stream m v a -> Stream m v b -> Stream m v c -> Stream m v d
611 {-# INLINE zipWith3 #-}
612 zipWith3 f = zipWith3M (\a b c -> return (f a b c))
613
614 zipWith4 :: Monad m => (a -> b -> c -> d -> e)
615 -> Stream m v a -> Stream m v b -> Stream m v c -> Stream m v d
616 -> Stream m v e
617 {-# INLINE zipWith4 #-}
618 zipWith4 f = zipWith4M (\a b c d -> return (f a b c d))
619
620 zipWith5 :: Monad m => (a -> b -> c -> d -> e -> f)
621 -> Stream m v a -> Stream m v b -> Stream m v c -> Stream m v d
622 -> Stream m v e -> Stream m v f
623 {-# INLINE zipWith5 #-}
624 zipWith5 f = zipWith5M (\a b c d e -> return (f a b c d e))
625
626 zipWith6 :: Monad m => (a -> b -> c -> d -> e -> f -> g)
627 -> Stream m v a -> Stream m v b -> Stream m v c -> Stream m v d
628 -> Stream m v e -> Stream m v f -> Stream m v g
629 {-# INLINE zipWith6 #-}
630 zipWith6 fn = zipWith6M (\a b c d e f -> return (fn a b c d e f))
631
632 zip :: Monad m => Stream m v a -> Stream m v b -> Stream m v (a,b)
633 {-# INLINE zip #-}
634 zip = zipWith (,)
635
636 zip3 :: Monad m => Stream m v a -> Stream m v b -> Stream m v c -> Stream m v (a,b,c)
637 {-# INLINE zip3 #-}
638 zip3 = zipWith3 (,,)
639
640 zip4 :: Monad m => Stream m v a -> Stream m v b -> Stream m v c -> Stream m v d
641 -> Stream m v (a,b,c,d)
642 {-# INLINE zip4 #-}
643 zip4 = zipWith4 (,,,)
644
645 zip5 :: Monad m => Stream m v a -> Stream m v b -> Stream m v c -> Stream m v d
646 -> Stream m v e -> Stream m v (a,b,c,d,e)
647 {-# INLINE zip5 #-}
648 zip5 = zipWith5 (,,,,)
649
650 zip6 :: Monad m => Stream m v a -> Stream m v b -> Stream m v c -> Stream m v d
651 -> Stream m v e -> Stream m v f -> Stream m v (a,b,c,d,e,f)
652 {-# INLINE zip6 #-}
653 zip6 = zipWith6 (,,,,,)
654
655 -- Filtering
656 -- ---------
657
658 -- | Drop elements which do not satisfy the predicate
659 filter :: Monad m => (a -> Bool) -> Stream m v a -> Stream m v a
660 {-# INLINE filter #-}
661 filter f = filterM (return . f)
662
663 -- | Drop elements which do not satisfy the monadic predicate
664 filterM :: Monad m => (a -> m Bool) -> Stream m v a -> Stream m v a
665 {-# INLINE_STREAM filterM #-}
666 filterM f (Stream (Unf step s) _ n) = simple step' s (toMax n)
667 where
668 {-# INLINE_INNER step' #-}
669 step' s = do
670 r <- step s
671 case r of
672 Yield x s' -> do
673 b <- f x
674 return $ if b then Yield x s'
675 else Skip s'
676 Skip s' -> return $ Skip s'
677 Done -> return $ Done
678
679 -- | Longest prefix of elements that satisfy the predicate
680 takeWhile :: Monad m => (a -> Bool) -> Stream m v a -> Stream m v a
681 {-# INLINE takeWhile #-}
682 takeWhile f = takeWhileM (return . f)
683
684 -- | Longest prefix of elements that satisfy the monadic predicate
685 takeWhileM :: Monad m => (a -> m Bool) -> Stream m v a -> Stream m v a
686 {-# INLINE_STREAM takeWhileM #-}
687 takeWhileM f (Stream (Unf step s) _ n) = simple step' s (toMax n)
688 where
689 {-# INLINE_INNER step' #-}
690 step' s = do
691 r <- step s
692 case r of
693 Yield x s' -> do
694 b <- f x
695 return $ if b then Yield x s' else Done
696 Skip s' -> return $ Skip s'
697 Done -> return $ Done
698
699 -- | Drop the longest prefix of elements that satisfy the predicate
700 dropWhile :: Monad m => (a -> Bool) -> Stream m v a -> Stream m v a
701 {-# INLINE dropWhile #-}
702 dropWhile f = dropWhileM (return . f)
703
704 data DropWhile s a = DropWhile_Drop s | DropWhile_Yield a s | DropWhile_Next s
705
706 -- | Drop the longest prefix of elements that satisfy the monadic predicate
707 dropWhileM :: Monad m => (a -> m Bool) -> Stream m v a -> Stream m v a
708 {-# INLINE_STREAM dropWhileM #-}
709 dropWhileM f (Stream (Unf step s) _ n) = simple step' (DropWhile_Drop s) (toMax n)
710 where
711 -- NOTE: we jump through hoops here to have only one Yield; local data
712 -- declarations would be nice!
713
714 {-# INLINE_INNER step' #-}
715 step' (DropWhile_Drop s)
716 = do
717 r <- step s
718 case r of
719 Yield x s' -> do
720 b <- f x
721 return $ if b then Skip (DropWhile_Drop s')
722 else Skip (DropWhile_Yield x s')
723 Skip s' -> return $ Skip (DropWhile_Drop s')
724 Done -> return $ Done
725
726 step' (DropWhile_Yield x s) = return $ Yield x (DropWhile_Next s)
727
728 step' (DropWhile_Next s)
729 = liftM (\r ->
730 case r of
731 Yield x s' -> Skip (DropWhile_Yield x s')
732 Skip s' -> Skip (DropWhile_Next s')
733 Done -> Done
734 ) (step s)
735
736 -- Searching
737 -- ---------
738
739 infix 4 `elem`
740 -- | Check whether the 'Stream' contains an element
741 elem :: (Monad m, Eq a) => a -> Stream m v a -> m Bool
742 {-# INLINE_STREAM elem #-}
743 elem x (Stream (Unf step s) _ _) = elem_loop SPEC s
744 where
745 elem_loop !sPEC s
746 = do
747 r <- step s
748 case r of
749 Yield y s' | x == y -> return True
750 | otherwise -> elem_loop SPEC s'
751 Skip s' -> elem_loop SPEC s'
752 Done -> return False
753
754 infix 4 `notElem`
755 -- | Inverse of `elem`
756 notElem :: (Monad m, Eq a) => a -> Stream m v a -> m Bool
757 {-# INLINE notElem #-}
758 notElem x s = liftM not (elem x s)
759
760 -- | Yield 'Just' the first element that satisfies the predicate or 'Nothing'
761 -- if no such element exists.
762 find :: Monad m => (a -> Bool) -> Stream m v a -> m (Maybe a)
763 {-# INLINE find #-}
764 find f = findM (return . f)
765
766 -- | Yield 'Just' the first element that satisfies the monadic predicate or
767 -- 'Nothing' if no such element exists.
768 findM :: Monad m => (a -> m Bool) -> Stream m v a -> m (Maybe a)
769 {-# INLINE_STREAM findM #-}
770 findM f (Stream (Unf step s) _ _) = find_loop SPEC s
771 where
772 find_loop !sPEC s
773 = do
774 r <- step s
775 case r of
776 Yield x s' -> do
777 b <- f x
778 if b then return $ Just x
779 else find_loop SPEC s'
780 Skip s' -> find_loop SPEC s'
781 Done -> return Nothing
782
783 -- | Yield 'Just' the index of the first element that satisfies the predicate
784 -- or 'Nothing' if no such element exists.
785 findIndex :: Monad m => (a -> Bool) -> Stream m v a -> m (Maybe Int)
786 {-# INLINE_STREAM findIndex #-}
787 findIndex f = findIndexM (return . f)
788
789 -- | Yield 'Just' the index of the first element that satisfies the monadic
790 -- predicate or 'Nothing' if no such element exists.
791 findIndexM :: Monad m => (a -> m Bool) -> Stream m v a -> m (Maybe Int)
792 {-# INLINE_STREAM findIndexM #-}
793 findIndexM f (Stream (Unf step s) _ _) = findIndex_loop SPEC s 0
794 where
795 findIndex_loop !sPEC s i
796 = do
797 r <- step s
798 case r of
799 Yield x s' -> do
800 b <- f x
801 if b then return $ Just i
802 else findIndex_loop SPEC s' (i+1)
803 Skip s' -> findIndex_loop SPEC s' i
804 Done -> return Nothing
805
806 -- Folding
807 -- -------
808
809 -- | Left fold
810 foldl :: Monad m => (a -> b -> a) -> a -> Stream m v b -> m a
811 {-# INLINE foldl #-}
812 foldl f = foldlM (\a b -> return (f a b))
813
814 -- | Left fold with a monadic operator
815 foldlM :: Monad m => (a -> b -> m a) -> a -> Stream m v b -> m a
816 {-# INLINE_STREAM foldlM #-}
817 foldlM m z (Stream (Unf step s) _ _) = foldlM_loop SPEC z s
818 where
819 foldlM_loop !sPEC z s
820 = do
821 r <- step s
822 case r of
823 Yield x s' -> do { z' <- m z x; foldlM_loop SPEC z' s' }
824 Skip s' -> foldlM_loop SPEC z s'
825 Done -> return z
826
827 -- | Same as 'foldlM'
828 foldM :: Monad m => (a -> b -> m a) -> a -> Stream m v b -> m a
829 {-# INLINE foldM #-}
830 foldM = foldlM
831
832 vfoldlM :: Monad m => (a -> b -> m a) -> (a -> v b -> m a) -> a -> Stream m v b -> m a
833 {-# INLINE_STREAM vfoldlM #-}
834 vfoldlM p q z (Stream _ (Unf step s) _) = vfoldlM_loop SPEC z s
835 where
836 vfoldlM_loop !sPEC z s
837 = do
838 r <- step s
839 case r of
840 Yield (Left x) s' -> do { z' <- p z x; vfoldlM_loop SPEC z' s' }
841 Yield (Right v) s' -> do { z' <- q z v; vfoldlM_loop SPEC z' s' }
842 Skip s' -> vfoldlM_loop SPEC z s'
843 Done -> return z
844
845 -- | Left fold over a non-empty 'Stream'
846 foldl1 :: Monad m => (a -> a -> a) -> Stream m v a -> m a
847 {-# INLINE foldl1 #-}
848 foldl1 f = foldl1M (\a b -> return (f a b))
849
850 -- | Left fold over a non-empty 'Stream' with a monadic operator
851 foldl1M :: Monad m => (a -> a -> m a) -> Stream m v a -> m a
852 {-# INLINE_STREAM foldl1M #-}
853 foldl1M f (Stream (Unf step s) _ sz) = foldl1M_loop SPEC s
854 where
855 foldl1M_loop !sPEC s
856 = do
857 r <- step s
858 case r of
859 Yield x s' -> foldlM f x (simple step s' (sz - 1))
860 Skip s' -> foldl1M_loop SPEC s'
861 Done -> EMPTY_STREAM "foldl1M"
862
863 -- | Same as 'foldl1M'
864 fold1M :: Monad m => (a -> a -> m a) -> Stream m v a -> m a
865 {-# INLINE fold1M #-}
866 fold1M = foldl1M
867
868 -- | Left fold with a strict accumulator
869 foldl' :: Monad m => (a -> b -> a) -> a -> Stream m v b -> m a
870 {-# INLINE foldl' #-}
871 foldl' f = foldlM' (\a b -> return (f a b))
872
873 -- | Left fold with a strict accumulator and a monadic operator
874 foldlM' :: Monad m => (a -> b -> m a) -> a -> Stream m v b -> m a
875 {-# INLINE_STREAM foldlM' #-}
876 foldlM' m z (Stream (Unf step s) _ _) = foldlM'_loop SPEC z s
877 where
878 foldlM'_loop !sPEC z s
879 = z `seq`
880 do
881 r <- step s
882 case r of
883 Yield x s' -> do { z' <- m z x; foldlM'_loop SPEC z' s' }
884 Skip s' -> foldlM'_loop SPEC z s'
885 Done -> return z
886
887 -- | Same as 'foldlM''
888 foldM' :: Monad m => (a -> b -> m a) -> a -> Stream m v b -> m a
889 {-# INLINE foldM' #-}
890 foldM' = foldlM'
891
892 vfoldlM' :: Monad m => (a -> b -> m a) -> (a -> v b -> m a) -> a -> Stream m v b -> m a
893 {-# INLINE_STREAM vfoldlM' #-}
894 vfoldlM' p q z (Stream _ (Unf step s) _) = vfoldlM'_loop SPEC z s
895 where
896 vfoldlM'_loop !sPEC z s
897 = z `seq` do
898 r <- step s
899 case r of
900 Yield (Left x) s' -> do { z' <- p z x; vfoldlM'_loop SPEC z' s' }
901 Yield (Right v) s' -> do { z' <- q z v; vfoldlM'_loop SPEC z' s' }
902 Skip s' -> vfoldlM'_loop SPEC z s'
903 Done -> return z
904
905 -- | Left fold over a non-empty 'Stream' with a strict accumulator
906 foldl1' :: Monad m => (a -> a -> a) -> Stream m v a -> m a
907 {-# INLINE foldl1' #-}
908 foldl1' f = foldl1M' (\a b -> return (f a b))
909
910 -- | Left fold over a non-empty 'Stream' with a strict accumulator and a
911 -- monadic operator
912 foldl1M' :: Monad m => (a -> a -> m a) -> Stream m v a -> m a
913 {-# INLINE_STREAM foldl1M' #-}
914 foldl1M' f (Stream (Unf step s) _ sz) = foldl1M'_loop SPEC s
915 where
916 foldl1M'_loop !sPEC s
917 = do
918 r <- step s
919 case r of
920 Yield x s' -> foldlM' f x (simple step s' (sz - 1))
921 Skip s' -> foldl1M'_loop SPEC s'
922 Done -> EMPTY_STREAM "foldl1M'"
923
924 -- | Same as 'foldl1M''
925 fold1M' :: Monad m => (a -> a -> m a) -> Stream m v a -> m a
926 {-# INLINE fold1M' #-}
927 fold1M' = foldl1M'
928
929 -- | Right fold
930 foldr :: Monad m => (a -> b -> b) -> b -> Stream m v a -> m b
931 {-# INLINE foldr #-}
932 foldr f = foldrM (\a b -> return (f a b))
933
934 -- | Right fold with a monadic operator
935 foldrM :: Monad m => (a -> b -> m b) -> b -> Stream m v a -> m b
936 {-# INLINE_STREAM foldrM #-}
937 foldrM f z (Stream (Unf step s) _ _) = foldrM_loop SPEC s
938 where
939 foldrM_loop !sPEC s
940 = do
941 r <- step s
942 case r of
943 Yield x s' -> f x =<< foldrM_loop SPEC s'
944 Skip s' -> foldrM_loop SPEC s'
945 Done -> return z
946
947 -- | Right fold over a non-empty stream
948 foldr1 :: Monad m => (a -> a -> a) -> Stream m v a -> m a
949 {-# INLINE foldr1 #-}
950 foldr1 f = foldr1M (\a b -> return (f a b))
951
952 -- | Right fold over a non-empty stream with a monadic operator
953 foldr1M :: Monad m => (a -> a -> m a) -> Stream m v a -> m a
954 {-# INLINE_STREAM foldr1M #-}
955 foldr1M f (Stream (Unf step s) _ _) = foldr1M_loop0 SPEC s
956 where
957 foldr1M_loop0 !sPEC s
958 = do
959 r <- step s
960 case r of
961 Yield x s' -> foldr1M_loop1 SPEC x s'
962 Skip s' -> foldr1M_loop0 SPEC s'
963 Done -> EMPTY_STREAM "foldr1M"
964
965 foldr1M_loop1 !sPEC x s
966 = do
967 r <- step s
968 case r of
969 Yield y s' -> f x =<< foldr1M_loop1 SPEC y s'
970 Skip s' -> foldr1M_loop1 SPEC x s'
971 Done -> return x
972
973 -- Specialised folds
974 -- -----------------
975
976 and :: Monad m => Stream m v Bool -> m Bool
977 {-# INLINE_STREAM and #-}
978 and (Stream (Unf step s) _ _) = and_loop SPEC s
979 where
980 and_loop !sPEC s
981 = do
982 r <- step s
983 case r of
984 Yield False _ -> return False
985 Yield True s' -> and_loop SPEC s'
986 Skip s' -> and_loop SPEC s'
987 Done -> return True
988
989 or :: Monad m => Stream m v Bool -> m Bool
990 {-# INLINE_STREAM or #-}
991 or (Stream (Unf step s) _ _) = or_loop SPEC s
992 where
993 or_loop !sPEC s
994 = do
995 r <- step s
996 case r of
997 Yield False s' -> or_loop SPEC s'
998 Yield True _ -> return True
999 Skip s' -> or_loop SPEC s'
1000 Done -> return False
1001
1002 concatMap :: Monad m => (a -> Stream m v b) -> Stream m v a -> Stream m v b
1003 {-# INLINE concatMap #-}
1004 concatMap f = concatMapM (return . f)
1005
1006 concatMapM :: Monad m => (a -> m (Stream m v b)) -> Stream m v a -> Stream m v b
1007 {-# INLINE_STREAM concatMapM #-}
1008 concatMapM f (Stream (Unf step s) _ _) = simple concatMap_go (Left s) Unknown
1009 where
1010 concatMap_go (Left s) = do
1011 r <- step s
1012 case r of
1013 Yield a s' -> do
1014 b_stream <- f a
1015 return $ Skip (Right (b_stream, s'))
1016 Skip s' -> return $ Skip (Left s')
1017 Done -> return Done
1018 concatMap_go (Right (Stream (Unf inner_step inner_s) _ sz, s)) = do
1019 r <- inner_step inner_s
1020 case r of
1021 Yield b inner_s' -> return $ Yield b (Right (simple inner_step inner_s' sz, s))
1022 Skip inner_s' -> return $ Skip (Right (simple inner_step inner_s' sz, s))
1023 Done -> return $ Skip (Left s)
1024
1025 -- | Create a 'Stream' of values from a 'Stream' of streamable things
1026 flatten :: Monad m => (a -> m s) -> (s -> m (Step s b)) -> Size
1027 -> Stream m v a -> Stream m v b
1028 {-# INLINE_STREAM flatten #-}
1029 flatten mk istep sz (Stream (Unf ostep t) _ _) = simple step (Left t) sz
1030 where
1031 {-# INLINE_INNER step #-}
1032 step (Left t) = do
1033 r <- ostep t
1034 case r of
1035 Yield a t' -> do
1036 s <- mk a
1037 s `seq` return (Skip (Right (s,t')))
1038 Skip t' -> return $ Skip (Left t')
1039 Done -> return $ Done
1040
1041
1042 step (Right (s,t)) = do
1043 r <- istep s
1044 case r of
1045 Yield x s' -> return $ Yield x (Right (s',t))
1046 Skip s' -> return $ Skip (Right (s',t))
1047 Done -> return $ Skip (Left t)
1048
1049 -- Unfolding
1050 -- ---------
1051
1052 -- | Unfold
1053 unfoldr :: Monad m => (s -> Maybe (a, s)) -> s -> Stream m u a
1054 {-# INLINE_STREAM unfoldr #-}
1055 unfoldr f = unfoldrM (return . f)
1056
1057 -- | Unfold with a monadic function
1058 unfoldrM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m u a
1059 {-# INLINE_STREAM unfoldrM #-}
1060 unfoldrM f s = simple step s Unknown
1061 where
1062 {-# INLINE_INNER step #-}
1063 step s = liftM (\r ->
1064 case r of
1065 Just (x, s') -> Yield x s'
1066 Nothing -> Done
1067 ) (f s)
1068
1069 -- | Unfold at most @n@ elements
1070 unfoldrN :: Monad m => Int -> (s -> Maybe (a, s)) -> s -> Stream m u a
1071 {-# INLINE_STREAM unfoldrN #-}
1072 unfoldrN n f = unfoldrNM n (return . f)
1073
1074 -- | Unfold at most @n@ elements with a monadic functions
1075 unfoldrNM :: Monad m => Int -> (s -> m (Maybe (a, s))) -> s -> Stream m u a
1076 {-# INLINE_STREAM unfoldrNM #-}
1077 unfoldrNM n f s = simple step (s,n) (Max (delay_inline max n 0))
1078 where
1079 {-# INLINE_INNER step #-}
1080 step (s,n) | n <= 0 = return Done
1081 | otherwise = liftM (\r ->
1082 case r of
1083 Just (x,s') -> Yield x (s',n-1)
1084 Nothing -> Done
1085 ) (f s)
1086
1087 -- | Apply monadic function n times to value. Zeroth element is original value.
1088 iterateNM :: Monad m => Int -> (a -> m a) -> a -> Stream m u a
1089 {-# INLINE_STREAM iterateNM #-}
1090 iterateNM n f x0 = simple step (x0,n) (Exact (delay_inline max n 0))
1091 where
1092 {-# INLINE_INNER step #-}
1093 step (x,i) | i <= 0 = return Done
1094 | i == n = return $ Yield x (x,i-1)
1095 | otherwise = do a <- f x
1096 return $ Yield a (a,i-1)
1097
1098 -- | Apply function n times to value. Zeroth element is original value.
1099 iterateN :: Monad m => Int -> (a -> a) -> a -> Stream m u a
1100 {-# INLINE_STREAM iterateN #-}
1101 iterateN n f x0 = iterateNM n (return . f) x0
1102
1103 -- Scans
1104 -- -----
1105
1106 -- | Prefix scan
1107 prescanl :: Monad m => (a -> b -> a) -> a -> Stream m v b -> Stream m v a
1108 {-# INLINE prescanl #-}
1109 prescanl f = prescanlM (\a b -> return (f a b))
1110
1111 -- | Prefix scan with a monadic operator
1112 prescanlM :: Monad m => (a -> b -> m a) -> a -> Stream m v b -> Stream m v a
1113 {-# INLINE_STREAM prescanlM #-}
1114 prescanlM f z (Stream (Unf step s) _ sz) = simple step' (s,z) sz
1115 where
1116 {-# INLINE_INNER step' #-}
1117 step' (s,x) = do
1118 r <- step s
1119 case r of
1120 Yield y s' -> do
1121 z <- f x y
1122 return $ Yield x (s', z)
1123 Skip s' -> return $ Skip (s', x)
1124 Done -> return Done
1125
1126 -- | Prefix scan with strict accumulator
1127 prescanl' :: Monad m => (a -> b -> a) -> a -> Stream m v b -> Stream m v a
1128 {-# INLINE prescanl' #-}
1129 prescanl' f = prescanlM' (\a b -> return (f a b))
1130
1131 -- | Prefix scan with strict accumulator and a monadic operator
1132 prescanlM' :: Monad m => (a -> b -> m a) -> a -> Stream m v b -> Stream m v a
1133 {-# INLINE_STREAM prescanlM' #-}
1134 prescanlM' f z (Stream (Unf step s) _ sz) = simple step' (s,z) sz
1135 where
1136 {-# INLINE_INNER step' #-}
1137 step' (s,x) = x `seq`
1138 do
1139 r <- step s
1140 case r of
1141 Yield y s' -> do
1142 z <- f x y
1143 return $ Yield x (s', z)
1144 Skip s' -> return $ Skip (s', x)
1145 Done -> return Done
1146
1147 -- | Suffix scan
1148 postscanl :: Monad m => (a -> b -> a) -> a -> Stream m v b -> Stream m v a
1149 {-# INLINE postscanl #-}
1150 postscanl f = postscanlM (\a b -> return (f a b))
1151
1152 -- | Suffix scan with a monadic operator
1153 postscanlM :: Monad m => (a -> b -> m a) -> a -> Stream m v b -> Stream m v a
1154 {-# INLINE_STREAM postscanlM #-}
1155 postscanlM f z (Stream (Unf step s) _ sz) = simple step' (s,z) sz
1156 where
1157 {-# INLINE_INNER step' #-}
1158 step' (s,x) = do
1159 r <- step s
1160 case r of
1161 Yield y s' -> do
1162 z <- f x y
1163 return $ Yield z (s',z)
1164 Skip s' -> return $ Skip (s',x)
1165 Done -> return Done
1166
1167 -- | Suffix scan with strict accumulator
1168 postscanl' :: Monad m => (a -> b -> a) -> a -> Stream m v b -> Stream m v a
1169 {-# INLINE postscanl' #-}
1170 postscanl' f = postscanlM' (\a b -> return (f a b))
1171
1172 -- | Suffix scan with strict acccumulator and a monadic operator
1173 postscanlM' :: Monad m => (a -> b -> m a) -> a -> Stream m v b -> Stream m v a
1174 {-# INLINE_STREAM postscanlM' #-}
1175 postscanlM' f z (Stream (Unf step s) _ sz) = z `seq` simple step' (s,z) sz
1176 where
1177 {-# INLINE_INNER step' #-}
1178 step' (s,x) = x `seq`
1179 do
1180 r <- step s
1181 case r of
1182 Yield y s' -> do
1183 z <- f x y
1184 z `seq` return (Yield z (s',z))
1185 Skip s' -> return $ Skip (s',x)
1186 Done -> return Done
1187
1188 -- | Haskell-style scan
1189 scanl :: Monad m => (a -> b -> a) -> a -> Stream m v b -> Stream m v a
1190 {-# INLINE scanl #-}
1191 scanl f = scanlM (\a b -> return (f a b))
1192
1193 -- | Haskell-style scan with a monadic operator
1194 scanlM :: Monad m => (a -> b -> m a) -> a -> Stream m v b -> Stream m v a
1195 {-# INLINE scanlM #-}
1196 scanlM f z s = z `cons` postscanlM f z s
1197
1198 -- | Haskell-style scan with strict accumulator
1199 scanl' :: Monad m => (a -> b -> a) -> a -> Stream m v b -> Stream m v a
1200 {-# INLINE scanl' #-}
1201 scanl' f = scanlM' (\a b -> return (f a b))
1202
1203 -- | Haskell-style scan with strict accumulator and a monadic operator
1204 scanlM' :: Monad m => (a -> b -> m a) -> a -> Stream m v b -> Stream m v a
1205 {-# INLINE scanlM' #-}
1206 scanlM' f z s = z `seq` (z `cons` postscanlM f z s)
1207
1208 -- | Scan over a non-empty 'Stream'
1209 scanl1 :: Monad m => (a -> a -> a) -> Stream m v a -> Stream m v a
1210 {-# INLINE scanl1 #-}
1211 scanl1 f = scanl1M (\x y -> return (f x y))
1212
1213 -- | Scan over a non-empty 'Stream' with a monadic operator
1214 scanl1M :: Monad m => (a -> a -> m a) -> Stream m v a -> Stream m v a
1215 {-# INLINE_STREAM scanl1M #-}
1216 scanl1M f (Stream (Unf step s) _ sz) = simple step' (s, Nothing) sz
1217 where
1218 {-# INLINE_INNER step' #-}
1219 step' (s, Nothing) = do
1220 r <- step s
1221 case r of
1222 Yield x s' -> return $ Yield x (s', Just x)
1223 Skip s' -> return $ Skip (s', Nothing)
1224 Done -> EMPTY_STREAM "scanl1M"
1225
1226 step' (s, Just x) = do
1227 r <- step s
1228 case r of
1229 Yield y s' -> do
1230 z <- f x y
1231 return $ Yield z (s', Just z)
1232 Skip s' -> return $ Skip (s', Just x)
1233 Done -> return Done
1234
1235 -- | Scan over a non-empty 'Stream' with a strict accumulator
1236 scanl1' :: Monad m => (a -> a -> a) -> Stream m v a -> Stream m v a
1237 {-# INLINE scanl1' #-}
1238 scanl1' f = scanl1M' (\x y -> return (f x y))
1239
1240 -- | Scan over a non-empty 'Stream' with a strict accumulator and a monadic
1241 -- operator
1242 scanl1M' :: Monad m => (a -> a -> m a) -> Stream m v a -> Stream m v a
1243 {-# INLINE_STREAM scanl1M' #-}
1244 scanl1M' f (Stream (Unf step s) _ sz) = simple step' (s, Nothing) sz
1245 where
1246 {-# INLINE_INNER step' #-}
1247 step' (s, Nothing) = do
1248 r <- step s
1249 case r of
1250 Yield x s' -> x `seq` return (Yield x (s', Just x))
1251 Skip s' -> return $ Skip (s', Nothing)
1252 Done -> EMPTY_STREAM "scanl1M"
1253
1254 step' (s, Just x) = x `seq`
1255 do
1256 r <- step s
1257 case r of
1258 Yield y s' -> do
1259 z <- f x y
1260 z `seq` return (Yield z (s', Just z))
1261 Skip s' -> return $ Skip (s', Just x)
1262 Done -> return Done
1263
1264 -- Enumerations
1265 -- ------------
1266
1267 -- The Enum class is broken for this, there just doesn't seem to be a
1268 -- way to implement this generically. We have to specialise for as many types
1269 -- as we can but this doesn't help in polymorphic loops.
1270
1271 -- | Yield a 'Stream' of the given length containing the values @x@, @x+y@,
1272 -- @x+y+y@ etc.
1273 enumFromStepN :: (Num a, Monad m) => a -> a -> Int -> Stream m v a
1274 {-# INLINE_STREAM enumFromStepN #-}
1275 enumFromStepN x y n = x `seq` y `seq` n `seq`
1276 simple step (x,n) (Exact (delay_inline max n 0))
1277 where
1278 {-# INLINE_INNER step #-}
1279 step (x,n) | n > 0 = return $ Yield x (x+y,n-1)
1280 | otherwise = return $ Done
1281
1282 -- | Enumerate values
1283 --
1284 -- /WARNING:/ This operation can be very inefficient. If at all possible, use
1285 -- 'enumFromStepN' instead.
1286 enumFromTo :: (Enum a, Monad m) => a -> a -> Stream m v a
1287 {-# INLINE_STREAM enumFromTo #-}
1288 enumFromTo x y = fromList [x .. y]
1289
1290 -- NOTE: We use (x+1) instead of (succ x) below because the latter checks for
1291 -- overflow which can't happen here.
1292
1293 -- FIXME: add "too large" test for Int
1294 enumFromTo_small :: (Integral a, Monad m) => a -> a -> Stream m v a
1295 {-# INLINE_STREAM enumFromTo_small #-}
1296 enumFromTo_small x y = x `seq` y `seq` simple step x (Exact n)
1297 where
1298 n = delay_inline max (fromIntegral y - fromIntegral x + 1) 0
1299
1300 {-# INLINE_INNER step #-}
1301 step x | x <= y = return $ Yield x (x+1)
1302 | otherwise = return $ Done
1303
1304 {-# RULES
1305
1306 "enumFromTo<Int8> [Stream]"
1307 enumFromTo = enumFromTo_small :: Monad m => Int8 -> Int8 -> Stream m v Int8
1308
1309 "enumFromTo<Int16> [Stream]"
1310 enumFromTo = enumFromTo_small :: Monad m => Int16 -> Int16 -> Stream m v Int16
1311
1312 "enumFromTo<Word8> [Stream]"
1313 enumFromTo = enumFromTo_small :: Monad m => Word8 -> Word8 -> Stream m v Word8
1314
1315 "enumFromTo<Word16> [Stream]"
1316 enumFromTo = enumFromTo_small :: Monad m => Word16 -> Word16 -> Stream m v Word16
1317
1318 #-}
1319
1320 #if WORD_SIZE_IN_BITS > 32
1321
1322 {-# RULES
1323
1324 "enumFromTo<Int32> [Stream]"
1325 enumFromTo = enumFromTo_small :: Monad m => Int32 -> Int32 -> Stream m v Int32
1326
1327 "enumFromTo<Word32> [Stream]"
1328 enumFromTo = enumFromTo_small :: Monad m => Word32 -> Word32 -> Stream m v Word32
1329
1330 #-}
1331
1332 #endif
1333
1334 -- NOTE: We could implement a generic "too large" test:
1335 --
1336 -- len x y | x > y = 0
1337 -- | n > 0 && n <= fromIntegral (maxBound :: Int) = fromIntegral n
1338 -- | otherwise = error
1339 -- where
1340 -- n = y-x+1
1341 --
1342 -- Alas, GHC won't eliminate unnecessary comparisons (such as n >= 0 for
1343 -- unsigned types). See http://hackage.haskell.org/trac/ghc/ticket/3744
1344 --
1345
1346 enumFromTo_int :: (Integral a, Monad m) => a -> a -> Stream m v a
1347 {-# INLINE_STREAM enumFromTo_int #-}
1348 enumFromTo_int x y = x `seq` y `seq` simple step x (Exact (len x y))
1349 where
1350 {-# INLINE [0] len #-}
1351 len x y | x > y = 0
1352 | otherwise = BOUNDS_CHECK(check) "enumFromTo" "vector too large"
1353 (n > 0)
1354 $ fromIntegral n
1355 where
1356 n = y-x+1
1357
1358 {-# INLINE_INNER step #-}
1359 step x | x <= y = return $ Yield x (x+1)
1360 | otherwise = return $ Done
1361
1362 {-# RULES
1363
1364 "enumFromTo<Int> [Stream]"
1365 enumFromTo = enumFromTo_int :: Monad m => Int -> Int -> Stream m v Int
1366
1367 #if WORD_SIZE_IN_BITS > 32
1368
1369 "enumFromTo<Int64> [Stream]"
1370 enumFromTo = enumFromTo_int :: Monad m => Int64 -> Int64 -> Stream m v Int64
1371
1372 #else
1373
1374 "enumFromTo<Int32> [Stream]"
1375 enumFromTo = enumFromTo_int :: Monad m => Int32 -> Int32 -> Stream m v Int32
1376
1377 #endif
1378
1379 #-}
1380
1381 enumFromTo_big_word :: (Integral a, Monad m) => a -> a -> Stream m v a
1382 {-# INLINE_STREAM enumFromTo_big_word #-}
1383 enumFromTo_big_word x y = x `seq` y `seq` simple step x (Exact (len x y))
1384 where
1385 {-# INLINE [0] len #-}
1386 len x y | x > y = 0
1387 | otherwise = BOUNDS_CHECK(check) "enumFromTo" "vector too large"
1388 (n < fromIntegral (maxBound :: Int))
1389 $ fromIntegral (n+1)
1390 where
1391 n = y-x
1392
1393 {-# INLINE_INNER step #-}
1394 step x | x <= y = return $ Yield x (x+1)
1395 | otherwise = return $ Done
1396
1397 {-# RULES
1398
1399 "enumFromTo<Word> [Stream]"
1400 enumFromTo = enumFromTo_big_word :: Monad m => Word -> Word -> Stream m v Word
1401
1402 "enumFromTo<Word64> [Stream]"
1403 enumFromTo = enumFromTo_big_word
1404 :: Monad m => Word64 -> Word64 -> Stream m v Word64
1405
1406 #if WORD_SIZE_IN_BITS == 32
1407
1408 "enumFromTo<Word32> [Stream]"
1409 enumFromTo = enumFromTo_big_word
1410 :: Monad m => Word32 -> Word32 -> Stream m v Word32
1411
1412 #endif
1413
1414 "enumFromTo<Integer> [Stream]"
1415 enumFromTo = enumFromTo_big_word
1416 :: Monad m => Integer -> Integer -> Stream m v Integer
1417
1418 #-}
1419
1420 -- FIXME: the "too large" test is totally wrong
1421 enumFromTo_big_int :: (Integral a, Monad m) => a -> a -> Stream m v a
1422 {-# INLINE_STREAM enumFromTo_big_int #-}
1423 enumFromTo_big_int x y = x `seq` y `seq` simple step x (Exact (len x y))
1424 where
1425 {-# INLINE [0] len #-}
1426 len x y | x > y = 0
1427 | otherwise = BOUNDS_CHECK(check) "enumFromTo" "vector too large"
1428 (n > 0 && n <= fromIntegral (maxBound :: Int))
1429 $ fromIntegral n
1430 where
1431 n = y-x+1
1432
1433 {-# INLINE_INNER step #-}
1434 step x | x <= y = return $ Yield x (x+1)
1435 | otherwise = return $ Done
1436
1437 #if WORD_SIZE_IN_BITS > 32
1438
1439 {-# RULES
1440
1441 "enumFromTo<Int64> [Stream]"
1442 enumFromTo = enumFromTo_big :: Monad m => Int64 -> Int64 -> Stream m v Int64
1443
1444 #-}
1445
1446 #endif
1447
1448 enumFromTo_char :: Monad m => Char -> Char -> Stream m v Char
1449 {-# INLINE_STREAM enumFromTo_char #-}
1450 enumFromTo_char x y = x `seq` y `seq` simple step xn (Exact n)
1451 where
1452 xn = ord x
1453 yn = ord y
1454
1455 n = delay_inline max 0 (yn - xn + 1)
1456
1457 {-# INLINE_INNER step #-}
1458 step xn | xn <= yn = return $ Yield (unsafeChr xn) (xn+1)
1459 | otherwise = return $ Done
1460
1461 {-# RULES
1462
1463 "enumFromTo<Char> [Stream]"
1464 enumFromTo = enumFromTo_char
1465
1466 #-}
1467
1468 ------------------------------------------------------------------------
1469
1470 -- Specialise enumFromTo for Float and Double.
1471 -- Also, try to do something about pairs?
1472
1473 enumFromTo_double :: (Monad m, Ord a, RealFrac a) => a -> a -> Stream m v a
1474 {-# INLINE_STREAM enumFromTo_double #-}
1475 enumFromTo_double n m = n `seq` m `seq` simple step n (Max (len n m))
1476 where
1477 lim = m + 1/2 -- important to float out
1478
1479 {-# INLINE [0] len #-}
1480 len x y | x > y = 0
1481 | otherwise = BOUNDS_CHECK(check) "enumFromTo" "vector too large"
1482 (n > 0)
1483 $ fromIntegral n
1484 where
1485 n = truncate (y-x)+2
1486
1487 {-# INLINE_INNER step #-}
1488 step x | x <= lim = return $ Yield x (x+1)
1489 | otherwise = return $ Done
1490
1491 {-# RULES
1492
1493 "enumFromTo<Double> [Stream]"
1494 enumFromTo = enumFromTo_double :: Monad m => Double -> Double -> Stream m v Double
1495
1496 "enumFromTo<Float> [Stream]"
1497 enumFromTo = enumFromTo_double :: Monad m => Float -> Float -> Stream m v Float
1498
1499 #-}
1500
1501 ------------------------------------------------------------------------
1502
1503 -- | Enumerate values with a given step.
1504 --
1505 -- /WARNING:/ This operation is very inefficient. If at all possible, use
1506 -- 'enumFromStepN' instead.
1507 enumFromThenTo :: (Enum a, Monad m) => a -> a -> a -> Stream m v a
1508 {-# INLINE_STREAM enumFromThenTo #-}
1509 enumFromThenTo x y z = fromList [x, y .. z]
1510
1511 -- FIXME: Specialise enumFromThenTo.
1512
1513 -- Conversions
1514 -- -----------
1515
1516 -- | Convert a 'Stream' to a list
1517 toList :: Monad m => Stream m v a -> m [a]
1518 {-# INLINE toList #-}
1519 toList = foldr (:) []
1520
1521 -- | Convert a list to a 'Stream'
1522 fromList :: Monad m => [a] -> Stream m v a
1523 {-# INLINE fromList #-}
1524 fromList xs = unsafeFromList Unknown xs
1525
1526 -- | Convert the first @n@ elements of a list to a 'Stream'
1527 fromListN :: Monad m => Int -> [a] -> Stream m v a
1528 {-# INLINE_STREAM fromListN #-}
1529 fromListN n xs = simple step (xs,n) (Max (delay_inline max n 0))
1530 where
1531 {-# INLINE_INNER step #-}
1532 step (xs,n) | n <= 0 = return Done
1533 step (x:xs,n) = return (Yield x (xs,n-1))
1534 step ([],n) = return Done
1535
1536 -- | Convert a list to a 'Stream' with the given 'Size' hint.
1537 unsafeFromList :: Monad m => Size -> [a] -> Stream m v a
1538 {-# INLINE_STREAM unsafeFromList #-}
1539 unsafeFromList sz xs = simple step xs sz
1540 where
1541 step (x:xs) = return (Yield x xs)
1542 step [] = return Done
1543
1544 fromVector :: (Monad m, Vector v a) => v a -> Stream m v a
1545 {-# INLINE_STREAM fromVector #-}
1546 fromVector v = v `seq` n `seq` Stream (Unf step 0) (Unf vstep True) (Exact n)
1547 where
1548 n = basicLength v
1549
1550 {-# INLINE step #-}
1551 step i | i >= n = return Done
1552 | otherwise = case basicUnsafeIndexM v i of
1553 Box x -> return $ Yield x (i+1)
1554
1555
1556 {-# INLINE vstep #-}
1557 vstep True = return (Yield (Right v) False)
1558 vstep False = return Done
1559
1560 reVector :: Monad m => Stream m u a -> Stream m v a
1561 {-# INLINE_STREAM reVector #-}
1562 reVector (Stream (Unf step s) _ n) = simple step s n
1563
1564 {-# RULES
1565
1566 "reVector [Vector]"
1567 reVector = id
1568
1569 "reVector/reVector [Vector]" forall s.
1570 reVector (reVector s) = s
1571
1572 #-}
1573
1574