Replace type families by GADTs for associating a monad with a mutable vector
[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 size, sized, unfold, toList, fromList,
9 empty, singleton, replicate, (++),
10 map, filter, zipWith,
11 foldr, foldl, foldl',
12 mapM_, foldM
13 ) where
14
15 import Data.Vector.Stream.Size
16
17 import Prelude hiding ( replicate, (++), map, filter, zipWith,
18 foldr, foldl,
19 mapM_ )
20
21 data Step s a = Yield a s
22 | Skip s
23 | Done
24
25 data Stream a = forall s. Stream (s -> Step s a) s Size
26
27 size :: Stream a -> Size
28 {-# INLINE size #-}
29 size (Stream _ _ sz) = sz
30
31 sized :: Stream a -> Size -> Stream a
32 {-# INLINE_STREAM sized #-}
33 sized (Stream step s _) sz = Stream step s sz
34
35 unfold :: (s -> Maybe (a, s)) -> s -> Stream a
36 {-# INLINE_STREAM unfold #-}
37 unfold f s = Stream step s Unknown
38 where
39 {-# INLINE step #-}
40 step s = case f s of
41 Just (x, s') -> Yield x s'
42 Nothing -> Done
43
44 toList :: Stream a -> [a]
45 {-# INLINE toList #-}
46 toList s = foldr (:) [] s
47
48 fromList :: [a] -> Stream a
49 {-# INLINE_STREAM fromList #-}
50 fromList xs = Stream step xs Unknown
51 where
52 step (x:xs) = Yield x xs
53 step [] = Done
54
55 empty :: Stream a
56 {-# INLINE_STREAM empty #-}
57 empty = Stream (const Done) () (Exact 0)
58
59 singleton :: a -> Stream a
60 {-# INLINE_STREAM singleton #-}
61 singleton x = Stream step True (Exact 1)
62 where
63 {-# INLINE step #-}
64 step True = Yield x False
65 step False = Done
66
67 replicate :: Int -> a -> Stream a
68 {-# INLINE_STREAM replicate #-}
69 replicate n x = Stream step n (Exact (max n 0))
70 where
71 {-# INLINE step #-}
72 step i | i > 0 = Yield x (i-1)
73 | otherwise = Done
74
75 infixr ++
76 (++) :: Stream a -> Stream a -> Stream a
77 {-# INLINE_STREAM (++) #-}
78 Stream stepa sa na ++ Stream stepb sb nb = Stream step (Left sa) (na + nb)
79 where
80 step (Left sa) = case stepa sa of
81 Yield x sa' -> Yield x (Left sa')
82 Skip sa' -> Skip (Left sa')
83 Done -> Skip (Right sb)
84 step (Right sb) = case stepb sb of
85 Yield x sb' -> Yield x (Right sb')
86 Skip sb' -> Skip (Right sb')
87 Done -> Done
88
89 map :: (a -> b) -> Stream a -> Stream b
90 {-# INLINE_STREAM map #-}
91 map f (Stream step s n) = Stream step' s n
92 where
93 {-# INLINE step' #-}
94 step' s = case step s of
95 Yield x s' -> Yield (f x) s'
96 Skip s' -> Skip s'
97 Done -> Done
98
99 filter :: (a -> Bool) -> Stream a -> Stream a
100 {-# INLINE_STREAM filter #-}
101 filter f (Stream step s n) = Stream step' s (toMax n)
102 where
103 {-# INLINE step' #-}
104 step' s = case step s of
105 Yield x s' | f x -> Yield x s'
106 | otherwise -> Skip s'
107 Skip s' -> Skip s'
108 Done -> Done
109
110 zipWith :: (a -> b -> c) -> Stream a -> Stream b -> Stream c
111 {-# INLINE_STREAM zipWith #-}
112 zipWith f (Stream stepa sa na) (Stream stepb sb nb)
113 = Stream step (sa, sb, Nothing) (smaller na nb)
114 where
115 {-# INLINE step #-}
116 step (sa, sb, Nothing) = case stepa sa of
117 Yield x sa' -> Skip (sa', sb, Just x)
118 Skip sa' -> Skip (sa', sb, Nothing)
119 Done -> Done
120
121 step (sa, sb, Just x) = case stepb sb of
122 Yield y sb' -> Yield (f x y) (sa, sb', Nothing)
123 Skip sb' -> Skip (sa, sb', Just x)
124 Done -> Done
125
126 foldl :: (a -> b -> a) -> a -> Stream b -> a
127 {-# INLINE_STREAM foldl #-}
128 foldl f z (Stream step s _) = foldl_go z s
129 where
130 foldl_go z s = case step s of
131 Yield x s' -> foldl_go (f z x) s'
132 Skip s' -> foldl_go z s'
133 Done -> z
134
135 foldl' :: (a -> b -> a) -> a -> Stream b -> a
136 {-# INLINE_STREAM foldl' #-}
137 foldl' f !z (Stream step s _) = foldl_go z s
138 where
139 foldl_go !z s = case step s of
140 Yield x s' -> foldl_go (f z x) s'
141 Skip s' -> foldl_go z s'
142 Done -> z
143
144 foldr :: (a -> b -> b) -> b -> Stream a -> b
145 {-# INLINE_STREAM foldr #-}
146 foldr f z (Stream step s _) = foldr_go s
147 where
148 foldr_go s = case step s of
149 Yield x s' -> f x (foldr_go s')
150 Skip s' -> foldr_go s'
151 Done -> z
152
153 mapM_ :: Monad m => (a -> m ()) -> Stream a -> m ()
154 {-# INLINE_STREAM mapM_ #-}
155 mapM_ m (Stream step s _) = mapM_go s
156 where
157 mapM_go s = case step s of
158 Yield x s' -> do { m x; mapM_go s' }
159 Skip s' -> mapM_go s'
160 Done -> return ()
161
162 foldM :: Monad m => (a -> b -> m a) -> a -> Stream b -> m a
163 {-# INLINE_STREAM foldM #-}
164 foldM m z (Stream step s _) = foldM_go z s
165 where
166 foldM_go z s = case step s of
167 Yield x s' -> do { z' <- m z x; foldM_go z' s' }
168 Skip s' -> foldM_go z s'
169 Done -> return z
170