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