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