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