Comments
[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 #include "phases.h"
16
17 module Data.Vector.Fusion.Stream (
18 -- * Types
19 Step(..), Stream, MStream,
20
21 -- * In-place markers
22 inplace, inplace',
23
24 -- * Size hints
25 size, sized,
26
27 -- * Length information
28 length, null,
29
30 -- * Construction
31 empty, singleton, cons, snoc, replicate, (++),
32
33 -- * Accessing individual elements
34 head, last, (!!),
35
36 -- * Substreams
37 extract, init, tail, take, drop,
38
39 -- * Mapping
40 map, concatMap,
41
42 -- * Zipping
43 zipWith, zipWith3,
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 -- * Conversions
67 toList, fromList, liftStream,
68
69 -- * Monadic combinators
70 mapM_, foldM, fold1M, foldM', fold1M'
71 ) where
72
73 import Data.Vector.Fusion.Stream.Size
74 import Data.Vector.Fusion.Util
75 import Data.Vector.Fusion.Stream.Monadic ( Step(..) )
76 import qualified Data.Vector.Fusion.Stream.Monadic as M
77
78 import Prelude hiding ( length, null,
79 replicate, (++),
80 head, last, (!!),
81 init, tail, take, drop,
82 map, concatMap,
83 zipWith, zipWith3,
84 filter, takeWhile, dropWhile,
85 elem, notElem,
86 foldl, foldl1, foldr, foldr1,
87 and, or,
88 scanl, scanl1,
89 mapM_ )
90
91 -- | The type of pure streams
92 type Stream = M.Stream Id
93
94 -- | Alternative name for monadic streams
95 type MStream = M.Stream
96
97 inplace :: (forall m. Monad m => M.Stream m a -> M.Stream m a)
98 -> Stream a -> Stream a
99 {-# INLINE_STREAM inplace #-}
100 inplace f s = f s
101
102 inplace' :: (forall m. Monad m => M.Stream m a -> M.Stream m b)
103 -> Stream a -> Stream b
104 {-# INLINE_STREAM inplace' #-}
105 inplace' f s = f s
106
107 {-# RULES
108
109 "inplace/inplace [Vector]"
110 forall (f :: forall m. Monad m => MStream m a -> MStream m a)
111 (g :: forall m. Monad m => MStream m a -> MStream m a)
112 s.
113 inplace f (inplace g s) = inplace (f . g) s
114
115 "inplace' [Vector]"
116 forall (f :: forall m. Monad m => MStream m a -> MStream m a).
117 inplace' f = inplace f
118
119 #-}
120
121 -- | Convert a pure stream to a monadic stream
122 liftStream :: Monad m => Stream a -> M.Stream m a
123 {-# INLINE_STREAM liftStream #-}
124 liftStream (M.Stream step s sz) = M.Stream (return . unId . step) s sz
125
126 -- | 'Size' hint of a 'Stream'
127 size :: Stream a -> Size
128 {-# INLINE size #-}
129 size = M.size
130
131 -- | Attach a 'Size' hint to a 'Stream'
132 sized :: Stream a -> Size -> Stream a
133 {-# INLINE sized #-}
134 sized = M.sized
135
136 -- Length
137 -- ------
138
139 -- | Length of a 'Stream'
140 length :: Stream a -> Int
141 {-# INLINE length #-}
142 length = unId . M.length
143
144 -- | Check if a 'Stream' is empty
145 null :: Stream a -> Bool
146 {-# INLINE null #-}
147 null = unId . M.null
148
149 -- Construction
150 -- ------------
151
152 -- | Empty 'Stream'
153 empty :: Stream a
154 {-# INLINE empty #-}
155 empty = M.empty
156
157 -- | Singleton 'Stream'
158 singleton :: a -> Stream a
159 {-# INLINE singleton #-}
160 singleton = M.singleton
161
162 -- | Replicate a value to a given length
163 replicate :: Int -> a -> Stream a
164 {-# INLINE replicate #-}
165 replicate = M.replicate
166
167 -- | Prepend an element
168 cons :: a -> Stream a -> Stream a
169 {-# INLINE cons #-}
170 cons = M.cons
171
172 -- | Append an element
173 snoc :: Stream a -> a -> Stream a
174 {-# INLINE snoc #-}
175 snoc = M.snoc
176
177 infixr 5 ++
178 -- | Concatenate two 'Stream's
179 (++) :: Stream a -> Stream a -> Stream a
180 {-# INLINE (++) #-}
181 (++) = (M.++)
182
183 -- Accessing elements
184 -- ------------------
185
186 -- | First element of the 'Stream' or error if empty
187 head :: Stream a -> a
188 {-# INLINE head #-}
189 head = unId . M.head
190
191 -- | Last element of the 'Stream' or error if empty
192 last :: Stream a -> a
193 {-# INLINE last #-}
194 last = unId . M.last
195
196 -- | Element at the given position
197 (!!) :: Stream a -> Int -> a
198 {-# INLINE (!!) #-}
199 s !! i = unId (s M.!! i)
200
201 -- Substreams
202 -- ----------
203
204 -- | Extract a substream of the given length starting at the given position.
205 extract :: Stream a -> Int -- ^ starting index
206 -> Int -- ^ length
207 -> Stream a
208 {-# INLINE extract #-}
209 extract = M.extract
210
211 -- | All but the last element
212 init :: Stream a -> Stream a
213 {-# INLINE init #-}
214 init = M.init
215
216 -- | All but the first element
217 tail :: Stream a -> Stream a
218 {-# INLINE tail #-}
219 tail = M.tail
220
221 -- | The first @n@ elements
222 take :: Int -> Stream a -> Stream a
223 {-# INLINE take #-}
224 take = M.take
225
226 -- | All but the first @n@ elements
227 drop :: Int -> Stream a -> Stream a
228 {-# INLINE drop #-}
229 drop = M.drop
230
231 -- Mapping
232 -- ---------------
233
234 -- | Map a function over a 'Stream'
235 map :: (a -> b) -> Stream a -> Stream b
236 {-# INLINE map #-}
237 map = M.map
238
239 concatMap :: (a -> Stream b) -> Stream a -> Stream b
240 {-# INLINE concatMap #-}
241 concatMap = M.concatMap
242
243 -- Zipping
244 -- -------
245
246 -- | Zip two 'Stream's with the given function
247 zipWith :: (a -> b -> c) -> Stream a -> Stream b -> Stream c
248 {-# INLINE zipWith #-}
249 zipWith = M.zipWith
250
251 -- | Zip three 'Stream's with the given function
252 zipWith3 :: (a -> b -> c -> d) -> Stream a -> Stream b -> Stream c -> Stream d
253 {-# INLINE zipWith3 #-}
254 zipWith3 = M.zipWith3
255
256 -- Filtering
257 -- ---------
258
259 -- | Drop elements which do not satisfy the predicate
260 filter :: (a -> Bool) -> Stream a -> Stream a
261 {-# INLINE filter #-}
262 filter = M.filter
263
264 -- | Longest prefix of elements that satisfy the predicate
265 takeWhile :: (a -> Bool) -> Stream a -> Stream a
266 {-# INLINE takeWhile #-}
267 takeWhile = M.takeWhile
268
269 -- | Drop the longest prefix of elements that satisfy the predicate
270 dropWhile :: (a -> Bool) -> Stream a -> Stream a
271 {-# INLINE dropWhile #-}
272 dropWhile = M.dropWhile
273
274 -- Searching
275 -- ---------
276
277 infix 4 `elem`
278 -- | Check whether the 'Stream' contains an element
279 elem :: Eq a => a -> Stream a -> Bool
280 {-# INLINE elem #-}
281 elem x = unId . M.elem x
282
283 infix 4 `notElem`
284 -- | Inverse of `elem`
285 notElem :: Eq a => a -> Stream a -> Bool
286 {-# INLINE notElem #-}
287 notElem x = unId . M.notElem x
288
289 -- | Yield 'Just' the first element matching the predicate or 'Nothing' if no
290 -- such element exists.
291 find :: (a -> Bool) -> Stream a -> Maybe a
292 {-# INLINE find #-}
293 find f = unId . M.find f
294
295 -- | Yield 'Just' the index of the first element matching the predicate or
296 -- 'Nothing' if no such element exists.
297 findIndex :: (a -> Bool) -> Stream a -> Maybe Int
298 {-# INLINE findIndex #-}
299 findIndex f = unId . M.findIndex f
300
301 -- Folding
302 -- -------
303
304 -- | Left fold
305 foldl :: (a -> b -> a) -> a -> Stream b -> a
306 {-# INLINE foldl #-}
307 foldl f z = unId . M.foldl f z
308
309 -- | Left fold on non-empty 'Stream's
310 foldl1 :: (a -> a -> a) -> Stream a -> a
311 {-# INLINE foldl1 #-}
312 foldl1 f = unId . M.foldl1 f
313
314 -- | Left fold with strict accumulator
315 foldl' :: (a -> b -> a) -> a -> Stream b -> a
316 {-# INLINE foldl' #-}
317 foldl' f z = unId . M.foldl' f z
318
319 -- | Left fold on non-empty 'Stream's with strict accumulator
320 foldl1' :: (a -> a -> a) -> Stream a -> a
321 {-# INLINE foldl1' #-}
322 foldl1' f = unId . M.foldl1' f
323
324 -- | Right fold
325 foldr :: (a -> b -> b) -> b -> Stream a -> b
326 {-# INLINE foldr #-}
327 foldr f z = unId . M.foldr f z
328
329 -- | Right fold on non-empty 'Stream's
330 foldr1 :: (a -> a -> a) -> Stream a -> a
331 {-# INLINE foldr1 #-}
332 foldr1 f = unId . M.foldr1 f
333
334 -- Specialised folds
335 -- -----------------
336
337 and :: Stream Bool -> Bool
338 {-# INLINE and #-}
339 and = unId . M.and
340
341 or :: Stream Bool -> Bool
342 {-# INLINE or #-}
343 or = unId . M.or
344
345 -- Unfolding
346 -- ---------
347
348 -- | Unfold
349 unfoldr :: (s -> Maybe (a, s)) -> s -> Stream a
350 {-# INLINE unfoldr #-}
351 unfoldr = M.unfoldr
352
353 -- Scans
354 -- -----
355
356 -- | Prefix scan
357 prescanl :: (a -> b -> a) -> a -> Stream b -> Stream a
358 {-# INLINE prescanl #-}
359 prescanl = M.prescanl
360
361 -- | Prefix scan with strict accumulator
362 prescanl' :: (a -> b -> a) -> a -> Stream b -> Stream a
363 {-# INLINE prescanl' #-}
364 prescanl' = M.prescanl'
365
366 -- | Suffix scan
367 postscanl :: (a -> b -> a) -> a -> Stream b -> Stream a
368 {-# INLINE postscanl #-}
369 postscanl = M.postscanl
370
371 -- | Suffix scan with strict accumulator
372 postscanl' :: (a -> b -> a) -> a -> Stream b -> Stream a
373 {-# INLINE postscanl' #-}
374 postscanl' = M.postscanl'
375
376 -- | Haskell-style scan
377 scanl :: (a -> b -> a) -> a -> Stream b -> Stream a
378 {-# INLINE scanl #-}
379 scanl = M.scanl
380
381 -- | Haskell-style scan with strict accumulator
382 scanl' :: (a -> b -> a) -> a -> Stream b -> Stream a
383 {-# INLINE scanl' #-}
384 scanl' = M.scanl'
385
386 -- | Scan over a non-empty 'Stream'
387 scanl1 :: (a -> a -> a) -> Stream a -> Stream a
388 {-# INLINE scanl1 #-}
389 scanl1 = M.scanl1
390
391 -- | Scan over a non-empty 'Stream' with a strict accumulator
392 scanl1' :: (a -> a -> a) -> Stream a -> Stream a
393 {-# INLINE scanl1' #-}
394 scanl1' = M.scanl1'
395
396
397 -- Comparisons
398 -- -----------
399
400 -- | Check if two 'Stream's are equal
401 eq :: Eq a => Stream a -> Stream a -> Bool
402 {-# INLINE_STREAM eq #-}
403 eq (M.Stream step1 s1 _) (M.Stream step2 s2 _) = eq_loop0 s1 s2
404 where
405 eq_loop0 s1 s2 = case unId (step1 s1) of
406 Yield x s1' -> eq_loop1 x s1' s2
407 Skip s1' -> eq_loop0 s1' s2
408 Done -> null (M.Stream step2 s2 Unknown)
409
410 eq_loop1 x s1 s2 = case unId (step2 s2) of
411 Yield y s2' -> x == y && eq_loop0 s1 s2'
412 Skip s2' -> eq_loop1 x s1 s2'
413 Done -> False
414
415 -- | Lexicographically compare two 'Stream's
416 cmp :: Ord a => Stream a -> Stream a -> Ordering
417 {-# INLINE_STREAM cmp #-}
418 cmp (M.Stream step1 s1 _) (M.Stream step2 s2 _) = cmp_loop0 s1 s2
419 where
420 cmp_loop0 s1 s2 = case unId (step1 s1) of
421 Yield x s1' -> cmp_loop1 x s1' s2
422 Skip s1' -> cmp_loop0 s1' s2
423 Done -> if null (M.Stream step2 s2 Unknown)
424 then EQ else LT
425
426 cmp_loop1 x s1 s2 = case unId (step2 s2) of
427 Yield y s2' -> case x `compare` y of
428 EQ -> cmp_loop0 s1 s2'
429 c -> c
430 Skip s2' -> cmp_loop1 x s1 s2'
431 Done -> GT
432
433 instance Eq a => Eq (M.Stream Id a) where
434 {-# INLINE (==) #-}
435 (==) = eq
436
437 instance Ord a => Ord (M.Stream Id a) where
438 {-# INLINE compare #-}
439 compare = cmp
440
441 -- Monadic combinators
442 -- -------------------
443
444 -- | Apply a monadic action to each element of the stream
445 mapM_ :: Monad m => (a -> m ()) -> Stream a -> m ()
446 {-# INLINE mapM_ #-}
447 mapM_ f = M.mapM_ f . liftStream
448
449 -- | Monadic fold
450 foldM :: Monad m => (a -> b -> m a) -> a -> Stream b -> m a
451 {-# INLINE foldM #-}
452 foldM m z = M.foldM m z . liftStream
453
454 -- | Monadic fold over non-empty stream
455 fold1M :: Monad m => (a -> a -> m a) -> Stream a -> m a
456 {-# INLINE fold1M #-}
457 fold1M m = M.fold1M m . liftStream
458
459 -- | Monadic fold with strict accumulator
460 foldM' :: Monad m => (a -> b -> m a) -> a -> Stream b -> m a
461 {-# INLINE foldM' #-}
462 foldM' m z = M.foldM' m z . liftStream
463
464 -- | Monad fold over non-empty stream with strict accumulator
465 fold1M' :: Monad m => (a -> a -> m a) -> Stream a -> m a
466 {-# INLINE fold1M' #-}
467 fold1M' m = M.fold1M' m . liftStream
468
469
470 -- Conversions
471 -- -----------
472
473 -- | Convert a 'Stream' to a list
474 toList :: Stream a -> [a]
475 {-# INLINE toList #-}
476 toList s = unId (M.toList s)
477
478 -- | Create a 'Stream' from a list
479 fromList :: [a] -> Stream a
480 {-# INLINE fromList #-}
481 fromList = M.fromList
482