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