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