661fe3321ccb54ba61f445c618b4e11b0a7c18b4
[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,
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,
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 and = unId . M.and
302
303 or :: Stream Bool -> Bool
304 or = unId . M.or
305
306 -- Unfolding
307 -- ---------
308
309 -- | Unfold
310 unfoldr :: (s -> Maybe (a, s)) -> s -> Stream a
311 {-# INLINE unfoldr #-}
312 unfoldr = M.unfoldr
313
314 -- Scans
315 -- -----
316
317 -- | Prefix scan
318 prescanl :: (a -> b -> a) -> a -> Stream b -> Stream a
319 {-# INLINE prescanl #-}
320 prescanl = M.prescanl
321
322 -- | Prefix scan with strict accumulator
323 prescanl' :: (a -> b -> a) -> a -> Stream b -> Stream a
324 {-# INLINE prescanl' #-}
325 prescanl' = M.prescanl'
326
327 -- Comparisons
328 -- -----------
329
330 -- | Check if two 'Stream's are equal
331 eq :: Eq a => Stream a -> Stream a -> Bool
332 {-# INLINE_STREAM eq #-}
333 eq (M.Stream step1 s1 _) (M.Stream step2 s2 _) = eq_loop0 s1 s2
334 where
335 eq_loop0 s1 s2 = case unId (step1 s1) of
336 Yield x s1' -> eq_loop1 x s1' s2
337 Skip s1' -> eq_loop0 s1' s2
338 Done -> null (M.Stream step2 s2 Unknown)
339
340 eq_loop1 x s1 s2 = case unId (step2 s2) of
341 Yield y s2' -> x == y && eq_loop0 s1 s2'
342 Skip s2' -> eq_loop1 x s1 s2'
343 Done -> False
344
345 -- | Lexicographically compare two 'Stream's
346 cmp :: Ord a => Stream a -> Stream a -> Ordering
347 {-# INLINE_STREAM cmp #-}
348 cmp (M.Stream step1 s1 _) (M.Stream step2 s2 _) = cmp_loop0 s1 s2
349 where
350 cmp_loop0 s1 s2 = case unId (step1 s1) of
351 Yield x s1' -> cmp_loop1 x s1' s2
352 Skip s1' -> cmp_loop0 s1' s2
353 Done -> if null (M.Stream step2 s2 Unknown)
354 then EQ else LT
355
356 cmp_loop1 x s1 s2 = case unId (step2 s2) of
357 Yield y s2' -> case x `compare` y of
358 EQ -> cmp_loop0 s1 s2'
359 c -> c
360 Skip s2' -> cmp_loop1 x s1 s2'
361 Done -> GT
362
363 instance Eq a => Eq (M.Stream Id a) where
364 {-# INLINE (==) #-}
365 (==) = eq
366
367 instance Ord a => Ord (M.Stream Id a) where
368 {-# INLINE compare #-}
369 compare = cmp
370
371 -- Monadic combinators
372 -- -------------------
373
374 -- | Apply a monadic action to each element of the stream
375 mapM_ :: Monad m => (a -> m ()) -> Stream a -> m ()
376 {-# INLINE_STREAM mapM_ #-}
377 mapM_ f = M.mapM_ f . liftStream
378
379 -- | Monadic fold
380 foldM :: Monad m => (a -> b -> m a) -> a -> Stream b -> m a
381 {-# INLINE_STREAM foldM #-}
382 foldM m z = M.foldM m z . liftStream
383
384 -- Conversions
385 -- -----------
386
387 -- | Convert a 'Stream' to a list
388 toList :: Stream a -> [a]
389 {-# INLINE toList #-}
390 toList s = unId (M.toList s)
391
392 -- | Create a 'Stream' from a list
393 fromList :: [a] -> Stream a
394 {-# INLINE fromList #-}
395 fromList = M.fromList
396