Haddock comments
[darcs-mirrors/vector.git] / Data / Vector / Stream.hs
1 {-# LANGUAGE ExistentialQuantification #-}
2
3 -- |
4 -- Module : Data.Vector.Stream.Size
5 -- Copyright : (c) Roman Leshchinskiy 2008
6 -- License : BSD-style
7 --
8 -- Maintainer : rl@cse.unsw.edu.au
9 -- Stability : experimental
10 -- Portability : non-portable
11 --
12 -- Fusible streams
13 --
14
15 #include "phases.h"
16
17 module Data.Vector.Stream (
18 -- * Types
19 Step(..), Stream(..),
20
21 -- * Size hints
22 size, sized,
23
24 -- * Length information
25 length, null,
26
27 -- * Construction
28 empty, singleton, cons, snoc, replicate, (++),
29
30 -- * Accessing individual elements
31 head, last, (!!),
32
33 -- * Substreams
34 init, tail, take, drop,
35
36 -- * Mapping and zipping
37 map, zipWith,
38
39 -- * Filtering
40 filter, takeWhile, dropWhile,
41
42 -- * Folding
43 foldl, foldl1, foldl', foldl1', foldr, foldr1,
44
45 -- * Unfolding
46 unfold,
47
48 -- * Conversion to/from lists
49 toList, fromList,
50
51 -- * Monadic combinators
52 mapM_, foldM
53 ) where
54
55 import Data.Vector.Stream.Size
56
57 import Prelude hiding ( length, null,
58 replicate, (++),
59 head, last, (!!),
60 init, tail, take, drop,
61 map, zipWith,
62 filter, takeWhile, dropWhile,
63 foldl, foldl1, foldr, foldr1,
64 mapM_ )
65
66 data Step s a = Yield a s
67 | Skip s
68 | Done
69
70 -- | The type of fusible streams
71 data Stream a = forall s. Stream (s -> Step s a) s Size
72
73 -- | 'Size' hint of a 'Stream'
74 size :: Stream a -> Size
75 {-# INLINE size #-}
76 size (Stream _ _ sz) = sz
77
78 -- | Attach a 'Size' hint to a 'Stream'
79 sized :: Stream a -> Size -> Stream a
80 {-# INLINE_STREAM sized #-}
81 sized (Stream step s _) sz = Stream step s sz
82
83 -- | Unfold
84 unfold :: (s -> Maybe (a, s)) -> s -> Stream a
85 {-# INLINE_STREAM unfold #-}
86 unfold f s = Stream step s Unknown
87 where
88 {-# INLINE step #-}
89 step s = case f s of
90 Just (x, s') -> Yield x s'
91 Nothing -> Done
92
93 -- | Convert a 'Stream' to a list
94 toList :: Stream a -> [a]
95 {-# INLINE toList #-}
96 toList s = foldr (:) [] s
97
98 -- | Create a 'Stream' from a list
99 fromList :: [a] -> Stream a
100 {-# INLINE_STREAM fromList #-}
101 fromList xs = Stream step xs Unknown
102 where
103 step (x:xs) = Yield x xs
104 step [] = Done
105
106 -- Length
107 -- ------
108
109 -- | Length of a 'Stream'
110 length :: Stream a -> Int
111 {-# INLINE_STREAM length #-}
112 length s = foldl' (\n _ -> n+1) 0 s
113
114 -- | Check if a 'Stream' is empty
115 null :: Stream a -> Bool
116 {-# INLINE_STREAM null #-}
117 null s = foldr (\_ _ -> False) True s
118
119 -- Construction
120 -- ------------
121
122 -- | Empty 'Stream'
123 empty :: Stream a
124 {-# INLINE_STREAM empty #-}
125 empty = Stream (const Done) () (Exact 0)
126
127 -- | Singleton 'Stream'
128 singleton :: a -> Stream a
129 {-# INLINE_STREAM singleton #-}
130 singleton x = Stream step True (Exact 1)
131 where
132 {-# INLINE step #-}
133 step True = Yield x False
134 step False = Done
135
136 -- | Replicate a value to a given length
137 replicate :: Int -> a -> Stream a
138 {-# INLINE_STREAM replicate #-}
139 replicate n x = Stream step n (Exact (max n 0))
140 where
141 {-# INLINE step #-}
142 step i | i > 0 = Yield x (i-1)
143 | otherwise = Done
144
145 -- | Prepend an element
146 cons :: a -> Stream a -> Stream a
147 {-# INLINE cons #-}
148 cons x s = singleton x ++ s
149
150 -- | Append an element
151 snoc :: Stream a -> a -> Stream a
152 {-# INLINE snoc #-}
153 snoc s x = s ++ singleton x
154
155 infixr 5 ++
156 -- | Concatenate two 'Stream's
157 (++) :: Stream a -> Stream a -> Stream a
158 {-# INLINE_STREAM (++) #-}
159 Stream stepa sa na ++ Stream stepb sb nb = Stream step (Left sa) (na + nb)
160 where
161 step (Left sa) = case stepa sa of
162 Yield x sa' -> Yield x (Left sa')
163 Skip sa' -> Skip (Left sa')
164 Done -> Skip (Right sb)
165 step (Right sb) = case stepb sb of
166 Yield x sb' -> Yield x (Right sb')
167 Skip sb' -> Skip (Right sb')
168 Done -> Done
169
170 -- Accessing elements
171 -- ------------------
172
173 -- | First element of the 'Stream' or error if empty
174 head :: Stream a -> a
175 {-# INLINE_STREAM head #-}
176 head (Stream step s _) = head_loop s
177 where
178 head_loop s = case step s of
179 Yield x _ -> x
180 Skip s' -> head_loop s'
181 Done -> error "Data.Vector.Stream.head: empty stream"
182
183 -- | Last element of the 'Stream' or error if empty
184 last :: Stream a -> a
185 {-# INLINE_STREAM last #-}
186 last (Stream step s _) = last_loop0 s
187 where
188 last_loop0 s = case step s of
189 Yield x s' -> last_loop1 x s'
190 Skip s' -> last_loop0 s'
191 Done -> error "Data.Vector.Stream.last: empty stream"
192
193 last_loop1 x s = case step s of
194 Yield y s' -> last_loop1 y s'
195 Skip s' -> last_loop1 x s'
196 Done -> x
197
198 -- | Element at the given position
199 (!!) :: Stream a -> Int -> a
200 {-# INLINE (!!) #-}
201 s !! i = head (drop i s)
202
203 -- Substreams
204 -- ----------
205
206 -- | All but the last element
207 init :: Stream a -> Stream a
208 {-# INLINE_STREAM init #-}
209 init (Stream step s sz) = Stream step' (Nothing, s) (sz - 1)
210 where
211 {-# INLINE step' #-}
212 step' (Nothing, s) = case step s of
213 Yield x s' -> Skip (Just x, s')
214 Skip s' -> Skip (Nothing, s')
215 Done -> Done -- FIXME: should be an error
216
217 step' (Just x, s) = case step s of
218 Yield y s' -> Yield x (Just y, s')
219 Skip s' -> Skip (Just x, s')
220 Done -> Done
221
222 -- | All but the first element
223 tail :: Stream a -> Stream a
224 {-# INLINE_STREAM tail #-}
225 tail (Stream step s sz) = Stream step' (Left s) (sz - 1)
226 where
227 {-# INLINE step' #-}
228 step' (Left s) = case step s of
229 Yield x s' -> Skip (Right s')
230 Skip s' -> Skip (Left s')
231 Done -> Done -- FIXME: should be error?
232
233 step' (Right s) = case step s of
234 Yield x s' -> Yield x (Right s')
235 Skip s' -> Skip (Right s')
236 Done -> Done
237
238 -- | The first @n@ elements
239 take :: Int -> Stream a -> Stream a
240 {-# INLINE_STREAM take #-}
241 take n (Stream step s sz) = Stream step' (s, 0) (smaller (Exact n) sz)
242 where
243 {-# INLINE step' #-}
244 step' (s, i) | i < n = case step s of
245 Yield x s' -> Yield x (s', i+1)
246 Skip s' -> Skip (s', i)
247 Done -> Done
248 step' (s, i) = Done
249
250 data Drop s = Drop_Drop s Int | Drop_Keep s
251
252 -- | All but the first @n@ elements
253 drop :: Int -> Stream a -> Stream a
254 {-# INLINE_STREAM drop #-}
255 drop n (Stream step s sz) = Stream step' (Drop_Drop s 0) (sz - Exact n)
256 where
257 {-# INLINE step' #-}
258 step' (Drop_Drop s i) | i < n = case step s of
259 Yield x s' -> Skip (Drop_Drop s' (i+1))
260 Skip s' -> Skip (Drop_Drop s' i)
261 Done -> Done
262 | otherwise = Skip (Drop_Keep s)
263
264 step' (Drop_Keep s) = case step s of
265 Yield x s' -> Yield x (Drop_Keep s')
266 Skip s' -> Skip (Drop_Keep s')
267 Done -> Done
268
269
270 -- Mapping/zipping
271 -- ---------------
272
273 instance Functor Stream where
274 {-# INLINE_STREAM fmap #-}
275 fmap = map
276
277 -- | Map a function over a 'Stream'
278 map :: (a -> b) -> Stream a -> Stream b
279 {-# INLINE_STREAM map #-}
280 map f (Stream step s n) = Stream step' s n
281 where
282 {-# INLINE step' #-}
283 step' s = case step s of
284 Yield x s' -> Yield (f x) s'
285 Skip s' -> Skip s'
286 Done -> Done
287
288 -- | Zip two 'Stream's with the given function
289 zipWith :: (a -> b -> c) -> Stream a -> Stream b -> Stream c
290 {-# INLINE_STREAM zipWith #-}
291 zipWith f (Stream stepa sa na) (Stream stepb sb nb)
292 = Stream step (sa, sb, Nothing) (smaller na nb)
293 where
294 {-# INLINE step #-}
295 step (sa, sb, Nothing) = case stepa sa of
296 Yield x sa' -> Skip (sa', sb, Just x)
297 Skip sa' -> Skip (sa', sb, Nothing)
298 Done -> Done
299
300 step (sa, sb, Just x) = case stepb sb of
301 Yield y sb' -> Yield (f x y) (sa, sb', Nothing)
302 Skip sb' -> Skip (sa, sb', Just x)
303 Done -> Done
304
305 -- Filtering
306 -- ---------
307
308 -- | Drop elements which do not satisfy the predicate
309 filter :: (a -> Bool) -> Stream a -> Stream a
310 {-# INLINE_STREAM filter #-}
311 filter f (Stream step s n) = Stream step' s (toMax n)
312 where
313 {-# INLINE step' #-}
314 step' s = case step s of
315 Yield x s' | f x -> Yield x s'
316 | otherwise -> Skip s'
317 Skip s' -> Skip s'
318 Done -> Done
319
320 -- | Longest prefix of elements that satisfy the predicate
321 takeWhile :: (a -> Bool) -> Stream a -> Stream a
322 {-# INLINE_STREAM takeWhile #-}
323 takeWhile f (Stream step s n) = Stream step' s (toMax n)
324 where
325 {-# INLINE step' #-}
326 step' s = case step s of
327 Yield x s' | f x -> Yield x s'
328 | otherwise -> Done
329 Skip s' -> Skip s'
330 Done -> Done
331
332
333 data DropWhile s a = DropWhile_Drop s | DropWhile_Yield a s | DropWhile_Next s
334
335 -- | Drop the longest prefix of elements that satisfy the predicate
336 dropWhile :: (a -> Bool) -> Stream a -> Stream a
337 {-# INLINE_STREAM dropWhile #-}
338 dropWhile f (Stream step s n) = Stream step' (DropWhile_Drop s) (toMax n)
339 where
340 -- NOTE: we jump through hoops here to have only one Yield; local data
341 -- declarations would be nice!
342
343 {-# INLINE step' #-}
344 step' (DropWhile_Drop s)
345 = case step s of
346 Yield x s' | f x -> Skip (DropWhile_Drop s')
347 | otherwise -> Skip (DropWhile_Yield x s')
348 Skip s' -> Skip (DropWhile_Drop s')
349 Done -> Done
350
351 step' (DropWhile_Yield x s) = Yield x (DropWhile_Next s)
352
353 step' (DropWhile_Next s) = case step s of
354 Yield x s' -> Skip (DropWhile_Yield x s')
355 Skip s' -> Skip (DropWhile_Next s')
356 Done -> Done
357
358 -- Folding
359 -- -------
360
361 -- | Left fold
362 foldl :: (a -> b -> a) -> a -> Stream b -> a
363 {-# INLINE_STREAM foldl #-}
364 foldl f z (Stream step s _) = foldl_go z s
365 where
366 foldl_go z s = case step s of
367 Yield x s' -> foldl_go (f z x) s'
368 Skip s' -> foldl_go z s'
369 Done -> z
370
371 -- | Left fold on non-empty 'Stream's
372 foldl1 :: (a -> a -> a) -> Stream a -> a
373 {-# INLINE_STREAM foldl1 #-}
374 foldl1 f (Stream step s sz) = foldl1_loop s
375 where
376 foldl1_loop s = case step s of
377 Yield x s' -> foldl f x (Stream step s' (sz - 1))
378 Skip s' -> foldl1_loop s'
379 Done -> error "Data.Vector.Stream.foldl1: empty stream"
380
381 -- | Left fold with strict accumulator
382 foldl' :: (a -> b -> a) -> a -> Stream b -> a
383 {-# INLINE_STREAM foldl' #-}
384 foldl' f z (Stream step s _) = foldl_go z s
385 where
386 foldl_go z s = z `seq`
387 case step s of
388 Yield x s' -> foldl_go (f z x) s'
389 Skip s' -> foldl_go z s'
390 Done -> z
391
392 -- | Left fold on non-empty 'Stream's with strict accumulator
393 foldl1' :: (a -> a -> a) -> Stream a -> a
394 {-# INLINE_STREAM foldl1' #-}
395 foldl1' f (Stream step s sz) = foldl1'_loop s
396 where
397 foldl1'_loop s = case step s of
398 Yield x s' -> foldl' f x (Stream step s' (sz - 1))
399 Skip s' -> foldl1'_loop s'
400 Done -> error "Data.Vector.Stream.foldl1': empty stream"
401
402 -- | Right fold
403 foldr :: (a -> b -> b) -> b -> Stream a -> b
404 {-# INLINE_STREAM foldr #-}
405 foldr f z (Stream step s _) = foldr_go s
406 where
407 foldr_go s = case step s of
408 Yield x s' -> f x (foldr_go s')
409 Skip s' -> foldr_go s'
410 Done -> z
411
412 -- | Right fold on non-empty 'Stream's
413 foldr1 :: (a -> a -> a) -> Stream a -> a
414 {-# INLINE_STREAM foldr1 #-}
415 foldr1 f (Stream step s sz) = foldr1_loop s
416 where
417 foldr1_loop s = case step s of
418 Yield x s' -> foldr f x (Stream step s' (sz - 1))
419 Skip s' -> foldr1_loop s'
420 Done -> error "Data.Vector.Stream.foldr1: empty stream"
421
422 -- | Apply a monadic action to each element of the stream
423 mapM_ :: Monad m => (a -> m ()) -> Stream a -> m ()
424 {-# INLINE_STREAM mapM_ #-}
425 mapM_ m (Stream step s _) = mapM_go s
426 where
427 mapM_go s = case step s of
428 Yield x s' -> do { m x; mapM_go s' }
429 Skip s' -> mapM_go s'
430 Done -> return ()
431
432 -- | Monadic fold
433 foldM :: Monad m => (a -> b -> m a) -> a -> Stream b -> m a
434 {-# INLINE_STREAM foldM #-}
435 foldM m z (Stream step s _) = foldM_go z s
436 where
437 foldM_go z s = case step s of
438 Yield x s' -> do { z' <- m z x; foldM_go z' s' }
439 Skip s' -> foldM_go z s'
440 Done -> return z
441