Use Stream.Step
[darcs-mirrors/vector.git] / Data / Vector / Fusion / MStream.hs
1 {-# LANGUAGE ExistentialQuantification #-}
2
3 #include "phases.h"
4
5 module Data.Vector.Fusion.MStream (
6 MStream,
7
8 sized,
9
10 unfoldM, foldM,
11
12 map, mapM, filter, filterM
13 ) where
14
15 import Data.Vector.Fusion.Stream.Step
16 import Data.Vector.Fusion.Stream.Size
17
18 import Control.Monad ( liftM )
19 import Prelude hiding ( map, mapM, filter )
20
21 data MStream m a = forall s. MStream (s -> m (Step s a)) s Size
22
23 sized :: MStream m a -> Size -> MStream m a
24 {-# INLINE_STREAM sized #-}
25 sized (MStream step s _) sz = MStream step s sz
26
27 unfoldM :: Monad m => (s -> m (Maybe (a, s))) -> s -> MStream m a
28 {-# INLINE_STREAM unfoldM #-}
29 unfoldM f s = MStream step s Unknown
30 where
31 {-# INLINE step #-}
32 step s = do
33 r <- f s
34 case r of
35 Just (x, s') -> return $ Yield x s'
36 Nothing -> return $ Done
37
38 map :: Monad m => (a -> b) -> MStream m a -> MStream m b
39 {-# INLINE map #-}
40 map f = mapM (return . f)
41
42 mapM :: Monad m => (a -> m b) -> MStream m a -> MStream m b
43 {-# INLINE_STREAM mapM #-}
44 mapM f (MStream step s n) = MStream step' s n
45 where
46 {-# INLINE step' #-}
47 step' s = do
48 r <- step s
49 case r of
50 Yield x s' -> liftM (`Yield` s') (f x)
51 Skip s' -> return (Skip s')
52 Done -> return Done
53
54 filter :: Monad m => (a -> Bool) -> MStream m a -> MStream m a
55 {-# INLINE filter #-}
56 filter f = filterM (return . f)
57
58 filterM :: Monad m => (a -> m Bool) -> MStream m a -> MStream m a
59 {-# INLINE_STREAM filterM #-}
60 filterM f (MStream step s n) = MStream step' s (toMax n)
61 where
62 {-# INLINE step' #-}
63 step' s = do
64 r <- step s
65 case r of
66 Yield x s' -> do
67 b <- f x
68 if b then return $ Yield x s'
69 else return $ Skip s'
70 Skip s' -> return $ Skip s'
71 Done -> return $ Done
72
73 foldl :: Monad m => (a -> b -> a) -> a -> MStream m b -> m a
74 {-# INLINE foldl #-}
75 foldl f = foldM (\a b -> return (f a b))
76
77 foldr :: Monad m => (a -> b -> b) -> b -> MStream m a -> m b
78 {-# INLINE foldr #-}
79 foldr f = foldrM (\a b -> return (f a b))
80
81 foldlM :: Monad m => (a -> b -> m a) -> a -> MStream m b -> m a
82 {-# INLINE_STREAM foldlM #-}
83 foldlM m z (MStream step s _) = foldlM_go z s
84 where
85 foldlM_go z s = do
86 r <- step s
87 case r of
88 Yield x s' -> do { z' <- m z x; foldlM_go z' s' }
89 Skip s' -> foldlM_go z s'
90 Done -> return z
91
92 foldlM' :: Monad m => (a -> b -> m a) -> a -> MStream m b -> m a
93 {-# INLINE_STREAM foldlM' #-}
94 foldlM' m z (MStream step s _) = foldlM'_go z s
95 where
96 foldlM'_go z s = z `seq`
97 do
98 r <- step s
99 case r of
100 Yield x s' -> do { z' <- m z x; foldlM'_go z' s' }
101 Skip s' -> foldlM'_go z s'
102 Done -> return z
103
104 foldM :: Monad m => (a -> b -> m a) -> a -> MStream m b -> m a
105 {-# INLINE foldM #-}
106 foldM = foldlM
107
108 foldM' :: Monad m => (a -> b -> m a) -> a -> MStream m b -> m a
109 {-# INLINE foldM' #-}
110 foldM' = foldlM'
111
112 foldrM :: Monad m => (a -> b -> m b) -> b -> MStream m a -> m b
113 {-# INLINE_STREAM foldrM #-}
114 foldrM f z (MStream step s _) = foldrM_go s
115 where
116 foldrM_go s = do
117 r <- step s
118 case r of
119 Yield x s' -> f x =<< foldrM_go s'
120 Skip s' -> foldrM_go s'
121 Done -> return z
122