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