Add the uninplace rule
[darcs-mirrors/vector.git] / Data / Vector / Fusion / Stream.hs
1 {-# LANGUAGE ExistentialQuantification, FlexibleInstances #-}
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, Id(..),
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 M
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 newtype Id a = Id { unId :: a }
74
75 instance Functor Id where
76 fmap f (Id x) = Id (f x)
77
78 instance Monad Id where
79 return = Id
80 Id x >>= f = f x
81
82 -- | The type of fusible streams
83 type Stream = M.Stream Id
84
85 type MStream = M.Stream
86
87 liftStream :: Monad m => Stream a -> M.Stream m a
88 {-# INLINE_STREAM liftStream #-}
89 liftStream (M.Stream step s sz) = M.Stream (return . unId . step) s sz
90
91 -- | 'Size' hint of a 'Stream'
92 size :: Stream a -> Size
93 {-# INLINE size #-}
94 size = M.size
95
96 -- | Attach a 'Size' hint to a 'Stream'
97 sized :: Stream a -> Size -> Stream a
98 {-# INLINE sized #-}
99 sized = M.sized
100
101 -- | Unfold
102 unfold :: (s -> Maybe (a, s)) -> s -> Stream a
103 {-# INLINE unfold #-}
104 unfold = M.unfold
105
106 -- | Convert a 'Stream' to a list
107 toList :: Stream a -> [a]
108 {-# INLINE toList #-}
109 toList s = unId (M.toList s)
110
111 -- | Create a 'Stream' from a list
112 fromList :: [a] -> Stream a
113 {-# INLINE fromList #-}
114 fromList = M.fromList
115
116 -- Length
117 -- ------
118
119 -- | Length of a 'Stream'
120 length :: Stream a -> Int
121 {-# INLINE length #-}
122 length = unId . M.length
123
124 -- | Check if a 'Stream' is empty
125 null :: Stream a -> Bool
126 {-# INLINE null #-}
127 null = unId . M.null
128
129 -- Construction
130 -- ------------
131
132 -- | Empty 'Stream'
133 empty :: Stream a
134 {-# INLINE empty #-}
135 empty = M.empty
136
137 -- | Singleton 'Stream'
138 singleton :: a -> Stream a
139 {-# INLINE singleton #-}
140 singleton = M.singleton
141
142 -- | Replicate a value to a given length
143 replicate :: Int -> a -> Stream a
144 {-# INLINE_STREAM replicate #-}
145 replicate = M.replicate
146
147 -- | Prepend an element
148 cons :: a -> Stream a -> Stream a
149 {-# INLINE cons #-}
150 cons = M.cons
151
152 -- | Append an element
153 snoc :: Stream a -> a -> Stream a
154 {-# INLINE snoc #-}
155 snoc = M.snoc
156
157 infixr 5 ++
158 -- | Concatenate two 'Stream's
159 (++) :: Stream a -> Stream a -> Stream a
160 {-# INLINE (++) #-}
161 (++) = (M.++)
162
163 -- Accessing elements
164 -- ------------------
165
166 -- | First element of the 'Stream' or error if empty
167 head :: Stream a -> a
168 {-# INLINE head #-}
169 head = unId . M.head
170
171 -- | Last element of the 'Stream' or error if empty
172 last :: Stream a -> a
173 {-# INLINE last #-}
174 last = unId . M.last
175
176 -- | Element at the given position
177 (!!) :: Stream a -> Int -> a
178 {-# INLINE (!!) #-}
179 s !! i = unId (s M.!! i)
180
181 -- Substreams
182 -- ----------
183
184 -- | Extract a substream of the given length starting at the given position.
185 extract :: Stream a -> Int -- ^ starting index
186 -> Int -- ^ length
187 -> Stream a
188 {-# INLINE extract #-}
189 extract = M.extract
190
191 -- | All but the last element
192 init :: Stream a -> Stream a
193 {-# INLINE init #-}
194 init = M.init
195
196 -- | All but the first element
197 tail :: Stream a -> Stream a
198 {-# INLINE tail #-}
199 tail = M.tail
200
201 -- | The first @n@ elements
202 take :: Int -> Stream a -> Stream a
203 {-# INLINE take #-}
204 take = M.take
205
206 -- | All but the first @n@ elements
207 drop :: Int -> Stream a -> Stream a
208 {-# INLINE drop #-}
209 drop = M.drop
210
211 -- Mapping/zipping
212 -- ---------------
213
214 -- | Map a function over a 'Stream'
215 map :: (a -> b) -> Stream a -> Stream b
216 {-# INLINE map #-}
217 map = M.map
218
219 -- | Zip two 'Stream's with the given function
220 zipWith :: (a -> b -> c) -> Stream a -> Stream b -> Stream c
221 {-# INLINE zipWith #-}
222 zipWith = M.zipWith
223
224 -- Filtering
225 -- ---------
226
227 -- | Drop elements which do not satisfy the predicate
228 filter :: (a -> Bool) -> Stream a -> Stream a
229 {-# INLINE filter #-}
230 filter = M.filter
231
232 -- | Longest prefix of elements that satisfy the predicate
233 takeWhile :: (a -> Bool) -> Stream a -> Stream a
234 {-# INLINE takeWhile #-}
235 takeWhile = M.takeWhile
236
237 -- | Drop the longest prefix of elements that satisfy the predicate
238 dropWhile :: (a -> Bool) -> Stream a -> Stream a
239 {-# INLINE dropWhile #-}
240 dropWhile = M.dropWhile
241
242 -- Searching
243 -- ---------
244
245 infix 4 `elem`
246 -- | Check whether the 'Stream' contains an element
247 elem :: Eq a => a -> Stream a -> Bool
248 {-# INLINE elem #-}
249 elem x = unId . M.elem x
250
251 infix 4 `notElem`
252 -- | Inverse of `elem`
253 notElem :: Eq a => a -> Stream a -> Bool
254 {-# INLINE notElem #-}
255 notElem x = unId . M.notElem x
256
257 -- | Yield 'Just' the first element matching the predicate or 'Nothing' if no
258 -- such element exists.
259 find :: (a -> Bool) -> Stream a -> Maybe a
260 {-# INLINE find #-}
261 find f = unId . M.find f
262
263 -- | Yield 'Just' the index of the first element matching the predicate or
264 -- 'Nothing' if no such element exists.
265 findIndex :: (a -> Bool) -> Stream a -> Maybe Int
266 {-# INLINE findIndex #-}
267 findIndex f = unId . M.findIndex f
268
269 -- Folding
270 -- -------
271
272 -- | Left fold
273 foldl :: (a -> b -> a) -> a -> Stream b -> a
274 {-# INLINE foldl #-}
275 foldl f z = unId . M.foldl f z
276
277 -- | Left fold on non-empty 'Stream's
278 foldl1 :: (a -> a -> a) -> Stream a -> a
279 {-# INLINE foldl1 #-}
280 foldl1 f = unId . M.foldl1 f
281
282 -- | Left fold with strict accumulator
283 foldl' :: (a -> b -> a) -> a -> Stream b -> a
284 {-# INLINE foldl' #-}
285 foldl' f z = unId . M.foldl' f z
286
287 -- | Left fold on non-empty 'Stream's with strict accumulator
288 foldl1' :: (a -> a -> a) -> Stream a -> a
289 {-# INLINE foldl1' #-}
290 foldl1' f = unId . M.foldl1' f
291
292 -- | Right fold
293 foldr :: (a -> b -> b) -> b -> Stream a -> b
294 {-# INLINE foldr #-}
295 foldr f z = unId . M.foldr f z
296
297 -- | Right fold on non-empty 'Stream's
298 foldr1 :: (a -> a -> a) -> Stream a -> a
299 {-# INLINE foldr1 #-}
300 foldr1 f = unId . M.foldr1 f
301
302 -- Comparisons
303 -- -----------
304
305 eq :: Eq a => Stream a -> Stream a -> Bool
306 {-# INLINE_STREAM eq #-}
307 eq (M.Stream step1 s1 _) (M.Stream step2 s2 _) = eq_loop0 s1 s2
308 where
309 eq_loop0 s1 s2 = case unId (step1 s1) of
310 Yield x s1' -> eq_loop1 x s1' s2
311 Skip s1' -> eq_loop0 s1' s2
312 Done -> null (M.Stream step2 s2 Unknown)
313
314 eq_loop1 x s1 s2 = case unId (step2 s2) of
315 Yield y s2' -> x == y && eq_loop0 s1 s2'
316 Skip s2' -> eq_loop1 x s1 s2'
317 Done -> False
318
319 cmp :: Ord a => Stream a -> Stream a -> Ordering
320 {-# INLINE_STREAM cmp #-}
321 cmp (M.Stream step1 s1 _) (M.Stream step2 s2 _) = cmp_loop0 s1 s2
322 where
323 cmp_loop0 s1 s2 = case unId (step1 s1) of
324 Yield x s1' -> cmp_loop1 x s1' s2
325 Skip s1' -> cmp_loop0 s1' s2
326 Done -> if null (M.Stream step2 s2 Unknown)
327 then EQ else LT
328
329 cmp_loop1 x s1 s2 = case unId (step2 s2) of
330 Yield y s2' -> case x `compare` y of
331 EQ -> cmp_loop0 s1 s2'
332 c -> c
333 Skip s2' -> cmp_loop1 x s1 s2'
334 Done -> GT
335
336 instance Eq a => Eq (M.Stream Id a) where
337 {-# INLINE (==) #-}
338 (==) = eq
339
340 instance Ord a => Ord (M.Stream Id a) where
341 {-# INLINE compare #-}
342 compare = cmp
343
344 -- Monadic combinators
345 -- -------------------
346
347 -- | Apply a monadic action to each element of the stream
348 mapM_ :: Monad m => (a -> m ()) -> Stream a -> m ()
349 {-# INLINE_STREAM mapM_ #-}
350 mapM_ f = M.mapM_ f . liftStream
351
352 -- | Monadic fold
353 foldM :: Monad m => (a -> b -> m a) -> a -> Stream b -> m a
354 {-# INLINE_STREAM foldM #-}
355 foldM m z = M.foldM m z . liftStream
356