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