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