Add findIndices, elemIndex, elemIndices
[darcs-mirrors/vector.git] / Data / Vector / Primitive.hs
1 {-# LANGUAGE FlexibleInstances, MultiParamTypeClasses, TypeFamilies #-}
2
3 -- |
4 -- Module : Data.Vector.Primitive
5 -- Copyright : (c) Roman Leshchinskiy 2008-2009
6 -- License : BSD-style
7 --
8 -- Maintainer : Roman Leshchinskiy <rl@cse.unsw.edu.au>
9 -- Stability : experimental
10 -- Portability : non-portable
11 --
12 -- Unboxed vectors of primitive types.
13 --
14
15 module Data.Vector.Primitive (
16 Vector, MVector(..), Prim,
17
18 -- * Length information
19 length, null,
20
21 -- * Construction
22 empty, singleton, cons, snoc, replicate, (++), copy,
23
24 -- * Accessing individual elements
25 (!), head, last,
26
27 -- * Subvectors
28 slice, init, tail, take, drop,
29 unsafeSlice,
30
31 -- * Permutations
32 accum, accumulate_, (//), update_, backpermute, reverse,
33
34 -- * Mapping
35 map, imap, concatMap,
36
37 -- * Zipping and unzipping
38 zipWith, zipWith3, zipWith4, zipWith5, zipWith6,
39 izipWith, izipWith3, izipWith4, izipWith5, izipWith6,
40
41 -- * Filtering
42 filter, ifilter, takeWhile, dropWhile,
43
44 -- * Searching
45 elem, notElem, find, findIndex, findIndices, elemIndex, elemIndices,
46
47 -- * Folding
48 foldl, foldl1, foldl', foldl1', foldr, foldr1,
49 ifoldl, ifoldl', ifoldr,
50
51 -- * Specialised folds
52 sum, product, maximum, minimum, minIndex, maxIndex,
53
54 -- * Unfolding
55 unfoldr,
56
57 -- * Scans
58 prescanl, prescanl',
59 postscanl, postscanl',
60 scanl, scanl', scanl1, scanl1',
61
62 -- * Enumeration
63 enumFromTo, enumFromThenTo,
64
65 -- * Conversion to/from lists
66 toList, fromList,
67
68 -- * Unsafe operations
69 unsafeIndex,
70 unsafeAccum, unsafeAccumulate_,
71 unsafeUpd, unsafeUpdate_
72 ) where
73
74 import qualified Data.Vector.Generic as G
75 import Data.Vector.Primitive.Mutable ( MVector(..) )
76 import Data.Primitive.ByteArray
77 import Data.Primitive ( Prim )
78
79 import Control.Monad ( liftM )
80
81 import Prelude hiding ( length, null,
82 replicate, (++),
83 head, last,
84 init, tail, take, drop, reverse,
85 map, concatMap,
86 zipWith, zipWith3, zip, zip3, unzip, unzip3,
87 filter, takeWhile, dropWhile,
88 elem, notElem,
89 foldl, foldl1, foldr, foldr1,
90 sum, product, minimum, maximum,
91 scanl, scanl1,
92 enumFromTo, enumFromThenTo )
93
94 import qualified Prelude
95
96 -- | Unboxed vectors of primitive types
97 data Vector a = Vector {-# UNPACK #-} !Int
98 {-# UNPACK #-} !Int
99 {-# UNPACK #-} !ByteArray
100
101 instance (Show a, Prim a) => Show (Vector a) where
102 show = (Prelude.++ " :: Data.Vector.Primitive.Vector") . ("fromList " Prelude.++) . show . toList
103
104 type instance G.Mutable Vector = MVector
105
106 instance Prim a => G.Vector Vector a where
107 {-# INLINE unsafeFreeze #-}
108 unsafeFreeze (MVector i n marr)
109 = Vector i n `liftM` unsafeFreezeByteArray marr
110
111 {-# INLINE basicLength #-}
112 basicLength (Vector _ n _) = n
113
114 {-# INLINE basicUnsafeSlice #-}
115 basicUnsafeSlice (Vector i _ arr) j n = Vector (i+j) n arr
116
117 {-# INLINE basicUnsafeIndexM #-}
118 basicUnsafeIndexM (Vector i _ arr) j = return (indexByteArray arr (i+j))
119
120 {-# INLINE elemseq #-}
121 elemseq _ = seq
122
123 instance (Prim a, Eq a) => Eq (Vector a) where
124 {-# INLINE (==) #-}
125 (==) = G.eq
126
127 instance (Prim a, Ord a) => Ord (Vector a) where
128 {-# INLINE compare #-}
129 compare = G.cmp
130
131 -- Length
132 -- ------
133
134 length :: Prim a => Vector a -> Int
135 {-# INLINE length #-}
136 length = G.length
137
138 null :: Prim a => Vector a -> Bool
139 {-# INLINE null #-}
140 null = G.null
141
142 -- Construction
143 -- ------------
144
145 -- | Empty vector
146 empty :: Prim a => Vector a
147 {-# INLINE empty #-}
148 empty = G.empty
149
150 -- | Vector with exaclty one element
151 singleton :: Prim a => a -> Vector a
152 {-# INLINE singleton #-}
153 singleton = G.singleton
154
155 -- | Vector of the given length with the given value in each position
156 replicate :: Prim a => Int -> a -> Vector a
157 {-# INLINE replicate #-}
158 replicate = G.replicate
159
160 -- | Prepend an element
161 cons :: Prim a => a -> Vector a -> Vector a
162 {-# INLINE cons #-}
163 cons = G.cons
164
165 -- | Append an element
166 snoc :: Prim a => Vector a -> a -> Vector a
167 {-# INLINE snoc #-}
168 snoc = G.snoc
169
170 infixr 5 ++
171 -- | Concatenate two vectors
172 (++) :: Prim a => Vector a -> Vector a -> Vector a
173 {-# INLINE (++) #-}
174 (++) = (G.++)
175
176 -- | Create a copy of a vector. Useful when dealing with slices.
177 copy :: Prim a => Vector a -> Vector a
178 {-# INLINE copy #-}
179 copy = G.copy
180
181 -- Accessing individual elements
182 -- -----------------------------
183
184 -- | Indexing
185 (!) :: Prim a => Vector a -> Int -> a
186 {-# INLINE (!) #-}
187 (!) = (G.!)
188
189 -- | Unsafe indexing without bounds checks
190 unsafeIndex :: Prim a => Vector a -> Int -> a
191 {-# INLINE unsafeIndex #-}
192 unsafeIndex = G.unsafeIndex
193
194 -- | First element
195 head :: Prim a => Vector a -> a
196 {-# INLINE head #-}
197 head = G.head
198
199 -- | Last element
200 last :: Prim a => Vector a -> a
201 {-# INLINE last #-}
202 last = G.last
203
204 -- Subarrays
205 -- ---------
206
207 -- | Yield a part of the vector without copying it. Safer version of
208 -- 'basicUnsafeSlice'.
209 slice :: Prim a => Vector a -> Int -- ^ starting index
210 -> Int -- ^ length
211 -> Vector a
212 {-# INLINE slice #-}
213 slice = G.slice
214
215 -- | Unsafely yield a part of the vector without copying it and without
216 -- performing bounds checks.
217 unsafeSlice :: Prim a => Vector a -> Int -- ^ starting index
218 -> Int -- ^ length
219 -> Vector a
220 {-# INLINE unsafeSlice #-}
221 unsafeSlice = G.unsafeSlice
222
223 -- | Yield all but the last element without copying.
224 init :: Prim a => Vector a -> Vector a
225 {-# INLINE init #-}
226 init = G.init
227
228 -- | All but the first element (without copying).
229 tail :: Prim a => Vector a -> Vector a
230 {-# INLINE tail #-}
231 tail = G.tail
232
233 -- | Yield the first @n@ elements without copying.
234 take :: Prim a => Int -> Vector a -> Vector a
235 {-# INLINE take #-}
236 take = G.take
237
238 -- | Yield all but the first @n@ elements without copying.
239 drop :: Prim a => Int -> Vector a -> Vector a
240 {-# INLINE drop #-}
241 drop = G.drop
242
243 -- Permutations
244 -- ------------
245
246 unsafeAccum :: Prim a => (a -> b -> a) -> Vector a -> [(Int,b)] -> Vector a
247 {-# INLINE unsafeAccum #-}
248 unsafeAccum = G.unsafeAccum
249
250 unsafeAccumulate_ :: (Prim a, Prim b) =>
251 (a -> b -> a) -> Vector a -> Vector Int -> Vector b -> Vector a
252 {-# INLINE unsafeAccumulate_ #-}
253 unsafeAccumulate_ = G.unsafeAccumulate_
254
255 accum :: Prim a => (a -> b -> a) -> Vector a -> [(Int,b)] -> Vector a
256 {-# INLINE accum #-}
257 accum = G.accum
258
259 accumulate_ :: (Prim a, Prim b) =>
260 (a -> b -> a) -> Vector a -> Vector Int -> Vector b -> Vector a
261 {-# INLINE accumulate_ #-}
262 accumulate_ = G.accumulate_
263
264 unsafeUpd :: Prim a => Vector a -> [(Int, a)] -> Vector a
265 {-# INLINE unsafeUpd #-}
266 unsafeUpd = G.unsafeUpd
267
268 unsafeUpdate_ :: Prim a => Vector a -> Vector Int -> Vector a -> Vector a
269 {-# INLINE unsafeUpdate_ #-}
270 unsafeUpdate_ = G.unsafeUpdate_
271
272 (//) :: Prim a => Vector a -> [(Int, a)] -> Vector a
273 {-# INLINE (//) #-}
274 (//) = (G.//)
275
276 update_ :: Prim a => Vector a -> Vector Int -> Vector a -> Vector a
277 {-# INLINE update_ #-}
278 update_ = G.update_
279
280 backpermute :: Prim a => Vector a -> Vector Int -> Vector a
281 {-# INLINE backpermute #-}
282 backpermute = G.backpermute
283
284 reverse :: Prim a => Vector a -> Vector a
285 {-# INLINE reverse #-}
286 reverse = G.reverse
287
288 -- Mapping
289 -- -------
290
291 -- | Map a function over a vector
292 map :: (Prim a, Prim b) => (a -> b) -> Vector a -> Vector b
293 {-# INLINE map #-}
294 map = G.map
295
296 -- | Apply a function to every index/value pair
297 imap :: (Prim a, Prim b) => (Int -> a -> b) -> Vector a -> Vector b
298 {-# INLINE imap #-}
299 imap = G.imap
300
301 concatMap :: (Prim a, Prim b) => (a -> Vector b) -> Vector a -> Vector b
302 {-# INLINE concatMap #-}
303 concatMap = G.concatMap
304
305 -- Zipping/unzipping
306 -- -----------------
307
308 -- | Zip two vectors with the given function.
309 zipWith :: (Prim a, Prim b, Prim c)
310 => (a -> b -> c) -> Vector a -> Vector b -> Vector c
311 {-# INLINE zipWith #-}
312 zipWith = G.zipWith
313
314 -- | Zip three vectors with the given function.
315 zipWith3 :: (Prim a, Prim b, Prim c, Prim d)
316 => (a -> b -> c -> d) -> Vector a -> Vector b -> Vector c -> Vector d
317 {-# INLINE zipWith3 #-}
318 zipWith3 = G.zipWith3
319
320 zipWith4 :: (Prim a, Prim b, Prim c, Prim d, Prim e)
321 => (a -> b -> c -> d -> e)
322 -> Vector a -> Vector b -> Vector c -> Vector d -> Vector e
323 {-# INLINE zipWith4 #-}
324 zipWith4 = G.zipWith4
325
326 zipWith5 :: (Prim a, Prim b, Prim c, Prim d, Prim e, Prim f)
327 => (a -> b -> c -> d -> e -> f)
328 -> Vector a -> Vector b -> Vector c -> Vector d -> Vector e
329 -> Vector f
330 {-# INLINE zipWith5 #-}
331 zipWith5 = G.zipWith5
332
333 zipWith6 :: (Prim a, Prim b, Prim c, Prim d, Prim e, Prim f, Prim g)
334 => (a -> b -> c -> d -> e -> f -> g)
335 -> Vector a -> Vector b -> Vector c -> Vector d -> Vector e
336 -> Vector f -> Vector g
337 {-# INLINE zipWith6 #-}
338 zipWith6 = G.zipWith6
339
340 -- | Zip two vectors and their indices with the given function.
341 izipWith :: (Prim a, Prim b, Prim c)
342 => (Int -> a -> b -> c) -> Vector a -> Vector b -> Vector c
343 {-# INLINE izipWith #-}
344 izipWith = G.izipWith
345
346 -- | Zip three vectors and their indices with the given function.
347 izipWith3 :: (Prim a, Prim b, Prim c, Prim d)
348 => (Int -> a -> b -> c -> d)
349 -> Vector a -> Vector b -> Vector c -> Vector d
350 {-# INLINE izipWith3 #-}
351 izipWith3 = G.izipWith3
352
353 izipWith4 :: (Prim a, Prim b, Prim c, Prim d, Prim e)
354 => (Int -> a -> b -> c -> d -> e)
355 -> Vector a -> Vector b -> Vector c -> Vector d -> Vector e
356 {-# INLINE izipWith4 #-}
357 izipWith4 = G.izipWith4
358
359 izipWith5 :: (Prim a, Prim b, Prim c, Prim d, Prim e, Prim f)
360 => (Int -> a -> b -> c -> d -> e -> f)
361 -> Vector a -> Vector b -> Vector c -> Vector d -> Vector e
362 -> Vector f
363 {-# INLINE izipWith5 #-}
364 izipWith5 = G.izipWith5
365
366 izipWith6 :: (Prim a, Prim b, Prim c, Prim d, Prim e, Prim f, Prim g)
367 => (Int -> a -> b -> c -> d -> e -> f -> g)
368 -> Vector a -> Vector b -> Vector c -> Vector d -> Vector e
369 -> Vector f -> Vector g
370 {-# INLINE izipWith6 #-}
371 izipWith6 = G.izipWith6
372
373 -- Filtering
374 -- ---------
375
376 -- | Drop elements which do not satisfy the predicate
377 filter :: Prim a => (a -> Bool) -> Vector a -> Vector a
378 {-# INLINE filter #-}
379 filter = G.filter
380
381 -- | Drop elements that do not satisfy the predicate (applied to values and
382 -- their indices)
383 ifilter :: Prim a => (Int -> a -> Bool) -> Vector a -> Vector a
384 {-# INLINE ifilter #-}
385 ifilter = G.ifilter
386
387 -- | Yield the longest prefix of elements satisfying the predicate.
388 takeWhile :: Prim a => (a -> Bool) -> Vector a -> Vector a
389 {-# INLINE takeWhile #-}
390 takeWhile = G.takeWhile
391
392 -- | Drop the longest prefix of elements that satisfy the predicate.
393 dropWhile :: Prim a => (a -> Bool) -> Vector a -> Vector a
394 {-# INLINE dropWhile #-}
395 dropWhile = G.dropWhile
396
397 -- Searching
398 -- ---------
399
400 infix 4 `elem`
401 -- | Check whether the vector contains an element
402 elem :: (Prim a, Eq a) => a -> Vector a -> Bool
403 {-# INLINE elem #-}
404 elem = G.elem
405
406 infix 4 `notElem`
407 -- | Inverse of `elem`
408 notElem :: (Prim a, Eq a) => a -> Vector a -> Bool
409 {-# INLINE notElem #-}
410 notElem = G.notElem
411
412 -- | Yield 'Just' the first element matching the predicate or 'Nothing' if no
413 -- such element exists.
414 find :: Prim a => (a -> Bool) -> Vector a -> Maybe a
415 {-# INLINE find #-}
416 find = G.find
417
418 -- | Yield 'Just' the index of the first element matching the predicate or
419 -- 'Nothing' if no such element exists.
420 findIndex :: Prim a => (a -> Bool) -> Vector a -> Maybe Int
421 {-# INLINE findIndex #-}
422 findIndex = G.findIndex
423
424 -- | Yield the indices of elements satisfying the predicate
425 findIndices :: Prim a => (a -> Bool) -> Vector a -> Vector Int
426 {-# INLINE findIndices #-}
427 findIndices = G.findIndices
428
429 -- | Yield 'Just' the index of the first occurence of the given element or
430 -- 'Nothing' if the vector does not contain the element
431 elemIndex :: (Prim a, Eq a) => a -> Vector a -> Maybe Int
432 {-# INLINE elemIndex #-}
433 elemIndex = G.elemIndex
434
435 -- | Yield the indices of all occurences of the given element
436 elemIndices :: (Prim a, Eq a) => a -> Vector a -> Vector Int
437 {-# INLINE elemIndices #-}
438 elemIndices = G.elemIndices
439
440 -- Folding
441 -- -------
442
443 -- | Left fold
444 foldl :: Prim b => (a -> b -> a) -> a -> Vector b -> a
445 {-# INLINE foldl #-}
446 foldl = G.foldl
447
448 -- | Lefgt fold on non-empty vectors
449 foldl1 :: Prim a => (a -> a -> a) -> Vector a -> a
450 {-# INLINE foldl1 #-}
451 foldl1 = G.foldl1
452
453 -- | Left fold with strict accumulator
454 foldl' :: Prim b => (a -> b -> a) -> a -> Vector b -> a
455 {-# INLINE foldl' #-}
456 foldl' = G.foldl'
457
458 -- | Left fold on non-empty vectors with strict accumulator
459 foldl1' :: Prim a => (a -> a -> a) -> Vector a -> a
460 {-# INLINE foldl1' #-}
461 foldl1' = G.foldl1'
462
463 -- | Right fold
464 foldr :: Prim a => (a -> b -> b) -> b -> Vector a -> b
465 {-# INLINE foldr #-}
466 foldr = G.foldr
467
468 -- | Right fold on non-empty vectors
469 foldr1 :: Prim a => (a -> a -> a) -> Vector a -> a
470 {-# INLINE foldr1 #-}
471 foldr1 = G.foldr1
472
473 -- | Left fold (function applied to each element and its index)
474 ifoldl :: Prim b => (a -> Int -> b -> a) -> a -> Vector b -> a
475 {-# INLINE ifoldl #-}
476 ifoldl = G.ifoldl
477
478 -- | Left fold with strict accumulator (function applied to each element and
479 -- its index)
480 ifoldl' :: Prim b => (a -> Int -> b -> a) -> a -> Vector b -> a
481 {-# INLINE ifoldl' #-}
482 ifoldl' = G.ifoldl'
483
484 -- | Right fold (function applied to each element and its index)
485 ifoldr :: Prim a => (Int -> a -> b -> b) -> b -> Vector a -> b
486 {-# INLINE ifoldr #-}
487 ifoldr = G.ifoldr
488
489 -- Specialised folds
490 -- -----------------
491
492 sum :: (Prim a, Num a) => Vector a -> a
493 {-# INLINE sum #-}
494 sum = G.sum
495
496 product :: (Prim a, Num a) => Vector a -> a
497 {-# INLINE product #-}
498 product = G.product
499
500 maximum :: (Prim a, Ord a) => Vector a -> a
501 {-# INLINE maximum #-}
502 maximum = G.maximum
503
504 minimum :: (Prim a, Ord a) => Vector a -> a
505 {-# INLINE minimum #-}
506 minimum = G.minimum
507
508 minIndex :: (Prim a, Ord a) => Vector a -> Int
509 {-# INLINE minIndex #-}
510 minIndex = G.minIndex
511
512 maxIndex :: (Prim a, Ord a) => Vector a -> Int
513 {-# INLINE maxIndex #-}
514 maxIndex = G.maxIndex
515
516 -- Unfolding
517 -- ---------
518
519 unfoldr :: Prim a => (b -> Maybe (a, b)) -> b -> Vector a
520 {-# INLINE unfoldr #-}
521 unfoldr = G.unfoldr
522
523 -- Scans
524 -- -----
525
526 -- | Prefix scan
527 prescanl :: (Prim a, Prim b) => (a -> b -> a) -> a -> Vector b -> Vector a
528 {-# INLINE prescanl #-}
529 prescanl = G.prescanl
530
531 -- | Prefix scan with strict accumulator
532 prescanl' :: (Prim a, Prim b) => (a -> b -> a) -> a -> Vector b -> Vector a
533 {-# INLINE prescanl' #-}
534 prescanl' = G.prescanl'
535
536 -- | Suffix scan
537 postscanl :: (Prim a, Prim b) => (a -> b -> a) -> a -> Vector b -> Vector a
538 {-# INLINE postscanl #-}
539 postscanl = G.postscanl
540
541 -- | Suffix scan with strict accumulator
542 postscanl' :: (Prim a, Prim b) => (a -> b -> a) -> a -> Vector b -> Vector a
543 {-# INLINE postscanl' #-}
544 postscanl' = G.postscanl'
545
546 -- | Haskell-style scan
547 scanl :: (Prim a, Prim b) => (a -> b -> a) -> a -> Vector b -> Vector a
548 {-# INLINE scanl #-}
549 scanl = G.scanl
550
551 -- | Haskell-style scan with strict accumulator
552 scanl' :: (Prim a, Prim b) => (a -> b -> a) -> a -> Vector b -> Vector a
553 {-# INLINE scanl' #-}
554 scanl' = G.scanl'
555
556 -- | Scan over a non-empty 'Vector'
557 scanl1 :: Prim a => (a -> a -> a) -> Vector a -> Vector a
558 {-# INLINE scanl1 #-}
559 scanl1 = G.scanl1
560
561 -- | Scan over a non-empty 'Vector' with a strict accumulator
562 scanl1' :: Prim a => (a -> a -> a) -> Vector a -> Vector a
563 {-# INLINE scanl1' #-}
564 scanl1' = G.scanl1'
565
566 -- Enumeration
567 -- -----------
568
569 enumFromTo :: (Prim a, Enum a) => a -> a -> Vector a
570 {-# INLINE enumFromTo #-}
571 enumFromTo = G.enumFromTo
572
573 enumFromThenTo :: (Prim a, Enum a) => a -> a -> a -> Vector a
574 {-# INLINE enumFromThenTo #-}
575 enumFromThenTo = G.enumFromThenTo
576
577 -- Conversion to/from lists
578 -- ------------------------
579
580 -- | Convert a vector to a list
581 toList :: Prim a => Vector a -> [a]
582 {-# INLINE toList #-}
583 toList = G.toList
584
585 -- | Convert a list to a vector
586 fromList :: Prim a => [a] -> Vector a
587 {-# INLINE fromList #-}
588 fromList = G.fromList
589