Include phases.h
[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 empty, singleton, replicate,
9 map, filter, zipWith,
10 foldr, foldl, foldl',
11 mapM_, foldM
12 ) where
13
14 import Prelude hiding ( replicate, map, filter, zipWith,
15 foldr, foldl,
16 mapM_ )
17
18 data Step s a = Yield a s
19 | Skip s
20 | Done
21
22 data Stream a = forall s. Stream (s -> Step s a) s Int
23
24 empty :: Stream a
25 {-# INLINE_STREAM empty #-}
26 empty = Stream (const Done) () 0
27
28 singleton :: a -> Stream a
29 {-# INLINE_STREAM singleton #-}
30 singleton x = Stream step True 1
31 where
32 {-# INLINE step #-}
33 step True = Yield x False
34 step False = Done
35
36 replicate :: Int -> a -> Stream a
37 {-# INLINE_STREAM replicate #-}
38 replicate n x = Stream step n (max n 0)
39 where
40 {-# INLINE step #-}
41 step i | i > 0 = Yield x (i-1)
42 | otherwise = Done
43
44 map :: (a -> b) -> Stream a -> Stream b
45 {-# INLINE_STREAM map #-}
46 map f (Stream step s n) = Stream step' s n
47 where
48 {-# INLINE step' #-}
49 step' s = case step s of
50 Yield x s' -> Yield (f x) s'
51 Skip s' -> Skip s'
52 Done -> Done
53
54 filter :: (a -> Bool) -> Stream a -> Stream a
55 {-# INLINE_STREAM filter #-}
56 filter f (Stream step s n) = Stream step' s n
57 where
58 {-# INLINE step' #-}
59 step' s = case step s of
60 Yield x s' | f x -> Yield x s'
61 | otherwise -> Skip s'
62 Skip s' -> Skip s'
63 Done -> Done
64
65 zipWith :: (a -> b -> c) -> Stream a -> Stream b -> Stream c
66 {-# INLINE_STREAM zipWith #-}
67 zipWith f (Stream stepa sa na) (Stream stepb sb nb)
68 = Stream step (sa, sb, Nothing) (min na nb)
69 where
70 {-# INLINE step #-}
71 step (sa, sb, Nothing) = case stepa sa of
72 Yield x sa' -> Skip (sa', sb, Just x)
73 Skip sa' -> Skip (sa', sb, Nothing)
74 Done -> Done
75
76 step (sa, sb, Just x) = case stepb sb of
77 Yield y sb' -> Yield (f x y) (sa, sb', Nothing)
78 Skip sb' -> Skip (sa, sb', Just x)
79 Done -> Done
80
81 foldl :: (a -> b -> b) -> b -> Stream a -> b
82 {-# INLINE_STREAM foldl #-}
83 foldl f z (Stream step s _) = foldl_go z s
84 where
85 foldl_go z s = case step s of
86 Yield x s' -> foldl_go (f x z) s'
87 Skip s' -> foldl_go z s'
88 Done -> z
89
90 foldl' :: (a -> b -> b) -> b -> Stream a -> b
91 {-# INLINE_STREAM foldl' #-}
92 foldl' f z (Stream step s _) = foldl_go z s
93 where
94 foldl_go !z s = case step s of
95 Yield x s' -> foldl_go (f x z) s'
96 Skip s' -> foldl_go z s'
97 Done -> z
98
99 foldr :: (a -> b -> b) -> b -> Stream a -> b
100 {-# INLINE_STREAM foldr #-}
101 foldr f z (Stream step s _) = foldr_go s
102 where
103 foldr_go s = case step s of
104 Yield x s' -> f x (foldr_go s')
105 Skip s' -> foldr_go s'
106 Done -> z
107
108 mapM_ :: Monad m => (a -> m ()) -> Stream a -> m ()
109 {-# INLINE_STREAM mapM_ #-}
110 mapM_ m (Stream step s _) = mapM_go s
111 where
112 mapM_go s = case step s of
113 Yield x s' -> do { m x; mapM_go s' }
114 Skip s' -> mapM_go s'
115 Done -> return ()
116
117 foldM :: Monad m => (a -> b -> m a) -> a -> Stream b -> m a
118 {-# INLINE_STREAM foldM #-}
119 foldM m z (Stream step s _) = foldM_go z s
120 where
121 foldM_go z s = case step s of
122 Yield x s' -> do { z' <- m z x; foldM_go z' s' }
123 Skip s' -> foldM_go z s'
124 Done -> return z
125