New combinators
[darcs-mirrors/vector.git] / Data / Vector / Stream.hs
1 {-# LANGUAGE ExistentialQuantification, BangPatterns, CPP #-}
2
3 #include "phases.h"
4
5 module Data.Vector.Stream (
6 Step(..), Stream(..),
7
8 size, sized, unfold, toList, fromList,
9 empty, singleton, replicate, (++),
10 take, drop,
11 map, zipWith,
12 filter, takeWhile, dropWhile,
13 foldr, foldl, foldl',
14 mapM_, foldM
15 ) where
16
17 import Data.Vector.Stream.Size
18
19 import Prelude hiding ( replicate, (++),
20 take, drop,
21 map, zipWith,
22 filter, takeWhile, dropWhile,
23 foldr, foldl,
24 mapM_ )
25
26 data Step s a = Yield a s
27 | Skip s
28 | Done
29
30 data Stream a = forall s. Stream (s -> Step s a) s Size
31
32 size :: Stream a -> Size
33 {-# INLINE size #-}
34 size (Stream _ _ sz) = sz
35
36 sized :: Stream a -> Size -> Stream a
37 {-# INLINE_STREAM sized #-}
38 sized (Stream step s _) sz = Stream step s sz
39
40 unfold :: (s -> Maybe (a, s)) -> s -> Stream a
41 {-# INLINE_STREAM unfold #-}
42 unfold f s = Stream step s Unknown
43 where
44 {-# INLINE step #-}
45 step s = case f s of
46 Just (x, s') -> Yield x s'
47 Nothing -> Done
48
49 toList :: Stream a -> [a]
50 {-# INLINE toList #-}
51 toList s = foldr (:) [] s
52
53 fromList :: [a] -> Stream a
54 {-# INLINE_STREAM fromList #-}
55 fromList xs = Stream step xs Unknown
56 where
57 step (x:xs) = Yield x xs
58 step [] = Done
59
60 -- Construction
61 -- ------------
62
63 empty :: Stream a
64 {-# INLINE_STREAM empty #-}
65 empty = Stream (const Done) () (Exact 0)
66
67 singleton :: a -> Stream a
68 {-# INLINE_STREAM singleton #-}
69 singleton x = Stream step True (Exact 1)
70 where
71 {-# INLINE step #-}
72 step True = Yield x False
73 step False = Done
74
75 replicate :: Int -> a -> Stream a
76 {-# INLINE_STREAM replicate #-}
77 replicate n x = Stream step n (Exact (max n 0))
78 where
79 {-# INLINE step #-}
80 step i | i > 0 = Yield x (i-1)
81 | otherwise = Done
82
83 infixr ++
84 (++) :: Stream a -> Stream a -> Stream a
85 {-# INLINE_STREAM (++) #-}
86 Stream stepa sa na ++ Stream stepb sb nb = Stream step (Left sa) (na + nb)
87 where
88 step (Left sa) = case stepa sa of
89 Yield x sa' -> Yield x (Left sa')
90 Skip sa' -> Skip (Left sa')
91 Done -> Skip (Right sb)
92 step (Right sb) = case stepb sb of
93 Yield x sb' -> Yield x (Right sb')
94 Skip sb' -> Skip (Right sb')
95 Done -> Done
96
97 -- Substreams
98 -- ----------
99
100 take :: Int -> Stream a -> Stream a
101 {-# INLINE_STREAM take #-}
102 take n (Stream step s sz) = Stream step' (s, 0) (smaller (Exact n) sz)
103 where
104 {-# INLINE step' #-}
105 step' (s, i) | i < n = case step s of
106 Yield x s' -> Yield x (s', i+1)
107 Skip s' -> Skip (s', i)
108 Done -> Done
109 step' (s, i) = Done
110
111 data Drop s = Drop_Drop s Int | Drop_Keep s
112
113 drop :: Int -> Stream a -> Stream a
114 {-# INLINE_STREAM drop #-}
115 drop n (Stream step s sz) = Stream step' (Drop_Drop s 0) (sz - Exact n)
116 where
117 {-# INLINE step' #-}
118 step' (Drop_Drop s i) | i < n = case step s of
119 Yield x s' -> Skip (Drop_Drop s' (i+1))
120 Skip s' -> Skip (Drop_Drop s' i)
121 Done -> Done
122 | otherwise = Skip (Drop_Keep s)
123
124 step' (Drop_Keep s) = case step s of
125 Yield x s' -> Yield x (Drop_Keep s')
126 Skip s' -> Skip (Drop_Keep s')
127 Done -> Done
128
129
130 -- Mapping/zipping
131 -- ---------------
132
133 instance Functor Stream where
134 {-# INLINE_STREAM fmap #-}
135 fmap = map
136
137 map :: (a -> b) -> Stream a -> Stream b
138 {-# INLINE_STREAM map #-}
139 map f (Stream step s n) = Stream step' s n
140 where
141 {-# INLINE step' #-}
142 step' s = case step s of
143 Yield x s' -> Yield (f x) s'
144 Skip s' -> Skip s'
145 Done -> Done
146
147 zipWith :: (a -> b -> c) -> Stream a -> Stream b -> Stream c
148 {-# INLINE_STREAM zipWith #-}
149 zipWith f (Stream stepa sa na) (Stream stepb sb nb)
150 = Stream step (sa, sb, Nothing) (smaller na nb)
151 where
152 {-# INLINE step #-}
153 step (sa, sb, Nothing) = case stepa sa of
154 Yield x sa' -> Skip (sa', sb, Just x)
155 Skip sa' -> Skip (sa', sb, Nothing)
156 Done -> Done
157
158 step (sa, sb, Just x) = case stepb sb of
159 Yield y sb' -> Yield (f x y) (sa, sb', Nothing)
160 Skip sb' -> Skip (sa, sb', Just x)
161 Done -> Done
162
163 -- Filtering
164 -- ---------
165
166 filter :: (a -> Bool) -> Stream a -> Stream a
167 {-# INLINE_STREAM filter #-}
168 filter f (Stream step s n) = Stream step' s (toMax n)
169 where
170 {-# INLINE step' #-}
171 step' s = case step s of
172 Yield x s' | f x -> Yield x s'
173 | otherwise -> Skip s'
174 Skip s' -> Skip s'
175 Done -> Done
176
177 takeWhile :: (a -> Bool) -> Stream a -> Stream a
178 {-# INLINE_STREAM takeWhile #-}
179 takeWhile f (Stream step s n) = Stream step' s (toMax n)
180 where
181 {-# INLINE step' #-}
182 step' s = case step s of
183 Yield x s' | f x -> Yield x s'
184 | otherwise -> Done
185 Skip s' -> Skip s'
186 Done -> Done
187
188
189 data DropWhile s a = DropWhile_Drop s | DropWhile_Yield a s | DropWhile_Next s
190
191 dropWhile :: (a -> Bool) -> Stream a -> Stream a
192 {-# INLINE_STREAM dropWhile #-}
193 dropWhile f (Stream step s n) = Stream step' (DropWhile_Drop s) (toMax n)
194 where
195 -- NOTE: we jump through hoops here to have only one Yield; local data
196 -- declarations would be nice!
197
198 {-# INLINE step' #-}
199 step' (DropWhile_Drop s)
200 = case step s of
201 Yield x s' | f x -> Skip (DropWhile_Drop s')
202 | otherwise -> Skip (DropWhile_Yield x s')
203 Skip s' -> Skip (DropWhile_Drop s')
204 Done -> Done
205
206 step' (DropWhile_Yield x s) = Yield x (DropWhile_Next s)
207
208 step' (DropWhile_Next s) = case step s of
209 Yield x s' -> Skip (DropWhile_Yield x s')
210 Skip s' -> Skip (DropWhile_Next s')
211 Done -> Done
212
213 -- Folding
214 -- -------
215
216 foldl :: (a -> b -> a) -> a -> Stream b -> a
217 {-# INLINE_STREAM foldl #-}
218 foldl f z (Stream step s _) = foldl_go z s
219 where
220 foldl_go z s = case step s of
221 Yield x s' -> foldl_go (f z x) s'
222 Skip s' -> foldl_go z s'
223 Done -> z
224
225 foldl' :: (a -> b -> a) -> a -> Stream b -> a
226 {-# INLINE_STREAM foldl' #-}
227 foldl' f !z (Stream step s _) = foldl_go z s
228 where
229 foldl_go !z s = case step s of
230 Yield x s' -> foldl_go (f z x) s'
231 Skip s' -> foldl_go z s'
232 Done -> z
233
234 foldr :: (a -> b -> b) -> b -> Stream a -> b
235 {-# INLINE_STREAM foldr #-}
236 foldr f z (Stream step s _) = foldr_go s
237 where
238 foldr_go s = case step s of
239 Yield x s' -> f x (foldr_go s')
240 Skip s' -> foldr_go s'
241 Done -> z
242
243 mapM_ :: Monad m => (a -> m ()) -> Stream a -> m ()
244 {-# INLINE_STREAM mapM_ #-}
245 mapM_ m (Stream step s _) = mapM_go s
246 where
247 mapM_go s = case step s of
248 Yield x s' -> do { m x; mapM_go s' }
249 Skip s' -> mapM_go s'
250 Done -> return ()
251
252 foldM :: Monad m => (a -> b -> m a) -> a -> Stream b -> m a
253 {-# INLINE_STREAM foldM #-}
254 foldM m z (Stream step s _) = foldM_go z s
255 where
256 foldM_go z s = case step s of
257 Yield x s' -> do { z' <- m z x; foldM_go z' s' }
258 Skip s' -> foldM_go z s'
259 Done -> return z
260