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