Add prescans
[darcs-mirrors/vector.git] / Data / Vector / Fusion / Stream / Monadic.hs
1 {-# LANGUAGE ExistentialQuantification #-}
2
3 #include "phases.h"
4
5 module Data.Vector.Fusion.Stream.Monadic (
6 Stream(..),
7
8 -- * Size hints
9 size, sized,
10
11 -- * Length
12 length, null,
13
14 -- * Construction
15 empty, singleton, cons, snoc, replicate, (++),
16
17 -- * Accessing elements
18 head, last, (!!),
19
20 -- * Substreams
21 extract, init, tail, take, drop,
22
23 -- * Mapping and zipping
24 map, mapM, mapM_, zipWith, zipWithM,
25
26 -- * Filtering
27 filter, filterM, takeWhile, takeWhileM, dropWhile, dropWhileM,
28
29 -- * Searching
30 elem, notElem, find, findM, findIndex, findIndexM,
31
32 -- * Folding
33 foldl, foldlM, foldM, foldl1, foldl1M,
34 foldl', foldlM', foldl1', foldl1M',
35 foldr, foldrM, foldr1, foldr1M,
36
37 -- * Unfolding
38 unfold, unfoldM,
39
40 -- * Scans
41 prescanl, prescanlM, prescanl', prescanlM',
42
43 toList, fromList
44 ) where
45
46 import Data.Vector.Fusion.Stream.Step
47 import Data.Vector.Fusion.Stream.Size
48
49 import Control.Monad ( liftM )
50 import Prelude hiding ( length, null,
51 replicate, (++),
52 head, last, (!!),
53 init, tail, take, drop,
54 map, mapM, mapM_, zipWith,
55 filter, takeWhile, dropWhile,
56 elem, notElem,
57 foldl, foldl1, foldr, foldr1 )
58 import qualified Prelude
59
60 data Stream m a = forall s. Stream (s -> m (Step s a)) s Size
61
62 -- | 'Size' hint of a 'Stream'
63 size :: Stream m a -> Size
64 {-# INLINE size #-}
65 size (Stream _ _ sz) = sz
66
67 -- | Attach a 'Size' hint to a 'Stream'
68 sized :: Stream m a -> Size -> Stream m a
69 {-# INLINE_STREAM sized #-}
70 sized (Stream step s _) sz = Stream step s sz
71
72 -- Length
73 -- ------
74
75 -- | Length of a 'Stream'
76 length :: Monad m => Stream m a -> m Int
77 {-# INLINE_STREAM length #-}
78 length s = foldl' (\n _ -> n+1) 0 s
79
80 -- | Check if a 'Stream' is empty
81 null :: Monad m => Stream m a -> m Bool
82 {-# INLINE_STREAM null #-}
83 null s = foldr (\_ _ -> False) True s
84
85
86 -- Construction
87 -- ------------
88
89 -- | Empty 'Stream'
90 empty :: Monad m => Stream m a
91 {-# INLINE_STREAM empty #-}
92 empty = Stream (const (return Done)) () (Exact 0)
93
94 -- | Singleton 'Stream'
95 singleton :: Monad m => a -> Stream m a
96 {-# INLINE_STREAM singleton #-}
97 singleton x = Stream (return . step) True (Exact 1)
98 where
99 {-# INLINE step #-}
100 step True = Yield x False
101 step False = Done
102
103 -- | Replicate a value to a given length
104 replicate :: Monad m => Int -> a -> Stream m a
105 {-# INLINE_STREAM replicate #-}
106 replicate n x = Stream (return . step) n (Exact (max n 0))
107 where
108 {-# INLINE step #-}
109 step i | i > 0 = Yield x (i-1)
110 | otherwise = Done
111
112 -- | Prepend an element
113 cons :: Monad m => a -> Stream m a -> Stream m a
114 {-# INLINE cons #-}
115 cons x s = singleton x ++ s
116
117 -- | Append an element
118 snoc :: Monad m => Stream m a -> a -> Stream m a
119 {-# INLINE snoc #-}
120 snoc s x = s ++ singleton x
121
122 infixr 5 ++
123 -- | Concatenate two 'Stream's
124 (++) :: Monad m => Stream m a -> Stream m a -> Stream m a
125 {-# INLINE_STREAM (++) #-}
126 Stream stepa sa na ++ Stream stepb sb nb = Stream step (Left sa) (na + nb)
127 where
128 step (Left sa) = do
129 r <- stepa sa
130 case r of
131 Yield x sa' -> return $ Yield x (Left sa')
132 Skip sa' -> return $ Skip (Left sa')
133 Done -> return $ Skip (Right sb)
134 step (Right sb) = do
135 r <- stepb sb
136 case r of
137 Yield x sb' -> return $ Yield x (Right sb')
138 Skip sb' -> return $ Skip (Right sb')
139 Done -> return $ Done
140
141 -- Accessing elements
142 -- ------------------
143
144 -- | First element of the 'Stream' or error if empty
145 head :: Monad m => Stream m a -> m a
146 {-# INLINE_STREAM head #-}
147 head (Stream step s _) = head_loop s
148 where
149 head_loop s = do
150 r <- step s
151 case r of
152 Yield x _ -> return x
153 Skip s' -> head_loop s'
154 Done -> errorEmptyStream "head"
155
156 -- | Last element of the 'Stream' or error if empty
157 last :: Monad m => Stream m a -> m a
158 {-# INLINE_STREAM last #-}
159 last (Stream step s _) = last_loop0 s
160 where
161 last_loop0 s = do
162 r <- step s
163 case r of
164 Yield x s' -> last_loop1 x s'
165 Skip s' -> last_loop0 s'
166 Done -> errorEmptyStream "last"
167
168 last_loop1 x s = do
169 r <- step s
170 case r of
171 Yield y s' -> last_loop1 y s'
172 Skip s' -> last_loop1 x s'
173 Done -> return x
174
175 -- | Element at the given position
176 (!!) :: Monad m => Stream m a -> Int -> m a
177 {-# INLINE (!!) #-}
178 s !! i = head (drop i s)
179
180 -- Substreams
181 -- ----------
182
183 -- | Extract a substream of the given length starting at the given position.
184 extract :: Monad m => Stream m a -> Int -- ^ starting index
185 -> Int -- ^ length
186 -> Stream m a
187 {-# INLINE extract #-}
188 extract s i n = take n (drop i s)
189
190 -- | All but the last element
191 init :: Monad m => Stream m a -> Stream m a
192 {-# INLINE_STREAM init #-}
193 init (Stream step s sz) = Stream step' (Nothing, s) (sz - 1)
194 where
195 {-# INLINE step' #-}
196 step' (Nothing, s) = liftM (\r ->
197 case r of
198 Yield x s' -> Skip (Just x, s')
199 Skip s' -> Skip (Nothing, s')
200 Done -> Done -- FIXME: should be an error
201 ) (step s)
202
203 step' (Just x, s) = liftM (\r ->
204 case r of
205 Yield y s' -> Yield x (Just y, s')
206 Skip s' -> Skip (Just x, s')
207 Done -> Done
208 ) (step s)
209
210 -- | All but the first element
211 tail :: Monad m => Stream m a -> Stream m a
212 {-# INLINE_STREAM tail #-}
213 tail (Stream step s sz) = Stream step' (Left s) (sz - 1)
214 where
215 {-# INLINE step' #-}
216 step' (Left s) = liftM (\r ->
217 case r of
218 Yield x s' -> Skip (Right s')
219 Skip s' -> Skip (Left s')
220 Done -> Done -- FIXME: should be error?
221 ) (step s)
222
223 step' (Right s) = liftM (\r ->
224 case r of
225 Yield x s' -> Yield x (Right s')
226 Skip s' -> Skip (Right s')
227 Done -> Done
228 ) (step s)
229
230 -- | The first @n@ elements
231 take :: Monad m => Int -> Stream m a -> Stream m a
232 {-# INLINE_STREAM take #-}
233 take n (Stream step s sz) = Stream step' (s, 0) (smaller (Exact n) sz)
234 where
235 {-# INLINE step' #-}
236 step' (s, i) | i < n = liftM (\r ->
237 case r of
238 Yield x s' -> Yield x (s', i+1)
239 Skip s' -> Skip (s', i)
240 Done -> Done
241 ) (step s)
242 step' (s, i) = return Done
243
244 data Drop s = Drop_Drop s Int | Drop_Keep s
245
246 -- | All but the first @n@ elements
247 drop :: Monad m => Int -> Stream m a -> Stream m a
248 {-# INLINE_STREAM drop #-}
249 drop n (Stream step s sz) = Stream step' (Drop_Drop s 0) (sz - Exact n)
250 where
251 {-# INLINE step' #-}
252 step' (Drop_Drop s i) | i < n = liftM (\r ->
253 case r of
254 Yield x s' -> Skip (Drop_Drop s' (i+1))
255 Skip s' -> Skip (Drop_Drop s' i)
256 Done -> Done
257 ) (step s)
258 | otherwise = return $ Skip (Drop_Keep s)
259
260 step' (Drop_Keep s) = liftM (\r ->
261 case r of
262 Yield x s' -> Yield x (Drop_Keep s')
263 Skip s' -> Skip (Drop_Keep s')
264 Done -> Done
265 ) (step s)
266
267
268 -- Mapping/zipping
269 -- ---------------
270
271 instance Monad m => Functor (Stream m) where
272 {-# INLINE fmap #-}
273 fmap = map
274
275 map :: Monad m => (a -> b) -> Stream m a -> Stream m b
276 {-# INLINE map #-}
277 map f = mapM (return . f)
278
279 mapM :: Monad m => (a -> m b) -> Stream m a -> Stream m b
280 {-# INLINE_STREAM mapM #-}
281 mapM f (Stream step s n) = Stream step' s n
282 where
283 {-# INLINE step' #-}
284 step' s = do
285 r <- step s
286 case r of
287 Yield x s' -> liftM (`Yield` s') (f x)
288 Skip s' -> return (Skip s')
289 Done -> return Done
290
291 mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()
292 {-# INLINE_STREAM mapM_ #-}
293 mapM_ m (Stream step s _) = mapM_go s
294 where
295 mapM_go s = do
296 r <- step s
297 case r of
298 Yield x s' -> do { m x; mapM_go s' }
299 Skip s' -> mapM_go s'
300 Done -> return ()
301
302 -- | Zip two 'Stream's with the given function
303 zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
304 {-# INLINE zipWith #-}
305 zipWith f = zipWithM (\a b -> return (f a b))
306
307 -- | Zip two 'Stream's with the given function
308 zipWithM :: Monad m => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
309 {-# INLINE_STREAM zipWithM #-}
310 zipWithM f (Stream stepa sa na) (Stream stepb sb nb)
311 = Stream step (sa, sb, Nothing) (smaller na nb)
312 where
313 {-# INLINE step #-}
314 step (sa, sb, Nothing) = liftM (\r ->
315 case r of
316 Yield x sa' -> Skip (sa', sb, Just x)
317 Skip sa' -> Skip (sa', sb, Nothing)
318 Done -> Done
319 ) (stepa sa)
320
321 step (sa, sb, Just x) = do
322 r <- stepb sb
323 case r of
324 Yield y sb' ->
325 do
326 z <- f x y
327 return $ Yield z (sa, sb', Nothing)
328 Skip sb' -> return $ Skip (sa, sb', Just x)
329 Done -> return $ Done
330
331 -- Filtering
332 -- ---------
333
334 -- | Drop elements which do not satisfy the predicate
335 filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
336 {-# INLINE filter #-}
337 filter f = filterM (return . f)
338
339 -- | Drop elements which do not satisfy the predicate
340 filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
341 {-# INLINE_STREAM filterM #-}
342 filterM f (Stream step s n) = Stream step' s (toMax n)
343 where
344 {-# INLINE step' #-}
345 step' s = do
346 r <- step s
347 case r of
348 Yield x s' -> do
349 b <- f x
350 return $ if b then Yield x s'
351 else Skip s'
352 Skip s' -> return $ Skip s'
353 Done -> return $ Done
354
355 -- | Longest prefix of elements that satisfy the predicate
356 takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
357 {-# INLINE takeWhile #-}
358 takeWhile f = takeWhileM (return . f)
359
360 -- | Longest prefix of elements that satisfy the predicate
361 takeWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
362 {-# INLINE_STREAM takeWhileM #-}
363 takeWhileM f (Stream step s n) = Stream step' s (toMax n)
364 where
365 {-# INLINE step' #-}
366 step' s = do
367 r <- step s
368 case r of
369 Yield x s' -> do
370 b <- f x
371 return $ if b then Yield x s' else Done
372 Skip s' -> return $ Skip s'
373 Done -> return $ Done
374
375
376 dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
377 {-# INLINE dropWhile #-}
378 dropWhile f = dropWhileM (return . f)
379
380 data DropWhile s a = DropWhile_Drop s | DropWhile_Yield a s | DropWhile_Next s
381
382 -- | Drop the longest prefix of elements that satisfy the predicate
383 dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
384 {-# INLINE_STREAM dropWhileM #-}
385 dropWhileM f (Stream step s n) = Stream step' (DropWhile_Drop s) (toMax n)
386 where
387 -- NOTE: we jump through hoops here to have only one Yield; local data
388 -- declarations would be nice!
389
390 {-# INLINE step' #-}
391 step' (DropWhile_Drop s)
392 = do
393 r <- step s
394 case r of
395 Yield x s' -> do
396 b <- f x
397 return $ if b then Skip (DropWhile_Drop s')
398 else Skip (DropWhile_Yield x s')
399 Skip s' -> return $ Skip (DropWhile_Drop s')
400 Done -> return $ Done
401
402 step' (DropWhile_Yield x s) = return $ Yield x (DropWhile_Next s)
403
404 step' (DropWhile_Next s)
405 = liftM (\r ->
406 case r of
407 Yield x s' -> Skip (DropWhile_Yield x s')
408 Skip s' -> Skip (DropWhile_Next s')
409 Done -> Done
410 ) (step s)
411
412 -- Searching
413 -- ---------
414
415 infix 4 `elem`
416 -- | Check whether the 'Stream' contains an element
417 elem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
418 {-# INLINE_STREAM elem #-}
419 elem x (Stream step s _) = elem_loop s
420 where
421 elem_loop s = do
422 r <- step s
423 case r of
424 Yield y s' | x == y -> return True
425 | otherwise -> elem_loop s'
426 Skip s' -> elem_loop s'
427 Done -> return False
428
429 infix 4 `notElem`
430 -- | Inverse of `elem`
431 notElem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
432 {-# INLINE notElem #-}
433 notElem x s = liftM not (elem x s)
434
435 find :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe a)
436 {-# INLINE find #-}
437 find f = findM (return . f)
438
439 -- | Yield 'Just' the first element matching the predicate or 'Nothing' if no
440 -- such element exists.
441 findM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe a)
442 {-# INLINE_STREAM findM #-}
443 findM f (Stream step s _) = find_loop s
444 where
445 find_loop s = do
446 r <- step s
447 case r of
448 Yield x s' -> do
449 b <- f x
450 if b then return $ Just x
451 else find_loop s'
452 Skip s' -> find_loop s'
453 Done -> return Nothing
454
455 findIndex :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe Int)
456 {-# INLINE_STREAM findIndex #-}
457 findIndex f = findIndexM (return . f)
458
459 -- | Yield 'Just' the index of the first element matching the predicate or
460 -- 'Nothing' if no such element exists.
461 findIndexM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe Int)
462 {-# INLINE_STREAM findIndexM #-}
463 findIndexM f (Stream step s _) = findIndex_loop s 0
464 where
465 findIndex_loop s i = do
466 r <- step s
467 case r of
468 Yield x s' -> do
469 b <- f x
470 if b then return $ Just i
471 else findIndex_loop s' (i+1)
472 Skip s' -> findIndex_loop s' i
473 Done -> return Nothing
474
475 -- Folding
476 -- -------
477
478 foldl :: Monad m => (a -> b -> a) -> a -> Stream m b -> m a
479 {-# INLINE foldl #-}
480 foldl f = foldlM (\a b -> return (f a b))
481
482 foldlM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> m a
483 {-# INLINE_STREAM foldlM #-}
484 foldlM m z (Stream step s _) = foldlM_go z s
485 where
486 foldlM_go z s = do
487 r <- step s
488 case r of
489 Yield x s' -> do { z' <- m z x; foldlM_go z' s' }
490 Skip s' -> foldlM_go z s'
491 Done -> return z
492
493 foldM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> m a
494 {-# INLINE foldM #-}
495 foldM = foldlM
496
497 foldl1 :: Monad m => (a -> a -> a) -> Stream m a -> m a
498 {-# INLINE foldl1 #-}
499 foldl1 f = foldl1M (\a b -> return (f a b))
500
501 foldl1M :: Monad m => (a -> a -> m a) -> Stream m a -> m a
502 {-# INLINE_STREAM foldl1M #-}
503 foldl1M f (Stream step s sz) = foldl1M_go s
504 where
505 foldl1M_go s = do
506 r <- step s
507 case r of
508 Yield x s' -> foldlM f x (Stream step s' (sz - 1))
509 Skip s' -> foldl1M_go s'
510 Done -> errorEmptyStream "foldl1M"
511
512 foldl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> m a
513 {-# INLINE foldl' #-}
514 foldl' f = foldlM' (\a b -> return (f a b))
515
516 foldlM' :: Monad m => (a -> b -> m a) -> a -> Stream m b -> m a
517 {-# INLINE_STREAM foldlM' #-}
518 foldlM' m z (Stream step s _) = foldlM'_go z s
519 where
520 foldlM'_go z s = z `seq`
521 do
522 r <- step s
523 case r of
524 Yield x s' -> do { z' <- m z x; foldlM'_go z' s' }
525 Skip s' -> foldlM'_go z s'
526 Done -> return z
527
528 foldl1' :: Monad m => (a -> a -> a) -> Stream m a -> m a
529 {-# INLINE foldl1' #-}
530 foldl1' f = foldl1M' (\a b -> return (f a b))
531
532 foldl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> m a
533 {-# INLINE_STREAM foldl1M' #-}
534 foldl1M' f (Stream step s sz) = foldl1M'_go s
535 where
536 foldl1M'_go s = do
537 r <- step s
538 case r of
539 Yield x s' -> foldlM' f x (Stream step s' (sz - 1))
540 Skip s' -> foldl1M'_go s'
541 Done -> errorEmptyStream "foldl1M'"
542
543 foldr :: Monad m => (a -> b -> b) -> b -> Stream m a -> m b
544 {-# INLINE foldr #-}
545 foldr f = foldrM (\a b -> return (f a b))
546
547 foldrM :: Monad m => (a -> b -> m b) -> b -> Stream m a -> m b
548 {-# INLINE_STREAM foldrM #-}
549 foldrM f z (Stream step s _) = foldrM_go s
550 where
551 foldrM_go s = do
552 r <- step s
553 case r of
554 Yield x s' -> f x =<< foldrM_go s'
555 Skip s' -> foldrM_go s'
556 Done -> return z
557
558 foldr1 :: Monad m => (a -> a -> a) -> Stream m a -> m a
559 {-# INLINE foldr1 #-}
560 foldr1 f = foldr1M (\a b -> return (f a b))
561
562 foldr1M :: Monad m => (a -> a -> m a) -> Stream m a -> m a
563 {-# INLINE_STREAM foldr1M #-}
564 foldr1M f (Stream step s _) = foldr1M_go0 s
565 where
566 foldr1M_go0 s = do
567 r <- step s
568 case r of
569 Yield x s' -> foldr1M_go1 x s'
570 Skip s' -> foldr1M_go0 s'
571 Done -> errorEmptyStream "foldr1M"
572
573 foldr1M_go1 x s = do
574 r <- step s
575 case r of
576 Yield y s' -> f x =<< foldr1M_go1 y s'
577 Skip s' -> foldr1M_go1 x s'
578 Done -> return x
579
580 -- Unfolding
581 -- ---------
582
583 -- | Unfold
584 unfold :: Monad m => (s -> Maybe (a, s)) -> s -> Stream m a
585 {-# INLINE_STREAM unfold #-}
586 unfold f = unfoldM (return . f)
587
588 unfoldM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m a
589 {-# INLINE_STREAM unfoldM #-}
590 unfoldM f s = Stream step s Unknown
591 where
592 {-# INLINE step #-}
593 step s = liftM (\r ->
594 case r of
595 Just (x, s') -> Yield x s'
596 Nothing -> Done
597 ) (f s)
598
599 -- Scans
600 -- -----
601
602 -- | Prefix scan
603 prescanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
604 {-# INLINE prescanl #-}
605 prescanl f = prescanlM (\a b -> return (f a b))
606
607 -- | Prefix scan
608 prescanlM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a
609 {-# INLINE_STREAM prescanlM #-}
610 prescanlM f z (Stream step s sz) = Stream step' (s,z) sz
611 where
612 {-# INLINE step' #-}
613 step' (s,x) = do
614 r <- step s
615 case r of
616 Yield y s' -> do
617 z <- f x y
618 return $ Yield x (s', z)
619 Skip s' -> return $ Skip (s', x)
620 Done -> return Done
621
622
623 -- | Prefix scan with strict accumulator
624 prescanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
625 {-# INLINE prescanl' #-}
626 prescanl' f = prescanlM' (\a b -> return (f a b))
627
628 -- | Prefix scan with strict accumulator
629 prescanlM' :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a
630 {-# INLINE_STREAM prescanlM' #-}
631 prescanlM' f z (Stream step s sz) = Stream step' (s,z) sz
632 where
633 {-# INLINE step' #-}
634 step' (s,x) = x `seq`
635 do
636 r <- step s
637 case r of
638 Yield y s' -> do
639 z <- f x y
640 return $ Yield x (s', z)
641 Skip s' -> return $ Skip (s', x)
642 Done -> return Done
643
644 -- Conversions
645 -- -----------
646
647 toList :: Monad m => Stream m a -> m [a]
648 {-# INLINE toList #-}
649 toList = foldr (:) []
650
651 fromList :: Monad m => [a] -> Stream m a
652 {-# INLINE_STREAM fromList #-}
653 fromList xs = Stream step xs Unknown
654 where
655 step (x:xs) = return (Yield x xs)
656 step [] = return Done
657
658
659 errorEmptyStream :: String -> a
660 errorEmptyStream s = error $ "Data.Vector.Fusion.Stream.Monadic."
661 Prelude.++ s Prelude.++ ": empty stream"
662