1 {-# LANGUAGE ExistentialQuantification #-}
3 #include "phases.h"
6 Stream,
8 sized,
10 unfoldM, foldM,
12 map, mapM, filter, filterM
13 ) where
15 import Data.Vector.Fusion.Stream.Step
16 import Data.Vector.Fusion.Stream.Size
18 import Control.Monad ( liftM )
19 import Prelude hiding ( map, mapM, filter )
21 data Stream m a = forall s. Stream (s -> m (Step s a)) s Size
23 sized :: Stream m a -> Size -> Stream m a
24 {-# INLINE_STREAM sized #-}
25 sized (Stream step s _) sz = Stream step s sz
27 unfoldM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m a
28 {-# INLINE_STREAM unfoldM #-}
29 unfoldM f s = Stream 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
38 map :: Monad m => (a -> b) -> Stream m a -> Stream m b
39 {-# INLINE map #-}
40 map f = mapM (return . f)
42 mapM :: Monad m => (a -> m b) -> Stream m a -> Stream m b
43 {-# INLINE_STREAM mapM #-}
44 mapM f (Stream step s n) = Stream 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
54 filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
55 {-# INLINE filter #-}
56 filter f = filterM (return . f)
58 filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
59 {-# INLINE_STREAM filterM #-}
60 filterM f (Stream step s n) = Stream 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
73 foldl :: Monad m => (a -> b -> a) -> a -> Stream m b -> m a
74 {-# INLINE foldl #-}
75 foldl f = foldM (\a b -> return (f a b))
77 foldr :: Monad m => (a -> b -> b) -> b -> Stream m a -> m b
78 {-# INLINE foldr #-}
79 foldr f = foldrM (\a b -> return (f a b))
81 foldlM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> m a
82 {-# INLINE_STREAM foldlM #-}
83 foldlM m z (Stream 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
92 foldlM' :: Monad m => (a -> b -> m a) -> a -> Stream m b -> m a
93 {-# INLINE_STREAM foldlM' #-}
94 foldlM' m z (Stream 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
104 foldM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> m a
105 {-# INLINE foldM #-}
106 foldM = foldlM
108 foldM' :: Monad m => (a -> b -> m a) -> a -> Stream m b -> m a
109 {-# INLINE foldM' #-}
110 foldM' = foldlM'
112 foldrM :: Monad m => (a -> b -> m b) -> b -> Stream m a -> m b
113 {-# INLINE_STREAM foldrM #-}
114 foldrM f z (Stream 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