Add the uninplace rule
[darcs-mirrors/vector.git] / Data / Vector / Fusion / Stream / Monadic.hs
1 {-# LANGUAGE ExistentialQuantification #-}
2
3 #include "phases.h"
4
5 module Data.Vector.Fusion.Stream.Monadic (
6 Stream(..),
7
8 -- * Size hints
9 size, sized,
10
11 -- * Length
12 length, null,
13
14 -- * Construction
15 empty, singleton, cons, snoc, replicate, (++),
16
17 -- * Accessing elements
18 head, last, (!!),
19
20 -- * Substreams
21 extract, init, tail, take, drop,
22
23 -- * Mapping and zipping
24 map, mapM, mapM_, zipWith, zipWithM,
25
26 -- * Filtering
27 filter, filterM, takeWhile, takeWhileM, dropWhile, dropWhileM,
28
29 -- * Searching
30 elem, notElem, find, findM, findIndex, findIndexM,
31
32 -- * Folding
33 foldl, foldlM, foldM, foldl1, foldl1M,
34 foldl', foldlM', foldl1', foldl1M',
35 foldr, foldrM, foldr1, foldr1M,
36
37 -- * Unfolding
38 unfold, unfoldM,
39
40 toList, fromList
41 ) where
42
43 import Data.Vector.Fusion.Stream.Step
44 import Data.Vector.Fusion.Stream.Size
45
46 import Control.Monad ( liftM )
47 import Prelude hiding ( length, null,
48 replicate, (++),
49 head, last, (!!),
50 init, tail, take, drop,
51 map, mapM, mapM_, zipWith,
52 filter, takeWhile, dropWhile,
53 elem, notElem,
54 foldl, foldl1, foldr, foldr1 )
55 import qualified Prelude
56
57 data Stream m a = forall s. Stream (s -> m (Step s a)) s Size
58
59 -- | 'Size' hint of a 'Stream'
60 size :: Stream m a -> Size
61 {-# INLINE size #-}
62 size (Stream _ _ sz) = sz
63
64 -- | Attach a 'Size' hint to a 'Stream'
65 sized :: Stream m a -> Size -> Stream m a
66 {-# INLINE_STREAM sized #-}
67 sized (Stream step s _) sz = Stream step s sz
68
69 -- Length
70 -- ------
71
72 -- | Length of a 'Stream'
73 length :: Monad m => Stream m a -> m Int
74 {-# INLINE_STREAM length #-}
75 length s = foldl' (\n _ -> n+1) 0 s
76
77 -- | Check if a 'Stream' is empty
78 null :: Monad m => Stream m a -> m Bool
79 {-# INLINE_STREAM null #-}
80 null s = foldr (\_ _ -> False) True s
81
82
83 -- Construction
84 -- ------------
85
86 -- | Empty 'Stream'
87 empty :: Monad m => Stream m a
88 {-# INLINE_STREAM empty #-}
89 empty = Stream (const (return Done)) () (Exact 0)
90
91 -- | Singleton 'Stream'
92 singleton :: Monad m => a -> Stream m a
93 {-# INLINE_STREAM singleton #-}
94 singleton x = Stream (return . step) True (Exact 1)
95 where
96 {-# INLINE step #-}
97 step True = Yield x False
98 step False = Done
99
100 -- | Replicate a value to a given length
101 replicate :: Monad m => Int -> a -> Stream m a
102 {-# INLINE_STREAM replicate #-}
103 replicate n x = Stream (return . step) n (Exact (max n 0))
104 where
105 {-# INLINE step #-}
106 step i | i > 0 = Yield x (i-1)
107 | otherwise = Done
108
109 -- | Prepend an element
110 cons :: Monad m => a -> Stream m a -> Stream m a
111 {-# INLINE cons #-}
112 cons x s = singleton x ++ s
113
114 -- | Append an element
115 snoc :: Monad m => Stream m a -> a -> Stream m a
116 {-# INLINE snoc #-}
117 snoc s x = s ++ singleton x
118
119 infixr 5 ++
120 -- | Concatenate two 'Stream's
121 (++) :: Monad m => Stream m a -> Stream m a -> Stream m a
122 {-# INLINE_STREAM (++) #-}
123 Stream stepa sa na ++ Stream stepb sb nb = Stream step (Left sa) (na + nb)
124 where
125 step (Left sa) = do
126 r <- stepa sa
127 case r of
128 Yield x sa' -> return $ Yield x (Left sa')
129 Skip sa' -> return $ Skip (Left sa')
130 Done -> return $ Skip (Right sb)
131 step (Right sb) = do
132 r <- stepb sb
133 case r of
134 Yield x sb' -> return $ Yield x (Right sb')
135 Skip sb' -> return $ Skip (Right sb')
136 Done -> return $ Done
137
138 -- Accessing elements
139 -- ------------------
140
141 -- | First element of the 'Stream' or error if empty
142 head :: Monad m => Stream m a -> m a
143 {-# INLINE_STREAM head #-}
144 head (Stream step s _) = head_loop s
145 where
146 head_loop s = do
147 r <- step s
148 case r of
149 Yield x _ -> return x
150 Skip s' -> head_loop s'
151 Done -> errorEmptyStream "head"
152
153 -- | Last element of the 'Stream' or error if empty
154 last :: Monad m => Stream m a -> m a
155 {-# INLINE_STREAM last #-}
156 last (Stream step s _) = last_loop0 s
157 where
158 last_loop0 s = do
159 r <- step s
160 case r of
161 Yield x s' -> last_loop1 x s'
162 Skip s' -> last_loop0 s'
163 Done -> errorEmptyStream "last"
164
165 last_loop1 x s = do
166 r <- step s
167 case r of
168 Yield y s' -> last_loop1 y s'
169 Skip s' -> last_loop1 x s'
170 Done -> return x
171
172 -- | Element at the given position
173 (!!) :: Monad m => Stream m a -> Int -> m a
174 {-# INLINE (!!) #-}
175 s !! i = head (drop i s)
176
177 -- Substreams
178 -- ----------
179
180 -- | Extract a substream of the given length starting at the given position.
181 extract :: Monad m => Stream m a -> Int -- ^ starting index
182 -> Int -- ^ length
183 -> Stream m a
184 {-# INLINE extract #-}
185 extract s i n = take n (drop i s)
186
187 -- | All but the last element
188 init :: Monad m => Stream m a -> Stream m a
189 {-# INLINE_STREAM init #-}
190 init (Stream step s sz) = Stream step' (Nothing, s) (sz - 1)
191 where
192 {-# INLINE step' #-}
193 step' (Nothing, s) = liftM (\r ->
194 case r of
195 Yield x s' -> Skip (Just x, s')
196 Skip s' -> Skip (Nothing, s')
197 Done -> Done -- FIXME: should be an error
198 ) (step s)
199
200 step' (Just x, s) = liftM (\r ->
201 case r of
202 Yield y s' -> Yield x (Just y, s')
203 Skip s' -> Skip (Just x, s')
204 Done -> Done
205 ) (step s)
206
207 -- | All but the first element
208 tail :: Monad m => Stream m a -> Stream m a
209 {-# INLINE_STREAM tail #-}
210 tail (Stream step s sz) = Stream step' (Left s) (sz - 1)
211 where
212 {-# INLINE step' #-}
213 step' (Left s) = liftM (\r ->
214 case r of
215 Yield x s' -> Skip (Right s')
216 Skip s' -> Skip (Left s')
217 Done -> Done -- FIXME: should be error?
218 ) (step s)
219
220 step' (Right s) = liftM (\r ->
221 case r of
222 Yield x s' -> Yield x (Right s')
223 Skip s' -> Skip (Right s')
224 Done -> Done
225 ) (step s)
226
227 -- | The first @n@ elements
228 take :: Monad m => Int -> Stream m a -> Stream m a
229 {-# INLINE_STREAM take #-}
230 take n (Stream step s sz) = Stream step' (s, 0) (smaller (Exact n) sz)
231 where
232 {-# INLINE step' #-}
233 step' (s, i) | i < n = liftM (\r ->
234 case r of
235 Yield x s' -> Yield x (s', i+1)
236 Skip s' -> Skip (s', i)
237 Done -> Done
238 ) (step s)
239 step' (s, i) = return Done
240
241 data Drop s = Drop_Drop s Int | Drop_Keep s
242
243 -- | All but the first @n@ elements
244 drop :: Monad m => Int -> Stream m a -> Stream m a
245 {-# INLINE_STREAM drop #-}
246 drop n (Stream step s sz) = Stream step' (Drop_Drop s 0) (sz - Exact n)
247 where
248 {-# INLINE step' #-}
249 step' (Drop_Drop s i) | i < n = liftM (\r ->
250 case r of
251 Yield x s' -> Skip (Drop_Drop s' (i+1))
252 Skip s' -> Skip (Drop_Drop s' i)
253 Done -> Done
254 ) (step s)
255 | otherwise = return $ Skip (Drop_Keep s)
256
257 step' (Drop_Keep s) = liftM (\r ->
258 case r of
259 Yield x s' -> Yield x (Drop_Keep s')
260 Skip s' -> Skip (Drop_Keep s')
261 Done -> Done
262 ) (step s)
263
264
265 -- Mapping/zipping
266 -- ---------------
267
268 instance Monad m => Functor (Stream m) where
269 {-# INLINE fmap #-}
270 fmap = map
271
272 map :: Monad m => (a -> b) -> Stream m a -> Stream m b
273 {-# INLINE map #-}
274 map f = mapM (return . f)
275
276 mapM :: Monad m => (a -> m b) -> Stream m a -> Stream m b
277 {-# INLINE_STREAM mapM #-}
278 mapM f (Stream step s n) = Stream step' s n
279 where
280 {-# INLINE step' #-}
281 step' s = do
282 r <- step s
283 case r of
284 Yield x s' -> liftM (`Yield` s') (f x)
285 Skip s' -> return (Skip s')
286 Done -> return Done
287
288 mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()
289 {-# INLINE_STREAM mapM_ #-}
290 mapM_ m (Stream step s _) = mapM_go s
291 where
292 mapM_go s = do
293 r <- step s
294 case r of
295 Yield x s' -> do { m x; mapM_go s' }
296 Skip s' -> mapM_go s'
297 Done -> return ()
298
299 -- | Zip two 'Stream's with the given function
300 zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
301 {-# INLINE zipWith #-}
302 zipWith f = zipWithM (\a b -> return (f a b))
303
304 -- | Zip two 'Stream's with the given function
305 zipWithM :: Monad m => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
306 {-# INLINE_STREAM zipWithM #-}
307 zipWithM f (Stream stepa sa na) (Stream stepb sb nb)
308 = Stream step (sa, sb, Nothing) (smaller na nb)
309 where
310 {-# INLINE step #-}
311 step (sa, sb, Nothing) = liftM (\r ->
312 case r of
313 Yield x sa' -> Skip (sa', sb, Just x)
314 Skip sa' -> Skip (sa', sb, Nothing)
315 Done -> Done
316 ) (stepa sa)
317
318 step (sa, sb, Just x) = do
319 r <- stepb sb
320 case r of
321 Yield y sb' ->
322 do
323 z <- f x y
324 return $ Yield z (sa, sb', Nothing)
325 Skip sb' -> return $ Skip (sa, sb', Just x)
326 Done -> return $ Done
327
328 -- Filtering
329 -- ---------
330
331 -- | Drop elements which do not satisfy the predicate
332 filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
333 {-# INLINE filter #-}
334 filter f = filterM (return . f)
335
336 -- | Drop elements which do not satisfy the predicate
337 filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
338 {-# INLINE_STREAM filterM #-}
339 filterM f (Stream step s n) = Stream step' s (toMax n)
340 where
341 {-# INLINE step' #-}
342 step' s = do
343 r <- step s
344 case r of
345 Yield x s' -> do
346 b <- f x
347 return $ if b then Yield x s'
348 else Skip s'
349 Skip s' -> return $ Skip s'
350 Done -> return $ Done
351
352 -- | Longest prefix of elements that satisfy the predicate
353 takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
354 {-# INLINE takeWhile #-}
355 takeWhile f = takeWhileM (return . f)
356
357 -- | Longest prefix of elements that satisfy the predicate
358 takeWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
359 {-# INLINE_STREAM takeWhileM #-}
360 takeWhileM f (Stream step s n) = Stream step' s (toMax n)
361 where
362 {-# INLINE step' #-}
363 step' s = do
364 r <- step s
365 case r of
366 Yield x s' -> do
367 b <- f x
368 return $ if b then Yield x s' else Done
369 Skip s' -> return $ Skip s'
370 Done -> return $ Done
371
372
373 dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
374 {-# INLINE dropWhile #-}
375 dropWhile f = dropWhileM (return . f)
376
377 data DropWhile s a = DropWhile_Drop s | DropWhile_Yield a s | DropWhile_Next s
378
379 -- | Drop the longest prefix of elements that satisfy the predicate
380 dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
381 {-# INLINE_STREAM dropWhileM #-}
382 dropWhileM f (Stream step s n) = Stream step' (DropWhile_Drop s) (toMax n)
383 where
384 -- NOTE: we jump through hoops here to have only one Yield; local data
385 -- declarations would be nice!
386
387 {-# INLINE step' #-}
388 step' (DropWhile_Drop s)
389 = do
390 r <- step s
391 case r of
392 Yield x s' -> do
393 b <- f x
394 return $ if b then Skip (DropWhile_Drop s')
395 else Skip (DropWhile_Yield x s')
396 Skip s' -> return $ Skip (DropWhile_Drop s')
397 Done -> return $ Done
398
399 step' (DropWhile_Yield x s) = return $ Yield x (DropWhile_Next s)
400
401 step' (DropWhile_Next s)
402 = liftM (\r ->
403 case r of
404 Yield x s' -> Skip (DropWhile_Yield x s')
405 Skip s' -> Skip (DropWhile_Next s')
406 Done -> Done
407 ) (step s)
408
409 -- Searching
410 -- ---------
411
412 infix 4 `elem`
413 -- | Check whether the 'Stream' contains an element
414 elem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
415 {-# INLINE_STREAM elem #-}
416 elem x (Stream step s _) = elem_loop s
417 where
418 elem_loop s = do
419 r <- step s
420 case r of
421 Yield y s' | x == y -> return True
422 | otherwise -> elem_loop s'
423 Skip s' -> elem_loop s'
424 Done -> return False
425
426 infix 4 `notElem`
427 -- | Inverse of `elem`
428 notElem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
429 {-# INLINE notElem #-}
430 notElem x s = liftM not (elem x s)
431
432 find :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe a)
433 {-# INLINE find #-}
434 find f = findM (return . f)
435
436 -- | Yield 'Just' the first element matching the predicate or 'Nothing' if no
437 -- such element exists.
438 findM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe a)
439 {-# INLINE_STREAM findM #-}
440 findM f (Stream step s _) = find_loop s
441 where
442 find_loop s = do
443 r <- step s
444 case r of
445 Yield x s' -> do
446 b <- f x
447 if b then return $ Just x
448 else find_loop s'
449 Skip s' -> find_loop s'
450 Done -> return Nothing
451
452 findIndex :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe Int)
453 {-# INLINE_STREAM findIndex #-}
454 findIndex f = findIndexM (return . f)
455
456 -- | Yield 'Just' the index of the first element matching the predicate or
457 -- 'Nothing' if no such element exists.
458 findIndexM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe Int)
459 {-# INLINE_STREAM findIndexM #-}
460 findIndexM f (Stream step s _) = findIndex_loop s 0
461 where
462 findIndex_loop s i = do
463 r <- step s
464 case r of
465 Yield x s' -> do
466 b <- f x
467 if b then return $ Just i
468 else findIndex_loop s' (i+1)
469 Skip s' -> findIndex_loop s' i
470 Done -> return Nothing
471
472 -- Folding
473 -- -------
474
475 foldl :: Monad m => (a -> b -> a) -> a -> Stream m b -> m a
476 {-# INLINE foldl #-}
477 foldl f = foldlM (\a b -> return (f a b))
478
479 foldlM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> m a
480 {-# INLINE_STREAM foldlM #-}
481 foldlM m z (Stream step s _) = foldlM_go z s
482 where
483 foldlM_go z s = do
484 r <- step s
485 case r of
486 Yield x s' -> do { z' <- m z x; foldlM_go z' s' }
487 Skip s' -> foldlM_go z s'
488 Done -> return z
489
490 foldM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> m a
491 {-# INLINE foldM #-}
492 foldM = foldlM
493
494 foldl1 :: Monad m => (a -> a -> a) -> Stream m a -> m a
495 {-# INLINE foldl1 #-}
496 foldl1 f = foldl1M (\a b -> return (f a b))
497
498 foldl1M :: Monad m => (a -> a -> m a) -> Stream m a -> m a
499 {-# INLINE_STREAM foldl1M #-}
500 foldl1M f (Stream step s sz) = foldl1M_go s
501 where
502 foldl1M_go s = do
503 r <- step s
504 case r of
505 Yield x s' -> foldlM f x (Stream step s' (sz - 1))
506 Skip s' -> foldl1M_go s'
507 Done -> errorEmptyStream "foldl1M"
508
509 foldl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> m a
510 {-# INLINE foldl' #-}
511 foldl' f = foldlM' (\a b -> return (f a b))
512
513 foldlM' :: Monad m => (a -> b -> m a) -> a -> Stream m b -> m a
514 {-# INLINE_STREAM foldlM' #-}
515 foldlM' m z (Stream step s _) = foldlM'_go z s
516 where
517 foldlM'_go z s = z `seq`
518 do
519 r <- step s
520 case r of
521 Yield x s' -> do { z' <- m z x; foldlM'_go z' s' }
522 Skip s' -> foldlM'_go z s'
523 Done -> return z
524
525 foldl1' :: Monad m => (a -> a -> a) -> Stream m a -> m a
526 {-# INLINE foldl1' #-}
527 foldl1' f = foldl1M' (\a b -> return (f a b))
528
529 foldl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> m a
530 {-# INLINE_STREAM foldl1M' #-}
531 foldl1M' f (Stream step s sz) = foldl1M'_go s
532 where
533 foldl1M'_go s = do
534 r <- step s
535 case r of
536 Yield x s' -> foldlM' f x (Stream step s' (sz - 1))
537 Skip s' -> foldl1M'_go s'
538 Done -> errorEmptyStream "foldl1M'"
539
540 foldr :: Monad m => (a -> b -> b) -> b -> Stream m a -> m b
541 {-# INLINE foldr #-}
542 foldr f = foldrM (\a b -> return (f a b))
543
544 foldrM :: Monad m => (a -> b -> m b) -> b -> Stream m a -> m b
545 {-# INLINE_STREAM foldrM #-}
546 foldrM f z (Stream step s _) = foldrM_go s
547 where
548 foldrM_go s = do
549 r <- step s
550 case r of
551 Yield x s' -> f x =<< foldrM_go s'
552 Skip s' -> foldrM_go s'
553 Done -> return z
554
555 foldr1 :: Monad m => (a -> a -> a) -> Stream m a -> m a
556 {-# INLINE foldr1 #-}
557 foldr1 f = foldr1M (\a b -> return (f a b))
558
559 foldr1M :: Monad m => (a -> a -> m a) -> Stream m a -> m a
560 {-# INLINE_STREAM foldr1M #-}
561 foldr1M f (Stream step s _) = foldr1M_go0 s
562 where
563 foldr1M_go0 s = do
564 r <- step s
565 case r of
566 Yield x s' -> foldr1M_go1 x s'
567 Skip s' -> foldr1M_go0 s'
568 Done -> errorEmptyStream "foldr1M"
569
570 foldr1M_go1 x s = do
571 r <- step s
572 case r of
573 Yield y s' -> f x =<< foldr1M_go1 y s'
574 Skip s' -> foldr1M_go1 x s'
575 Done -> return x
576
577 -- Unfolding
578 -- ---------
579
580 -- | Unfold
581 unfold :: Monad m => (s -> Maybe (a, s)) -> s -> Stream m a
582 {-# INLINE_STREAM unfold #-}
583 unfold f = unfoldM (return . f)
584
585 unfoldM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m a
586 {-# INLINE_STREAM unfoldM #-}
587 unfoldM f s = Stream step s Unknown
588 where
589 {-# INLINE step #-}
590 step s = liftM (\r ->
591 case r of
592 Just (x, s') -> Yield x s'
593 Nothing -> Done
594 ) (f s)
595
596
597 -- Conversions
598 -- -----------
599
600 toList :: Monad m => Stream m a -> m [a]
601 {-# INLINE toList #-}
602 toList = foldr (:) []
603
604 fromList :: Monad m => [a] -> Stream m a
605 {-# INLINE_STREAM fromList #-}
606 fromList xs = Stream step xs Unknown
607 where
608 step (x:xs) = return (Yield x xs)
609 step [] = return Done
610
611
612 errorEmptyStream :: String -> a
613 errorEmptyStream s = error $ "Data.Vector.Fusion.Stream.Monadic."
614 Prelude.++ s Prelude.++ ": empty stream"
615