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