Implement and test zipWith3, zip3, unzip, unzip3
[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 : Roman Leshchinskiy <rl@cse.unsw.edu.au>
9 -- Stability : experimental
10 -- Portability : non-portable
11 --
12 -- Streams for stream fusion
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
37 map, concatMap,
38
39 -- * Zipping
40 zipWith, zipWith3,
41
42 -- * Filtering
43 filter, takeWhile, dropWhile,
44
45 -- * Searching
46 elem, notElem, find, findIndex,
47
48 -- * Folding
49 foldl, foldl1, foldl', foldl1', foldr, foldr1,
50
51 -- * Specialised folds
52 and, or,
53
54 -- * Unfolding
55 unfoldr,
56
57 -- * Scans
58 prescanl, prescanl',
59
60 -- * Conversions
61 toList, fromList, liftStream,
62
63 -- * Monadic combinators
64 mapM_, foldM
65 ) where
66
67 import Data.Vector.Fusion.Stream.Size
68 import Data.Vector.Fusion.Stream.Monadic ( Step(..) )
69 import qualified Data.Vector.Fusion.Stream.Monadic as M
70
71 import Prelude hiding ( length, null,
72 replicate, (++),
73 head, last, (!!),
74 init, tail, take, drop,
75 map, concatMap,
76 zipWith, zipWith3,
77 filter, takeWhile, dropWhile,
78 elem, notElem,
79 foldl, foldl1, foldr, foldr1,
80 and, or,
81 mapM_ )
82
83
84 -- | Identity monad
85 newtype Id a = Id { unId :: a }
86
87 instance Functor Id where
88 fmap f (Id x) = Id (f x)
89
90 instance Monad Id where
91 return = Id
92 Id x >>= f = f x
93
94 -- | The type of pure streams
95 type Stream = M.Stream Id
96
97 -- | Alternative name for monadic streams
98 type MStream = M.Stream
99
100 -- | Convert a pure stream to a monadic stream
101 liftStream :: Monad m => Stream a -> M.Stream m a
102 {-# INLINE_STREAM liftStream #-}
103 liftStream (M.Stream step s sz) = M.Stream (return . unId . step) s sz
104
105 -- | 'Size' hint of a 'Stream'
106 size :: Stream a -> Size
107 {-# INLINE size #-}
108 size = M.size
109
110 -- | Attach a 'Size' hint to a 'Stream'
111 sized :: Stream a -> Size -> Stream a
112 {-# INLINE sized #-}
113 sized = M.sized
114
115 -- Length
116 -- ------
117
118 -- | Length of a 'Stream'
119 length :: Stream a -> Int
120 {-# INLINE length #-}
121 length = unId . M.length
122
123 -- | Check if a 'Stream' is empty
124 null :: Stream a -> Bool
125 {-# INLINE null #-}
126 null = unId . M.null
127
128 -- Construction
129 -- ------------
130
131 -- | Empty 'Stream'
132 empty :: Stream a
133 {-# INLINE empty #-}
134 empty = M.empty
135
136 -- | Singleton 'Stream'
137 singleton :: a -> Stream a
138 {-# INLINE singleton #-}
139 singleton = M.singleton
140
141 -- | Replicate a value to a given length
142 replicate :: Int -> a -> Stream a
143 {-# INLINE_STREAM replicate #-}
144 replicate = M.replicate
145
146 -- | Prepend an element
147 cons :: a -> Stream a -> Stream a
148 {-# INLINE cons #-}
149 cons = M.cons
150
151 -- | Append an element
152 snoc :: Stream a -> a -> Stream a
153 {-# INLINE snoc #-}
154 snoc = M.snoc
155
156 infixr 5 ++
157 -- | Concatenate two 'Stream's
158 (++) :: Stream a -> Stream a -> Stream a
159 {-# INLINE (++) #-}
160 (++) = (M.++)
161
162 -- Accessing elements
163 -- ------------------
164
165 -- | First element of the 'Stream' or error if empty
166 head :: Stream a -> a
167 {-# INLINE head #-}
168 head = unId . M.head
169
170 -- | Last element of the 'Stream' or error if empty
171 last :: Stream a -> a
172 {-# INLINE last #-}
173 last = unId . M.last
174
175 -- | Element at the given position
176 (!!) :: Stream a -> Int -> a
177 {-# INLINE (!!) #-}
178 s !! i = unId (s M.!! i)
179
180 -- Substreams
181 -- ----------
182
183 -- | Extract a substream of the given length starting at the given position.
184 extract :: Stream a -> Int -- ^ starting index
185 -> Int -- ^ length
186 -> Stream a
187 {-# INLINE extract #-}
188 extract = M.extract
189
190 -- | All but the last element
191 init :: Stream a -> Stream a
192 {-# INLINE init #-}
193 init = M.init
194
195 -- | All but the first element
196 tail :: Stream a -> Stream a
197 {-# INLINE tail #-}
198 tail = M.tail
199
200 -- | The first @n@ elements
201 take :: Int -> Stream a -> Stream a
202 {-# INLINE take #-}
203 take = M.take
204
205 -- | All but the first @n@ elements
206 drop :: Int -> Stream a -> Stream a
207 {-# INLINE drop #-}
208 drop = M.drop
209
210 -- Mapping
211 -- ---------------
212
213 -- | Map a function over a 'Stream'
214 map :: (a -> b) -> Stream a -> Stream b
215 {-# INLINE map #-}
216 map = M.map
217
218 concatMap :: (a -> Stream b) -> Stream a -> Stream b
219 {-# INLINE concatMap #-}
220 concatMap = M.concatMap
221
222 -- Zipping
223 -- -------
224
225 -- | Zip two 'Stream's with the given function
226 zipWith :: (a -> b -> c) -> Stream a -> Stream b -> Stream c
227 {-# INLINE zipWith #-}
228 zipWith = M.zipWith
229
230 -- | Zip three 'Stream's with the given function
231 zipWith3 :: (a -> b -> c -> d) -> Stream a -> Stream b -> Stream c -> Stream d
232 {-# INLINE zipWith3 #-}
233 zipWith3 = M.zipWith3
234
235 -- Filtering
236 -- ---------
237
238 -- | Drop elements which do not satisfy the predicate
239 filter :: (a -> Bool) -> Stream a -> Stream a
240 {-# INLINE filter #-}
241 filter = M.filter
242
243 -- | Longest prefix of elements that satisfy the predicate
244 takeWhile :: (a -> Bool) -> Stream a -> Stream a
245 {-# INLINE takeWhile #-}
246 takeWhile = M.takeWhile
247
248 -- | Drop the longest prefix of elements that satisfy the predicate
249 dropWhile :: (a -> Bool) -> Stream a -> Stream a
250 {-# INLINE dropWhile #-}
251 dropWhile = M.dropWhile
252
253 -- Searching
254 -- ---------
255
256 infix 4 `elem`
257 -- | Check whether the 'Stream' contains an element
258 elem :: Eq a => a -> Stream a -> Bool
259 {-# INLINE elem #-}
260 elem x = unId . M.elem x
261
262 infix 4 `notElem`
263 -- | Inverse of `elem`
264 notElem :: Eq a => a -> Stream a -> Bool
265 {-# INLINE notElem #-}
266 notElem x = unId . M.notElem x
267
268 -- | Yield 'Just' the first element matching the predicate or 'Nothing' if no
269 -- such element exists.
270 find :: (a -> Bool) -> Stream a -> Maybe a
271 {-# INLINE find #-}
272 find f = unId . M.find f
273
274 -- | Yield 'Just' the index of the first element matching the predicate or
275 -- 'Nothing' if no such element exists.
276 findIndex :: (a -> Bool) -> Stream a -> Maybe Int
277 {-# INLINE findIndex #-}
278 findIndex f = unId . M.findIndex f
279
280 -- Folding
281 -- -------
282
283 -- | Left fold
284 foldl :: (a -> b -> a) -> a -> Stream b -> a
285 {-# INLINE foldl #-}
286 foldl f z = unId . M.foldl f z
287
288 -- | Left fold on non-empty 'Stream's
289 foldl1 :: (a -> a -> a) -> Stream a -> a
290 {-# INLINE foldl1 #-}
291 foldl1 f = unId . M.foldl1 f
292
293 -- | Left fold with strict accumulator
294 foldl' :: (a -> b -> a) -> a -> Stream b -> a
295 {-# INLINE foldl' #-}
296 foldl' f z = unId . M.foldl' f z
297
298 -- | Left fold on non-empty 'Stream's with strict accumulator
299 foldl1' :: (a -> a -> a) -> Stream a -> a
300 {-# INLINE foldl1' #-}
301 foldl1' f = unId . M.foldl1' f
302
303 -- | Right fold
304 foldr :: (a -> b -> b) -> b -> Stream a -> b
305 {-# INLINE foldr #-}
306 foldr f z = unId . M.foldr f z
307
308 -- | Right fold on non-empty 'Stream's
309 foldr1 :: (a -> a -> a) -> Stream a -> a
310 {-# INLINE foldr1 #-}
311 foldr1 f = unId . M.foldr1 f
312
313 -- Specialised folds
314 -- -----------------
315
316 and :: Stream Bool -> Bool
317 {-# INLINE and #-}
318 and = unId . M.and
319
320 or :: Stream Bool -> Bool
321 {-# INLINE or #-}
322 or = unId . M.or
323
324 -- Unfolding
325 -- ---------
326
327 -- | Unfold
328 unfoldr :: (s -> Maybe (a, s)) -> s -> Stream a
329 {-# INLINE unfoldr #-}
330 unfoldr = M.unfoldr
331
332 -- Scans
333 -- -----
334
335 -- | Prefix scan
336 prescanl :: (a -> b -> a) -> a -> Stream b -> Stream a
337 {-# INLINE prescanl #-}
338 prescanl = M.prescanl
339
340 -- | Prefix scan with strict accumulator
341 prescanl' :: (a -> b -> a) -> a -> Stream b -> Stream a
342 {-# INLINE prescanl' #-}
343 prescanl' = M.prescanl'
344
345 -- Comparisons
346 -- -----------
347
348 -- | Check if two 'Stream's are equal
349 eq :: Eq a => Stream a -> Stream a -> Bool
350 {-# INLINE_STREAM eq #-}
351 eq (M.Stream step1 s1 _) (M.Stream step2 s2 _) = eq_loop0 s1 s2
352 where
353 eq_loop0 s1 s2 = case unId (step1 s1) of
354 Yield x s1' -> eq_loop1 x s1' s2
355 Skip s1' -> eq_loop0 s1' s2
356 Done -> null (M.Stream step2 s2 Unknown)
357
358 eq_loop1 x s1 s2 = case unId (step2 s2) of
359 Yield y s2' -> x == y && eq_loop0 s1 s2'
360 Skip s2' -> eq_loop1 x s1 s2'
361 Done -> False
362
363 -- | Lexicographically compare two 'Stream's
364 cmp :: Ord a => Stream a -> Stream a -> Ordering
365 {-# INLINE_STREAM cmp #-}
366 cmp (M.Stream step1 s1 _) (M.Stream step2 s2 _) = cmp_loop0 s1 s2
367 where
368 cmp_loop0 s1 s2 = case unId (step1 s1) of
369 Yield x s1' -> cmp_loop1 x s1' s2
370 Skip s1' -> cmp_loop0 s1' s2
371 Done -> if null (M.Stream step2 s2 Unknown)
372 then EQ else LT
373
374 cmp_loop1 x s1 s2 = case unId (step2 s2) of
375 Yield y s2' -> case x `compare` y of
376 EQ -> cmp_loop0 s1 s2'
377 c -> c
378 Skip s2' -> cmp_loop1 x s1 s2'
379 Done -> GT
380
381 instance Eq a => Eq (M.Stream Id a) where
382 {-# INLINE (==) #-}
383 (==) = eq
384
385 instance Ord a => Ord (M.Stream Id a) where
386 {-# INLINE compare #-}
387 compare = cmp
388
389 -- Monadic combinators
390 -- -------------------
391
392 -- | Apply a monadic action to each element of the stream
393 mapM_ :: Monad m => (a -> m ()) -> Stream a -> m ()
394 {-# INLINE_STREAM mapM_ #-}
395 mapM_ f = M.mapM_ f . liftStream
396
397 -- | Monadic fold
398 foldM :: Monad m => (a -> b -> m a) -> a -> Stream b -> m a
399 {-# INLINE_STREAM foldM #-}
400 foldM m z = M.foldM m z . liftStream
401
402 -- Conversions
403 -- -----------
404
405 -- | Convert a 'Stream' to a list
406 toList :: Stream a -> [a]
407 {-# INLINE toList #-}
408 toList s = unId (M.toList s)
409
410 -- | Create a 'Stream' from a list
411 fromList :: [a] -> Stream a
412 {-# INLINE fromList #-}
413 fromList = M.fromList
414