Switch the order of argument to slice and unsafeSlice
[darcs-mirrors/vector.git] / Data / Vector / Fusion / Stream / Monadic.hs
1 {-# LANGUAGE ExistentialQuantification, Rank2Types #-}
2
3 -- |
4 -- Module : Data.Vector.Fusion.Stream.Monadic
5 -- Copyright : (c) Roman Leshchinskiy 2008-2009
6 -- License : BSD-style
7 --
8 -- Maintainer : Roman Leshchinskiy <rl@cse.unsw.edu.au>
9 -- Stability : experimental
10 -- Portability : non-portable
11 --
12 -- Monadic streams
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, 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,
35
36 -- * Zipping
37 indexed,
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
59 -- * Scans
60 prescanl, prescanlM, prescanl', prescanlM',
61 postscanl, postscanlM, postscanl', postscanlM',
62 scanl, scanlM, scanl', scanlM',
63 scanl1, scanl1M, scanl1', scanl1M',
64
65 -- * Enumerations
66 enumFromTo, enumFromThenTo,
67
68 -- * Conversions
69 toList, fromList
70 ) where
71
72 import Data.Vector.Fusion.Stream.Size
73 import Data.Vector.Fusion.Util ( Box(..), delay_inline )
74
75 import Data.Char ( ord )
76 import GHC.Base ( unsafeChr )
77 import Control.Monad ( liftM )
78 import Prelude hiding ( length, null,
79 replicate, (++),
80 head, last, (!!),
81 init, tail, take, drop,
82 map, mapM, mapM_, concatMap,
83 zipWith, zipWith3, zip, zip3,
84 filter, takeWhile, dropWhile,
85 elem, notElem,
86 foldl, foldl1, foldr, foldr1,
87 and, or,
88 scanl, scanl1,
89 enumFromTo, enumFromThenTo )
90 import qualified Prelude
91
92 import Data.Int ( Int8, Int16, Int32, Int64 )
93 import Data.Word ( Word8, Word16, Word32, Word, Word64 )
94
95 #if __GLASGOW_HASKELL__ >= 613
96 import SpecConstr ( SpecConstrAnnotation(..) )
97 #endif
98
99 #include "vector.h"
100
101 data SPEC = SPEC | SPEC2
102 #if __GLASGOW_HASKELL__ >= 613
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_STREAM replicate #-}
159 -- NOTE: We delay inlining max here because GHC will create a join point for
160 -- the call to newArray# otherwise which is not really nice.
161 replicate n x = Stream (return . step) n (Exact (delay_inline max n 0))
162 where
163 {-# INLINE_INNER step #-}
164 step i | i > 0 = Yield x (i-1)
165 | otherwise = Done
166
167 generate :: Monad m => Int -> (Int -> a) -> Stream m a
168 {-# INLINE generate #-}
169 generate n f = generateM n (return . f)
170
171 -- | Generate a stream from its indices
172 generateM :: Monad m => Int -> (Int -> m a) -> Stream m a
173 {-# INLINE_STREAM generateM #-}
174 generateM n f = n `seq` Stream step 0 (Exact (delay_inline max n 0))
175 where
176 {-# INLINE_INNER step #-}
177 step i | i < n = do
178 x <- f i
179 return $ Yield x (i+1)
180 | otherwise = return Done
181
182 -- | Prepend an element
183 cons :: Monad m => a -> Stream m a -> Stream m a
184 {-# INLINE cons #-}
185 cons x s = singleton x ++ s
186
187 -- | Append an element
188 snoc :: Monad m => Stream m a -> a -> Stream m a
189 {-# INLINE snoc #-}
190 snoc s x = s ++ singleton x
191
192 infixr 5 ++
193 -- | Concatenate two 'Stream's
194 (++) :: Monad m => Stream m a -> Stream m a -> Stream m a
195 {-# INLINE_STREAM (++) #-}
196 Stream stepa sa na ++ Stream stepb sb nb = Stream step (Left sa) (na + nb)
197 where
198 {-# INLINE_INNER step #-}
199 step (Left sa) = do
200 r <- stepa sa
201 case r of
202 Yield x sa' -> return $ Yield x (Left sa')
203 Skip sa' -> return $ Skip (Left sa')
204 Done -> return $ Skip (Right sb)
205 step (Right sb) = do
206 r <- stepb sb
207 case r of
208 Yield x sb' -> return $ Yield x (Right sb')
209 Skip sb' -> return $ Skip (Right sb')
210 Done -> return $ Done
211
212 -- Accessing elements
213 -- ------------------
214
215 -- | First element of the 'Stream' or error if empty
216 head :: Monad m => Stream m a -> m a
217 {-# INLINE_STREAM head #-}
218 head (Stream step s _) = head_loop SPEC s
219 where
220 head_loop SPEC s
221 = do
222 r <- step s
223 case r of
224 Yield x _ -> return x
225 Skip s' -> head_loop SPEC s'
226 Done -> BOUNDS_ERROR(emptyStream) "head"
227
228
229
230 -- | Last element of the 'Stream' or error if empty
231 last :: Monad m => Stream m a -> m a
232 {-# INLINE_STREAM last #-}
233 last (Stream step s _) = last_loop0 SPEC s
234 where
235 last_loop0 SPEC s
236 = do
237 r <- step s
238 case r of
239 Yield x s' -> last_loop1 SPEC x s'
240 Skip s' -> last_loop0 SPEC s'
241 Done -> BOUNDS_ERROR(emptyStream) "last"
242
243 last_loop1 SPEC x s
244 = do
245 r <- step s
246 case r of
247 Yield y s' -> last_loop1 SPEC y s'
248 Skip s' -> last_loop1 SPEC x s'
249 Done -> return x
250
251 -- | Element at the given position
252 (!!) :: Monad m => Stream m a -> Int -> m a
253 {-# INLINE (!!) #-}
254 Stream step s _ !! i | i < 0 = BOUNDS_ERROR(error) "!!" "negative index"
255 | otherwise = index_loop SPEC s i
256 where
257 index_loop SPEC s i
258 = i `seq`
259 do
260 r <- step s
261 case r of
262 Yield x s' | i == 0 -> return x
263 | otherwise -> index_loop SPEC s' (i-1)
264 Skip s' -> index_loop SPEC s' i
265 Done -> BOUNDS_ERROR(emptyStream) "!!"
266
267 -- Substreams
268 -- ----------
269
270 -- | Extract a substream of the given length starting at the given position.
271 slice :: Monad m => Int -- ^ starting index
272 -> Int -- ^ length
273 -> Stream m a
274 -> Stream m a
275 {-# INLINE slice #-}
276 slice i n s = take n (drop i s)
277
278 -- | All but the last element
279 init :: Monad m => Stream m a -> Stream m a
280 {-# INLINE_STREAM init #-}
281 init (Stream step s sz) = Stream step' (Nothing, s) (sz - 1)
282 where
283 {-# INLINE_INNER step' #-}
284 step' (Nothing, s) = liftM (\r ->
285 case r of
286 Yield x s' -> Skip (Just x, s')
287 Skip s' -> Skip (Nothing, s')
288 Done -> BOUNDS_ERROR(emptyStream) "init"
289 ) (step s)
290
291 step' (Just x, s) = liftM (\r ->
292 case r of
293 Yield y s' -> Yield x (Just y, s')
294 Skip s' -> Skip (Just x, s')
295 Done -> Done
296 ) (step s)
297
298 -- | All but the first element
299 tail :: Monad m => Stream m a -> Stream m a
300 {-# INLINE_STREAM tail #-}
301 tail (Stream step s sz) = Stream step' (Left s) (sz - 1)
302 where
303 {-# INLINE_INNER step' #-}
304 step' (Left s) = liftM (\r ->
305 case r of
306 Yield x s' -> Skip (Right s')
307 Skip s' -> Skip (Left s')
308 Done -> BOUNDS_ERROR(emptyStream) "tail"
309 ) (step s)
310
311 step' (Right s) = liftM (\r ->
312 case r of
313 Yield x s' -> Yield x (Right s')
314 Skip s' -> Skip (Right s')
315 Done -> Done
316 ) (step s)
317
318 -- | The first @n@ elements
319 take :: Monad m => Int -> Stream m a -> Stream m a
320 {-# INLINE_STREAM take #-}
321 take n (Stream step s sz) = Stream step' (s, 0) (smaller (Exact n) sz)
322 where
323 {-# INLINE_INNER step' #-}
324 step' (s, i) | i < n = liftM (\r ->
325 case r of
326 Yield x s' -> Yield x (s', i+1)
327 Skip s' -> Skip (s', i)
328 Done -> Done
329 ) (step s)
330 step' (s, i) = return Done
331
332 -- | All but the first @n@ elements
333 drop :: Monad m => Int -> Stream m a -> Stream m a
334 {-# INLINE_STREAM drop #-}
335 drop n (Stream step s sz) = Stream step' (s, Just n) (sz - Exact n)
336 where
337 {-# INLINE_INNER step' #-}
338 step' (s, Just i) | i > 0 = liftM (\r ->
339 case r of
340 Yield x s' -> Skip (s', Just (i-1))
341 Skip s' -> Skip (s', Just i)
342 Done -> Done
343 ) (step s)
344 | otherwise = return $ Skip (s, Nothing)
345
346 step' (s, Nothing) = liftM (\r ->
347 case r of
348 Yield x s' -> Yield x (s', Nothing)
349 Skip s' -> Skip (s', Nothing)
350 Done -> Done
351 ) (step s)
352
353
354 -- Mapping
355 -- -------
356
357 instance Monad m => Functor (Stream m) where
358 {-# INLINE fmap #-}
359 fmap = map
360
361 -- | Map a function over a 'Stream'
362 map :: Monad m => (a -> b) -> Stream m a -> Stream m b
363 {-# INLINE map #-}
364 map f = mapM (return . f)
365
366
367 -- | Map a monadic function over a 'Stream'
368 mapM :: Monad m => (a -> m b) -> Stream m a -> Stream m b
369 {-# INLINE_STREAM mapM #-}
370 mapM f (Stream step s n) = Stream step' s n
371 where
372 {-# INLINE_INNER step' #-}
373 step' s = do
374 r <- step s
375 case r of
376 Yield x s' -> liftM (`Yield` s') (f x)
377 Skip s' -> return (Skip s')
378 Done -> return Done
379
380 -- | Execute a monadic action for each element of the 'Stream'
381 mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()
382 {-# INLINE_STREAM mapM_ #-}
383 mapM_ m (Stream step s _) = mapM_loop SPEC s
384 where
385 mapM_loop SPEC s
386 = do
387 r <- step s
388 case r of
389 Yield x s' -> do { m x; mapM_loop SPEC s' }
390 Skip s' -> mapM_loop SPEC s'
391 Done -> return ()
392
393 -- | Transform a 'Stream' to use a different monad
394 trans :: (Monad m, Monad m') => (forall a. m a -> m' a)
395 -> Stream m a -> Stream m' a
396 {-# INLINE_STREAM trans #-}
397 trans f (Stream step s n) = Stream (f . step) s n
398
399 unbox :: Monad m => Stream m (Box a) -> Stream m a
400 {-# INLINE_STREAM unbox #-}
401 unbox (Stream step s n) = Stream step' s n
402 where
403 {-# INLINE_INNER step' #-}
404 step' s = do
405 r <- step s
406 case r of
407 Yield (Box x) s' -> return $ Yield x s'
408 Skip s' -> return $ Skip s'
409 Done -> return $ Done
410
411 -- Zipping
412 -- -------
413
414 -- | Pair each element in a 'Stream' with its index
415 indexed :: Monad m => Stream m a -> Stream m (Int,a)
416 {-# INLINE_STREAM indexed #-}
417 indexed (Stream step s n) = Stream step' (s,0) n
418 where
419 {-# INLINE_INNER step' #-}
420 step' (s,i) = i `seq`
421 do
422 r <- step s
423 case r of
424 Yield x s' -> return $ Yield (i,x) (s', i+1)
425 Skip s' -> return $ Skip (s', i)
426 Done -> return Done
427
428 -- | Zip two 'Stream's with the given monadic function
429 zipWithM :: Monad m => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
430 {-# INLINE_STREAM zipWithM #-}
431 zipWithM f (Stream stepa sa na) (Stream stepb sb nb)
432 = Stream step (sa, sb, Nothing) (smaller na nb)
433 where
434 {-# INLINE_INNER step #-}
435 step (sa, sb, Nothing) = liftM (\r ->
436 case r of
437 Yield x sa' -> Skip (sa', sb, Just x)
438 Skip sa' -> Skip (sa', sb, Nothing)
439 Done -> Done
440 ) (stepa sa)
441
442 step (sa, sb, Just x) = do
443 r <- stepb sb
444 case r of
445 Yield y sb' ->
446 do
447 z <- f x y
448 return $ Yield z (sa, sb', Nothing)
449 Skip sb' -> return $ Skip (sa, sb', Just x)
450 Done -> return $ Done
451
452 -- FIXME: This might expose an opportunity for inplace execution.
453 {-# RULES
454
455 "zipWithM xs xs [Vector.Stream]" forall f xs.
456 zipWithM f xs xs = mapM (\x -> f x x) xs
457
458 #-}
459
460 zipWith3M :: Monad m => (a -> b -> c -> m d) -> Stream m a -> Stream m b -> Stream m c -> Stream m d
461 {-# INLINE_STREAM zipWith3M #-}
462 zipWith3M f (Stream stepa sa na) (Stream stepb sb nb) (Stream stepc sc nc)
463 = Stream step (sa, sb, sc, Nothing) (smaller na (smaller nb nc))
464 where
465 {-# INLINE_INNER step #-}
466 step (sa, sb, sc, Nothing) = do
467 r <- stepa sa
468 return $ case r of
469 Yield x sa' -> Skip (sa', sb, sc, Just (x, Nothing))
470 Skip sa' -> Skip (sa', sb, sc, Nothing)
471 Done -> Done
472
473 step (sa, sb, sc, Just (x, Nothing)) = do
474 r <- stepb sb
475 return $ case r of
476 Yield y sb' -> Skip (sa, sb', sc, Just (x, Just y))
477 Skip sb' -> Skip (sa, sb', sc, Just (x, Nothing))
478 Done -> Done
479
480 step (sa, sb, sc, Just (x, Just y)) = do
481 r <- stepc sc
482 case r of
483 Yield z sc' -> f x y z >>= (\res -> return $ Yield res (sa, sb, sc', Nothing))
484 Skip sc' -> return $ Skip (sa, sb, sc', Just (x, Just y))
485 Done -> return $ Done
486
487 zipWith4M :: Monad m => (a -> b -> c -> d -> m e)
488 -> Stream m a -> Stream m b -> Stream m c -> Stream m d
489 -> Stream m e
490 {-# INLINE zipWith4M #-}
491 zipWith4M f sa sb sc sd
492 = zipWithM (\(a,b) (c,d) -> f a b c d) (zip sa sb) (zip sc sd)
493
494 zipWith5M :: Monad m => (a -> b -> c -> d -> e -> m f)
495 -> Stream m a -> Stream m b -> Stream m c -> Stream m d
496 -> Stream m e -> Stream m f
497 {-# INLINE zipWith5M #-}
498 zipWith5M f sa sb sc sd se
499 = zipWithM (\(a,b,c) (d,e) -> f a b c d e) (zip3 sa sb sc) (zip sd se)
500
501 zipWith6M :: Monad m => (a -> b -> c -> d -> e -> f -> m g)
502 -> Stream m a -> Stream m b -> Stream m c -> Stream m d
503 -> Stream m e -> Stream m f -> Stream m g
504 {-# INLINE zipWith6M #-}
505 zipWith6M fn sa sb sc sd se sf
506 = zipWithM (\(a,b,c) (d,e,f) -> fn a b c d e f) (zip3 sa sb sc)
507 (zip3 sd se sf)
508
509 zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
510 {-# INLINE zipWith #-}
511 zipWith f = zipWithM (\a b -> return (f a b))
512
513 zipWith3 :: Monad m => (a -> b -> c -> d)
514 -> Stream m a -> Stream m b -> Stream m c -> Stream m d
515 {-# INLINE zipWith3 #-}
516 zipWith3 f = zipWith3M (\a b c -> return (f a b c))
517
518 zipWith4 :: Monad m => (a -> b -> c -> d -> e)
519 -> Stream m a -> Stream m b -> Stream m c -> Stream m d
520 -> Stream m e
521 {-# INLINE zipWith4 #-}
522 zipWith4 f = zipWith4M (\a b c d -> return (f a b c d))
523
524 zipWith5 :: Monad m => (a -> b -> c -> d -> e -> f)
525 -> Stream m a -> Stream m b -> Stream m c -> Stream m d
526 -> Stream m e -> Stream m f
527 {-# INLINE zipWith5 #-}
528 zipWith5 f = zipWith5M (\a b c d e -> return (f a b c d e))
529
530 zipWith6 :: Monad m => (a -> b -> c -> d -> e -> f -> g)
531 -> Stream m a -> Stream m b -> Stream m c -> Stream m d
532 -> Stream m e -> Stream m f -> Stream m g
533 {-# INLINE zipWith6 #-}
534 zipWith6 fn = zipWith6M (\a b c d e f -> return (fn a b c d e f))
535
536 zip :: Monad m => Stream m a -> Stream m b -> Stream m (a,b)
537 {-# INLINE zip #-}
538 zip = zipWith (,)
539
540 zip3 :: Monad m => Stream m a -> Stream m b -> Stream m c -> Stream m (a,b,c)
541 {-# INLINE zip3 #-}
542 zip3 = zipWith3 (,,)
543
544 zip4 :: Monad m => Stream m a -> Stream m b -> Stream m c -> Stream m d
545 -> Stream m (a,b,c,d)
546 {-# INLINE zip4 #-}
547 zip4 = zipWith4 (,,,)
548
549 zip5 :: Monad m => Stream m a -> Stream m b -> Stream m c -> Stream m d
550 -> Stream m e -> Stream m (a,b,c,d,e)
551 {-# INLINE zip5 #-}
552 zip5 = zipWith5 (,,,,)
553
554 zip6 :: Monad m => Stream m a -> Stream m b -> Stream m c -> Stream m d
555 -> Stream m e -> Stream m f -> Stream m (a,b,c,d,e,f)
556 {-# INLINE zip6 #-}
557 zip6 = zipWith6 (,,,,,)
558
559 -- Filtering
560 -- ---------
561
562 -- | Drop elements which do not satisfy the predicate
563 filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
564 {-# INLINE filter #-}
565 filter f = filterM (return . f)
566
567 -- | Drop elements which do not satisfy the monadic predicate
568 filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
569 {-# INLINE_STREAM filterM #-}
570 filterM f (Stream step s n) = Stream step' s (toMax n)
571 where
572 {-# INLINE_INNER step' #-}
573 step' s = do
574 r <- step s
575 case r of
576 Yield x s' -> do
577 b <- f x
578 return $ if b then Yield x s'
579 else Skip s'
580 Skip s' -> return $ Skip s'
581 Done -> return $ Done
582
583 -- | Longest prefix of elements that satisfy the predicate
584 takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
585 {-# INLINE takeWhile #-}
586 takeWhile f = takeWhileM (return . f)
587
588 -- | Longest prefix of elements that satisfy the monadic predicate
589 takeWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
590 {-# INLINE_STREAM takeWhileM #-}
591 takeWhileM f (Stream step s n) = Stream step' s (toMax n)
592 where
593 {-# INLINE_INNER step' #-}
594 step' s = do
595 r <- step s
596 case r of
597 Yield x s' -> do
598 b <- f x
599 return $ if b then Yield x s' else Done
600 Skip s' -> return $ Skip s'
601 Done -> return $ Done
602
603 -- | Drop the longest prefix of elements that satisfy the predicate
604 dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
605 {-# INLINE dropWhile #-}
606 dropWhile f = dropWhileM (return . f)
607
608 data DropWhile s a = DropWhile_Drop s | DropWhile_Yield a s | DropWhile_Next s
609
610 -- | Drop the longest prefix of elements that satisfy the monadic predicate
611 dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
612 {-# INLINE_STREAM dropWhileM #-}
613 dropWhileM f (Stream step s n) = Stream step' (DropWhile_Drop s) (toMax n)
614 where
615 -- NOTE: we jump through hoops here to have only one Yield; local data
616 -- declarations would be nice!
617
618 {-# INLINE_INNER step' #-}
619 step' (DropWhile_Drop s)
620 = do
621 r <- step s
622 case r of
623 Yield x s' -> do
624 b <- f x
625 return $ if b then Skip (DropWhile_Drop s')
626 else Skip (DropWhile_Yield x s')
627 Skip s' -> return $ Skip (DropWhile_Drop s')
628 Done -> return $ Done
629
630 step' (DropWhile_Yield x s) = return $ Yield x (DropWhile_Next s)
631
632 step' (DropWhile_Next s)
633 = liftM (\r ->
634 case r of
635 Yield x s' -> Skip (DropWhile_Yield x s')
636 Skip s' -> Skip (DropWhile_Next s')
637 Done -> Done
638 ) (step s)
639
640 -- Searching
641 -- ---------
642
643 infix 4 `elem`
644 -- | Check whether the 'Stream' contains an element
645 elem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
646 {-# INLINE_STREAM elem #-}
647 elem x (Stream step s _) = elem_loop SPEC s
648 where
649 elem_loop SPEC s
650 = do
651 r <- step s
652 case r of
653 Yield y s' | x == y -> return True
654 | otherwise -> elem_loop SPEC s'
655 Skip s' -> elem_loop SPEC s'
656 Done -> return False
657
658 infix 4 `notElem`
659 -- | Inverse of `elem`
660 notElem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
661 {-# INLINE notElem #-}
662 notElem x s = liftM not (elem x s)
663
664 -- | Yield 'Just' the first element that satisfies the predicate or 'Nothing'
665 -- if no such element exists.
666 find :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe a)
667 {-# INLINE find #-}
668 find f = findM (return . f)
669
670 -- | Yield 'Just' the first element that satisfies the monadic predicate or
671 -- 'Nothing' if no such element exists.
672 findM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe a)
673 {-# INLINE_STREAM findM #-}
674 findM f (Stream step s _) = find_loop SPEC s
675 where
676 find_loop SPEC s
677 = do
678 r <- step s
679 case r of
680 Yield x s' -> do
681 b <- f x
682 if b then return $ Just x
683 else find_loop SPEC s'
684 Skip s' -> find_loop SPEC s'
685 Done -> return Nothing
686
687 -- | Yield 'Just' the index of the first element that satisfies the predicate
688 -- or 'Nothing' if no such element exists.
689 findIndex :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe Int)
690 {-# INLINE_STREAM findIndex #-}
691 findIndex f = findIndexM (return . f)
692
693 -- | Yield 'Just' the index of the first element that satisfies the monadic
694 -- predicate or 'Nothing' if no such element exists.
695 findIndexM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe Int)
696 {-# INLINE_STREAM findIndexM #-}
697 findIndexM f (Stream step s _) = findIndex_loop SPEC s 0
698 where
699 findIndex_loop SPEC s i
700 = do
701 r <- step s
702 case r of
703 Yield x s' -> do
704 b <- f x
705 if b then return $ Just i
706 else findIndex_loop SPEC s' (i+1)
707 Skip s' -> findIndex_loop SPEC s' i
708 Done -> return Nothing
709
710 -- Folding
711 -- -------
712
713 -- | Left fold
714 foldl :: Monad m => (a -> b -> a) -> a -> Stream m b -> m a
715 {-# INLINE foldl #-}
716 foldl f = foldlM (\a b -> return (f a b))
717
718 -- | Left fold with a monadic operator
719 foldlM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> m a
720 {-# INLINE_STREAM foldlM #-}
721 foldlM m z (Stream step s _) = foldlM_loop SPEC z s
722 where
723 foldlM_loop SPEC z s
724 = do
725 r <- step s
726 case r of
727 Yield x s' -> do { z' <- m z x; foldlM_loop SPEC z' s' }
728 Skip s' -> foldlM_loop SPEC z s'
729 Done -> return z
730
731 -- | Same as 'foldlM'
732 foldM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> m a
733 {-# INLINE foldM #-}
734 foldM = foldlM
735
736 -- | Left fold over a non-empty 'Stream'
737 foldl1 :: Monad m => (a -> a -> a) -> Stream m a -> m a
738 {-# INLINE foldl1 #-}
739 foldl1 f = foldl1M (\a b -> return (f a b))
740
741 -- | Left fold over a non-empty 'Stream' with a monadic operator
742 foldl1M :: Monad m => (a -> a -> m a) -> Stream m a -> m a
743 {-# INLINE_STREAM foldl1M #-}
744 foldl1M f (Stream step s sz) = foldl1M_loop SPEC s
745 where
746 foldl1M_loop SPEC s
747 = do
748 r <- step s
749 case r of
750 Yield x s' -> foldlM f x (Stream step s' (sz - 1))
751 Skip s' -> foldl1M_loop SPEC s'
752 Done -> BOUNDS_ERROR(emptyStream) "foldl1M"
753
754 -- | Same as 'foldl1M'
755 fold1M :: Monad m => (a -> a -> m a) -> Stream m a -> m a
756 {-# INLINE fold1M #-}
757 fold1M = foldl1M
758
759 -- | Left fold with a strict accumulator
760 foldl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> m a
761 {-# INLINE foldl' #-}
762 foldl' f = foldlM' (\a b -> return (f a b))
763
764 -- | Left fold with a strict accumulator and a monadic operator
765 foldlM' :: Monad m => (a -> b -> m a) -> a -> Stream m b -> m a
766 {-# INLINE_STREAM foldlM' #-}
767 foldlM' m z (Stream step s _) = foldlM'_loop SPEC z s
768 where
769 foldlM'_loop SPEC z s
770 = z `seq`
771 do
772 r <- step s
773 case r of
774 Yield x s' -> do { z' <- m z x; foldlM'_loop SPEC z' s' }
775 Skip s' -> foldlM'_loop SPEC z s'
776 Done -> return z
777
778 -- | Same as 'foldlM''
779 foldM' :: Monad m => (a -> b -> m a) -> a -> Stream m b -> m a
780 {-# INLINE foldM' #-}
781 foldM' = foldlM'
782
783 -- | Left fold over a non-empty 'Stream' with a strict accumulator
784 foldl1' :: Monad m => (a -> a -> a) -> Stream m a -> m a
785 {-# INLINE foldl1' #-}
786 foldl1' f = foldl1M' (\a b -> return (f a b))
787
788 -- | Left fold over a non-empty 'Stream' with a strict accumulator and a
789 -- monadic operator
790 foldl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> m a
791 {-# INLINE_STREAM foldl1M' #-}
792 foldl1M' f (Stream step s sz) = foldl1M'_loop SPEC s
793 where
794 foldl1M'_loop SPEC s
795 = do
796 r <- step s
797 case r of
798 Yield x s' -> foldlM' f x (Stream step s' (sz - 1))
799 Skip s' -> foldl1M'_loop SPEC s'
800 Done -> BOUNDS_ERROR(emptyStream) "foldl1M'"
801
802 -- | Same as 'foldl1M''
803 fold1M' :: Monad m => (a -> a -> m a) -> Stream m a -> m a
804 {-# INLINE fold1M' #-}
805 fold1M' = foldl1M'
806
807 -- | Right fold
808 foldr :: Monad m => (a -> b -> b) -> b -> Stream m a -> m b
809 {-# INLINE foldr #-}
810 foldr f = foldrM (\a b -> return (f a b))
811
812 -- | Right fold with a monadic operator
813 foldrM :: Monad m => (a -> b -> m b) -> b -> Stream m a -> m b
814 {-# INLINE_STREAM foldrM #-}
815 foldrM f z (Stream step s _) = foldrM_loop SPEC s
816 where
817 foldrM_loop SPEC s
818 = do
819 r <- step s
820 case r of
821 Yield x s' -> f x =<< foldrM_loop SPEC s'
822 Skip s' -> foldrM_loop SPEC s'
823 Done -> return z
824
825 -- | Right fold over a non-empty stream
826 foldr1 :: Monad m => (a -> a -> a) -> Stream m a -> m a
827 {-# INLINE foldr1 #-}
828 foldr1 f = foldr1M (\a b -> return (f a b))
829
830 -- | Right fold over a non-empty stream with a monadic operator
831 foldr1M :: Monad m => (a -> a -> m a) -> Stream m a -> m a
832 {-# INLINE_STREAM foldr1M #-}
833 foldr1M f (Stream step s _) = foldr1M_loop0 SPEC s
834 where
835 foldr1M_loop0 SPEC s
836 = do
837 r <- step s
838 case r of
839 Yield x s' -> foldr1M_loop1 SPEC x s'
840 Skip s' -> foldr1M_loop0 SPEC s'
841 Done -> BOUNDS_ERROR(emptyStream) "foldr1M"
842
843 foldr1M_loop1 SPEC x s
844 = do
845 r <- step s
846 case r of
847 Yield y s' -> f x =<< foldr1M_loop1 SPEC y s'
848 Skip s' -> foldr1M_loop1 SPEC x s'
849 Done -> return x
850
851 -- Specialised folds
852 -- -----------------
853
854 and :: Monad m => Stream m Bool -> m Bool
855 {-# INLINE_STREAM and #-}
856 and (Stream step s _) = and_loop SPEC s
857 where
858 and_loop SPEC s
859 = do
860 r <- step s
861 case r of
862 Yield False _ -> return False
863 Yield True s' -> and_loop SPEC s'
864 Skip s' -> and_loop SPEC s'
865 Done -> return True
866
867 or :: Monad m => Stream m Bool -> m Bool
868 {-# INLINE_STREAM or #-}
869 or (Stream step s _) = or_loop SPEC s
870 where
871 or_loop SPEC s
872 = do
873 r <- step s
874 case r of
875 Yield False s' -> or_loop SPEC s'
876 Yield True _ -> return True
877 Skip s' -> or_loop SPEC s'
878 Done -> return False
879
880 concatMap :: Monad m => (a -> Stream m b) -> Stream m a -> Stream m b
881 {-# INLINE concatMap #-}
882 concatMap f = concatMapM (return . f)
883
884 concatMapM :: Monad m => (a -> m (Stream m b)) -> Stream m a -> Stream m b
885 {-# INLINE_STREAM concatMapM #-}
886 concatMapM f (Stream step s _) = Stream concatMap_go (Left s) Unknown
887 where
888 concatMap_go (Left s) = do
889 r <- step s
890 case r of
891 Yield a s' -> do
892 b_stream <- f a
893 return $ Skip (Right (b_stream, s'))
894 Skip s' -> return $ Skip (Left s')
895 Done -> return Done
896 concatMap_go (Right (Stream inner_step inner_s sz, s)) = do
897 r <- inner_step inner_s
898 case r of
899 Yield b inner_s' -> return $ Yield b (Right (Stream inner_step inner_s' sz, s))
900 Skip inner_s' -> return $ Skip (Right (Stream inner_step inner_s' sz, s))
901 Done -> return $ Skip (Left s)
902
903 -- Unfolding
904 -- ---------
905
906 -- | Unfold
907 unfoldr :: Monad m => (s -> Maybe (a, s)) -> s -> Stream m a
908 {-# INLINE_STREAM unfoldr #-}
909 unfoldr f = unfoldrM (return . f)
910
911 -- | Unfold with a monadic function
912 unfoldrM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m a
913 {-# INLINE_STREAM unfoldrM #-}
914 unfoldrM f s = Stream step s Unknown
915 where
916 {-# INLINE_INNER step #-}
917 step s = liftM (\r ->
918 case r of
919 Just (x, s') -> Yield x s'
920 Nothing -> Done
921 ) (f s)
922
923 -- Scans
924 -- -----
925
926 -- | Prefix scan
927 prescanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
928 {-# INLINE prescanl #-}
929 prescanl f = prescanlM (\a b -> return (f a b))
930
931 -- | Prefix scan with a monadic operator
932 prescanlM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a
933 {-# INLINE_STREAM prescanlM #-}
934 prescanlM f z (Stream step s sz) = Stream step' (s,z) sz
935 where
936 {-# INLINE_INNER step' #-}
937 step' (s,x) = do
938 r <- step s
939 case r of
940 Yield y s' -> do
941 z <- f x y
942 return $ Yield x (s', z)
943 Skip s' -> return $ Skip (s', x)
944 Done -> return Done
945
946 -- | Prefix scan with strict accumulator
947 prescanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
948 {-# INLINE prescanl' #-}
949 prescanl' f = prescanlM' (\a b -> return (f a b))
950
951 -- | Prefix scan with strict accumulator and a monadic operator
952 prescanlM' :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a
953 {-# INLINE_STREAM prescanlM' #-}
954 prescanlM' f z (Stream step s sz) = Stream step' (s,z) sz
955 where
956 {-# INLINE_INNER step' #-}
957 step' (s,x) = x `seq`
958 do
959 r <- step s
960 case r of
961 Yield y s' -> do
962 z <- f x y
963 return $ Yield x (s', z)
964 Skip s' -> return $ Skip (s', x)
965 Done -> return Done
966
967 -- | Suffix scan
968 postscanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
969 {-# INLINE postscanl #-}
970 postscanl f = postscanlM (\a b -> return (f a b))
971
972 -- | Suffix scan with a monadic operator
973 postscanlM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a
974 {-# INLINE_STREAM postscanlM #-}
975 postscanlM f z (Stream step s sz) = Stream step' (s,z) sz
976 where
977 {-# INLINE_INNER step' #-}
978 step' (s,x) = do
979 r <- step s
980 case r of
981 Yield y s' -> do
982 z <- f x y
983 return $ Yield z (s',z)
984 Skip s' -> return $ Skip (s',x)
985 Done -> return Done
986
987 -- | Suffix scan with strict accumulator
988 postscanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
989 {-# INLINE postscanl' #-}
990 postscanl' f = postscanlM' (\a b -> return (f a b))
991
992 -- | Suffix scan with strict acccumulator and a monadic operator
993 postscanlM' :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a
994 {-# INLINE_STREAM postscanlM' #-}
995 postscanlM' f z (Stream step s sz) = z `seq` Stream step' (s,z) sz
996 where
997 {-# INLINE_INNER step' #-}
998 step' (s,x) = x `seq`
999 do
1000 r <- step s
1001 case r of
1002 Yield y s' -> do
1003 z <- f x y
1004 z `seq` return (Yield z (s',z))
1005 Skip s' -> return $ Skip (s',x)
1006 Done -> return Done
1007
1008 -- | Haskell-style scan
1009 scanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
1010 {-# INLINE scanl #-}
1011 scanl f = scanlM (\a b -> return (f a b))
1012
1013 -- | Haskell-style scan with a monadic operator
1014 scanlM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a
1015 {-# INLINE scanlM #-}
1016 scanlM f z s = z `cons` postscanlM f z s
1017
1018 -- | Haskell-style scan with strict accumulator
1019 scanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
1020 {-# INLINE scanl' #-}
1021 scanl' f = scanlM' (\a b -> return (f a b))
1022
1023 -- | Haskell-style scan with strict accumulator and a monadic operator
1024 scanlM' :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a
1025 {-# INLINE scanlM' #-}
1026 scanlM' f z s = z `seq` (z `cons` postscanlM f z s)
1027
1028 -- | Scan over a non-empty 'Stream'
1029 scanl1 :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
1030 {-# INLINE scanl1 #-}
1031 scanl1 f = scanl1M (\x y -> return (f x y))
1032
1033 -- | Scan over a non-empty 'Stream' with a monadic operator
1034 scanl1M :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
1035 {-# INLINE_STREAM scanl1M #-}
1036 scanl1M f (Stream step s sz) = Stream step' (s, Nothing) sz
1037 where
1038 {-# INLINE_INNER step' #-}
1039 step' (s, Nothing) = do
1040 r <- step s
1041 case r of
1042 Yield x s' -> return $ Yield x (s', Just x)
1043 Skip s' -> return $ Skip (s', Nothing)
1044 Done -> BOUNDS_ERROR(emptyStream) "scanl1M"
1045
1046 step' (s, Just x) = do
1047 r <- step s
1048 case r of
1049 Yield y s' -> do
1050 z <- f x y
1051 return $ Yield z (s', Just z)
1052 Skip s' -> return $ Skip (s', Just x)
1053 Done -> return Done
1054
1055 -- | Scan over a non-empty 'Stream' with a strict accumulator
1056 scanl1' :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
1057 {-# INLINE scanl1' #-}
1058 scanl1' f = scanl1M' (\x y -> return (f x y))
1059
1060 -- | Scan over a non-empty 'Stream' with a strict accumulator and a monadic
1061 -- operator
1062 scanl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
1063 {-# INLINE_STREAM scanl1M' #-}
1064 scanl1M' f (Stream step s sz) = Stream step' (s, Nothing) sz
1065 where
1066 {-# INLINE_INNER step' #-}
1067 step' (s, Nothing) = do
1068 r <- step s
1069 case r of
1070 Yield x s' -> x `seq` return (Yield x (s', Just x))
1071 Skip s' -> return $ Skip (s', Nothing)
1072 Done -> BOUNDS_ERROR(emptyStream) "scanl1M"
1073
1074 step' (s, Just x) = x `seq`
1075 do
1076 r <- step s
1077 case r of
1078 Yield y s' -> do
1079 z <- f x y
1080 z `seq` return (Yield z (s', Just z))
1081 Skip s' -> return $ Skip (s', Just x)
1082 Done -> return Done
1083
1084 -- Enumerations
1085 -- ------------
1086
1087 -- The Enum class is broken for this, there just doesn't seem to be a
1088 -- way to implement this generically. We have specialise for as many types as
1089 -- we can but this doesn't help in polymorphic loops.
1090
1091 -- | Enumerate values from @x@ to @y@
1092 enumFromTo :: (Enum a, Monad m) => a -> a -> Stream m a
1093 {-# INLINE_STREAM enumFromTo #-}
1094 enumFromTo x y = fromList [x .. y]
1095
1096 -- NOTE: We use (x+1) instead of (succ x) below because the latter checks for
1097 -- overflow which can't happen here.
1098
1099 -- FIXME: add "too large" test for Int
1100 enumFromTo_small :: (Integral a, Monad m) => a -> a -> Stream m a
1101 {-# INLINE_STREAM enumFromTo_small #-}
1102 enumFromTo_small x y = Stream step x (Exact n)
1103 where
1104 n = max (fromIntegral y - fromIntegral x + 1) 0
1105
1106 {-# INLINE_INNER step #-}
1107 step x | x <= y = return $ Yield x (x+1)
1108 | otherwise = return $ Done
1109
1110 {-# RULES
1111
1112 "enumFromTo<Int> [Stream]"
1113 enumFromTo = enumFromTo_small :: Monad m => Int -> Int -> Stream m Int
1114
1115 "enumFromTo<Int8> [Stream]"
1116 enumFromTo = enumFromTo_small :: Monad m => Int8 -> Int8 -> Stream m Int8
1117
1118 "enumFromTo<Int16> [Stream]"
1119 enumFromTo = enumFromTo_small :: Monad m => Int16 -> Int16 -> Stream m Int16
1120
1121 "enumFromTo<Int32> [Stream]"
1122 enumFromTo = enumFromTo_small :: Monad m => Int32 -> Int32 -> Stream m Int32
1123
1124 "enumFromTo<Word8> [Stream]"
1125 enumFromTo = enumFromTo_small :: Monad m => Word8 -> Word8 -> Stream m Word8
1126
1127 "enumFromTo<Word16> [Stream]"
1128 enumFromTo = enumFromTo_small :: Monad m => Word16 -> Word16 -> Stream m Word16
1129
1130 #-}
1131
1132 -- FIXME: the "too large" test is totally wrong
1133 enumFromTo_big :: (Integral a, Monad m) => a -> a -> Stream m a
1134 {-# INLINE_STREAM enumFromTo_big #-}
1135 enumFromTo_big x y = Stream step x (Exact n)
1136 where
1137 n | x > y = 0
1138 | y - x < fromIntegral (maxBound :: Int) = fromIntegral (y-x+1)
1139 | otherwise = error $ "vector.enumFromTo_big: Array too large"
1140
1141 {-# INLINE_INNER step #-}
1142 step x | x <= y = return $ Yield x (x+1)
1143 | otherwise = return $ Done
1144
1145 {-# RULES
1146
1147 "enumFromTo<Int64> [Stream]"
1148 enumFromTo = enumFromTo_big :: Monad m => Int64 -> Int64 -> Stream m Int64
1149
1150 "enumFromTo<Word32> [Stream]"
1151 enumFromTo = enumFromTo_big :: Monad m => Word32 -> Word32 -> Stream m Word32
1152
1153 "enumFromTo<Word64> [Stream]"
1154 enumFromTo = enumFromTo_big :: Monad m => Word64 -> Word64 -> Stream m Word64
1155
1156 "enumFromTo<Integer> [Stream]"
1157 enumFromTo = enumFromTo_big :: Monad m => Integer -> Integer -> Stream m Integer
1158
1159 #-}
1160
1161 enumFromTo_char :: Monad m => Char -> Char -> Stream m Char
1162 {-# INLINE_STREAM enumFromTo_char #-}
1163 enumFromTo_char x y = Stream step xn (Exact n)
1164 where
1165 xn = ord x
1166 yn = ord y
1167
1168 n | xn > yn = 0
1169 | otherwise = yn - xn + 1
1170
1171 {-# INLINE_INNER step #-}
1172 step xn | xn <= yn = return $ Yield (unsafeChr xn) (xn+1)
1173 | otherwise = return $ Done
1174
1175 {-# RULES
1176
1177 "enumFromTo<Char> [Stream]"
1178 enumFromTo = enumFromTo_char
1179
1180 #-}
1181
1182 -- | Enumerate values from @x@ to @y@
1183 enumFromThenTo :: (Enum a, Monad m) => a -> a -> a -> Stream m a
1184 {-# INLINE_STREAM enumFromThenTo #-}
1185 enumFromThenTo x y z = fromList [x, y .. z]
1186
1187 -- Conversions
1188 -- -----------
1189
1190 -- | Convert a 'Stream' to a list
1191 toList :: Monad m => Stream m a -> m [a]
1192 {-# INLINE toList #-}
1193 toList = foldr (:) []
1194
1195 -- | Convert a list to a 'Stream'
1196 fromList :: Monad m => [a] -> Stream m a
1197 {-# INLINE_STREAM fromList #-}
1198 fromList xs = Stream step xs Unknown
1199 where
1200 step (x:xs) = return (Yield x xs)
1201 step [] = return Done
1202