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