Fix type of Stream.mapM_
[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, 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, fromListN, 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 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 -- Substreams
208 -- ----------
209
210 -- | Extract a substream of the given length starting at the given position.
211 slice :: Int -- ^ starting index
212 -> Int -- ^ length
213 -> Stream a
214 -> Stream a
215 {-# INLINE slice #-}
216 slice = M.slice
217
218 -- | All but the last element
219 init :: Stream a -> Stream a
220 {-# INLINE init #-}
221 init = M.init
222
223 -- | All but the first element
224 tail :: Stream a -> Stream a
225 {-# INLINE tail #-}
226 tail = M.tail
227
228 -- | The first @n@ elements
229 take :: Int -> Stream a -> Stream a
230 {-# INLINE take #-}
231 take = M.take
232
233 -- | All but the first @n@ elements
234 drop :: Int -> Stream a -> Stream a
235 {-# INLINE drop #-}
236 drop = M.drop
237
238 -- Mapping
239 -- ---------------
240
241 -- | Map a function over a 'Stream'
242 map :: (a -> b) -> Stream a -> Stream b
243 {-# INLINE map #-}
244 map = M.map
245
246 unbox :: Stream (Box a) -> Stream a
247 {-# INLINE unbox #-}
248 unbox = M.unbox
249
250 concatMap :: (a -> Stream b) -> Stream a -> Stream b
251 {-# INLINE concatMap #-}
252 concatMap = M.concatMap
253
254 -- Zipping
255 -- -------
256
257 -- | Pair each element in a 'Stream' with its index
258 indexed :: Stream a -> Stream (Int,a)
259 {-# INLINE indexed #-}
260 indexed = M.indexed
261
262 -- | Pair each element in a 'Stream' with its index, starting from the right
263 -- and counting down
264 indexedR :: Int -> Stream a -> Stream (Int,a)
265 {-# INLINE_STREAM indexedR #-}
266 indexedR = M.indexedR
267
268 -- | Zip two 'Stream's with the given function
269 zipWith :: (a -> b -> c) -> Stream a -> Stream b -> Stream c
270 {-# INLINE zipWith #-}
271 zipWith = M.zipWith
272
273 -- | Zip three 'Stream's with the given function
274 zipWith3 :: (a -> b -> c -> d) -> Stream a -> Stream b -> Stream c -> Stream d
275 {-# INLINE zipWith3 #-}
276 zipWith3 = M.zipWith3
277
278 zipWith4 :: (a -> b -> c -> d -> e)
279 -> Stream a -> Stream b -> Stream c -> Stream d
280 -> Stream e
281 {-# INLINE zipWith4 #-}
282 zipWith4 = M.zipWith4
283
284 zipWith5 :: (a -> b -> c -> d -> e -> f)
285 -> Stream a -> Stream b -> Stream c -> Stream d
286 -> Stream e -> Stream f
287 {-# INLINE zipWith5 #-}
288 zipWith5 = M.zipWith5
289
290 zipWith6 :: (a -> b -> c -> d -> e -> f -> g)
291 -> Stream a -> Stream b -> Stream c -> Stream d
292 -> Stream e -> Stream f -> Stream g
293 {-# INLINE zipWith6 #-}
294 zipWith6 = M.zipWith6
295
296 zip :: Stream a -> Stream b -> Stream (a,b)
297 {-# INLINE zip #-}
298 zip = M.zip
299
300 zip3 :: Stream a -> Stream b -> Stream c -> Stream (a,b,c)
301 {-# INLINE zip3 #-}
302 zip3 = M.zip3
303
304 zip4 :: Stream a -> Stream b -> Stream c -> Stream d
305 -> Stream (a,b,c,d)
306 {-# INLINE zip4 #-}
307 zip4 = M.zip4
308
309 zip5 :: Stream a -> Stream b -> Stream c -> Stream d
310 -> Stream e -> Stream (a,b,c,d,e)
311 {-# INLINE zip5 #-}
312 zip5 = M.zip5
313
314 zip6 :: Stream a -> Stream b -> Stream c -> Stream d
315 -> Stream e -> Stream f -> Stream (a,b,c,d,e,f)
316 {-# INLINE zip6 #-}
317 zip6 = M.zip6
318
319 -- Filtering
320 -- ---------
321
322 -- | Drop elements which do not satisfy the predicate
323 filter :: (a -> Bool) -> Stream a -> Stream a
324 {-# INLINE filter #-}
325 filter = M.filter
326
327 -- | Longest prefix of elements that satisfy the predicate
328 takeWhile :: (a -> Bool) -> Stream a -> Stream a
329 {-# INLINE takeWhile #-}
330 takeWhile = M.takeWhile
331
332 -- | Drop the longest prefix of elements that satisfy the predicate
333 dropWhile :: (a -> Bool) -> Stream a -> Stream a
334 {-# INLINE dropWhile #-}
335 dropWhile = M.dropWhile
336
337 -- Searching
338 -- ---------
339
340 infix 4 `elem`
341 -- | Check whether the 'Stream' contains an element
342 elem :: Eq a => a -> Stream a -> Bool
343 {-# INLINE elem #-}
344 elem x = unId . M.elem x
345
346 infix 4 `notElem`
347 -- | Inverse of `elem`
348 notElem :: Eq a => a -> Stream a -> Bool
349 {-# INLINE notElem #-}
350 notElem x = unId . M.notElem x
351
352 -- | Yield 'Just' the first element matching the predicate or 'Nothing' if no
353 -- such element exists.
354 find :: (a -> Bool) -> Stream a -> Maybe a
355 {-# INLINE find #-}
356 find f = unId . M.find f
357
358 -- | Yield 'Just' the index of the first element matching the predicate or
359 -- 'Nothing' if no such element exists.
360 findIndex :: (a -> Bool) -> Stream a -> Maybe Int
361 {-# INLINE findIndex #-}
362 findIndex f = unId . M.findIndex f
363
364 -- Folding
365 -- -------
366
367 -- | Left fold
368 foldl :: (a -> b -> a) -> a -> Stream b -> a
369 {-# INLINE foldl #-}
370 foldl f z = unId . M.foldl f z
371
372 -- | Left fold on non-empty 'Stream's
373 foldl1 :: (a -> a -> a) -> Stream a -> a
374 {-# INLINE foldl1 #-}
375 foldl1 f = unId . M.foldl1 f
376
377 -- | Left fold with strict accumulator
378 foldl' :: (a -> b -> a) -> a -> Stream b -> a
379 {-# INLINE foldl' #-}
380 foldl' f z = unId . M.foldl' f z
381
382 -- | Left fold on non-empty 'Stream's with strict accumulator
383 foldl1' :: (a -> a -> a) -> Stream a -> a
384 {-# INLINE foldl1' #-}
385 foldl1' f = unId . M.foldl1' f
386
387 -- | Right fold
388 foldr :: (a -> b -> b) -> b -> Stream a -> b
389 {-# INLINE foldr #-}
390 foldr f z = unId . M.foldr f z
391
392 -- | Right fold on non-empty 'Stream's
393 foldr1 :: (a -> a -> a) -> Stream a -> a
394 {-# INLINE foldr1 #-}
395 foldr1 f = unId . M.foldr1 f
396
397 -- Specialised folds
398 -- -----------------
399
400 and :: Stream Bool -> Bool
401 {-# INLINE and #-}
402 and = unId . M.and
403
404 or :: Stream Bool -> Bool
405 {-# INLINE or #-}
406 or = unId . M.or
407
408 -- Unfolding
409 -- ---------
410
411 -- | Unfold
412 unfoldr :: (s -> Maybe (a, s)) -> s -> Stream a
413 {-# INLINE unfoldr #-}
414 unfoldr = M.unfoldr
415
416 -- | Unfold at most @n@ elements
417 unfoldrN :: Int -> (s -> Maybe (a, s)) -> s -> Stream a
418 {-# INLINE unfoldrN #-}
419 unfoldrN = M.unfoldrN
420
421 -- Scans
422 -- -----
423
424 -- | Prefix scan
425 prescanl :: (a -> b -> a) -> a -> Stream b -> Stream a
426 {-# INLINE prescanl #-}
427 prescanl = M.prescanl
428
429 -- | Prefix scan with strict accumulator
430 prescanl' :: (a -> b -> a) -> a -> Stream b -> Stream a
431 {-# INLINE prescanl' #-}
432 prescanl' = M.prescanl'
433
434 -- | Suffix scan
435 postscanl :: (a -> b -> a) -> a -> Stream b -> Stream a
436 {-# INLINE postscanl #-}
437 postscanl = M.postscanl
438
439 -- | Suffix scan with strict accumulator
440 postscanl' :: (a -> b -> a) -> a -> Stream b -> Stream a
441 {-# INLINE postscanl' #-}
442 postscanl' = M.postscanl'
443
444 -- | Haskell-style scan
445 scanl :: (a -> b -> a) -> a -> Stream b -> Stream a
446 {-# INLINE scanl #-}
447 scanl = M.scanl
448
449 -- | Haskell-style scan with strict accumulator
450 scanl' :: (a -> b -> a) -> a -> Stream b -> Stream a
451 {-# INLINE scanl' #-}
452 scanl' = M.scanl'
453
454 -- | Scan over a non-empty 'Stream'
455 scanl1 :: (a -> a -> a) -> Stream a -> Stream a
456 {-# INLINE scanl1 #-}
457 scanl1 = M.scanl1
458
459 -- | Scan over a non-empty 'Stream' with a strict accumulator
460 scanl1' :: (a -> a -> a) -> Stream a -> Stream a
461 {-# INLINE scanl1' #-}
462 scanl1' = M.scanl1'
463
464
465 -- Comparisons
466 -- -----------
467
468 -- | Check if two 'Stream's are equal
469 eq :: Eq a => Stream a -> Stream a -> Bool
470 {-# INLINE_STREAM eq #-}
471 eq (M.Stream step1 s1 _) (M.Stream step2 s2 _) = eq_loop0 s1 s2
472 where
473 eq_loop0 s1 s2 = case unId (step1 s1) of
474 Yield x s1' -> eq_loop1 x s1' s2
475 Skip s1' -> eq_loop0 s1' s2
476 Done -> null (M.Stream step2 s2 Unknown)
477
478 eq_loop1 x s1 s2 = case unId (step2 s2) of
479 Yield y s2' -> x == y && eq_loop0 s1 s2'
480 Skip s2' -> eq_loop1 x s1 s2'
481 Done -> False
482
483 -- | Lexicographically compare two 'Stream's
484 cmp :: Ord a => Stream a -> Stream a -> Ordering
485 {-# INLINE_STREAM cmp #-}
486 cmp (M.Stream step1 s1 _) (M.Stream step2 s2 _) = cmp_loop0 s1 s2
487 where
488 cmp_loop0 s1 s2 = case unId (step1 s1) of
489 Yield x s1' -> cmp_loop1 x s1' s2
490 Skip s1' -> cmp_loop0 s1' s2
491 Done -> if null (M.Stream step2 s2 Unknown)
492 then EQ else LT
493
494 cmp_loop1 x s1 s2 = case unId (step2 s2) of
495 Yield y s2' -> case x `compare` y of
496 EQ -> cmp_loop0 s1 s2'
497 c -> c
498 Skip s2' -> cmp_loop1 x s1 s2'
499 Done -> GT
500
501 instance Eq a => Eq (M.Stream Id a) where
502 {-# INLINE (==) #-}
503 (==) = eq
504
505 instance Ord a => Ord (M.Stream Id a) where
506 {-# INLINE compare #-}
507 compare = cmp
508
509 -- Monadic combinators
510 -- -------------------
511
512 -- | Apply a monadic action to each element of the stream
513 mapM_ :: Monad m => (a -> m b) -> Stream a -> m ()
514 {-# INLINE mapM_ #-}
515 mapM_ f = M.mapM_ f . liftStream
516
517 -- | Monadic fold
518 foldM :: Monad m => (a -> b -> m a) -> a -> Stream b -> m a
519 {-# INLINE foldM #-}
520 foldM m z = M.foldM m z . liftStream
521
522 -- | Monadic fold over non-empty stream
523 fold1M :: Monad m => (a -> a -> m a) -> Stream a -> m a
524 {-# INLINE fold1M #-}
525 fold1M m = M.fold1M m . liftStream
526
527 -- | Monadic fold with strict accumulator
528 foldM' :: Monad m => (a -> b -> m a) -> a -> Stream b -> m a
529 {-# INLINE foldM' #-}
530 foldM' m z = M.foldM' m z . liftStream
531
532 -- | Monad fold over non-empty stream with strict accumulator
533 fold1M' :: Monad m => (a -> a -> m a) -> Stream a -> m a
534 {-# INLINE fold1M' #-}
535 fold1M' m = M.fold1M' m . liftStream
536
537 -- Enumerations
538 -- ------------
539
540 -- | Yield a 'Stream' of the given length containing the values @x@, @x+y@,
541 -- @x+y+y@ etc.
542 enumFromStepN :: Num a => a -> a -> Int -> Stream a
543 {-# INLINE enumFromStepN #-}
544 enumFromStepN = M.enumFromStepN
545
546 -- | Enumerate values
547 --
548 -- /WARNING:/ This operations can be very inefficient. If at all possible, use
549 -- 'enumFromStepN' instead.
550 enumFromTo :: Enum a => a -> a -> Stream a
551 {-# INLINE enumFromTo #-}
552 enumFromTo = M.enumFromTo
553
554 -- | Enumerate values with a given step.
555 --
556 -- /WARNING:/ This operations is very inefficient. If at all possible, use
557 -- 'enumFromStepN' instead.
558 enumFromThenTo :: Enum a => a -> a -> a -> Stream a
559 {-# INLINE enumFromThenTo #-}
560 enumFromThenTo = M.enumFromThenTo
561
562 -- Conversions
563 -- -----------
564
565 -- | Convert a 'Stream' to a list
566 toList :: Stream a -> [a]
567 {-# INLINE toList #-}
568 -- toList s = unId (M.toList s)
569 toList s = build (\c n -> toListFB c n s)
570
571 -- This supports foldr/build list fusion that GHC implements
572 toListFB :: (a -> b -> b) -> b -> Stream a -> b
573 {-# INLINE [0] toListFB #-}
574 toListFB c n (M.Stream step s _) = go s
575 where
576 go s = case unId (step s) of
577 Yield x s' -> x `c` go s'
578 Skip s' -> go s'
579 Done -> n
580
581 -- | Create a 'Stream' from a list
582 fromList :: [a] -> Stream a
583 {-# INLINE fromList #-}
584 fromList = M.fromList
585
586 -- | Create a 'Stream' from the first @n@ elements of a list
587 --
588 -- > fromListN n xs = fromList (take n xs)
589 fromListN :: Int -> [a] -> Stream a
590 {-# INLINE fromListN #-}
591 fromListN = M.fromListN
592