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