18bbabf8d7db8b428ce46193e0c15ea54017fb7e
[darcs-mirrors/vector.git] / Data / Vector / Fusion / Stream.hs
1 {-# LANGUAGE FlexibleInstances, Rank2Types #-}
2
3 -- |
4 -- Module : Data.Vector.Fusion.Stream
5 -- Copyright : (c) Roman Leshchinskiy 2008-2010
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, flatten, 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, iterateN,
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, fromListN, unsafeFromList, liftStream,
71
72 -- * Monadic combinators
73 mapM, mapM_, zipWithM, zipWithM_, filterM, 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, mapM_ )
96
97 import GHC.Base ( build )
98
99 #include "vector.h"
100
101 -- | The type of pure streams
102 type Stream = M.Stream Id
103
104 -- | Alternative name for monadic streams
105 type MStream = M.Stream
106
107 inplace :: (forall m. Monad m => M.Stream m a -> M.Stream m b)
108 -> Stream a -> Stream b
109 {-# INLINE_STREAM inplace #-}
110 inplace f s = s `seq` f s
111
112 {-# RULES
113
114 "inplace/inplace [Vector]"
115 forall (f :: forall m. Monad m => MStream m a -> MStream m a)
116 (g :: forall m. Monad m => MStream m a -> MStream m a)
117 s.
118 inplace f (inplace g s) = inplace (f . g) s
119
120 #-}
121
122 -- | Convert a pure stream to a monadic stream
123 liftStream :: Monad m => Stream a -> M.Stream m a
124 {-# INLINE_STREAM liftStream #-}
125 liftStream (M.Stream step s sz) = M.Stream (return . unId . step) s sz
126
127 -- | 'Size' hint of a 'Stream'
128 size :: Stream a -> Size
129 {-# INLINE size #-}
130 size = M.size
131
132 -- | Attach a 'Size' hint to a 'Stream'
133 sized :: Stream a -> Size -> Stream a
134 {-# INLINE sized #-}
135 sized = M.sized
136
137 -- Length
138 -- ------
139
140 -- | Length of a 'Stream'
141 length :: Stream a -> Int
142 {-# INLINE length #-}
143 length = unId . M.length
144
145 -- | Check if a 'Stream' is empty
146 null :: Stream a -> Bool
147 {-# INLINE null #-}
148 null = unId . M.null
149
150 -- Construction
151 -- ------------
152
153 -- | Empty 'Stream'
154 empty :: Stream a
155 {-# INLINE empty #-}
156 empty = M.empty
157
158 -- | Singleton 'Stream'
159 singleton :: a -> Stream a
160 {-# INLINE singleton #-}
161 singleton = M.singleton
162
163 -- | Replicate a value to a given length
164 replicate :: Int -> a -> Stream a
165 {-# INLINE replicate #-}
166 replicate = M.replicate
167
168 -- | Generate a stream from its indices
169 generate :: Int -> (Int -> a) -> Stream a
170 {-# INLINE generate #-}
171 generate = M.generate
172
173 -- | Prepend an element
174 cons :: a -> Stream a -> Stream a
175 {-# INLINE cons #-}
176 cons = M.cons
177
178 -- | Append an element
179 snoc :: Stream a -> a -> Stream a
180 {-# INLINE snoc #-}
181 snoc = M.snoc
182
183 infixr 5 ++
184 -- | Concatenate two 'Stream's
185 (++) :: Stream a -> Stream a -> Stream a
186 {-# INLINE (++) #-}
187 (++) = (M.++)
188
189 -- Accessing elements
190 -- ------------------
191
192 -- | First element of the 'Stream' or error if empty
193 head :: Stream a -> a
194 {-# INLINE head #-}
195 head = unId . M.head
196
197 -- | Last element of the 'Stream' or error if empty
198 last :: Stream a -> a
199 {-# INLINE last #-}
200 last = unId . M.last
201
202 -- | Element at the given position
203 (!!) :: Stream a -> Int -> a
204 {-# INLINE (!!) #-}
205 s !! i = unId (s M.!! i)
206
207 infixl 9 !?
208 -- | Element at the given position or 'Nothing' if out of bounds
209 (!?) :: Stream a -> Int -> Maybe a
210 {-# INLINE (!?) #-}
211 s !? i = unId (s M.!? i)
212
213 -- Substreams
214 -- ----------
215
216 -- | Extract a substream of the given length starting at the given position.
217 slice :: Int -- ^ starting index
218 -> Int -- ^ length
219 -> Stream a
220 -> Stream a
221 {-# INLINE slice #-}
222 slice = M.slice
223
224 -- | All but the last element
225 init :: Stream a -> Stream a
226 {-# INLINE init #-}
227 init = M.init
228
229 -- | All but the first element
230 tail :: Stream a -> Stream a
231 {-# INLINE tail #-}
232 tail = M.tail
233
234 -- | The first @n@ elements
235 take :: Int -> Stream a -> Stream a
236 {-# INLINE take #-}
237 take = M.take
238
239 -- | All but the first @n@ elements
240 drop :: Int -> Stream a -> Stream a
241 {-# INLINE drop #-}
242 drop = M.drop
243
244 -- Mapping
245 -- ---------------
246
247 -- | Map a function over a 'Stream'
248 map :: (a -> b) -> Stream a -> Stream b
249 {-# INLINE map #-}
250 map = M.map
251
252 unbox :: Stream (Box a) -> Stream a
253 {-# INLINE unbox #-}
254 unbox = M.unbox
255
256 concatMap :: (a -> Stream b) -> Stream a -> Stream b
257 {-# INLINE concatMap #-}
258 concatMap = M.concatMap
259
260 -- Zipping
261 -- -------
262
263 -- | Pair each element in a 'Stream' with its index
264 indexed :: Stream a -> Stream (Int,a)
265 {-# INLINE indexed #-}
266 indexed = M.indexed
267
268 -- | Pair each element in a 'Stream' with its index, starting from the right
269 -- and counting down
270 indexedR :: Int -> Stream a -> Stream (Int,a)
271 {-# INLINE_STREAM indexedR #-}
272 indexedR = M.indexedR
273
274 -- | Zip two 'Stream's with the given function
275 zipWith :: (a -> b -> c) -> Stream a -> Stream b -> Stream c
276 {-# INLINE zipWith #-}
277 zipWith = M.zipWith
278
279 -- | Zip three 'Stream's with the given function
280 zipWith3 :: (a -> b -> c -> d) -> Stream a -> Stream b -> Stream c -> Stream d
281 {-# INLINE zipWith3 #-}
282 zipWith3 = M.zipWith3
283
284 zipWith4 :: (a -> b -> c -> d -> e)
285 -> Stream a -> Stream b -> Stream c -> Stream d
286 -> Stream e
287 {-# INLINE zipWith4 #-}
288 zipWith4 = M.zipWith4
289
290 zipWith5 :: (a -> b -> c -> d -> e -> f)
291 -> Stream a -> Stream b -> Stream c -> Stream d
292 -> Stream e -> Stream f
293 {-# INLINE zipWith5 #-}
294 zipWith5 = M.zipWith5
295
296 zipWith6 :: (a -> b -> c -> d -> e -> f -> g)
297 -> Stream a -> Stream b -> Stream c -> Stream d
298 -> Stream e -> Stream f -> Stream g
299 {-# INLINE zipWith6 #-}
300 zipWith6 = M.zipWith6
301
302 zip :: Stream a -> Stream b -> Stream (a,b)
303 {-# INLINE zip #-}
304 zip = M.zip
305
306 zip3 :: Stream a -> Stream b -> Stream c -> Stream (a,b,c)
307 {-# INLINE zip3 #-}
308 zip3 = M.zip3
309
310 zip4 :: Stream a -> Stream b -> Stream c -> Stream d
311 -> Stream (a,b,c,d)
312 {-# INLINE zip4 #-}
313 zip4 = M.zip4
314
315 zip5 :: Stream a -> Stream b -> Stream c -> Stream d
316 -> Stream e -> Stream (a,b,c,d,e)
317 {-# INLINE zip5 #-}
318 zip5 = M.zip5
319
320 zip6 :: Stream a -> Stream b -> Stream c -> Stream d
321 -> Stream e -> Stream f -> Stream (a,b,c,d,e,f)
322 {-# INLINE zip6 #-}
323 zip6 = M.zip6
324
325 -- Filtering
326 -- ---------
327
328 -- | Drop elements which do not satisfy the predicate
329 filter :: (a -> Bool) -> Stream a -> Stream a
330 {-# INLINE filter #-}
331 filter = M.filter
332
333 -- | Longest prefix of elements that satisfy the predicate
334 takeWhile :: (a -> Bool) -> Stream a -> Stream a
335 {-# INLINE takeWhile #-}
336 takeWhile = M.takeWhile
337
338 -- | Drop the longest prefix of elements that satisfy the predicate
339 dropWhile :: (a -> Bool) -> Stream a -> Stream a
340 {-# INLINE dropWhile #-}
341 dropWhile = M.dropWhile
342
343 -- Searching
344 -- ---------
345
346 infix 4 `elem`
347 -- | Check whether the 'Stream' contains an element
348 elem :: Eq a => a -> Stream a -> Bool
349 {-# INLINE elem #-}
350 elem x = unId . M.elem x
351
352 infix 4 `notElem`
353 -- | Inverse of `elem`
354 notElem :: Eq a => a -> Stream a -> Bool
355 {-# INLINE notElem #-}
356 notElem x = unId . M.notElem x
357
358 -- | Yield 'Just' the first element matching the predicate or 'Nothing' if no
359 -- such element exists.
360 find :: (a -> Bool) -> Stream a -> Maybe a
361 {-# INLINE find #-}
362 find f = unId . M.find f
363
364 -- | Yield 'Just' the index of the first element matching the predicate or
365 -- 'Nothing' if no such element exists.
366 findIndex :: (a -> Bool) -> Stream a -> Maybe Int
367 {-# INLINE findIndex #-}
368 findIndex f = unId . M.findIndex f
369
370 -- Folding
371 -- -------
372
373 -- | Left fold
374 foldl :: (a -> b -> a) -> a -> Stream b -> a
375 {-# INLINE foldl #-}
376 foldl f z = unId . M.foldl f z
377
378 -- | Left fold on non-empty 'Stream's
379 foldl1 :: (a -> a -> a) -> Stream a -> a
380 {-# INLINE foldl1 #-}
381 foldl1 f = unId . M.foldl1 f
382
383 -- | Left fold with strict accumulator
384 foldl' :: (a -> b -> a) -> a -> Stream b -> a
385 {-# INLINE foldl' #-}
386 foldl' f z = unId . M.foldl' f z
387
388 -- | Left fold on non-empty 'Stream's with strict accumulator
389 foldl1' :: (a -> a -> a) -> Stream a -> a
390 {-# INLINE foldl1' #-}
391 foldl1' f = unId . M.foldl1' f
392
393 -- | Right fold
394 foldr :: (a -> b -> b) -> b -> Stream a -> b
395 {-# INLINE foldr #-}
396 foldr f z = unId . M.foldr f z
397
398 -- | Right fold on non-empty 'Stream's
399 foldr1 :: (a -> a -> a) -> Stream a -> a
400 {-# INLINE foldr1 #-}
401 foldr1 f = unId . M.foldr1 f
402
403 -- Specialised folds
404 -- -----------------
405
406 and :: Stream Bool -> Bool
407 {-# INLINE and #-}
408 and = unId . M.and
409
410 or :: Stream Bool -> Bool
411 {-# INLINE or #-}
412 or = unId . M.or
413
414 -- Unfolding
415 -- ---------
416
417 -- | Unfold
418 unfoldr :: (s -> Maybe (a, s)) -> s -> Stream a
419 {-# INLINE unfoldr #-}
420 unfoldr = M.unfoldr
421
422 -- | Unfold at most @n@ elements
423 unfoldrN :: Int -> (s -> Maybe (a, s)) -> s -> Stream a
424 {-# INLINE unfoldrN #-}
425 unfoldrN = M.unfoldrN
426
427 -- | Apply function n-1 times to value. Zeroth element is original value.
428 iterateN :: Int -> (a -> a) -> a -> Stream a
429 {-# INLINE iterateN #-}
430 iterateN = M.iterateN
431
432 -- Scans
433 -- -----
434
435 -- | Prefix scan
436 prescanl :: (a -> b -> a) -> a -> Stream b -> Stream a
437 {-# INLINE prescanl #-}
438 prescanl = M.prescanl
439
440 -- | Prefix scan with strict accumulator
441 prescanl' :: (a -> b -> a) -> a -> Stream b -> Stream a
442 {-# INLINE prescanl' #-}
443 prescanl' = M.prescanl'
444
445 -- | Suffix scan
446 postscanl :: (a -> b -> a) -> a -> Stream b -> Stream a
447 {-# INLINE postscanl #-}
448 postscanl = M.postscanl
449
450 -- | Suffix scan with strict accumulator
451 postscanl' :: (a -> b -> a) -> a -> Stream b -> Stream a
452 {-# INLINE postscanl' #-}
453 postscanl' = M.postscanl'
454
455 -- | Haskell-style scan
456 scanl :: (a -> b -> a) -> a -> Stream b -> Stream a
457 {-# INLINE scanl #-}
458 scanl = M.scanl
459
460 -- | Haskell-style scan with strict accumulator
461 scanl' :: (a -> b -> a) -> a -> Stream b -> Stream a
462 {-# INLINE scanl' #-}
463 scanl' = M.scanl'
464
465 -- | Scan over a non-empty 'Stream'
466 scanl1 :: (a -> a -> a) -> Stream a -> Stream a
467 {-# INLINE scanl1 #-}
468 scanl1 = M.scanl1
469
470 -- | Scan over a non-empty 'Stream' with a strict accumulator
471 scanl1' :: (a -> a -> a) -> Stream a -> Stream a
472 {-# INLINE scanl1' #-}
473 scanl1' = M.scanl1'
474
475
476 -- Comparisons
477 -- -----------
478
479 -- | Check if two 'Stream's are equal
480 eq :: Eq a => Stream a -> Stream a -> Bool
481 {-# INLINE_STREAM eq #-}
482 eq (M.Stream step1 s1 _) (M.Stream step2 s2 _) = eq_loop0 s1 s2
483 where
484 eq_loop0 s1 s2 = case unId (step1 s1) of
485 Yield x s1' -> eq_loop1 x s1' s2
486 Skip s1' -> eq_loop0 s1' s2
487 Done -> null (M.Stream step2 s2 Unknown)
488
489 eq_loop1 x s1 s2 = case unId (step2 s2) of
490 Yield y s2' -> x == y && eq_loop0 s1 s2'
491 Skip s2' -> eq_loop1 x s1 s2'
492 Done -> False
493
494 -- | Lexicographically compare two 'Stream's
495 cmp :: Ord a => Stream a -> Stream a -> Ordering
496 {-# INLINE_STREAM cmp #-}
497 cmp (M.Stream step1 s1 _) (M.Stream step2 s2 _) = cmp_loop0 s1 s2
498 where
499 cmp_loop0 s1 s2 = case unId (step1 s1) of
500 Yield x s1' -> cmp_loop1 x s1' s2
501 Skip s1' -> cmp_loop0 s1' s2
502 Done -> if null (M.Stream step2 s2 Unknown)
503 then EQ else LT
504
505 cmp_loop1 x s1 s2 = case unId (step2 s2) of
506 Yield y s2' -> case x `compare` y of
507 EQ -> cmp_loop0 s1 s2'
508 c -> c
509 Skip s2' -> cmp_loop1 x s1 s2'
510 Done -> GT
511
512 instance Eq a => Eq (M.Stream Id a) where
513 {-# INLINE (==) #-}
514 (==) = eq
515
516 instance Ord a => Ord (M.Stream Id a) where
517 {-# INLINE compare #-}
518 compare = cmp
519
520 -- Monadic combinators
521 -- -------------------
522
523 -- | Apply a monadic action to each element of the stream, producing a monadic
524 -- stream of results
525 mapM :: Monad m => (a -> m b) -> Stream a -> M.Stream m b
526 {-# INLINE mapM #-}
527 mapM f = M.mapM f . liftStream
528
529 -- | Apply a monadic action to each element of the stream
530 mapM_ :: Monad m => (a -> m b) -> Stream a -> m ()
531 {-# INLINE mapM_ #-}
532 mapM_ f = M.mapM_ f . liftStream
533
534 zipWithM :: Monad m => (a -> b -> m c) -> Stream a -> Stream b -> M.Stream m c
535 {-# INLINE zipWithM #-}
536 zipWithM f as bs = M.zipWithM f (liftStream as) (liftStream bs)
537
538 zipWithM_ :: Monad m => (a -> b -> m c) -> Stream a -> Stream b -> m ()
539 {-# INLINE zipWithM_ #-}
540 zipWithM_ f as bs = M.zipWithM_ f (liftStream as) (liftStream bs)
541
542 -- | Yield a monadic stream of elements that satisfy the monadic predicate
543 filterM :: Monad m => (a -> m Bool) -> Stream a -> M.Stream m a
544 {-# INLINE filterM #-}
545 filterM f = M.filterM f . liftStream
546
547 -- | Monadic fold
548 foldM :: Monad m => (a -> b -> m a) -> a -> Stream b -> m a
549 {-# INLINE foldM #-}
550 foldM m z = M.foldM m z . liftStream
551
552 -- | Monadic fold over non-empty stream
553 fold1M :: Monad m => (a -> a -> m a) -> Stream a -> m a
554 {-# INLINE fold1M #-}
555 fold1M m = M.fold1M m . liftStream
556
557 -- | Monadic fold with strict accumulator
558 foldM' :: Monad m => (a -> b -> m a) -> a -> Stream b -> m a
559 {-# INLINE foldM' #-}
560 foldM' m z = M.foldM' m z . liftStream
561
562 -- | Monad fold over non-empty stream with strict accumulator
563 fold1M' :: Monad m => (a -> a -> m a) -> Stream a -> m a
564 {-# INLINE fold1M' #-}
565 fold1M' m = M.fold1M' m . liftStream
566
567 -- Enumerations
568 -- ------------
569
570 -- | Yield a 'Stream' of the given length containing the values @x@, @x+y@,
571 -- @x+y+y@ etc.
572 enumFromStepN :: Num a => a -> a -> Int -> Stream a
573 {-# INLINE enumFromStepN #-}
574 enumFromStepN = M.enumFromStepN
575
576 -- | Enumerate values
577 --
578 -- /WARNING:/ This operations can be very inefficient. If at all possible, use
579 -- 'enumFromStepN' instead.
580 enumFromTo :: Enum a => a -> a -> Stream a
581 {-# INLINE enumFromTo #-}
582 enumFromTo = M.enumFromTo
583
584 -- | Enumerate values with a given step.
585 --
586 -- /WARNING:/ This operations is very inefficient. If at all possible, use
587 -- 'enumFromStepN' instead.
588 enumFromThenTo :: Enum a => a -> a -> a -> Stream a
589 {-# INLINE enumFromThenTo #-}
590 enumFromThenTo = M.enumFromThenTo
591
592 -- Conversions
593 -- -----------
594
595 -- | Convert a 'Stream' to a list
596 toList :: Stream a -> [a]
597 {-# INLINE toList #-}
598 -- toList s = unId (M.toList s)
599 toList s = build (\c n -> toListFB c n s)
600
601 -- This supports foldr/build list fusion that GHC implements
602 toListFB :: (a -> b -> b) -> b -> Stream a -> b
603 {-# INLINE [0] toListFB #-}
604 toListFB c n (M.Stream step s _) = go s
605 where
606 go s = case unId (step s) of
607 Yield x s' -> x `c` go s'
608 Skip s' -> go s'
609 Done -> n
610
611 -- | Create a 'Stream' from a list
612 fromList :: [a] -> Stream a
613 {-# INLINE fromList #-}
614 fromList = M.fromList
615
616 -- | Create a 'Stream' from the first @n@ elements of a list
617 --
618 -- > fromListN n xs = fromList (take n xs)
619 fromListN :: Int -> [a] -> Stream a
620 {-# INLINE fromListN #-}
621 fromListN = M.fromListN
622
623 unsafeFromList :: Size -> [a] -> Stream a
624 {-# INLINE unsafeFromList #-}
625 unsafeFromList = M.unsafeFromList
626
627 -- | Create a 'Stream' of values from a 'Stream' of streamable things
628 flatten :: (a -> s) -> (s -> Step s b) -> Size -> Stream a -> Stream b
629 {-# INLINE_STREAM flatten #-}
630 flatten mk istep sz = M.flatten (return . mk) (return . istep) sz . liftStream
631