Add comment
[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 -- FIXME: slicing doesn't work with the inplace stuff at the moment
272
273 -- | Yield a part of the vector without copying it. Safer version of
274 -- 'unsafeSlice'.
275 slice :: IVector v a => v a -> Int -- ^ starting index
276 -> Int -- ^ length
277 -> v a
278 {-# INLINE_STREAM slice #-}
279 slice v i n = assert (i >= 0 && n >= 0 && i+n <= length v)
280 $ unsafeSlice v i n
281
282 -- | Copy @n@ elements starting at the given position to a new vector.
283 extract :: IVector v a => v a -> Int -- ^ starting index
284 -> Int -- ^ length
285 -> v a
286 {-# INLINE extract #-}
287 extract v i n = unstream (Stream.extract (stream v) i n)
288
289 -- | Yield the first @n@ elements without copying.
290 takeSlice :: IVector v a => Int -> v a -> v a
291 {-# INLINE takeSlice #-}
292 takeSlice n v = slice v 0 n
293
294 -- | Copy the first @n@ elements to a new vector.
295 take :: IVector v a => Int -> v a -> v a
296 {-# INLINE take #-}
297 take n = unstream . Stream.take n . stream
298
299 -- | Yield all but the first @n@ elements without copying.
300 dropSlice :: IVector v a => Int -> v a -> v a
301 {-# INLINE dropSlice #-}
302 dropSlice n v = slice v n (length v - n)
303
304 -- | Copy all but the first @n@ elements to a new vectors.
305 drop :: IVector v a => Int -> v a -> v a
306 {-# INLINE drop #-}
307 drop n = unstream . Stream.drop n . stream
308
309 {-# RULES
310
311 "slice/unstream [IVector]" forall v i n s.
312 slice (new' v (New.unstream s)) i n = extract (new' v (New.unstream s)) i n
313
314 #-}
315
316 -- Permutations
317 -- ------------
318
319 (//) :: IVector v a => v a -> [(Int, a)] -> v a
320 {-# INLINE (//) #-}
321 v // us = new (New.update (New.unstream (stream v))
322 (Stream.fromList us))
323
324 update :: (IVector v a, IVector v (Int, a)) => v a -> v (Int, a) -> v a
325 {-# INLINE update #-}
326 update v w = new (New.update (New.unstream (stream v)) (stream w))
327
328 bpermute :: (IVector v a, IVector v Int) => v a -> v Int -> v a
329 {-# INLINE bpermute #-}
330 bpermute v is = v `seq` map (v!) is
331
332 -- Mapping/zipping
333 -- ---------------
334
335 -- | Map a function over a vector
336 map :: (IVector v a, IVector v b) => (a -> b) -> v a -> v b
337 {-# INLINE map #-}
338 map f = unstream . Stream.map f . stream
339
340 inplace_map :: IVector v a => (a -> a) -> v a -> v a
341 {-# INLINE inplace_map #-}
342 inplace_map f = unstream . inplace (MStream.map f) . stream
343
344 {-# RULES
345
346 "map->inplace_map [IVector]" map = inplace_map
347
348 #-}
349
350 -- | Zip two vectors with the given function.
351 zipWith :: (IVector v a, IVector v b, IVector v c) => (a -> b -> c) -> v a -> v b -> v c
352 {-# INLINE zipWith #-}
353 zipWith f xs ys = unstream (Stream.zipWith f (stream xs) (stream ys))
354
355 zip :: (IVector v a, IVector v b, IVector v (a,b)) => v a -> v b -> v (a, b)
356 {-# INLINE zip #-}
357 zip = zipWith (,)
358
359 -- Comparisons
360 -- -----------
361
362 eq :: (IVector v a, Eq a) => v a -> v a -> Bool
363 {-# INLINE eq #-}
364 xs `eq` ys = stream xs == stream ys
365
366 cmp :: (IVector v a, Ord a) => v a -> v a -> Ordering
367 {-# INLINE cmp #-}
368 cmp xs ys = compare (stream xs) (stream ys)
369
370 -- Filtering
371 -- ---------
372
373 -- | Drop elements which do not satisfy the predicate
374 filter :: IVector v a => (a -> Bool) -> v a -> v a
375 {-# INLINE filter #-}
376 filter f = unstream . inplace (MStream.filter f) . stream
377
378 -- | Yield the longest prefix of elements satisfying the predicate without
379 -- copying.
380 takeWhileSlice :: IVector v a => (a -> Bool) -> v a -> v a
381 {-# INLINE_STREAM takeWhileSlice #-}
382 takeWhileSlice f v = case findIndex (not . f) v of
383 Just n -> takeSlice n v
384 Nothing -> v
385
386 -- | Copy the longest prefix of elements satisfying the predicate to a new
387 -- vector
388 takeWhile :: IVector v a => (a -> Bool) -> v a -> v a
389 {-# INLINE takeWhile #-}
390 takeWhile f = unstream . Stream.takeWhile f . stream
391
392 -- | Drop the longest prefix of elements that satisfy the predicate without
393 -- copying
394 dropWhileSlice :: IVector v a => (a -> Bool) -> v a -> v a
395 {-# INLINE_STREAM dropWhileSlice #-}
396 dropWhileSlice f v = case findIndex (not . f) v of
397 Just n -> dropSlice n v
398 Nothing -> v
399
400 -- | Drop the longest prefix of elements that satisfy the predicate and copy
401 -- the rest to a new vector.
402 dropWhile :: IVector v a => (a -> Bool) -> v a -> v a
403 {-# INLINE dropWhile #-}
404 dropWhile f = unstream . Stream.dropWhile f . stream
405
406 {-# RULES
407
408 "takeWhileSlice/unstream" forall v f s.
409 takeWhileSlice f (new' v (New.unstream s)) = takeWhile f (new' v (New.unstream s))
410
411 "dropWhileSlice/unstream" forall v f s.
412 dropWhileSlice f (new' v (New.unstream s)) = dropWhile f (new' v (New.unstream s))
413
414 #-}
415
416 -- Searching
417 -- ---------
418
419 infix 4 `elem`
420 -- | Check whether the vector contains an element
421 elem :: (IVector v a, Eq a) => a -> v a -> Bool
422 {-# INLINE elem #-}
423 elem x = Stream.elem x . stream
424
425 infix 4 `notElem`
426 -- | Inverse of `elem`
427 notElem :: (IVector v a, Eq a) => a -> v a -> Bool
428 {-# INLINE notElem #-}
429 notElem x = Stream.notElem x . stream
430
431 -- | Yield 'Just' the first element matching the predicate or 'Nothing' if no
432 -- such element exists.
433 find :: IVector v a => (a -> Bool) -> v a -> Maybe a
434 {-# INLINE find #-}
435 find f = Stream.find f . stream
436
437 -- | Yield 'Just' the index of the first element matching the predicate or
438 -- 'Nothing' if no such element exists.
439 findIndex :: IVector v a => (a -> Bool) -> v a -> Maybe Int
440 {-# INLINE findIndex #-}
441 findIndex f = Stream.findIndex f . stream
442
443 -- Folding
444 -- -------
445
446 -- | Left fold
447 foldl :: IVector v b => (a -> b -> a) -> a -> v b -> a
448 {-# INLINE foldl #-}
449 foldl f z = Stream.foldl f z . stream
450
451 -- | Lefgt fold on non-empty vectors
452 foldl1 :: IVector v a => (a -> a -> a) -> v a -> a
453 {-# INLINE foldl1 #-}
454 foldl1 f = Stream.foldl1 f . stream
455
456 -- | Left fold with strict accumulator
457 foldl' :: IVector v b => (a -> b -> a) -> a -> v b -> a
458 {-# INLINE foldl' #-}
459 foldl' f z = Stream.foldl' f z . stream
460
461 -- | Left fold on non-empty vectors with strict accumulator
462 foldl1' :: IVector v a => (a -> a -> a) -> v a -> a
463 {-# INLINE foldl1' #-}
464 foldl1' f = Stream.foldl1' f . stream
465
466 -- | Right fold
467 foldr :: IVector v a => (a -> b -> b) -> b -> v a -> b
468 {-# INLINE foldr #-}
469 foldr f z = Stream.foldr f z . stream
470
471 -- | Right fold on non-empty vectors
472 foldr1 :: IVector v a => (a -> a -> a) -> v a -> a
473 {-# INLINE foldr1 #-}
474 foldr1 f = Stream.foldr1 f . stream
475
476 -- Scans
477 -- -----
478
479 -- | Prefix scan
480 prescanl :: (IVector v a, IVector v b) => (a -> b -> a) -> a -> v b -> v a
481 {-# INLINE prescanl #-}
482 prescanl f z = unstream . Stream.prescanl f z . stream
483
484 inplace_prescanl :: IVector v a => (a -> a -> a) -> a -> v a -> v a
485 {-# INLINE inplace_prescanl #-}
486 inplace_prescanl f z = unstream . inplace (MStream.prescanl f z) . stream
487
488 {-# RULES
489
490 "prescanl -> inplace_prescanl [IVector]" prescanl = inplace_prescanl
491
492 #-}
493
494 -- | Prefix scan with strict accumulator
495 prescanl' :: (IVector v a, IVector v b) => (a -> b -> a) -> a -> v b -> v a
496 {-# INLINE prescanl' #-}
497 prescanl' f z = unstream . Stream.prescanl' f z . stream
498
499 inplace_prescanl' :: IVector v a => (a -> a -> a) -> a -> v a -> v a
500 {-# INLINE inplace_prescanl' #-}
501 inplace_prescanl' f z = unstream . inplace (MStream.prescanl' f z) . stream
502
503 {-# RULES
504
505 "prescanl' -> inplace_prescanl' [IVector]" prescanl' = inplace_prescanl'
506
507 #-}
508
509
510 -- | Convert a vector to a list
511 toList :: IVector v a => v a -> [a]
512 {-# INLINE toList #-}
513 toList = Stream.toList . stream
514
515 -- | Convert a list to a vector
516 fromList :: IVector v a => [a] -> v a
517 {-# INLINE fromList #-}
518 fromList = unstream . Stream.fromList
519