62070ef3eb1524c2dae3ebf7b2beb2af759e6eaa
[darcs-mirrors/vector.git] / Data / Vector / Fusion / Stream.hs
1 {-# LANGUAGE ExistentialQuantification, FlexibleInstances #-}
2
3 -- |
4 -- Module : Data.Vector.Fusion.Stream
5 -- Copyright : (c) Roman Leshchinskiy 2008
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, Id(..),
20
21 -- * Size hints
22 size, sized,
23
24 -- * Length information
25 length, null,
26
27 -- * Construction
28 empty, singleton, cons, snoc, replicate, (++),
29
30 -- * Accessing individual elements
31 head, last, (!!),
32
33 -- * Substreams
34 extract, init, tail, take, drop,
35
36 -- * Mapping and zipping
37 map, zipWith,
38
39 -- * Filtering
40 filter, takeWhile, dropWhile,
41
42 -- * Searching
43 elem, notElem, find, findIndex,
44
45 -- * Folding
46 foldl, foldl1, foldl', foldl1', foldr, foldr1,
47
48 -- * Specialised folds
49 and, or, concatMap,
50
51 -- * Unfolding
52 unfoldr,
53
54 -- * Scans
55 prescanl, prescanl',
56
57 -- * Conversions
58 toList, fromList, liftStream,
59
60 -- * Monadic combinators
61 mapM_, foldM
62 ) where
63
64 import Data.Vector.Fusion.Stream.Size
65 import Data.Vector.Fusion.Stream.Monadic ( Step(..) )
66 import qualified Data.Vector.Fusion.Stream.Monadic as M
67
68 import Prelude hiding ( length, null,
69 replicate, (++),
70 head, last, (!!),
71 init, tail, take, drop,
72 map, zipWith,
73 filter, takeWhile, dropWhile,
74 elem, notElem,
75 foldl, foldl1, foldr, foldr1,
76 and, or, concatMap,
77 mapM_ )
78
79
80 -- | Identity monad
81 newtype Id a = Id { unId :: a }
82
83 instance Functor Id where
84 fmap f (Id x) = Id (f x)
85
86 instance Monad Id where
87 return = Id
88 Id x >>= f = f x
89
90 -- | The type of pure streams
91 type Stream = M.Stream Id
92
93 -- | Alternative name for monadic streams
94 type MStream = M.Stream
95
96 -- | Convert a pure stream to a monadic stream
97 liftStream :: Monad m => Stream a -> M.Stream m a
98 {-# INLINE_STREAM liftStream #-}
99 liftStream (M.Stream step s sz) = M.Stream (return . unId . step) s sz
100
101 -- | 'Size' hint of a 'Stream'
102 size :: Stream a -> Size
103 {-# INLINE size #-}
104 size = M.size
105
106 -- | Attach a 'Size' hint to a 'Stream'
107 sized :: Stream a -> Size -> Stream a
108 {-# INLINE sized #-}
109 sized = M.sized
110
111 -- Length
112 -- ------
113
114 -- | Length of a 'Stream'
115 length :: Stream a -> Int
116 {-# INLINE length #-}
117 length = unId . M.length
118
119 -- | Check if a 'Stream' is empty
120 null :: Stream a -> Bool
121 {-# INLINE null #-}
122 null = unId . M.null
123
124 -- Construction
125 -- ------------
126
127 -- | Empty 'Stream'
128 empty :: Stream a
129 {-# INLINE empty #-}
130 empty = M.empty
131
132 -- | Singleton 'Stream'
133 singleton :: a -> Stream a
134 {-# INLINE singleton #-}
135 singleton = M.singleton
136
137 -- | Replicate a value to a given length
138 replicate :: Int -> a -> Stream a
139 {-# INLINE_STREAM replicate #-}
140 replicate = M.replicate
141
142 -- | Prepend an element
143 cons :: a -> Stream a -> Stream a
144 {-# INLINE cons #-}
145 cons = M.cons
146
147 -- | Append an element
148 snoc :: Stream a -> a -> Stream a
149 {-# INLINE snoc #-}
150 snoc = M.snoc
151
152 infixr 5 ++
153 -- | Concatenate two 'Stream's
154 (++) :: Stream a -> Stream a -> Stream a
155 {-# INLINE (++) #-}
156 (++) = (M.++)
157
158 -- Accessing elements
159 -- ------------------
160
161 -- | First element of the 'Stream' or error if empty
162 head :: Stream a -> a
163 {-# INLINE head #-}
164 head = unId . M.head
165
166 -- | Last element of the 'Stream' or error if empty
167 last :: Stream a -> a
168 {-# INLINE last #-}
169 last = unId . M.last
170
171 -- | Element at the given position
172 (!!) :: Stream a -> Int -> a
173 {-# INLINE (!!) #-}
174 s !! i = unId (s M.!! i)
175
176 -- Substreams
177 -- ----------
178
179 -- | Extract a substream of the given length starting at the given position.
180 extract :: Stream a -> Int -- ^ starting index
181 -> Int -- ^ length
182 -> Stream a
183 {-# INLINE extract #-}
184 extract = M.extract
185
186 -- | All but the last element
187 init :: Stream a -> Stream a
188 {-# INLINE init #-}
189 init = M.init
190
191 -- | All but the first element
192 tail :: Stream a -> Stream a
193 {-# INLINE tail #-}
194 tail = M.tail
195
196 -- | The first @n@ elements
197 take :: Int -> Stream a -> Stream a
198 {-# INLINE take #-}
199 take = M.take
200
201 -- | All but the first @n@ elements
202 drop :: Int -> Stream a -> Stream a
203 {-# INLINE drop #-}
204 drop = M.drop
205
206 -- Mapping/zipping
207 -- ---------------
208
209 -- | Map a function over a 'Stream'
210 map :: (a -> b) -> Stream a -> Stream b
211 {-# INLINE map #-}
212 map = M.map
213
214 -- | Zip two 'Stream's with the given function
215 zipWith :: (a -> b -> c) -> Stream a -> Stream b -> Stream c
216 {-# INLINE zipWith #-}
217 zipWith = M.zipWith
218
219 -- Filtering
220 -- ---------
221
222 -- | Drop elements which do not satisfy the predicate
223 filter :: (a -> Bool) -> Stream a -> Stream a
224 {-# INLINE filter #-}
225 filter = M.filter
226
227 -- | Longest prefix of elements that satisfy the predicate
228 takeWhile :: (a -> Bool) -> Stream a -> Stream a
229 {-# INLINE takeWhile #-}
230 takeWhile = M.takeWhile
231
232 -- | Drop the longest prefix of elements that satisfy the predicate
233 dropWhile :: (a -> Bool) -> Stream a -> Stream a
234 {-# INLINE dropWhile #-}
235 dropWhile = M.dropWhile
236
237 -- Searching
238 -- ---------
239
240 infix 4 `elem`
241 -- | Check whether the 'Stream' contains an element
242 elem :: Eq a => a -> Stream a -> Bool
243 {-# INLINE elem #-}
244 elem x = unId . M.elem x
245
246 infix 4 `notElem`
247 -- | Inverse of `elem`
248 notElem :: Eq a => a -> Stream a -> Bool
249 {-# INLINE notElem #-}
250 notElem x = unId . M.notElem x
251
252 -- | Yield 'Just' the first element matching the predicate or 'Nothing' if no
253 -- such element exists.
254 find :: (a -> Bool) -> Stream a -> Maybe a
255 {-# INLINE find #-}
256 find f = unId . M.find f
257
258 -- | Yield 'Just' the index of the first element matching the predicate or
259 -- 'Nothing' if no such element exists.
260 findIndex :: (a -> Bool) -> Stream a -> Maybe Int
261 {-# INLINE findIndex #-}
262 findIndex f = unId . M.findIndex f
263
264 -- Folding
265 -- -------
266
267 -- | Left fold
268 foldl :: (a -> b -> a) -> a -> Stream b -> a
269 {-# INLINE foldl #-}
270 foldl f z = unId . M.foldl f z
271
272 -- | Left fold on non-empty 'Stream's
273 foldl1 :: (a -> a -> a) -> Stream a -> a
274 {-# INLINE foldl1 #-}
275 foldl1 f = unId . M.foldl1 f
276
277 -- | Left fold with strict accumulator
278 foldl' :: (a -> b -> a) -> a -> Stream b -> a
279 {-# INLINE foldl' #-}
280 foldl' f z = unId . M.foldl' f z
281
282 -- | Left fold on non-empty 'Stream's with strict accumulator
283 foldl1' :: (a -> a -> a) -> Stream a -> a
284 {-# INLINE foldl1' #-}
285 foldl1' f = unId . M.foldl1' f
286
287 -- | Right fold
288 foldr :: (a -> b -> b) -> b -> Stream a -> b
289 {-# INLINE foldr #-}
290 foldr f z = unId . M.foldr f z
291
292 -- | Right fold on non-empty 'Stream's
293 foldr1 :: (a -> a -> a) -> Stream a -> a
294 {-# INLINE foldr1 #-}
295 foldr1 f = unId . M.foldr1 f
296
297 -- Specialised folds
298 -- -----------------
299
300 and :: Stream Bool -> Bool
301 {-# INLINE and #-}
302 and = unId . M.and
303
304 or :: Stream Bool -> Bool
305 {-# INLINE or #-}
306 or = unId . M.or
307
308 concatMap :: (a -> Stream b) -> Stream a -> Stream b
309 {-# INLINE concatMap #-}
310 concatMap = M.concatMap
311
312 -- Unfolding
313 -- ---------
314
315 -- | Unfold
316 unfoldr :: (s -> Maybe (a, s)) -> s -> Stream a
317 {-# INLINE unfoldr #-}
318 unfoldr = M.unfoldr
319
320 -- Scans
321 -- -----
322
323 -- | Prefix scan
324 prescanl :: (a -> b -> a) -> a -> Stream b -> Stream a
325 {-# INLINE prescanl #-}
326 prescanl = M.prescanl
327
328 -- | Prefix scan with strict accumulator
329 prescanl' :: (a -> b -> a) -> a -> Stream b -> Stream a
330 {-# INLINE prescanl' #-}
331 prescanl' = M.prescanl'
332
333 -- Comparisons
334 -- -----------
335
336 -- | Check if two 'Stream's are equal
337 eq :: Eq a => Stream a -> Stream a -> Bool
338 {-# INLINE_STREAM eq #-}
339 eq (M.Stream step1 s1 _) (M.Stream step2 s2 _) = eq_loop0 s1 s2
340 where
341 eq_loop0 s1 s2 = case unId (step1 s1) of
342 Yield x s1' -> eq_loop1 x s1' s2
343 Skip s1' -> eq_loop0 s1' s2
344 Done -> null (M.Stream step2 s2 Unknown)
345
346 eq_loop1 x s1 s2 = case unId (step2 s2) of
347 Yield y s2' -> x == y && eq_loop0 s1 s2'
348 Skip s2' -> eq_loop1 x s1 s2'
349 Done -> False
350
351 -- | Lexicographically compare two 'Stream's
352 cmp :: Ord a => Stream a -> Stream a -> Ordering
353 {-# INLINE_STREAM cmp #-}
354 cmp (M.Stream step1 s1 _) (M.Stream step2 s2 _) = cmp_loop0 s1 s2
355 where
356 cmp_loop0 s1 s2 = case unId (step1 s1) of
357 Yield x s1' -> cmp_loop1 x s1' s2
358 Skip s1' -> cmp_loop0 s1' s2
359 Done -> if null (M.Stream step2 s2 Unknown)
360 then EQ else LT
361
362 cmp_loop1 x s1 s2 = case unId (step2 s2) of
363 Yield y s2' -> case x `compare` y of
364 EQ -> cmp_loop0 s1 s2'
365 c -> c
366 Skip s2' -> cmp_loop1 x s1 s2'
367 Done -> GT
368
369 instance Eq a => Eq (M.Stream Id a) where
370 {-# INLINE (==) #-}
371 (==) = eq
372
373 instance Ord a => Ord (M.Stream Id a) where
374 {-# INLINE compare #-}
375 compare = cmp
376
377 -- Monadic combinators
378 -- -------------------
379
380 -- | Apply a monadic action to each element of the stream
381 mapM_ :: Monad m => (a -> m ()) -> Stream a -> m ()
382 {-# INLINE_STREAM mapM_ #-}
383 mapM_ f = M.mapM_ f . liftStream
384
385 -- | Monadic fold
386 foldM :: Monad m => (a -> b -> m a) -> a -> Stream b -> m a
387 {-# INLINE_STREAM foldM #-}
388 foldM m z = M.foldM m z . liftStream
389
390 -- Conversions
391 -- -----------
392
393 -- | Convert a 'Stream' to a list
394 toList :: Stream a -> [a]
395 {-# INLINE toList #-}
396 toList s = unId (M.toList s)
397
398 -- | Create a 'Stream' from a list
399 fromList :: [a] -> Stream a
400 {-# INLINE fromList #-}
401 fromList = M.fromList
402