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