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