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