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