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