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