Add lots of index-related collective operations
[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, (++),
30
31 -- * Accessing individual elements
32 head, last, (!!),
33
34 -- * Substreams
35 extract, init, tail, take, drop,
36
37 -- * Mapping
38 map, concatMap, unbox,
39
40 -- * Zipping
41 indexed,
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 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 -- | Prepend an element
165 cons :: a -> Stream a -> Stream a
166 {-# INLINE cons #-}
167 cons = M.cons
168
169 -- | Append an element
170 snoc :: Stream a -> a -> Stream a
171 {-# INLINE snoc #-}
172 snoc = M.snoc
173
174 infixr 5 ++
175 -- | Concatenate two 'Stream's
176 (++) :: Stream a -> Stream a -> Stream a
177 {-# INLINE (++) #-}
178 (++) = (M.++)
179
180 -- Accessing elements
181 -- ------------------
182
183 -- | First element of the 'Stream' or error if empty
184 head :: Stream a -> a
185 {-# INLINE head #-}
186 head = unId . M.head
187
188 -- | Last element of the 'Stream' or error if empty
189 last :: Stream a -> a
190 {-# INLINE last #-}
191 last = unId . M.last
192
193 -- | Element at the given position
194 (!!) :: Stream a -> Int -> a
195 {-# INLINE (!!) #-}
196 s !! i = unId (s M.!! i)
197
198 -- Substreams
199 -- ----------
200
201 -- | Extract a substream of the given length starting at the given position.
202 extract :: Stream a -> Int -- ^ starting index
203 -> Int -- ^ length
204 -> Stream a
205 {-# INLINE extract #-}
206 extract = M.extract
207
208 -- | All but the last element
209 init :: Stream a -> Stream a
210 {-# INLINE init #-}
211 init = M.init
212
213 -- | All but the first element
214 tail :: Stream a -> Stream a
215 {-# INLINE tail #-}
216 tail = M.tail
217
218 -- | The first @n@ elements
219 take :: Int -> Stream a -> Stream a
220 {-# INLINE take #-}
221 take = M.take
222
223 -- | All but the first @n@ elements
224 drop :: Int -> Stream a -> Stream a
225 {-# INLINE drop #-}
226 drop = M.drop
227
228 -- Mapping
229 -- ---------------
230
231 -- | Map a function over a 'Stream'
232 map :: (a -> b) -> Stream a -> Stream b
233 {-# INLINE map #-}
234 map = M.map
235
236 unbox :: Stream (Box a) -> Stream a
237 {-# INLINE unbox #-}
238 unbox = M.unbox
239
240 concatMap :: (a -> Stream b) -> Stream a -> Stream b
241 {-# INLINE concatMap #-}
242 concatMap = M.concatMap
243
244 -- Zipping
245 -- -------
246
247 -- | Associate all 'Stream' elements with their indices
248 indexed :: Stream a -> Stream (Int,a)
249 {-# INLINE indexed #-}
250 indexed = M.indexed
251
252 -- | Zip two 'Stream's with the given function
253 zipWith :: (a -> b -> c) -> Stream a -> Stream b -> Stream c
254 {-# INLINE zipWith #-}
255 zipWith = M.zipWith
256
257 -- | Zip three 'Stream's with the given function
258 zipWith3 :: (a -> b -> c -> d) -> Stream a -> Stream b -> Stream c -> Stream d
259 {-# INLINE zipWith3 #-}
260 zipWith3 = M.zipWith3
261
262 zipWith4 :: (a -> b -> c -> d -> e)
263 -> Stream a -> Stream b -> Stream c -> Stream d
264 -> Stream e
265 {-# INLINE zipWith4 #-}
266 zipWith4 = M.zipWith4
267
268 zipWith5 :: (a -> b -> c -> d -> e -> f)
269 -> Stream a -> Stream b -> Stream c -> Stream d
270 -> Stream e -> Stream f
271 {-# INLINE zipWith5 #-}
272 zipWith5 = M.zipWith5
273
274 zipWith6 :: (a -> b -> c -> d -> e -> f -> g)
275 -> Stream a -> Stream b -> Stream c -> Stream d
276 -> Stream e -> Stream f -> Stream g
277 {-# INLINE zipWith6 #-}
278 zipWith6 = M.zipWith6
279
280 zip :: Stream a -> Stream b -> Stream (a,b)
281 {-# INLINE zip #-}
282 zip = M.zip
283
284 zip3 :: Stream a -> Stream b -> Stream c -> Stream (a,b,c)
285 {-# INLINE zip3 #-}
286 zip3 = M.zip3
287
288 zip4 :: Stream a -> Stream b -> Stream c -> Stream d
289 -> Stream (a,b,c,d)
290 {-# INLINE zip4 #-}
291 zip4 = M.zip4
292
293 zip5 :: Stream a -> Stream b -> Stream c -> Stream d
294 -> Stream e -> Stream (a,b,c,d,e)
295 {-# INLINE zip5 #-}
296 zip5 = M.zip5
297
298 zip6 :: Stream a -> Stream b -> Stream c -> Stream d
299 -> Stream e -> Stream f -> Stream (a,b,c,d,e,f)
300 {-# INLINE zip6 #-}
301 zip6 = M.zip6
302
303 -- Filtering
304 -- ---------
305
306 -- | Drop elements which do not satisfy the predicate
307 filter :: (a -> Bool) -> Stream a -> Stream a
308 {-# INLINE filter #-}
309 filter = M.filter
310
311 -- | Longest prefix of elements that satisfy the predicate
312 takeWhile :: (a -> Bool) -> Stream a -> Stream a
313 {-# INLINE takeWhile #-}
314 takeWhile = M.takeWhile
315
316 -- | Drop the longest prefix of elements that satisfy the predicate
317 dropWhile :: (a -> Bool) -> Stream a -> Stream a
318 {-# INLINE dropWhile #-}
319 dropWhile = M.dropWhile
320
321 -- Searching
322 -- ---------
323
324 infix 4 `elem`
325 -- | Check whether the 'Stream' contains an element
326 elem :: Eq a => a -> Stream a -> Bool
327 {-# INLINE elem #-}
328 elem x = unId . M.elem x
329
330 infix 4 `notElem`
331 -- | Inverse of `elem`
332 notElem :: Eq a => a -> Stream a -> Bool
333 {-# INLINE notElem #-}
334 notElem x = unId . M.notElem x
335
336 -- | Yield 'Just' the first element matching the predicate or 'Nothing' if no
337 -- such element exists.
338 find :: (a -> Bool) -> Stream a -> Maybe a
339 {-# INLINE find #-}
340 find f = unId . M.find f
341
342 -- | Yield 'Just' the index of the first element matching the predicate or
343 -- 'Nothing' if no such element exists.
344 findIndex :: (a -> Bool) -> Stream a -> Maybe Int
345 {-# INLINE findIndex #-}
346 findIndex f = unId . M.findIndex f
347
348 -- Folding
349 -- -------
350
351 -- | Left fold
352 foldl :: (a -> b -> a) -> a -> Stream b -> a
353 {-# INLINE foldl #-}
354 foldl f z = unId . M.foldl f z
355
356 -- | Left fold on non-empty 'Stream's
357 foldl1 :: (a -> a -> a) -> Stream a -> a
358 {-# INLINE foldl1 #-}
359 foldl1 f = unId . M.foldl1 f
360
361 -- | Left fold with strict accumulator
362 foldl' :: (a -> b -> a) -> a -> Stream b -> a
363 {-# INLINE foldl' #-}
364 foldl' f z = unId . M.foldl' f z
365
366 -- | Left fold on non-empty 'Stream's with strict accumulator
367 foldl1' :: (a -> a -> a) -> Stream a -> a
368 {-# INLINE foldl1' #-}
369 foldl1' f = unId . M.foldl1' f
370
371 -- | Right fold
372 foldr :: (a -> b -> b) -> b -> Stream a -> b
373 {-# INLINE foldr #-}
374 foldr f z = unId . M.foldr f z
375
376 -- | Right fold on non-empty 'Stream's
377 foldr1 :: (a -> a -> a) -> Stream a -> a
378 {-# INLINE foldr1 #-}
379 foldr1 f = unId . M.foldr1 f
380
381 -- Specialised folds
382 -- -----------------
383
384 and :: Stream Bool -> Bool
385 {-# INLINE and #-}
386 and = unId . M.and
387
388 or :: Stream Bool -> Bool
389 {-# INLINE or #-}
390 or = unId . M.or
391
392 -- Unfolding
393 -- ---------
394
395 -- | Unfold
396 unfoldr :: (s -> Maybe (a, s)) -> s -> Stream a
397 {-# INLINE unfoldr #-}
398 unfoldr = M.unfoldr
399
400 -- Scans
401 -- -----
402
403 -- | Prefix scan
404 prescanl :: (a -> b -> a) -> a -> Stream b -> Stream a
405 {-# INLINE prescanl #-}
406 prescanl = M.prescanl
407
408 -- | Prefix scan with strict accumulator
409 prescanl' :: (a -> b -> a) -> a -> Stream b -> Stream a
410 {-# INLINE prescanl' #-}
411 prescanl' = M.prescanl'
412
413 -- | Suffix scan
414 postscanl :: (a -> b -> a) -> a -> Stream b -> Stream a
415 {-# INLINE postscanl #-}
416 postscanl = M.postscanl
417
418 -- | Suffix scan with strict accumulator
419 postscanl' :: (a -> b -> a) -> a -> Stream b -> Stream a
420 {-# INLINE postscanl' #-}
421 postscanl' = M.postscanl'
422
423 -- | Haskell-style scan
424 scanl :: (a -> b -> a) -> a -> Stream b -> Stream a
425 {-# INLINE scanl #-}
426 scanl = M.scanl
427
428 -- | Haskell-style scan with strict accumulator
429 scanl' :: (a -> b -> a) -> a -> Stream b -> Stream a
430 {-# INLINE scanl' #-}
431 scanl' = M.scanl'
432
433 -- | Scan over a non-empty 'Stream'
434 scanl1 :: (a -> a -> a) -> Stream a -> Stream a
435 {-# INLINE scanl1 #-}
436 scanl1 = M.scanl1
437
438 -- | Scan over a non-empty 'Stream' with a strict accumulator
439 scanl1' :: (a -> a -> a) -> Stream a -> Stream a
440 {-# INLINE scanl1' #-}
441 scanl1' = M.scanl1'
442
443
444 -- Comparisons
445 -- -----------
446
447 -- | Check if two 'Stream's are equal
448 eq :: Eq a => Stream a -> Stream a -> Bool
449 {-# INLINE_STREAM eq #-}
450 eq (M.Stream step1 s1 _) (M.Stream step2 s2 _) = eq_loop0 s1 s2
451 where
452 eq_loop0 s1 s2 = case unId (step1 s1) of
453 Yield x s1' -> eq_loop1 x s1' s2
454 Skip s1' -> eq_loop0 s1' s2
455 Done -> null (M.Stream step2 s2 Unknown)
456
457 eq_loop1 x s1 s2 = case unId (step2 s2) of
458 Yield y s2' -> x == y && eq_loop0 s1 s2'
459 Skip s2' -> eq_loop1 x s1 s2'
460 Done -> False
461
462 -- | Lexicographically compare two 'Stream's
463 cmp :: Ord a => Stream a -> Stream a -> Ordering
464 {-# INLINE_STREAM cmp #-}
465 cmp (M.Stream step1 s1 _) (M.Stream step2 s2 _) = cmp_loop0 s1 s2
466 where
467 cmp_loop0 s1 s2 = case unId (step1 s1) of
468 Yield x s1' -> cmp_loop1 x s1' s2
469 Skip s1' -> cmp_loop0 s1' s2
470 Done -> if null (M.Stream step2 s2 Unknown)
471 then EQ else LT
472
473 cmp_loop1 x s1 s2 = case unId (step2 s2) of
474 Yield y s2' -> case x `compare` y of
475 EQ -> cmp_loop0 s1 s2'
476 c -> c
477 Skip s2' -> cmp_loop1 x s1 s2'
478 Done -> GT
479
480 instance Eq a => Eq (M.Stream Id a) where
481 {-# INLINE (==) #-}
482 (==) = eq
483
484 instance Ord a => Ord (M.Stream Id a) where
485 {-# INLINE compare #-}
486 compare = cmp
487
488 -- Monadic combinators
489 -- -------------------
490
491 -- | Apply a monadic action to each element of the stream
492 mapM_ :: Monad m => (a -> m ()) -> Stream a -> m ()
493 {-# INLINE mapM_ #-}
494 mapM_ f = M.mapM_ f . liftStream
495
496 -- | Monadic fold
497 foldM :: Monad m => (a -> b -> m a) -> a -> Stream b -> m a
498 {-# INLINE foldM #-}
499 foldM m z = M.foldM m z . liftStream
500
501 -- | Monadic fold over non-empty stream
502 fold1M :: Monad m => (a -> a -> m a) -> Stream a -> m a
503 {-# INLINE fold1M #-}
504 fold1M m = M.fold1M m . liftStream
505
506 -- | Monadic fold with strict accumulator
507 foldM' :: Monad m => (a -> b -> m a) -> a -> Stream b -> m a
508 {-# INLINE foldM' #-}
509 foldM' m z = M.foldM' m z . liftStream
510
511 -- | Monad fold over non-empty stream with strict accumulator
512 fold1M' :: Monad m => (a -> a -> m a) -> Stream a -> m a
513 {-# INLINE fold1M' #-}
514 fold1M' m = M.fold1M' m . liftStream
515
516 -- Enumerations
517 -- ------------
518
519 enumFromTo :: Enum a => a -> a -> Stream a
520 {-# INLINE enumFromTo #-}
521 enumFromTo = M.enumFromTo
522
523 enumFromThenTo :: Enum a => a -> a -> a -> Stream a
524 {-# INLINE enumFromThenTo #-}
525 enumFromThenTo = M.enumFromThenTo
526
527 -- Conversions
528 -- -----------
529
530 -- | Convert a 'Stream' to a list
531 toList :: Stream a -> [a]
532 {-# INLINE toList #-}
533 toList s = unId (M.toList s)
534
535 -- | Create a 'Stream' from a list
536 fromList :: [a] -> Stream a
537 {-# INLINE fromList #-}
538 fromList = M.fromList
539