Add prescans
[darcs-mirrors/vector.git] / Data / Vector / IVector.hs
1 {-# LANGUAGE Rank2Types, MultiParamTypeClasses, FlexibleContexts,
2 ScopedTypeVariables #-}
3 -- |
4 -- Module : Data.Vector.IVector
5 -- Copyright : (c) Roman Leshchinskiy 2008
6 -- License : BSD-style
7 --
8 -- Maintainer : rl@cse.unsw.edu.au
9 -- Stability : experimental
10 -- Portability : non-portable
11 --
12 -- Generic interface to pure vectors
13 --
14
15 #include "phases.h"
16
17 module Data.Vector.IVector (
18 -- * Immutable vectors
19 IVector,
20
21 -- * Length information
22 length,
23
24 -- * Construction
25 empty, singleton, cons, snoc, replicate, (++),
26
27 -- * Accessing individual elements
28 (!), head, last,
29
30 -- * Subvectors
31 slice, extract, takeSlice, take, dropSlice, drop,
32
33 -- * Permutations
34 (//), update, bpermute,
35
36 -- * Mapping and zipping
37 map, zipWith, zip,
38
39 -- * Comparisons
40 eq, cmp,
41
42 -- * Filtering
43 filter, takeWhileSlice, takeWhile, dropWhileSlice, dropWhile,
44
45 -- * Searching
46 elem, notElem, find, findIndex,
47
48 -- * Folding
49 foldl, foldl1, foldl', foldl1', foldr, foldr1,
50
51 -- * Scans
52 prescanl, prescanl',
53
54 -- * Conversion to/from lists
55 toList, fromList,
56
57 -- * Conversion to/from Streams
58 stream, unstream,
59
60 -- * MVector-based initialisation
61 new,
62
63 -- * Unsafe functions
64 unsafeSlice, unsafeIndex,
65
66 -- * Utility functions
67 vlength, vnew
68 ) where
69
70 import qualified Data.Vector.MVector as MVector
71 import Data.Vector.MVector ( MVector )
72
73 import qualified Data.Vector.MVector.New as New
74 import Data.Vector.MVector.New ( New )
75
76 import qualified Data.Vector.Fusion.Stream as Stream
77 import Data.Vector.Fusion.Stream ( Stream, MStream )
78 import qualified Data.Vector.Fusion.Stream.Monadic as MStream
79 import Data.Vector.Fusion.Stream.Size
80
81 import Control.Exception ( assert )
82
83 import Prelude hiding ( length,
84 replicate, (++),
85 head, last,
86 init, tail, take, drop,
87 map, zipWith, zip,
88 filter, takeWhile, dropWhile,
89 elem, notElem,
90 foldl, foldl1, foldr, foldr1 )
91
92 -- | Class of immutable vectors.
93 --
94 class IVector v a where
95 -- | Construct a pure vector from a monadic initialiser (not fusible!)
96 vnew :: (forall mv m. MVector mv m a => m (mv a)) -> v a
97
98 -- | Length of the vector (not fusible!)
99 vlength :: v a -> Int
100
101 -- | Yield a part of the vector without copying it. No range checks!
102 unsafeSlice :: v a -> Int -> Int -> v a
103
104 -- | Apply the given function to the element at the given position. This
105 -- interface prevents us from being too lazy. Suppose we had
106 --
107 -- > unsafeIndex' :: v a -> Int -> a
108 --
109 -- instead. Now, if we wanted to copy a vector, we'd do something like
110 --
111 -- > copy mv v ... = ... unsafeWrite mv i (unsafeIndex' v i) ...
112 --
113 -- For lazy vectors, the indexing would not be evaluated which means that we
114 -- would retain a reference to the original vector in each element we write.
115 -- This would be bad!
116 --
117 -- With 'unsafeIndex', we can do
118 --
119 -- > copy mv v ... = ... unsafeIndex v i (unsafeWrite mv i) ...
120 --
121 -- which does not have this problem.
122 --
123 unsafeIndex :: v a -> Int -> (a -> b) -> b
124
125 -- Fusion
126 -- ------
127
128 -- | Construct a pure vector from a monadic initialiser
129 new :: IVector v a => New a -> v a
130 {-# INLINE new #-}
131 new m = new' undefined m
132
133 -- | Same as 'new' but with a dummy argument necessary for correctly typing
134 -- the rule @uninplace@.
135 new' :: IVector v a => v a -> New a -> v a
136 {-# INLINE_STREAM new' #-}
137 new' _ m = vnew (New.run m)
138
139 -- | Convert a vector to a 'Stream'
140 stream :: IVector v a => v a -> Stream a
141 {-# INLINE_STREAM stream #-}
142 stream v = v `seq` (Stream.unfold get 0 `Stream.sized` Exact n)
143 where
144 n = length v
145
146 {-# INLINE get #-}
147 get i | i < n = unsafeIndex v i $ \x -> Just (x, i+1)
148 | otherwise = Nothing
149
150 -- | Create a vector from a 'Stream'
151 unstream :: IVector v a => Stream a -> v a
152 {-# INLINE unstream #-}
153 unstream s = new (New.unstream s)
154
155 {-# RULES
156
157 "stream/unstream [IVector]" forall v s.
158 stream (new' v (New.unstream s)) = s
159
160 "New.unstream/stream/new [IVector]" forall v p.
161 New.unstream (stream (new' v p)) = p
162
163 #-}
164
165 inplace :: (forall m. Monad m => MStream m a -> MStream m a)
166 -> Stream a -> Stream a
167 {-# INLINE_STREAM inplace #-}
168 inplace f s = f s
169
170 {-# RULES
171
172 "inplace [IVector]"
173 forall (f :: forall m. Monad m => MStream m a -> MStream m a) v m.
174 New.unstream (inplace f (stream (new' v m))) = New.transform f m
175
176 "uninplace [IVector]"
177 forall (f :: forall m. Monad m => MStream m a -> MStream m a) v m.
178 stream (new' v (New.transform f m)) = inplace f (stream (new' v m))
179
180 "inplace/inplace [IVector]"
181 forall (f :: forall m. Monad m => MStream m a -> MStream m a)
182 (g :: forall m. Monad m => MStream m a -> MStream m a)
183 s.
184 inplace f (inplace g s) = inplace (f . g) s
185
186 #-}
187
188 -- Length
189 -- ------
190
191 length :: IVector v a => v a -> Int
192 {-# INLINE_STREAM length #-}
193 length v = vlength v
194
195 {-# RULES
196
197 "length/unstream [IVector]" forall v s.
198 length (new' v (New.unstream s)) = Stream.length s
199
200 #-}
201
202 -- Construction
203 -- ------------
204
205 -- | Empty vector
206 empty :: IVector v a => v a
207 {-# INLINE empty #-}
208 empty = unstream Stream.empty
209
210 -- | Vector with exaclty one element
211 singleton :: IVector v a => a -> v a
212 {-# INLINE singleton #-}
213 singleton x = unstream (Stream.singleton x)
214
215 -- | Vector of the given length with the given value in each position
216 replicate :: IVector v a => Int -> a -> v a
217 {-# INLINE replicate #-}
218 replicate n = unstream . Stream.replicate n
219
220 -- | Prepend an element
221 cons :: IVector v a => a -> v a -> v a
222 {-# INLINE cons #-}
223 cons x = unstream . Stream.cons x . stream
224
225 -- | Append an element
226 snoc :: IVector v a => v a -> a -> v a
227 {-# INLINE snoc #-}
228 snoc v = unstream . Stream.snoc (stream v)
229
230 infixr 5 ++
231 -- | Concatenate two vectors
232 (++) :: IVector v a => v a -> v a -> v a
233 {-# INLINE (++) #-}
234 v ++ w = unstream (stream v Stream.++ stream w)
235
236 -- Accessing individual elements
237 -- -----------------------------
238
239 -- | Indexing
240 (!) :: IVector v a => v a -> Int -> a
241 {-# INLINE_STREAM (!) #-}
242 v ! i = assert (i >= 0 && i < length v)
243 $ unsafeIndex v i id
244
245 -- | First element
246 head :: IVector v a => v a -> a
247 {-# INLINE_STREAM head #-}
248 head v = v ! 0
249
250 -- | Last element
251 last :: IVector v a => v a -> a
252 {-# INLINE_STREAM last #-}
253 last v = v ! (length v - 1)
254
255 {-# RULES
256
257 "(!)/unstream [IVector]" forall v i s.
258 new' v (New.unstream s) ! i = s Stream.!! i
259
260 "head/unstream [IVector]" forall v s.
261 head (new' v (New.unstream s)) = Stream.head s
262
263 "last/unstream [IVector]" forall v s.
264 last (new' v (New.unstream s)) = Stream.last s
265
266 #-}
267
268 -- Subarrays
269 -- ---------
270
271 -- | Yield a part of the vector without copying it. Safer version of
272 -- 'unsafeSlice'.
273 slice :: IVector v a => v a -> Int -- ^ starting index
274 -> Int -- ^ length
275 -> v a
276 {-# INLINE slice #-}
277 slice v i n = assert (i >= 0 && n >= 0 && i+n <= length v)
278 $ unsafeSlice v i n
279
280 -- | Copy @n@ elements starting at the given position to a new vector.
281 extract :: IVector v a => v a -> Int -- ^ starting index
282 -> Int -- ^ length
283 -> v a
284 {-# INLINE extract #-}
285 extract v i n = unstream (Stream.extract (stream v) i n)
286
287 -- | Yield the first @n@ elements without copying.
288 takeSlice :: IVector v a => Int -> v a -> v a
289 {-# INLINE takeSlice #-}
290 takeSlice n v = slice v 0 n
291
292 -- | Copy the first @n@ elements to a new vector.
293 take :: IVector v a => Int -> v a -> v a
294 {-# INLINE take #-}
295 take n = unstream . Stream.take n . stream
296
297 -- | Yield all but the first @n@ elements without copying.
298 dropSlice :: IVector v a => Int -> v a -> v a
299 {-# INLINE dropSlice #-}
300 dropSlice n v = slice v n (length v - n)
301
302 -- | Copy all but the first @n@ elements to a new vectors.
303 drop :: IVector v a => Int -> v a -> v a
304 {-# INLINE drop #-}
305 drop n = unstream . Stream.drop n . stream
306
307 {-# RULES
308
309 "slice/extract [IVector]" forall v i n s.
310 slice (new' v (New.unstream s)) i n = extract (new' v (New.unstream s)) i n
311
312 "takeSlice/unstream [IVector]" forall v n s.
313 takeSlice n (new' v (New.unstream s)) = take n (new' v (New.unstream s))
314
315 "dropSlice/unstream [IVector]" forall v n s.
316 dropSlice n (new' v (New.unstream s)) = drop n (new' v (New.unstream s))
317
318 #-}
319
320 -- Permutations
321 -- ------------
322
323 (//) :: IVector v a => v a -> [(Int, a)] -> v a
324 {-# INLINE (//) #-}
325 v // us = new (New.update (New.unstream (stream v))
326 (Stream.fromList us))
327
328 update :: (IVector v a, IVector v (Int, a)) => v a -> v (Int, a) -> v a
329 {-# INLINE update #-}
330 update v w = new (New.update (New.unstream (stream v)) (stream w))
331
332 bpermute :: (IVector v a, IVector v Int) => v a -> v Int -> v a
333 {-# INLINE bpermute #-}
334 bpermute v is = is `seq` map (v!) is
335
336 -- Mapping/zipping
337 -- ---------------
338
339 -- | Map a function over a vector
340 map :: (IVector v a, IVector v b) => (a -> b) -> v a -> v b
341 {-# INLINE map #-}
342 map f = unstream . Stream.map f . stream
343
344 inplace_map :: IVector v a => (a -> a) -> v a -> v a
345 {-# INLINE inplace_map #-}
346 inplace_map f = unstream . inplace (MStream.map f) . stream
347
348 {-# RULES
349
350 "map->inplace_map [IVector]" map = inplace_map
351
352 #-}
353
354 -- | Zip two vectors with the given function.
355 zipWith :: (IVector v a, IVector v b, IVector v c) => (a -> b -> c) -> v a -> v b -> v c
356 {-# INLINE zipWith #-}
357 zipWith f xs ys = unstream (Stream.zipWith f (stream xs) (stream ys))
358
359 zip :: (IVector v a, IVector v b, IVector v (a,b)) => v a -> v b -> v (a, b)
360 {-# INLINE zip #-}
361 zip = zipWith (,)
362
363 -- Comparisons
364 -- -----------
365
366 eq :: (IVector v a, Eq a) => v a -> v a -> Bool
367 {-# INLINE eq #-}
368 xs `eq` ys = stream xs == stream ys
369
370 cmp :: (IVector v a, Ord a) => v a -> v a -> Ordering
371 {-# INLINE cmp #-}
372 cmp xs ys = compare (stream xs) (stream ys)
373
374 -- Filtering
375 -- ---------
376
377 -- | Drop elements which do not satisfy the predicate
378 filter :: IVector v a => (a -> Bool) -> v a -> v a
379 {-# INLINE filter #-}
380 filter f = unstream . inplace (MStream.filter f) . stream
381
382 -- | Yield the longest prefix of elements satisfying the predicate without
383 -- copying.
384 takeWhileSlice :: IVector v a => (a -> Bool) -> v a -> v a
385 {-# INLINE takeWhileSlice #-}
386 takeWhileSlice f v = case findIndex (not . f) v of
387 Just n -> takeSlice n v
388 Nothing -> v
389
390 -- | Copy the longest prefix of elements satisfying the predicate to a new
391 -- vector
392 takeWhile :: IVector v a => (a -> Bool) -> v a -> v a
393 {-# INLINE takeWhile #-}
394 takeWhile f = unstream . Stream.takeWhile f . stream
395
396 -- | Drop the longest prefix of elements that satisfy the predicate without
397 -- copying
398 dropWhileSlice :: IVector v a => (a -> Bool) -> v a -> v a
399 {-# INLINE dropWhileSlice #-}
400 dropWhileSlice f v = case findIndex (not . f) v of
401 Just n -> dropSlice n v
402 Nothing -> v
403
404 -- | Drop the longest prefix of elements that satisfy the predicate and copy
405 -- the rest to a new vector.
406 dropWhile :: IVector v a => (a -> Bool) -> v a -> v a
407 {-# INLINE dropWhile #-}
408 dropWhile f = unstream . Stream.dropWhile f . stream
409
410 {-# RULES
411
412 "takeWhileSlice/unstream" forall v f s.
413 takeWhileSlice f (new' v (New.unstream s)) = takeWhile f (new' v (New.unstream s))
414
415 "dropWhileSlice/unstream" forall v f s.
416 dropWhileSlice f (new' v (New.unstream s)) = dropWhile f (new' v (New.unstream s))
417
418 #-}
419
420 -- Searching
421 -- ---------
422
423 infix 4 `elem`
424 -- | Check whether the vector contains an element
425 elem :: (IVector v a, Eq a) => a -> v a -> Bool
426 {-# INLINE elem #-}
427 elem x = Stream.elem x . stream
428
429 infix 4 `notElem`
430 -- | Inverse of `elem`
431 notElem :: (IVector v a, Eq a) => a -> v a -> Bool
432 {-# INLINE notElem #-}
433 notElem x = Stream.notElem x . stream
434
435 -- | Yield 'Just' the first element matching the predicate or 'Nothing' if no
436 -- such element exists.
437 find :: IVector v a => (a -> Bool) -> v a -> Maybe a
438 {-# INLINE find #-}
439 find f = Stream.find f . stream
440
441 -- | Yield 'Just' the index of the first element matching the predicate or
442 -- 'Nothing' if no such element exists.
443 findIndex :: IVector v a => (a -> Bool) -> v a -> Maybe Int
444 {-# INLINE findIndex #-}
445 findIndex f = Stream.findIndex f . stream
446
447 -- Folding
448 -- -------
449
450 -- | Left fold
451 foldl :: IVector v b => (a -> b -> a) -> a -> v b -> a
452 {-# INLINE foldl #-}
453 foldl f z = Stream.foldl f z . stream
454
455 -- | Lefgt fold on non-empty vectors
456 foldl1 :: IVector v a => (a -> a -> a) -> v a -> a
457 {-# INLINE foldl1 #-}
458 foldl1 f = Stream.foldl1 f . stream
459
460 -- | Left fold with strict accumulator
461 foldl' :: IVector v b => (a -> b -> a) -> a -> v b -> a
462 {-# INLINE foldl' #-}
463 foldl' f z = Stream.foldl' f z . stream
464
465 -- | Left fold on non-empty vectors with strict accumulator
466 foldl1' :: IVector v a => (a -> a -> a) -> v a -> a
467 {-# INLINE foldl1' #-}
468 foldl1' f = Stream.foldl1' f . stream
469
470 -- | Right fold
471 foldr :: IVector v a => (a -> b -> b) -> b -> v a -> b
472 {-# INLINE foldr #-}
473 foldr f z = Stream.foldr f z . stream
474
475 -- | Right fold on non-empty vectors
476 foldr1 :: IVector v a => (a -> a -> a) -> v a -> a
477 {-# INLINE foldr1 #-}
478 foldr1 f = Stream.foldr1 f . stream
479
480 -- Scans
481 -- -----
482
483 -- | Prefix scan
484 prescanl :: (IVector v a, IVector v b) => (a -> b -> a) -> a -> v b -> v a
485 {-# INLINE prescanl #-}
486 prescanl f z = unstream . Stream.prescanl f z . stream
487
488 inplace_prescanl :: IVector v a => (a -> a -> a) -> a -> v a -> v a
489 {-# INLINE inplace_prescanl #-}
490 inplace_prescanl f z = unstream . inplace (MStream.prescanl f z) . stream
491
492 {-# RULES
493
494 "prescanl -> inplace_prescanl [IVector]" prescanl = inplace_prescanl
495
496 #-}
497
498 -- | Prefix scan with strict accumulator
499 prescanl' :: (IVector v a, IVector v b) => (a -> b -> a) -> a -> v b -> v a
500 {-# INLINE prescanl' #-}
501 prescanl' f z = unstream . Stream.prescanl' f z . stream
502
503 inplace_prescanl' :: IVector v a => (a -> a -> a) -> a -> v a -> v a
504 {-# INLINE inplace_prescanl' #-}
505 inplace_prescanl' f z = unstream . inplace (MStream.prescanl' f z) . stream
506
507 {-# RULES
508
509 "prescanl' -> inplace_prescanl' [IVector]" prescanl' = inplace_prescanl'
510
511 #-}
512
513
514 -- | Convert a vector to a list
515 toList :: IVector v a => v a -> [a]
516 {-# INLINE toList #-}
517 toList = Stream.toList . stream
518
519 -- | Convert a list to a vector
520 fromList :: IVector v a => [a] -> v a
521 {-# INLINE fromList #-}
522 fromList = unstream . Stream.fromList
523