9820dfed6f4734bfa6ae802525a54cabff2f2d82
[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 : Roman Leshchinskiy <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, null,
23
24 -- * Construction
25 empty, singleton, cons, snoc, replicate, (++), copy,
26
27 -- * Accessing individual elements
28 (!), head, last,
29
30 -- * Subvectors
31 slice, init, tail, take, drop,
32
33 -- * Permutations
34 accum, (//), update, bpermute,
35
36 -- * Mapping and zipping
37 map, zipWith, zip,
38
39 -- * Comparisons
40 eq, cmp,
41
42 -- * Filtering
43 filter, takeWhile, 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, null,
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 --
136 -- See http://hackage.haskell.org/trac/ghc/ticket/2600
137 new' :: IVector v a => v a -> New a -> v a
138 {-# INLINE_STREAM new' #-}
139 new' _ m = vnew (New.run m)
140
141 -- | Convert a vector to a 'Stream'
142 stream :: IVector v a => v a -> Stream a
143 {-# INLINE_STREAM stream #-}
144 stream v = v `seq` (Stream.unfold get 0 `Stream.sized` Exact n)
145 where
146 n = length v
147
148 {-# INLINE get #-}
149 get i | i < n = unsafeIndex v i $ \x -> Just (x, i+1)
150 | otherwise = Nothing
151
152 -- | Create a vector from a 'Stream'
153 unstream :: IVector v a => Stream a -> v a
154 {-# INLINE unstream #-}
155 unstream s = new (New.unstream s)
156
157 {-# RULES
158
159 "stream/unstream [IVector]" forall v s.
160 stream (new' v (New.unstream s)) = s
161
162 "New.unstream/stream/new [IVector]" forall v p.
163 New.unstream (stream (new' v p)) = p
164
165 #-}
166
167 inplace :: (forall m. Monad m => MStream m a -> MStream m a)
168 -> Stream a -> Stream a
169 {-# INLINE_STREAM inplace #-}
170 inplace f s = f s
171
172 {-# RULES
173
174 "inplace [IVector]"
175 forall (f :: forall m. Monad m => MStream m a -> MStream m a) v m.
176 New.unstream (inplace f (stream (new' v m))) = New.transform f m
177
178 "uninplace [IVector]"
179 forall (f :: forall m. Monad m => MStream m a -> MStream m a) v m.
180 stream (new' v (New.transform f m)) = inplace f (stream (new' v m))
181
182 "inplace/inplace [IVector]"
183 forall (f :: forall m. Monad m => MStream m a -> MStream m a)
184 (g :: forall m. Monad m => MStream m a -> MStream m a)
185 s.
186 inplace f (inplace g s) = inplace (f . g) s
187
188 #-}
189
190 -- Length
191 -- ------
192
193 length :: IVector v a => v a -> Int
194 {-# INLINE_STREAM length #-}
195 length v = vlength v
196
197 {-# RULES
198
199 "length/unstream [IVector]" forall v s.
200 length (new' v (New.unstream s)) = Stream.length s
201
202 #-}
203
204 null :: IVector v a => v a -> Bool
205 {-# INLINE_STREAM null #-}
206 null v = vlength v == 0
207
208 {-# RULES
209
210 "null/unstream [IVector]" forall v s.
211 null (new' v (New.unstream s)) = Stream.null s
212
213 #-}
214
215 -- Construction
216 -- ------------
217
218 -- | Empty vector
219 empty :: IVector v a => v a
220 {-# INLINE empty #-}
221 empty = unstream Stream.empty
222
223 -- | Vector with exaclty one element
224 singleton :: IVector v a => a -> v a
225 {-# INLINE singleton #-}
226 singleton x = unstream (Stream.singleton x)
227
228 -- | Vector of the given length with the given value in each position
229 replicate :: IVector v a => Int -> a -> v a
230 {-# INLINE replicate #-}
231 replicate n = unstream . Stream.replicate n
232
233 -- | Prepend an element
234 cons :: IVector v a => a -> v a -> v a
235 {-# INLINE cons #-}
236 cons x = unstream . Stream.cons x . stream
237
238 -- | Append an element
239 snoc :: IVector v a => v a -> a -> v a
240 {-# INLINE snoc #-}
241 snoc v = unstream . Stream.snoc (stream v)
242
243 infixr 5 ++
244 -- | Concatenate two vectors
245 (++) :: IVector v a => v a -> v a -> v a
246 {-# INLINE (++) #-}
247 v ++ w = unstream (stream v Stream.++ stream w)
248
249 -- | Create a copy of a vector. Useful when dealing with slices.
250 copy :: IVector v a => v a -> v a
251 {-# INLINE_STREAM copy #-}
252 copy = unstream . stream
253
254 {-# RULES
255
256 "copy/unstream [IVector]" forall v s.
257 copy (new' v (New.unstream s)) = new' v (New.unstream s)
258
259 #-}
260
261 -- Accessing individual elements
262 -- -----------------------------
263
264 -- | Indexing
265 (!) :: IVector v a => v a -> Int -> a
266 {-# INLINE_STREAM (!) #-}
267 v ! i = assert (i >= 0 && i < length v)
268 $ unsafeIndex v i id
269
270 -- | First element
271 head :: IVector v a => v a -> a
272 {-# INLINE_STREAM head #-}
273 head v = v ! 0
274
275 -- | Last element
276 last :: IVector v a => v a -> a
277 {-# INLINE_STREAM last #-}
278 last v = v ! (length v - 1)
279
280 {-# RULES
281
282 "(!)/unstream [IVector]" forall v i s.
283 new' v (New.unstream s) ! i = s Stream.!! i
284
285 "head/unstream [IVector]" forall v s.
286 head (new' v (New.unstream s)) = Stream.head s
287
288 "last/unstream [IVector]" forall v s.
289 last (new' v (New.unstream s)) = Stream.last s
290
291 #-}
292
293 -- Subarrays
294 -- ---------
295
296 -- FIXME: slicing doesn't work with the inplace stuff at the moment
297
298 -- | Yield a part of the vector without copying it. Safer version of
299 -- 'unsafeSlice'.
300 slice :: IVector v a => v a -> Int -- ^ starting index
301 -> Int -- ^ length
302 -> v a
303 {-# INLINE_STREAM slice #-}
304 slice v i n = assert (i >= 0 && n >= 0 && i+n <= length v)
305 $ unsafeSlice v i n
306
307 -- | Yield all but the last element without copying.
308 init :: IVector v a => v a -> v a
309 {-# INLINE_STREAM init #-}
310 init v = slice v 0 (length v - 1)
311
312 -- | All but the first element (without copying).
313 tail :: IVector v a => v a -> v a
314 {-# INLINE_STREAM tail #-}
315 tail v = slice v 1 (length v - 1)
316
317 -- | Yield the first @n@ elements without copying.
318 take :: IVector v a => Int -> v a -> v a
319 {-# INLINE_STREAM take #-}
320 take n v = slice v 0 (min n (length v))
321
322 -- | Yield all but the first @n@ elements without copying.
323 drop :: IVector v a => Int -> v a -> v a
324 {-# INLINE_STREAM drop #-}
325 drop n v = slice v n (max 0 (length v - n))
326
327 {-# RULES
328
329 "slice/new [IVector]" forall v p i n.
330 slice (new' v p) i n = new' v (New.slice p i n)
331
332 "init/new [IVector]" forall v p.
333 init (new' v p) = new' v (New.init p)
334
335 "tail/new [IVector]" forall v p.
336 tail (new' v p) = new' v (New.tail p)
337
338 "take/new [IVector]" forall n v p.
339 take n (new' v p) = new' v (New.take n p)
340
341 "drop/new [IVector]" forall n v p.
342 drop n (new' v p) = new' v (New.drop n p)
343
344 #-}
345
346 -- Permutations
347 -- ------------
348
349 accum :: IVector v a => (a -> b -> a) -> v a -> [(Int,b)] -> v a
350 {-# INLINE accum #-}
351 accum f v us = new (New.accum f (New.unstream (stream v))
352 (Stream.fromList us))
353
354 (//) :: IVector v a => v a -> [(Int, a)] -> v a
355 {-# INLINE (//) #-}
356 v // us = new (New.update (New.unstream (stream v))
357 (Stream.fromList us))
358
359 update :: (IVector v a, IVector v (Int, a)) => v a -> v (Int, a) -> v a
360 {-# INLINE update #-}
361 update v w = new (New.update (New.unstream (stream v)) (stream w))
362
363 bpermute :: (IVector v a, IVector v Int) => v a -> v Int -> v a
364 {-# INLINE bpermute #-}
365 bpermute v is = v `seq` map (v!) is
366
367 -- Mapping/zipping
368 -- ---------------
369
370 -- | Map a function over a vector
371 map :: (IVector v a, IVector v b) => (a -> b) -> v a -> v b
372 {-# INLINE map #-}
373 map f = unstream . Stream.map f . stream
374
375 inplace_map :: IVector v a => (a -> a) -> v a -> v a
376 {-# INLINE inplace_map #-}
377 inplace_map f = unstream . inplace (MStream.map f) . stream
378
379 {-# RULES
380
381 "map->inplace_map [IVector]" map = inplace_map
382
383 #-}
384
385 -- | Zip two vectors with the given function.
386 zipWith :: (IVector v a, IVector v b, IVector v c) => (a -> b -> c) -> v a -> v b -> v c
387 {-# INLINE zipWith #-}
388 zipWith f xs ys = unstream (Stream.zipWith f (stream xs) (stream ys))
389
390 zip :: (IVector v a, IVector v b, IVector v (a,b)) => v a -> v b -> v (a, b)
391 {-# INLINE zip #-}
392 zip = zipWith (,)
393
394 -- Comparisons
395 -- -----------
396
397 eq :: (IVector v a, Eq a) => v a -> v a -> Bool
398 {-# INLINE eq #-}
399 xs `eq` ys = stream xs == stream ys
400
401 cmp :: (IVector v a, Ord a) => v a -> v a -> Ordering
402 {-# INLINE cmp #-}
403 cmp xs ys = compare (stream xs) (stream ys)
404
405 -- Filtering
406 -- ---------
407
408 -- | Drop elements which do not satisfy the predicate
409 filter :: IVector v a => (a -> Bool) -> v a -> v a
410 {-# INLINE filter #-}
411 filter f = unstream . inplace (MStream.filter f) . stream
412
413 -- | Yield the longest prefix of elements satisfying the predicate.
414 takeWhile :: IVector v a => (a -> Bool) -> v a -> v a
415 {-# INLINE takeWhile #-}
416 takeWhile f = unstream . Stream.takeWhile f . stream
417
418 -- | Drop the longest prefix of elements that satisfy the predicate.
419 dropWhile :: IVector v a => (a -> Bool) -> v a -> v a
420 {-# INLINE dropWhile #-}
421 dropWhile f = unstream . Stream.dropWhile f . stream
422
423 -- Searching
424 -- ---------
425
426 infix 4 `elem`
427 -- | Check whether the vector contains an element
428 elem :: (IVector v a, Eq a) => a -> v a -> Bool
429 {-# INLINE elem #-}
430 elem x = Stream.elem x . stream
431
432 infix 4 `notElem`
433 -- | Inverse of `elem`
434 notElem :: (IVector v a, Eq a) => a -> v a -> Bool
435 {-# INLINE notElem #-}
436 notElem x = Stream.notElem x . stream
437
438 -- | Yield 'Just' the first element matching the predicate or 'Nothing' if no
439 -- such element exists.
440 find :: IVector v a => (a -> Bool) -> v a -> Maybe a
441 {-# INLINE find #-}
442 find f = Stream.find f . stream
443
444 -- | Yield 'Just' the index of the first element matching the predicate or
445 -- 'Nothing' if no such element exists.
446 findIndex :: IVector v a => (a -> Bool) -> v a -> Maybe Int
447 {-# INLINE findIndex #-}
448 findIndex f = Stream.findIndex f . stream
449
450 -- Folding
451 -- -------
452
453 -- | Left fold
454 foldl :: IVector v b => (a -> b -> a) -> a -> v b -> a
455 {-# INLINE foldl #-}
456 foldl f z = Stream.foldl f z . stream
457
458 -- | Lefgt fold on non-empty vectors
459 foldl1 :: IVector v a => (a -> a -> a) -> v a -> a
460 {-# INLINE foldl1 #-}
461 foldl1 f = Stream.foldl1 f . stream
462
463 -- | Left fold with strict accumulator
464 foldl' :: IVector v b => (a -> b -> a) -> a -> v b -> a
465 {-# INLINE foldl' #-}
466 foldl' f z = Stream.foldl' f z . stream
467
468 -- | Left fold on non-empty vectors with strict accumulator
469 foldl1' :: IVector v a => (a -> a -> a) -> v a -> a
470 {-# INLINE foldl1' #-}
471 foldl1' f = Stream.foldl1' f . stream
472
473 -- | Right fold
474 foldr :: IVector v a => (a -> b -> b) -> b -> v a -> b
475 {-# INLINE foldr #-}
476 foldr f z = Stream.foldr f z . stream
477
478 -- | Right fold on non-empty vectors
479 foldr1 :: IVector v a => (a -> a -> a) -> v a -> a
480 {-# INLINE foldr1 #-}
481 foldr1 f = Stream.foldr1 f . stream
482
483 -- Scans
484 -- -----
485
486 -- | Prefix scan
487 prescanl :: (IVector v a, IVector v b) => (a -> b -> a) -> a -> v b -> v a
488 {-# INLINE prescanl #-}
489 prescanl f z = unstream . Stream.prescanl f z . stream
490
491 inplace_prescanl :: IVector v a => (a -> a -> a) -> a -> v a -> v a
492 {-# INLINE inplace_prescanl #-}
493 inplace_prescanl f z = unstream . inplace (MStream.prescanl f z) . stream
494
495 {-# RULES
496
497 "prescanl -> inplace_prescanl [IVector]" prescanl = inplace_prescanl
498
499 #-}
500
501 -- | Prefix scan with strict accumulator
502 prescanl' :: (IVector v a, IVector v b) => (a -> b -> a) -> a -> v b -> v a
503 {-# INLINE prescanl' #-}
504 prescanl' f z = unstream . Stream.prescanl' f z . stream
505
506 inplace_prescanl' :: IVector v a => (a -> a -> a) -> a -> v a -> v a
507 {-# INLINE inplace_prescanl' #-}
508 inplace_prescanl' f z = unstream . inplace (MStream.prescanl' f z) . stream
509
510 {-# RULES
511
512 "prescanl' -> inplace_prescanl' [IVector]" prescanl' = inplace_prescanl'
513
514 #-}
515
516
517 -- | Convert a vector to a list
518 toList :: IVector v a => v a -> [a]
519 {-# INLINE toList #-}
520 toList = Stream.toList . stream
521
522 -- | Convert a list to a vector
523 fromList :: IVector v a => [a] -> v a
524 {-# INLINE fromList #-}
525 fromList = unstream . Stream.fromList
526