Rewrite Stream.drop
[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 -- | All but the first @n@ elements
245 drop :: Monad m => Int -> Stream m a -> Stream m a
246 {-# INLINE_STREAM drop #-}
247 drop n (Stream step s sz) = Stream step' (s, Just n) (sz - Exact n)
248 where
249 {-# INLINE step' #-}
250 step' (s, Just i) | i > 0 = liftM (\r ->
251 case r of
252 Yield x s' -> Skip (s', Just (i-1))
253 Skip s' -> Skip (s', Just i)
254 Done -> Done
255 ) (step s)
256 | otherwise = return $ Skip (s, Nothing)
257
258 step' (s, Nothing) = liftM (\r ->
259 case r of
260 Yield x s' -> Yield x (s', Nothing)
261 Skip s' -> Skip (s', Nothing)
262 Done -> Done
263 ) (step s)
264
265
266 -- Mapping/zipping
267 -- ---------------
268
269 instance Monad m => Functor (Stream m) where
270 {-# INLINE fmap #-}
271 fmap = map
272
273 map :: Monad m => (a -> b) -> Stream m a -> Stream m b
274 {-# INLINE map #-}
275 map f = mapM (return . f)
276
277 mapM :: Monad m => (a -> m b) -> Stream m a -> Stream m b
278 {-# INLINE_STREAM mapM #-}
279 mapM f (Stream step s n) = Stream step' s n
280 where
281 {-# INLINE step' #-}
282 step' s = do
283 r <- step s
284 case r of
285 Yield x s' -> liftM (`Yield` s') (f x)
286 Skip s' -> return (Skip s')
287 Done -> return Done
288
289 mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()
290 {-# INLINE_STREAM mapM_ #-}
291 mapM_ m (Stream step s _) = mapM_go s
292 where
293 mapM_go s = do
294 r <- step s
295 case r of
296 Yield x s' -> do { m x; mapM_go s' }
297 Skip s' -> mapM_go s'
298 Done -> return ()
299
300 -- | Zip two 'Stream's with the given function
301 zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
302 {-# INLINE zipWith #-}
303 zipWith f = zipWithM (\a b -> return (f a b))
304
305 -- | Zip two 'Stream's with the given function
306 zipWithM :: Monad m => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
307 {-# INLINE_STREAM zipWithM #-}
308 zipWithM f (Stream stepa sa na) (Stream stepb sb nb)
309 = Stream step (sa, sb, Nothing) (smaller na nb)
310 where
311 {-# INLINE step #-}
312 step (sa, sb, Nothing) = liftM (\r ->
313 case r of
314 Yield x sa' -> Skip (sa', sb, Just x)
315 Skip sa' -> Skip (sa', sb, Nothing)
316 Done -> Done
317 ) (stepa sa)
318
319 step (sa, sb, Just x) = do
320 r <- stepb sb
321 case r of
322 Yield y sb' ->
323 do
324 z <- f x y
325 return $ Yield z (sa, sb', Nothing)
326 Skip sb' -> return $ Skip (sa, sb', Just x)
327 Done -> return $ Done
328
329 -- Filtering
330 -- ---------
331
332 -- | Drop elements which do not satisfy the predicate
333 filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
334 {-# INLINE filter #-}
335 filter f = filterM (return . f)
336
337 -- | Drop elements which do not satisfy the predicate
338 filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
339 {-# INLINE_STREAM filterM #-}
340 filterM f (Stream step s n) = Stream step' s (toMax n)
341 where
342 {-# INLINE step' #-}
343 step' s = do
344 r <- step s
345 case r of
346 Yield x s' -> do
347 b <- f x
348 return $ if b then Yield x s'
349 else Skip s'
350 Skip s' -> return $ Skip s'
351 Done -> return $ Done
352
353 -- | Longest prefix of elements that satisfy the predicate
354 takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
355 {-# INLINE takeWhile #-}
356 takeWhile f = takeWhileM (return . f)
357
358 -- | Longest prefix of elements that satisfy the predicate
359 takeWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
360 {-# INLINE_STREAM takeWhileM #-}
361 takeWhileM f (Stream step s n) = Stream step' s (toMax n)
362 where
363 {-# INLINE step' #-}
364 step' s = do
365 r <- step s
366 case r of
367 Yield x s' -> do
368 b <- f x
369 return $ if b then Yield x s' else Done
370 Skip s' -> return $ Skip s'
371 Done -> return $ Done
372
373
374 dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
375 {-# INLINE dropWhile #-}
376 dropWhile f = dropWhileM (return . f)
377
378 data DropWhile s a = DropWhile_Drop s | DropWhile_Yield a s | DropWhile_Next s
379
380 -- | Drop the longest prefix of elements that satisfy the predicate
381 dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
382 {-# INLINE_STREAM dropWhileM #-}
383 dropWhileM f (Stream step s n) = Stream step' (DropWhile_Drop s) (toMax n)
384 where
385 -- NOTE: we jump through hoops here to have only one Yield; local data
386 -- declarations would be nice!
387
388 {-# INLINE step' #-}
389 step' (DropWhile_Drop s)
390 = do
391 r <- step s
392 case r of
393 Yield x s' -> do
394 b <- f x
395 return $ if b then Skip (DropWhile_Drop s')
396 else Skip (DropWhile_Yield x s')
397 Skip s' -> return $ Skip (DropWhile_Drop s')
398 Done -> return $ Done
399
400 step' (DropWhile_Yield x s) = return $ Yield x (DropWhile_Next s)
401
402 step' (DropWhile_Next s)
403 = liftM (\r ->
404 case r of
405 Yield x s' -> Skip (DropWhile_Yield x s')
406 Skip s' -> Skip (DropWhile_Next s')
407 Done -> Done
408 ) (step s)
409
410 -- Searching
411 -- ---------
412
413 infix 4 `elem`
414 -- | Check whether the 'Stream' contains an element
415 elem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
416 {-# INLINE_STREAM elem #-}
417 elem x (Stream step s _) = elem_loop s
418 where
419 elem_loop s = do
420 r <- step s
421 case r of
422 Yield y s' | x == y -> return True
423 | otherwise -> elem_loop s'
424 Skip s' -> elem_loop s'
425 Done -> return False
426
427 infix 4 `notElem`
428 -- | Inverse of `elem`
429 notElem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
430 {-# INLINE notElem #-}
431 notElem x s = liftM not (elem x s)
432
433 find :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe a)
434 {-# INLINE find #-}
435 find f = findM (return . f)
436
437 -- | Yield 'Just' the first element matching the predicate or 'Nothing' if no
438 -- such element exists.
439 findM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe a)
440 {-# INLINE_STREAM findM #-}
441 findM f (Stream step s _) = find_loop s
442 where
443 find_loop s = do
444 r <- step s
445 case r of
446 Yield x s' -> do
447 b <- f x
448 if b then return $ Just x
449 else find_loop s'
450 Skip s' -> find_loop s'
451 Done -> return Nothing
452
453 findIndex :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe Int)
454 {-# INLINE_STREAM findIndex #-}
455 findIndex f = findIndexM (return . f)
456
457 -- | Yield 'Just' the index of the first element matching the predicate or
458 -- 'Nothing' if no such element exists.
459 findIndexM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe Int)
460 {-# INLINE_STREAM findIndexM #-}
461 findIndexM f (Stream step s _) = findIndex_loop s 0
462 where
463 findIndex_loop s i = do
464 r <- step s
465 case r of
466 Yield x s' -> do
467 b <- f x
468 if b then return $ Just i
469 else findIndex_loop s' (i+1)
470 Skip s' -> findIndex_loop s' i
471 Done -> return Nothing
472
473 -- Folding
474 -- -------
475
476 foldl :: Monad m => (a -> b -> a) -> a -> Stream m b -> m a
477 {-# INLINE foldl #-}
478 foldl f = foldlM (\a b -> return (f a b))
479
480 foldlM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> m a
481 {-# INLINE_STREAM foldlM #-}
482 foldlM m z (Stream step s _) = foldlM_go z s
483 where
484 foldlM_go z s = do
485 r <- step s
486 case r of
487 Yield x s' -> do { z' <- m z x; foldlM_go z' s' }
488 Skip s' -> foldlM_go z s'
489 Done -> return z
490
491 foldM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> m a
492 {-# INLINE foldM #-}
493 foldM = foldlM
494
495 foldl1 :: Monad m => (a -> a -> a) -> Stream m a -> m a
496 {-# INLINE foldl1 #-}
497 foldl1 f = foldl1M (\a b -> return (f a b))
498
499 foldl1M :: Monad m => (a -> a -> m a) -> Stream m a -> m a
500 {-# INLINE_STREAM foldl1M #-}
501 foldl1M f (Stream step s sz) = foldl1M_go s
502 where
503 foldl1M_go s = do
504 r <- step s
505 case r of
506 Yield x s' -> foldlM f x (Stream step s' (sz - 1))
507 Skip s' -> foldl1M_go s'
508 Done -> errorEmptyStream "foldl1M"
509
510 foldl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> m a
511 {-# INLINE foldl' #-}
512 foldl' f = foldlM' (\a b -> return (f a b))
513
514 foldlM' :: Monad m => (a -> b -> m a) -> a -> Stream m b -> m a
515 {-# INLINE_STREAM foldlM' #-}
516 foldlM' m z (Stream step s _) = foldlM'_go z s
517 where
518 foldlM'_go z s = z `seq`
519 do
520 r <- step s
521 case r of
522 Yield x s' -> do { z' <- m z x; foldlM'_go z' s' }
523 Skip s' -> foldlM'_go z s'
524 Done -> return z
525
526 foldl1' :: Monad m => (a -> a -> a) -> Stream m a -> m a
527 {-# INLINE foldl1' #-}
528 foldl1' f = foldl1M' (\a b -> return (f a b))
529
530 foldl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> m a
531 {-# INLINE_STREAM foldl1M' #-}
532 foldl1M' f (Stream step s sz) = foldl1M'_go s
533 where
534 foldl1M'_go s = do
535 r <- step s
536 case r of
537 Yield x s' -> foldlM' f x (Stream step s' (sz - 1))
538 Skip s' -> foldl1M'_go s'
539 Done -> errorEmptyStream "foldl1M'"
540
541 foldr :: Monad m => (a -> b -> b) -> b -> Stream m a -> m b
542 {-# INLINE foldr #-}
543 foldr f = foldrM (\a b -> return (f a b))
544
545 foldrM :: Monad m => (a -> b -> m b) -> b -> Stream m a -> m b
546 {-# INLINE_STREAM foldrM #-}
547 foldrM f z (Stream step s _) = foldrM_go s
548 where
549 foldrM_go s = do
550 r <- step s
551 case r of
552 Yield x s' -> f x =<< foldrM_go s'
553 Skip s' -> foldrM_go s'
554 Done -> return z
555
556 foldr1 :: Monad m => (a -> a -> a) -> Stream m a -> m a
557 {-# INLINE foldr1 #-}
558 foldr1 f = foldr1M (\a b -> return (f a b))
559
560 foldr1M :: Monad m => (a -> a -> m a) -> Stream m a -> m a
561 {-# INLINE_STREAM foldr1M #-}
562 foldr1M f (Stream step s _) = foldr1M_go0 s
563 where
564 foldr1M_go0 s = do
565 r <- step s
566 case r of
567 Yield x s' -> foldr1M_go1 x s'
568 Skip s' -> foldr1M_go0 s'
569 Done -> errorEmptyStream "foldr1M"
570
571 foldr1M_go1 x s = do
572 r <- step s
573 case r of
574 Yield y s' -> f x =<< foldr1M_go1 y s'
575 Skip s' -> foldr1M_go1 x s'
576 Done -> return x
577
578 -- Unfolding
579 -- ---------
580
581 -- | Unfold
582 unfold :: Monad m => (s -> Maybe (a, s)) -> s -> Stream m a
583 {-# INLINE_STREAM unfold #-}
584 unfold f = unfoldM (return . f)
585
586 unfoldM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m a
587 {-# INLINE_STREAM unfoldM #-}
588 unfoldM f s = Stream step s Unknown
589 where
590 {-# INLINE step #-}
591 step s = liftM (\r ->
592 case r of
593 Just (x, s') -> Yield x s'
594 Nothing -> Done
595 ) (f s)
596
597 -- Scans
598 -- -----
599
600 -- | Prefix scan
601 prescanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
602 {-# INLINE prescanl #-}
603 prescanl f = prescanlM (\a b -> return (f a b))
604
605 -- | Prefix scan
606 prescanlM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a
607 {-# INLINE_STREAM prescanlM #-}
608 prescanlM f z (Stream step s sz) = Stream step' (s,z) sz
609 where
610 {-# INLINE step' #-}
611 step' (s,x) = do
612 r <- step s
613 case r of
614 Yield y s' -> do
615 z <- f x y
616 return $ Yield x (s', z)
617 Skip s' -> return $ Skip (s', x)
618 Done -> return Done
619
620
621 -- | Prefix scan with strict accumulator
622 prescanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
623 {-# INLINE prescanl' #-}
624 prescanl' f = prescanlM' (\a b -> return (f a b))
625
626 -- | Prefix scan with strict accumulator
627 prescanlM' :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a
628 {-# INLINE_STREAM prescanlM' #-}
629 prescanlM' f z (Stream step s sz) = Stream step' (s,z) sz
630 where
631 {-# INLINE step' #-}
632 step' (s,x) = x `seq`
633 do
634 r <- step s
635 case r of
636 Yield y s' -> do
637 z <- f x y
638 return $ Yield x (s', z)
639 Skip s' -> return $ Skip (s', x)
640 Done -> return Done
641
642 -- Conversions
643 -- -----------
644
645 toList :: Monad m => Stream m a -> m [a]
646 {-# INLINE toList #-}
647 toList = foldr (:) []
648
649 fromList :: Monad m => [a] -> Stream m a
650 {-# INLINE_STREAM fromList #-}
651 fromList xs = Stream step xs Unknown
652 where
653 step (x:xs) = return (Yield x xs)
654 step [] = return Done
655
656
657 errorEmptyStream :: String -> a
658 errorEmptyStream s = error $ "Data.Vector.Fusion.Stream.Monadic."
659 Prelude.++ s Prelude.++ ": empty stream"
660