phases.h -> vector.h
[darcs-mirrors/vector.git] / Data / Vector / Fusion / Stream.hs
1 {-# LANGUAGE ExistentialQuantification, FlexibleInstances, Rank2Types #-}
2
3 -- |
4 -- Module : Data.Vector.Fusion.Stream
5 -- Copyright : (c) Roman Leshchinskiy 2008-2009
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 "vector.h"
16
17 module Data.Vector.Fusion.Stream (
18 -- * Types
19 Step(..), Stream, MStream,
20
21 -- * In-place markers
22 inplace, inplace',
23
24 -- * Size hints
25 size, sized,
26
27 -- * Length information
28 length, null,
29
30 -- * Construction
31 empty, singleton, cons, snoc, replicate, (++),
32
33 -- * Accessing individual elements
34 head, last, (!!),
35
36 -- * Substreams
37 extract, init, tail, take, drop,
38
39 -- * Mapping
40 map, concatMap, unbox,
41
42 -- * Zipping
43 zipWith, zipWith3,
44
45 -- * Filtering
46 filter, takeWhile, dropWhile,
47
48 -- * Searching
49 elem, notElem, find, findIndex,
50
51 -- * Folding
52 foldl, foldl1, foldl', foldl1', foldr, foldr1,
53
54 -- * Specialised folds
55 and, or,
56
57 -- * Unfolding
58 unfoldr,
59
60 -- * Scans
61 prescanl, prescanl',
62 postscanl, postscanl',
63 scanl, scanl',
64 scanl1, scanl1',
65
66 -- * Enumerations
67 enumFromTo,
68
69 -- * Conversions
70 toList, fromList, liftStream,
71
72 -- * Monadic combinators
73 mapM_, foldM, fold1M, foldM', fold1M'
74 ) where
75
76 import Data.Vector.Fusion.Stream.Size
77 import Data.Vector.Fusion.Util
78 import Data.Vector.Fusion.Stream.Monadic ( Step(..) )
79 import qualified Data.Vector.Fusion.Stream.Monadic as M
80
81 import Prelude hiding ( length, null,
82 replicate, (++),
83 head, last, (!!),
84 init, tail, take, drop,
85 map, concatMap,
86 zipWith, zipWith3,
87 filter, takeWhile, dropWhile,
88 elem, notElem,
89 foldl, foldl1, foldr, foldr1,
90 and, or,
91 scanl, scanl1,
92 enumFromTo, enumFromThenTo,
93 mapM_ )
94
95 -- | The type of pure streams
96 type Stream = M.Stream Id
97
98 -- | Alternative name for monadic streams
99 type MStream = M.Stream
100
101 inplace :: (forall m. Monad m => M.Stream m a -> M.Stream m a)
102 -> Stream a -> Stream a
103 {-# INLINE_STREAM inplace #-}
104 inplace f s = f s
105
106 {-# RULES
107
108 "inplace/inplace [Vector]"
109 forall (f :: forall m. Monad m => MStream m a -> MStream m a)
110 (g :: forall m. Monad m => MStream m a -> MStream m a)
111 s.
112 inplace f (inplace g s) = inplace (f . g) s
113
114 #-}
115
116 inplace' :: (forall m. Monad m => M.Stream m a -> M.Stream m b)
117 -> Stream a -> Stream b
118 {-# INLINE_STREAM inplace' #-}
119 inplace' f s = f s
120
121 -- FIXME: We'd like to have this
122 {- RULES
123
124 "inplace' [Vector]" inplace' = inplace
125 -}
126 -- but it's only available in 6.13
127 -- (see http://hackage.haskell.org/trac/ghc/ticket/3670)
128
129 {-# RULES
130
131 "inplace' [Vector]"
132 forall (f :: forall m. Monad m => MStream m a -> MStream m a).
133 inplace' f = inplace f
134
135 #-}
136
137 -- | Convert a pure stream to a monadic stream
138 liftStream :: Monad m => Stream a -> M.Stream m a
139 {-# INLINE_STREAM liftStream #-}
140 liftStream (M.Stream step s sz) = M.Stream (return . unId . step) s sz
141
142 -- | 'Size' hint of a 'Stream'
143 size :: Stream a -> Size
144 {-# INLINE size #-}
145 size = M.size
146
147 -- | Attach a 'Size' hint to a 'Stream'
148 sized :: Stream a -> Size -> Stream a
149 {-# INLINE sized #-}
150 sized = M.sized
151
152 -- Length
153 -- ------
154
155 -- | Length of a 'Stream'
156 length :: Stream a -> Int
157 {-# INLINE length #-}
158 length = unId . M.length
159
160 -- | Check if a 'Stream' is empty
161 null :: Stream a -> Bool
162 {-# INLINE null #-}
163 null = unId . M.null
164
165 -- Construction
166 -- ------------
167
168 -- | Empty 'Stream'
169 empty :: Stream a
170 {-# INLINE empty #-}
171 empty = M.empty
172
173 -- | Singleton 'Stream'
174 singleton :: a -> Stream a
175 {-# INLINE singleton #-}
176 singleton = M.singleton
177
178 -- | Replicate a value to a given length
179 replicate :: Int -> a -> Stream a
180 {-# INLINE replicate #-}
181 replicate = M.replicate
182
183 -- | Prepend an element
184 cons :: a -> Stream a -> Stream a
185 {-# INLINE cons #-}
186 cons = M.cons
187
188 -- | Append an element
189 snoc :: Stream a -> a -> Stream a
190 {-# INLINE snoc #-}
191 snoc = M.snoc
192
193 infixr 5 ++
194 -- | Concatenate two 'Stream's
195 (++) :: Stream a -> Stream a -> Stream a
196 {-# INLINE (++) #-}
197 (++) = (M.++)
198
199 -- Accessing elements
200 -- ------------------
201
202 -- | First element of the 'Stream' or error if empty
203 head :: Stream a -> a
204 {-# INLINE head #-}
205 head = unId . M.head
206
207 -- | Last element of the 'Stream' or error if empty
208 last :: Stream a -> a
209 {-# INLINE last #-}
210 last = unId . M.last
211
212 -- | Element at the given position
213 (!!) :: Stream a -> Int -> a
214 {-# INLINE (!!) #-}
215 s !! i = unId (s M.!! i)
216
217 -- Substreams
218 -- ----------
219
220 -- | Extract a substream of the given length starting at the given position.
221 extract :: Stream a -> Int -- ^ starting index
222 -> Int -- ^ length
223 -> Stream a
224 {-# INLINE extract #-}
225 extract = M.extract
226
227 -- | All but the last element
228 init :: Stream a -> Stream a
229 {-# INLINE init #-}
230 init = M.init
231
232 -- | All but the first element
233 tail :: Stream a -> Stream a
234 {-# INLINE tail #-}
235 tail = M.tail
236
237 -- | The first @n@ elements
238 take :: Int -> Stream a -> Stream a
239 {-# INLINE take #-}
240 take = M.take
241
242 -- | All but the first @n@ elements
243 drop :: Int -> Stream a -> Stream a
244 {-# INLINE drop #-}
245 drop = M.drop
246
247 -- Mapping
248 -- ---------------
249
250 -- | Map a function over a 'Stream'
251 map :: (a -> b) -> Stream a -> Stream b
252 {-# INLINE map #-}
253 map = M.map
254
255 unbox :: Stream (Box a) -> Stream a
256 {-# INLINE unbox #-}
257 unbox = M.unbox
258
259 concatMap :: (a -> Stream b) -> Stream a -> Stream b
260 {-# INLINE concatMap #-}
261 concatMap = M.concatMap
262
263 -- Zipping
264 -- -------
265
266 -- | Zip two 'Stream's with the given function
267 zipWith :: (a -> b -> c) -> Stream a -> Stream b -> Stream c
268 {-# INLINE zipWith #-}
269 zipWith = M.zipWith
270
271 -- | Zip three 'Stream's with the given function
272 zipWith3 :: (a -> b -> c -> d) -> Stream a -> Stream b -> Stream c -> Stream d
273 {-# INLINE zipWith3 #-}
274 zipWith3 = M.zipWith3
275
276 -- Filtering
277 -- ---------
278
279 -- | Drop elements which do not satisfy the predicate
280 filter :: (a -> Bool) -> Stream a -> Stream a
281 {-# INLINE filter #-}
282 filter = M.filter
283
284 -- | Longest prefix of elements that satisfy the predicate
285 takeWhile :: (a -> Bool) -> Stream a -> Stream a
286 {-# INLINE takeWhile #-}
287 takeWhile = M.takeWhile
288
289 -- | Drop the longest prefix of elements that satisfy the predicate
290 dropWhile :: (a -> Bool) -> Stream a -> Stream a
291 {-# INLINE dropWhile #-}
292 dropWhile = M.dropWhile
293
294 -- Searching
295 -- ---------
296
297 infix 4 `elem`
298 -- | Check whether the 'Stream' contains an element
299 elem :: Eq a => a -> Stream a -> Bool
300 {-# INLINE elem #-}
301 elem x = unId . M.elem x
302
303 infix 4 `notElem`
304 -- | Inverse of `elem`
305 notElem :: Eq a => a -> Stream a -> Bool
306 {-# INLINE notElem #-}
307 notElem x = unId . M.notElem x
308
309 -- | Yield 'Just' the first element matching the predicate or 'Nothing' if no
310 -- such element exists.
311 find :: (a -> Bool) -> Stream a -> Maybe a
312 {-# INLINE find #-}
313 find f = unId . M.find f
314
315 -- | Yield 'Just' the index of the first element matching the predicate or
316 -- 'Nothing' if no such element exists.
317 findIndex :: (a -> Bool) -> Stream a -> Maybe Int
318 {-# INLINE findIndex #-}
319 findIndex f = unId . M.findIndex f
320
321 -- Folding
322 -- -------
323
324 -- | Left fold
325 foldl :: (a -> b -> a) -> a -> Stream b -> a
326 {-# INLINE foldl #-}
327 foldl f z = unId . M.foldl f z
328
329 -- | Left fold on non-empty 'Stream's
330 foldl1 :: (a -> a -> a) -> Stream a -> a
331 {-# INLINE foldl1 #-}
332 foldl1 f = unId . M.foldl1 f
333
334 -- | Left fold with strict accumulator
335 foldl' :: (a -> b -> a) -> a -> Stream b -> a
336 {-# INLINE foldl' #-}
337 foldl' f z = unId . M.foldl' f z
338
339 -- | Left fold on non-empty 'Stream's with strict accumulator
340 foldl1' :: (a -> a -> a) -> Stream a -> a
341 {-# INLINE foldl1' #-}
342 foldl1' f = unId . M.foldl1' f
343
344 -- | Right fold
345 foldr :: (a -> b -> b) -> b -> Stream a -> b
346 {-# INLINE foldr #-}
347 foldr f z = unId . M.foldr f z
348
349 -- | Right fold on non-empty 'Stream's
350 foldr1 :: (a -> a -> a) -> Stream a -> a
351 {-# INLINE foldr1 #-}
352 foldr1 f = unId . M.foldr1 f
353
354 -- Specialised folds
355 -- -----------------
356
357 and :: Stream Bool -> Bool
358 {-# INLINE and #-}
359 and = unId . M.and
360
361 or :: Stream Bool -> Bool
362 {-# INLINE or #-}
363 or = unId . M.or
364
365 -- Unfolding
366 -- ---------
367
368 -- | Unfold
369 unfoldr :: (s -> Maybe (a, s)) -> s -> Stream a
370 {-# INLINE unfoldr #-}
371 unfoldr = M.unfoldr
372
373 -- Scans
374 -- -----
375
376 -- | Prefix scan
377 prescanl :: (a -> b -> a) -> a -> Stream b -> Stream a
378 {-# INLINE prescanl #-}
379 prescanl = M.prescanl
380
381 -- | Prefix scan with strict accumulator
382 prescanl' :: (a -> b -> a) -> a -> Stream b -> Stream a
383 {-# INLINE prescanl' #-}
384 prescanl' = M.prescanl'
385
386 -- | Suffix scan
387 postscanl :: (a -> b -> a) -> a -> Stream b -> Stream a
388 {-# INLINE postscanl #-}
389 postscanl = M.postscanl
390
391 -- | Suffix scan with strict accumulator
392 postscanl' :: (a -> b -> a) -> a -> Stream b -> Stream a
393 {-# INLINE postscanl' #-}
394 postscanl' = M.postscanl'
395
396 -- | Haskell-style scan
397 scanl :: (a -> b -> a) -> a -> Stream b -> Stream a
398 {-# INLINE scanl #-}
399 scanl = M.scanl
400
401 -- | Haskell-style scan with strict accumulator
402 scanl' :: (a -> b -> a) -> a -> Stream b -> Stream a
403 {-# INLINE scanl' #-}
404 scanl' = M.scanl'
405
406 -- | Scan over a non-empty 'Stream'
407 scanl1 :: (a -> a -> a) -> Stream a -> Stream a
408 {-# INLINE scanl1 #-}
409 scanl1 = M.scanl1
410
411 -- | Scan over a non-empty 'Stream' with a strict accumulator
412 scanl1' :: (a -> a -> a) -> Stream a -> Stream a
413 {-# INLINE scanl1' #-}
414 scanl1' = M.scanl1'
415
416
417 -- Comparisons
418 -- -----------
419
420 -- | Check if two 'Stream's are equal
421 eq :: Eq a => Stream a -> Stream a -> Bool
422 {-# INLINE_STREAM eq #-}
423 eq (M.Stream step1 s1 _) (M.Stream step2 s2 _) = eq_loop0 s1 s2
424 where
425 eq_loop0 s1 s2 = case unId (step1 s1) of
426 Yield x s1' -> eq_loop1 x s1' s2
427 Skip s1' -> eq_loop0 s1' s2
428 Done -> null (M.Stream step2 s2 Unknown)
429
430 eq_loop1 x s1 s2 = case unId (step2 s2) of
431 Yield y s2' -> x == y && eq_loop0 s1 s2'
432 Skip s2' -> eq_loop1 x s1 s2'
433 Done -> False
434
435 -- | Lexicographically compare two 'Stream's
436 cmp :: Ord a => Stream a -> Stream a -> Ordering
437 {-# INLINE_STREAM cmp #-}
438 cmp (M.Stream step1 s1 _) (M.Stream step2 s2 _) = cmp_loop0 s1 s2
439 where
440 cmp_loop0 s1 s2 = case unId (step1 s1) of
441 Yield x s1' -> cmp_loop1 x s1' s2
442 Skip s1' -> cmp_loop0 s1' s2
443 Done -> if null (M.Stream step2 s2 Unknown)
444 then EQ else LT
445
446 cmp_loop1 x s1 s2 = case unId (step2 s2) of
447 Yield y s2' -> case x `compare` y of
448 EQ -> cmp_loop0 s1 s2'
449 c -> c
450 Skip s2' -> cmp_loop1 x s1 s2'
451 Done -> GT
452
453 instance Eq a => Eq (M.Stream Id a) where
454 {-# INLINE (==) #-}
455 (==) = eq
456
457 instance Ord a => Ord (M.Stream Id a) where
458 {-# INLINE compare #-}
459 compare = cmp
460
461 -- Monadic combinators
462 -- -------------------
463
464 -- | Apply a monadic action to each element of the stream
465 mapM_ :: Monad m => (a -> m ()) -> Stream a -> m ()
466 {-# INLINE mapM_ #-}
467 mapM_ f = M.mapM_ f . liftStream
468
469 -- | Monadic fold
470 foldM :: Monad m => (a -> b -> m a) -> a -> Stream b -> m a
471 {-# INLINE foldM #-}
472 foldM m z = M.foldM m z . liftStream
473
474 -- | Monadic fold over non-empty stream
475 fold1M :: Monad m => (a -> a -> m a) -> Stream a -> m a
476 {-# INLINE fold1M #-}
477 fold1M m = M.fold1M m . liftStream
478
479 -- | Monadic fold with strict accumulator
480 foldM' :: Monad m => (a -> b -> m a) -> a -> Stream b -> m a
481 {-# INLINE foldM' #-}
482 foldM' m z = M.foldM' m z . liftStream
483
484 -- | Monad fold over non-empty stream with strict accumulator
485 fold1M' :: Monad m => (a -> a -> m a) -> Stream a -> m a
486 {-# INLINE fold1M' #-}
487 fold1M' m = M.fold1M' m . liftStream
488
489 -- Enumerations
490 -- ------------
491
492 enumFromTo :: Enum a => a -> a -> Stream a
493 {-# INLINE enumFromTo #-}
494 enumFromTo = M.enumFromTo
495
496 -- Conversions
497 -- -----------
498
499 -- | Convert a 'Stream' to a list
500 toList :: Stream a -> [a]
501 {-# INLINE toList #-}
502 toList s = unId (M.toList s)
503
504 -- | Create a 'Stream' from a list
505 fromList :: [a] -> Stream a
506 {-# INLINE fromList #-}
507 fromList = M.fromList
508