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