Rename classes and modules
[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 length, null,
10 empty, singleton, cons, snoc, replicate, (++),
11 head, last, (!!),
12 init, tail, take, drop,
13 map, zipWith,
14 filter, takeWhile, dropWhile,
15 foldl, foldl1, foldl', foldl1', foldr, foldr1,
16 mapM_, foldM
17 ) where
18
19 import Data.Vector.Stream.Size
20
21 import Prelude hiding ( length, null,
22 replicate, (++),
23 head, last, (!!),
24 init, tail, take, drop,
25 map, zipWith,
26 filter, takeWhile, dropWhile,
27 foldl, foldl1, foldr, foldr1,
28 mapM_ )
29
30 data Step s a = Yield a s
31 | Skip s
32 | Done
33
34 data Stream a = forall s. Stream (s -> Step s a) s Size
35
36 size :: Stream a -> Size
37 {-# INLINE size #-}
38 size (Stream _ _ sz) = sz
39
40 sized :: Stream a -> Size -> Stream a
41 {-# INLINE_STREAM sized #-}
42 sized (Stream step s _) sz = Stream step s sz
43
44 unfold :: (s -> Maybe (a, s)) -> s -> Stream a
45 {-# INLINE_STREAM unfold #-}
46 unfold f s = Stream step s Unknown
47 where
48 {-# INLINE step #-}
49 step s = case f s of
50 Just (x, s') -> Yield x s'
51 Nothing -> Done
52
53 toList :: Stream a -> [a]
54 {-# INLINE toList #-}
55 toList s = foldr (:) [] s
56
57 fromList :: [a] -> Stream a
58 {-# INLINE_STREAM fromList #-}
59 fromList xs = Stream step xs Unknown
60 where
61 step (x:xs) = Yield x xs
62 step [] = Done
63
64 -- Length
65 -- ------
66
67 length :: Stream a -> Int
68 {-# INLINE_STREAM length #-}
69 length s = foldl' (\n _ -> n+1) 0 s
70
71 null :: Stream a -> Bool
72 {-# INLINE_STREAM null #-}
73 null s = foldr (\_ _ -> False) True s
74
75 -- Construction
76 -- ------------
77
78 empty :: Stream a
79 {-# INLINE_STREAM empty #-}
80 empty = Stream (const Done) () (Exact 0)
81
82 singleton :: a -> Stream a
83 {-# INLINE_STREAM singleton #-}
84 singleton x = Stream step True (Exact 1)
85 where
86 {-# INLINE step #-}
87 step True = Yield x False
88 step False = Done
89
90 replicate :: Int -> a -> Stream a
91 {-# INLINE_STREAM replicate #-}
92 replicate n x = Stream step n (Exact (max n 0))
93 where
94 {-# INLINE step #-}
95 step i | i > 0 = Yield x (i-1)
96 | otherwise = Done
97
98 cons :: a -> Stream a -> Stream a
99 {-# INLINE cons #-}
100 cons x s = singleton x ++ s
101
102 snoc :: Stream a -> a -> Stream a
103 {-# INLINE snoc #-}
104 snoc s x = s ++ singleton x
105
106 infixr 5 ++
107 (++) :: Stream a -> Stream a -> Stream a
108 {-# INLINE_STREAM (++) #-}
109 Stream stepa sa na ++ Stream stepb sb nb = Stream step (Left sa) (na + nb)
110 where
111 step (Left sa) = case stepa sa of
112 Yield x sa' -> Yield x (Left sa')
113 Skip sa' -> Skip (Left sa')
114 Done -> Skip (Right sb)
115 step (Right sb) = case stepb sb of
116 Yield x sb' -> Yield x (Right sb')
117 Skip sb' -> Skip (Right sb')
118 Done -> Done
119
120 -- Accessing elements
121 -- ------------------
122
123 head :: Stream a -> a
124 {-# INLINE_STREAM head #-}
125 head (Stream step s _) = head_loop s
126 where
127 head_loop s = case step s of
128 Yield x _ -> x
129 Skip s' -> head_loop s'
130 Done -> error "Data.Vector.Stream.head: empty stream"
131
132 last :: Stream a -> a
133 {-# INLINE_STREAM last #-}
134 last (Stream step s _) = last_loop0 s
135 where
136 last_loop0 s = case step s of
137 Yield x s' -> last_loop1 x s'
138 Skip s' -> last_loop0 s'
139 Done -> error "Data.Vector.Stream.last: empty stream"
140
141 last_loop1 x s = case step s of
142 Yield y s' -> last_loop1 y s'
143 Skip s' -> last_loop1 x s'
144 Done -> x
145
146 (!!) :: Stream a -> Int -> a
147 {-# INLINE (!!) #-}
148 s !! i = head (drop i s)
149
150 -- Substreams
151 -- ----------
152
153 init :: Stream a -> Stream a
154 {-# INLINE_STREAM init #-}
155 init (Stream step s sz) = Stream step' (Nothing, s) (sz - 1)
156 where
157 {-# INLINE step' #-}
158 step' (Nothing, s) = case step s of
159 Yield x s' -> Skip (Just x, s')
160 Skip s' -> Skip (Nothing, s')
161 Done -> Done -- FIXME: should be an error
162
163 step' (Just x, s) = case step s of
164 Yield y s' -> Yield x (Just y, s')
165 Skip s' -> Skip (Just x, s')
166 Done -> Done
167
168 tail :: Stream a -> Stream a
169 {-# INLINE_STREAM tail #-}
170 tail (Stream step s sz) = Stream step' (Left s) (sz - 1)
171 where
172 {-# INLINE step' #-}
173 step' (Left s) = case step s of
174 Yield x s' -> Skip (Right s')
175 Skip s' -> Skip (Left s')
176 Done -> Done -- FIXME: should be error?
177
178 step' (Right s) = case step s of
179 Yield x s' -> Yield x (Right s')
180 Skip s' -> Skip (Right s')
181 Done -> Done
182
183 take :: Int -> Stream a -> Stream a
184 {-# INLINE_STREAM take #-}
185 take n (Stream step s sz) = Stream step' (s, 0) (smaller (Exact n) sz)
186 where
187 {-# INLINE step' #-}
188 step' (s, i) | i < n = case step s of
189 Yield x s' -> Yield x (s', i+1)
190 Skip s' -> Skip (s', i)
191 Done -> Done
192 step' (s, i) = Done
193
194 data Drop s = Drop_Drop s Int | Drop_Keep s
195
196 drop :: Int -> Stream a -> Stream a
197 {-# INLINE_STREAM drop #-}
198 drop n (Stream step s sz) = Stream step' (Drop_Drop s 0) (sz - Exact n)
199 where
200 {-# INLINE step' #-}
201 step' (Drop_Drop s i) | i < n = case step s of
202 Yield x s' -> Skip (Drop_Drop s' (i+1))
203 Skip s' -> Skip (Drop_Drop s' i)
204 Done -> Done
205 | otherwise = Skip (Drop_Keep s)
206
207 step' (Drop_Keep s) = case step s of
208 Yield x s' -> Yield x (Drop_Keep s')
209 Skip s' -> Skip (Drop_Keep s')
210 Done -> Done
211
212
213 -- Mapping/zipping
214 -- ---------------
215
216 instance Functor Stream where
217 {-# INLINE_STREAM fmap #-}
218 fmap = map
219
220 map :: (a -> b) -> Stream a -> Stream b
221 {-# INLINE_STREAM map #-}
222 map f (Stream step s n) = Stream step' s n
223 where
224 {-# INLINE step' #-}
225 step' s = case step s of
226 Yield x s' -> Yield (f x) s'
227 Skip s' -> Skip s'
228 Done -> Done
229
230 zipWith :: (a -> b -> c) -> Stream a -> Stream b -> Stream c
231 {-# INLINE_STREAM zipWith #-}
232 zipWith f (Stream stepa sa na) (Stream stepb sb nb)
233 = Stream step (sa, sb, Nothing) (smaller na nb)
234 where
235 {-# INLINE step #-}
236 step (sa, sb, Nothing) = case stepa sa of
237 Yield x sa' -> Skip (sa', sb, Just x)
238 Skip sa' -> Skip (sa', sb, Nothing)
239 Done -> Done
240
241 step (sa, sb, Just x) = case stepb sb of
242 Yield y sb' -> Yield (f x y) (sa, sb', Nothing)
243 Skip sb' -> Skip (sa, sb', Just x)
244 Done -> Done
245
246 -- Filtering
247 -- ---------
248
249 filter :: (a -> Bool) -> Stream a -> Stream a
250 {-# INLINE_STREAM filter #-}
251 filter f (Stream step s n) = Stream step' s (toMax n)
252 where
253 {-# INLINE step' #-}
254 step' s = case step s of
255 Yield x s' | f x -> Yield x s'
256 | otherwise -> Skip s'
257 Skip s' -> Skip s'
258 Done -> Done
259
260 takeWhile :: (a -> Bool) -> Stream a -> Stream a
261 {-# INLINE_STREAM takeWhile #-}
262 takeWhile f (Stream step s n) = Stream step' s (toMax n)
263 where
264 {-# INLINE step' #-}
265 step' s = case step s of
266 Yield x s' | f x -> Yield x s'
267 | otherwise -> Done
268 Skip s' -> Skip s'
269 Done -> Done
270
271
272 data DropWhile s a = DropWhile_Drop s | DropWhile_Yield a s | DropWhile_Next s
273
274 dropWhile :: (a -> Bool) -> Stream a -> Stream a
275 {-# INLINE_STREAM dropWhile #-}
276 dropWhile f (Stream step s n) = Stream step' (DropWhile_Drop s) (toMax n)
277 where
278 -- NOTE: we jump through hoops here to have only one Yield; local data
279 -- declarations would be nice!
280
281 {-# INLINE step' #-}
282 step' (DropWhile_Drop s)
283 = case step s of
284 Yield x s' | f x -> Skip (DropWhile_Drop s')
285 | otherwise -> Skip (DropWhile_Yield x s')
286 Skip s' -> Skip (DropWhile_Drop s')
287 Done -> Done
288
289 step' (DropWhile_Yield x s) = Yield x (DropWhile_Next s)
290
291 step' (DropWhile_Next s) = case step s of
292 Yield x s' -> Skip (DropWhile_Yield x s')
293 Skip s' -> Skip (DropWhile_Next s')
294 Done -> Done
295
296 -- Folding
297 -- -------
298
299 foldl :: (a -> b -> a) -> a -> Stream b -> a
300 {-# INLINE_STREAM foldl #-}
301 foldl f z (Stream step s _) = foldl_go z s
302 where
303 foldl_go z s = case step s of
304 Yield x s' -> foldl_go (f z x) s'
305 Skip s' -> foldl_go z s'
306 Done -> z
307
308 foldl1 :: (a -> a -> a) -> Stream a -> a
309 {-# INLINE_STREAM foldl1 #-}
310 foldl1 f (Stream step s sz) = foldl1_loop s
311 where
312 foldl1_loop s = case step s of
313 Yield x s' -> foldl f x (Stream step s' (sz - 1))
314 Skip s' -> foldl1_loop s'
315 Done -> error "Data.Vector.Stream.foldl1: empty stream"
316
317 foldl' :: (a -> b -> a) -> a -> Stream b -> a
318 {-# INLINE_STREAM foldl' #-}
319 foldl' f z (Stream step s _) = foldl_go z s
320 where
321 foldl_go z s = z `seq`
322 case step s of
323 Yield x s' -> foldl_go (f z x) s'
324 Skip s' -> foldl_go z s'
325 Done -> z
326
327 foldl1' :: (a -> a -> a) -> Stream a -> a
328 {-# INLINE_STREAM foldl1' #-}
329 foldl1' f (Stream step s sz) = foldl1'_loop s
330 where
331 foldl1'_loop s = case step s of
332 Yield x s' -> foldl' f x (Stream step s' (sz - 1))
333 Skip s' -> foldl1'_loop s'
334 Done -> error "Data.Vector.Stream.foldl1': empty stream"
335
336
337 foldr :: (a -> b -> b) -> b -> Stream a -> b
338 {-# INLINE_STREAM foldr #-}
339 foldr f z (Stream step s _) = foldr_go s
340 where
341 foldr_go s = case step s of
342 Yield x s' -> f x (foldr_go s')
343 Skip s' -> foldr_go s'
344 Done -> z
345
346 foldr1 :: (a -> a -> a) -> Stream a -> a
347 {-# INLINE_STREAM foldr1 #-}
348 foldr1 f (Stream step s sz) = foldr1_loop s
349 where
350 foldr1_loop s = case step s of
351 Yield x s' -> foldr f x (Stream step s' (sz - 1))
352 Skip s' -> foldr1_loop s'
353 Done -> error "Data.Vector.Stream.foldr1: empty stream"
354
355 mapM_ :: Monad m => (a -> m ()) -> Stream a -> m ()
356 {-# INLINE_STREAM mapM_ #-}
357 mapM_ m (Stream step s _) = mapM_go s
358 where
359 mapM_go s = case step s of
360 Yield x s' -> do { m x; mapM_go s' }
361 Skip s' -> mapM_go s'
362 Done -> return ()
363
364 foldM :: Monad m => (a -> b -> m a) -> a -> Stream b -> m a
365 {-# INLINE_STREAM foldM #-}
366 foldM m z (Stream step s _) = foldM_go z s
367 where
368 foldM_go z s = case step s of
369 Yield x s' -> do { z' <- m z x; foldM_go z' s' }
370 Skip s' -> foldM_go z s'
371 Done -> return z
372