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