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