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