Documentation changes
[darcs-mirrors/vector.git] / Data / Vector / Generic.hs
1 {-# LANGUAGE Rank2Types, MultiParamTypeClasses, FlexibleContexts,
2 TypeFamilies, ScopedTypeVariables #-}
3 -- |
4 -- Module : Data.Vector.Generic
5 -- Copyright : (c) Roman Leshchinskiy 2008-2010
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 module Data.Vector.Generic (
16 -- * Immutable vectors
17 Vector(..), Mutable,
18
19 -- * Accessors
20
21 -- ** Length information
22 length, null,
23
24 -- ** Indexing
25 (!), head, last,
26 unsafeIndex, unsafeHead, unsafeLast,
27
28 -- ** Monadic indexing
29 indexM, headM, lastM,
30 unsafeIndexM, unsafeHeadM, unsafeLastM,
31
32 -- ** Extracting subvectors (slicing)
33 slice, init, tail, take, drop,
34 unsafeSlice, unsafeInit, unsafeTail, unsafeTake, unsafeDrop,
35
36 -- * Construction
37
38 -- ** Initialisation
39 empty, singleton, replicate, generate,
40
41 -- ** Monadic initialisation
42 replicateM, create,
43
44 -- ** Unfolding
45 unfoldr, unfoldrN,
46
47 -- ** Enumeration
48 enumFromN, enumFromStepN, enumFromTo, enumFromThenTo,
49
50 -- ** Concatenation
51 cons, snoc, (++),
52
53 -- ** Restricting memory usage
54 force,
55
56 -- * Modifying vectors
57
58 -- ** Bulk updates
59 (//), update, update_,
60 unsafeUpd, unsafeUpdate, unsafeUpdate_,
61
62 -- ** Accumulations
63 accum, accumulate, accumulate_,
64 unsafeAccum, unsafeAccumulate, unsafeAccumulate_,
65
66 -- ** Permutations
67 reverse, backpermute, unsafeBackpermute,
68
69 -- ** Safe destructive updates
70 modify,
71
72 -- * Elementwise operations
73
74 -- ** Mapping
75 map, imap, concatMap,
76
77 -- ** Monadic mapping
78 mapM, mapM_, forM, forM_,
79
80 -- ** Zipping
81 zipWith, zipWith3, zipWith4, zipWith5, zipWith6,
82 izipWith, izipWith3, izipWith4, izipWith5, izipWith6,
83 zip, zip3, zip4, zip5, zip6,
84
85 -- ** Monadic zipping
86 zipWithM, zipWithM_,
87
88 -- ** Unzipping
89 unzip, unzip3, unzip4, unzip5, unzip6,
90
91 -- * Working with predicates
92
93 -- ** Filtering
94 filter, ifilter, filterM,
95 takeWhile, dropWhile,
96
97 -- ** Partitioning
98 partition, unstablePartition, span, break,
99
100 -- ** Searching
101 elem, notElem, find, findIndex, findIndices, elemIndex, elemIndices,
102
103 -- * Folding
104 foldl, foldl1, foldl', foldl1', foldr, foldr1, foldr', foldr1',
105 ifoldl, ifoldl', ifoldr, ifoldr',
106
107 -- ** Specialised folds
108 all, any, and, or,
109 sum, product,
110 maximum, maximumBy, minimum, minimumBy,
111 minIndex, minIndexBy, maxIndex, maxIndexBy,
112
113 -- ** Monadic folds
114 foldM, foldM', fold1M, fold1M',
115
116 -- * Prefix sums (scans)
117 prescanl, prescanl',
118 postscanl, postscanl',
119 scanl, scanl', scanl1, scanl1',
120 prescanr, prescanr',
121 postscanr, postscanr',
122 scanr, scanr', scanr1, scanr1',
123
124 -- * Conversions
125
126 -- ** Lists
127 toList, fromList, fromListN,
128
129 -- ** Mutable vectors
130 copy, unsafeCopy,
131
132 -- * Fusion support
133
134 -- ** Conversion to/from Streams
135 stream, unstream, streamR, unstreamR,
136
137 -- ** Recycling support
138 new, clone,
139
140 -- * Utilities
141
142 -- ** Comparisons
143 eq, cmp,
144
145 -- ** @Data@ and @Typeable@
146 gfoldl, dataCast, mkType
147 ) where
148
149 import Data.Vector.Generic.Base
150
151 import Data.Vector.Generic.Mutable ( MVector )
152 import qualified Data.Vector.Generic.Mutable as M
153
154 import qualified Data.Vector.Generic.New as New
155 import Data.Vector.Generic.New ( New )
156
157 import qualified Data.Vector.Fusion.Stream as Stream
158 import Data.Vector.Fusion.Stream ( Stream, MStream, inplace )
159 import qualified Data.Vector.Fusion.Stream.Monadic as MStream
160 import Data.Vector.Fusion.Stream.Size
161 import Data.Vector.Fusion.Util
162
163 import Control.Monad.ST ( ST, runST )
164 import Control.Monad.Primitive
165 import qualified Control.Monad as Monad
166 import Prelude hiding ( length, null,
167 replicate, (++),
168 head, last,
169 init, tail, take, drop, reverse,
170 map, concatMap,
171 zipWith, zipWith3, zip, zip3, unzip, unzip3,
172 filter, takeWhile, dropWhile, span, break,
173 elem, notElem,
174 foldl, foldl1, foldr, foldr1,
175 all, any, and, or, sum, product, maximum, minimum,
176 scanl, scanl1, scanr, scanr1,
177 enumFromTo, enumFromThenTo,
178 mapM, mapM_ )
179
180 import Data.Typeable ( Typeable1, gcast1 )
181 import Data.Data ( Data, DataType, mkNorepType )
182
183 #include "vector.h"
184
185 -- Length information
186 -- ------------------
187
188 -- | Yield the length of the vector
189 length :: Vector v a => v a -> Int
190 {-# INLINE_STREAM length #-}
191 length v = basicLength v
192
193 {-# RULES
194
195 "length/unstream [Vector]" forall s.
196 length (new (New.unstream s)) = Stream.length s
197
198 #-}
199
200 -- | Test whether a vector if empty
201 null :: Vector v a => v a -> Bool
202 {-# INLINE_STREAM null #-}
203 null v = basicLength v == 0
204
205 {-# RULES
206
207 "null/unstream [Vector]" forall s.
208 null (new (New.unstream s)) = Stream.null s
209
210 #-}
211
212 -- Indexing
213 -- --------
214
215 -- | Indexing
216 (!) :: Vector v a => v a -> Int -> a
217 {-# INLINE_STREAM (!) #-}
218 v ! i = BOUNDS_CHECK(checkIndex) "(!)" i (length v)
219 $ unId (basicUnsafeIndexM v i)
220
221 -- | First element
222 head :: Vector v a => v a -> a
223 {-# INLINE_STREAM head #-}
224 head v = v ! 0
225
226 -- | Last element
227 last :: Vector v a => v a -> a
228 {-# INLINE_STREAM last #-}
229 last v = v ! (length v - 1)
230
231 -- | Unsafe indexing without bounds checking
232 unsafeIndex :: Vector v a => v a -> Int -> a
233 {-# INLINE_STREAM unsafeIndex #-}
234 unsafeIndex v i = UNSAFE_CHECK(checkIndex) "unsafeIndex" i (length v)
235 $ unId (basicUnsafeIndexM v i)
236
237 -- | Yield the first element of a vector without checking if the vector is
238 -- empty
239 unsafeHead :: Vector v a => v a -> a
240 {-# INLINE_STREAM unsafeHead #-}
241 unsafeHead v = unsafeIndex v 0
242
243 -- | Yield the last element of a vector without checking if the vector is
244 -- empty
245 unsafeLast :: Vector v a => v a -> a
246 {-# INLINE_STREAM unsafeLast #-}
247 unsafeLast v = unsafeIndex v (length v - 1)
248
249 {-# RULES
250
251 "(!)/unstream [Vector]" forall i s.
252 new (New.unstream s) ! i = s Stream.!! i
253
254 "head/unstream [Vector]" forall s.
255 head (new (New.unstream s)) = Stream.head s
256
257 "last/unstream [Vector]" forall s.
258 last (new (New.unstream s)) = Stream.last s
259
260 "unsafeIndex/unstream [Vector]" forall i s.
261 unsafeIndex (new (New.unstream s)) i = s Stream.!! i
262
263 "unsafeHead/unstream [Vector]" forall s.
264 unsafeHead (new (New.unstream s)) = Stream.head s
265
266 "unsafeLast/unstream [Vector]" forall s.
267 unsafeLast (new (New.unstream s)) = Stream.last s
268
269 #-}
270
271 -- Monadic indexing
272 -- ----------------
273
274 -- | Indexing in a monad.
275 --
276 -- The monad allows operations to be strict in the vector when necessary.
277 -- Suppose vector copying is implemented like this:
278 --
279 -- > copy mv v = ... write mv i (v ! i) ...
280 --
281 -- For lazy vectors, @v ! i@ would not be evaluated which means that @mv@
282 -- would unnecessarily retain a reference to @v@ in each element written.
283 --
284 -- With 'indexM', copying can be implemented like this instead:
285 --
286 -- > copy mv v = ... do
287 -- > x <- indexM v i
288 -- > write mv i x
289 --
290 -- Here, no references to @v@ are retained because indexing (but /not/ the
291 -- elements) is evaluated eagerly.
292 --
293 indexM :: (Vector v a, Monad m) => v a -> Int -> m a
294 {-# INLINE_STREAM indexM #-}
295 indexM v i = BOUNDS_CHECK(checkIndex) "indexM" i (length v)
296 $ basicUnsafeIndexM v i
297
298 -- | Yield the first element of a vector in a monad. See 'indexM' for an
299 -- explanation of why this is useful.
300 headM :: (Vector v a, Monad m) => v a -> m a
301 {-# INLINE_STREAM headM #-}
302 headM v = indexM v 0
303
304 -- | Yield the last element of a vector in a monad. See 'indexM' for an
305 -- explanation of why this is useful.
306 lastM :: (Vector v a, Monad m) => v a -> m a
307 {-# INLINE_STREAM lastM #-}
308 lastM v = indexM v (length v - 1)
309
310 -- | Indexing in a monad without bounds checks. See 'indexM' for an
311 -- explanation of why this is useful.
312 unsafeIndexM :: (Vector v a, Monad m) => v a -> Int -> m a
313 {-# INLINE_STREAM unsafeIndexM #-}
314 unsafeIndexM v i = UNSAFE_CHECK(checkIndex) "unsafeIndexM" i (length v)
315 $ basicUnsafeIndexM v i
316
317 -- | Yield the first element in a monad without checking for empty vectors.
318 -- See 'indexM' for an explanation of why this is useful.
319 unsafeHeadM :: (Vector v a, Monad m) => v a -> m a
320 {-# INLINE_STREAM unsafeHeadM #-}
321 unsafeHeadM v = unsafeIndexM v 0
322
323 -- | Yield the last element in a monad without checking for empty vectors.
324 -- See 'indexM' for an explanation of why this is useful.
325 unsafeLastM :: (Vector v a, Monad m) => v a -> m a
326 {-# INLINE_STREAM unsafeLastM #-}
327 unsafeLastM v = unsafeIndexM v (length v - 1)
328
329 -- FIXME: the rhs of these rules are lazy in the stream which is WRONG
330 {- RULES
331
332 "indexM/unstream [Vector]" forall v i s.
333 indexM (new' v (New.unstream s)) i = return (s Stream.!! i)
334
335 "headM/unstream [Vector]" forall v s.
336 headM (new' v (New.unstream s)) = return (Stream.head s)
337
338 "lastM/unstream [Vector]" forall v s.
339 lastM (new' v (New.unstream s)) = return (Stream.last s)
340
341 -}
342
343 -- Extracting subvectors (slicing)
344 -- -------------------------------
345
346 -- | Yield a slice of the vector without copying it. The vector must contain
347 -- at least (starting index + length) elements.
348 slice :: Vector v a => Int -- ^ starting index
349 -> Int -- ^ length
350 -> v a
351 -> v a
352 {-# INLINE_STREAM slice #-}
353 slice i n v = BOUNDS_CHECK(checkSlice) "slice" i n (length v)
354 $ basicUnsafeSlice i n v
355
356 -- | Yield all but the last element without copying. The vector may not be
357 -- empty.
358 init :: Vector v a => v a -> v a
359 {-# INLINE_STREAM init #-}
360 init v = slice 0 (length v - 1) v
361
362 -- | Yield all but the first element without copying. The vector may not be
363 -- empty.
364 tail :: Vector v a => v a -> v a
365 {-# INLINE_STREAM tail #-}
366 tail v = slice 1 (length v - 1) v
367
368 -- | Yield at the first @n@ elements without copying. The vector may contain
369 -- less than @n@ elements in which case the operation is an identity.
370 take :: Vector v a => Int -> v a -> v a
371 {-# INLINE_STREAM take #-}
372 take n v = unsafeSlice 0 (delay_inline min n' (length v)) v
373 where n' = max n 0
374
375 -- | Yield all but the first @n@ elements without copying. The vector may
376 -- contain less than @n@ elements in which case an empty vector is returned.
377 drop :: Vector v a => Int -> v a -> v a
378 {-# INLINE_STREAM drop #-}
379 drop n v = unsafeSlice (delay_inline min n' len)
380 (delay_inline max 0 (len - n')) v
381 where n' = max n 0
382 len = length v
383
384 -- | Yield a slice of the vector without copying. The vector must contain
385 -- at least (starting index + length) elements but this is not checked.
386 unsafeSlice :: Vector v a => Int -- ^ starting index
387 -> Int -- ^ length
388 -> v a
389 -> v a
390 {-# INLINE_STREAM unsafeSlice #-}
391 unsafeSlice i n v = UNSAFE_CHECK(checkSlice) "unsafeSlice" i n (length v)
392 $ basicUnsafeSlice i n v
393
394 -- | Yield all but the last element without copying. The vector may not be
395 -- empty but this is not checked.
396 unsafeInit :: Vector v a => v a -> v a
397 {-# INLINE_STREAM unsafeInit #-}
398 unsafeInit v = unsafeSlice 0 (length v - 1) v
399
400 -- | Yield all but the first element without copying. The vector may not be
401 -- empty but this is not checked.
402 unsafeTail :: Vector v a => v a -> v a
403 {-# INLINE_STREAM unsafeTail #-}
404 unsafeTail v = unsafeSlice 1 (length v - 1) v
405
406 -- | Yield the first @n@ elements without copying. The vector must contain at
407 -- least @n@ elements but this is not checked.
408 unsafeTake :: Vector v a => Int -> v a -> v a
409 {-# INLINE unsafeTake #-}
410 unsafeTake n v = unsafeSlice 0 n v
411
412 -- | Yield all but the first @n@ elements without copying. The vector must
413 -- contain at least @n@ elements but this is not checked.
414 unsafeDrop :: Vector v a => Int -> v a -> v a
415 {-# INLINE unsafeDrop #-}
416 unsafeDrop n v = unsafeSlice n (length v - n) v
417
418 {-# RULES
419
420 "slice/new [Vector]" forall i n p.
421 slice i n (new p) = new (New.slice i n p)
422
423 "init/new [Vector]" forall p.
424 init (new p) = new (New.init p)
425
426 "tail/new [Vector]" forall p.
427 tail (new p) = new (New.tail p)
428
429 "take/new [Vector]" forall n p.
430 take n (new p) = new (New.take n p)
431
432 "drop/new [Vector]" forall n p.
433 drop n (new p) = new (New.drop n p)
434
435 "unsafeSlice/new [Vector]" forall i n p.
436 unsafeSlice i n (new p) = new (New.unsafeSlice i n p)
437
438 "unsafeInit/new [Vector]" forall p.
439 unsafeInit (new p) = new (New.unsafeInit p)
440
441 "unsafeTail/new [Vector]" forall p.
442 unsafeTail (new p) = new (New.unsafeTail p)
443
444 #-}
445
446 -- Initialisation
447 -- --------------
448
449 -- | Empty vector
450 empty :: Vector v a => v a
451 {-# INLINE empty #-}
452 empty = unstream Stream.empty
453
454 -- | Vector with exactly one element
455 singleton :: forall v a. Vector v a => a -> v a
456 {-# INLINE singleton #-}
457 singleton x = elemseq (undefined :: v a) x
458 $ unstream (Stream.singleton x)
459
460 -- | Vector of the given length with the same value in each position
461 replicate :: forall v a. Vector v a => Int -> a -> v a
462 {-# INLINE replicate #-}
463 replicate n x = elemseq (undefined :: v a) x
464 $ unstream
465 $ Stream.replicate n x
466
467 -- | Generate a vector of the given length by applying the function to each
468 -- index
469 generate :: Vector v a => Int -> (Int -> a) -> v a
470 {-# INLINE generate #-}
471 generate n f = unstream (Stream.generate n f)
472
473 -- Unfolding
474 -- ---------
475
476 -- | Unfold
477 unfoldr :: Vector v a => (b -> Maybe (a, b)) -> b -> v a
478 {-# INLINE unfoldr #-}
479 unfoldr f = unstream . Stream.unfoldr f
480
481 -- | Unfoldr at most @n@ elements.
482 unfoldrN :: Vector v a => Int -> (b -> Maybe (a, b)) -> b -> v a
483 {-# INLINE unfoldrN #-}
484 unfoldrN n f = unstream . Stream.unfoldrN n f
485
486 -- Enumeration
487 -- -----------
488
489 -- | Yield a vector of the given length containing the values @x@, @x+1@ etc.
490 -- This operation is usually more efficient than 'enumFromTo'.
491 enumFromN :: (Vector v a, Num a) => a -> Int -> v a
492 {-# INLINE enumFromN #-}
493 enumFromN x n = enumFromStepN x 1 n
494
495 -- | Yield a vector of the given length containing the values @x@, @x+y@,
496 -- @x+y+y@ etc. This operations is usually more efficient than
497 -- 'enumFromThenTo'.
498 enumFromStepN :: forall v a. (Vector v a, Num a) => a -> a -> Int -> v a
499 {-# INLINE enumFromStepN #-}
500 enumFromStepN x y n = elemseq (undefined :: v a) x
501 $ elemseq (undefined :: v a) y
502 $ unstream
503 $ Stream.enumFromStepN x y n
504
505 -- | Enumerate values from @x@ to @y@.
506 --
507 -- /WARNING:/ This operation can be very inefficient. If at all possible, use
508 -- 'enumFromN' instead.
509 enumFromTo :: (Vector v a, Enum a) => a -> a -> v a
510 {-# INLINE enumFromTo #-}
511 enumFromTo x y = unstream (Stream.enumFromTo x y)
512
513 -- | Enumerate values from @x@ to @y@ with a specific step @z@.
514 --
515 -- /WARNING:/ This operation can be very inefficient. If at all possible, use
516 -- 'enumFromStepN' instead.
517 enumFromThenTo :: (Vector v a, Enum a) => a -> a -> a -> v a
518 {-# INLINE enumFromThenTo #-}
519 enumFromThenTo x y z = unstream (Stream.enumFromThenTo x y z)
520
521 -- Concatenation
522 -- -------------
523
524 -- | Prepend an element
525 cons :: forall v a. Vector v a => a -> v a -> v a
526 {-# INLINE cons #-}
527 cons x v = elemseq (undefined :: v a) x
528 $ unstream
529 $ Stream.cons x
530 $ stream v
531
532 -- | Append an element
533 snoc :: forall v a. Vector v a => v a -> a -> v a
534 {-# INLINE snoc #-}
535 snoc v x = elemseq (undefined :: v a) x
536 $ unstream
537 $ Stream.snoc (stream v) x
538
539 infixr 5 ++
540 -- | Concatenate two vectors
541 (++) :: Vector v a => v a -> v a -> v a
542 {-# INLINE (++) #-}
543 v ++ w = unstream (stream v Stream.++ stream w)
544
545 -- Monadic initialisation
546 -- ----------------------
547
548 -- | Perform the monadic action the given number of times and store the
549 -- results in a vector.
550 replicateM :: (Monad m, Vector v a) => Int -> m a -> m (v a)
551 -- FIXME: specialise for ST and IO?
552 {-# INLINE replicateM #-}
553 replicateM n m = fromListN n `Monad.liftM` Monad.replicateM n m
554
555 -- | Execute the monadic action and freeze the resulting vector.
556 create :: Vector v a => (forall s. ST s (Mutable v s a)) -> v a
557 {-# INLINE create #-}
558 create p = new (New.create p)
559
560 -- Restricting memory usage
561 -- ------------------------
562
563 -- | Yields its argument but forces it not to retain any extra memory.
564 --
565 -- This is especially useful when dealing with slices. For example:
566 --
567 -- > force (slice 0 2 <huge vector>)
568 --
569 -- Here, the slice retains a reference to the huge vector. Forcing it creates
570 -- a copy of just the elements that belong to the slice and allows the huge
571 -- vector to be garbage collected.
572 force :: Vector v a => v a -> v a
573 -- FIXME: we probably ought to inline this later as the rules still might fire
574 -- otherwise
575 {-# INLINE_STREAM force #-}
576 force v = new (clone v)
577
578 -- Bulk updates
579 -- ------------
580
581 -- | For each pair @(i,a)@ from the list, replace the vector element at
582 -- position @i@ by @a@.
583 --
584 -- > <5,9,2,7> // [(2,1),(0,3),(2,8)] = <3,9,8,7>
585 --
586 (//) :: Vector v a => v a -> [(Int, a)] -> v a
587 {-# INLINE (//) #-}
588 v // us = update_stream v (Stream.fromList us)
589
590 -- | For each pair @(i,a)@ from the vector of index/value pairs, replace the
591 -- vector element at position @i@ by @a@.
592 --
593 -- > update <5,9,2,7> <(2,1),(0,3),(2,8)> = <3,9,8,7>
594 --
595 update :: (Vector v a, Vector v (Int, a)) => v a -> v (Int, a) -> v a
596 {-# INLINE update #-}
597 update v w = update_stream v (stream w)
598
599 -- | For each index @i@ from the index vector and the corresponding value @a@
600 -- from the update vector, replace the element of the value vector at position
601 -- @i@ by @a@.
602 --
603 -- > update_ <5,9,2,7> <2,0,2> <1,3,8> = <3,9,8,7>
604 --
605 -- This function is useful for instances of 'Vector' that cannot store pairs.
606 -- Otherwise, 'update' is probably more convenient.
607 --
608 -- > update_ xs is ys = update xs (zip is ys)
609 --
610 update_ :: (Vector v a, Vector v Int) => v a -- ^ value vector
611 -> v Int -- ^ index vector
612 -> v a -- ^ update vector
613 -> v a
614 {-# INLINE update_ #-}
615 update_ v is w = update_stream v (Stream.zipWith (,) (stream is) (stream w))
616
617 update_stream :: Vector v a => v a -> Stream (Int,a) -> v a
618 {-# INLINE update_stream #-}
619 update_stream = modifyWithStream M.update
620
621 -- | For each pair @(i,a)@ from the list, replace the vector element at
622 -- position @i@ by @a@. The indices are not checked. The safe version of this
623 -- function is ('//').
624 --
625 -- > unsafeUpd <5,9,2,7> [(2,1),(0,3),(2,8)] = <3,9,8,7>
626 --
627 unsafeUpd :: Vector v a => v a -> [(Int, a)] -> v a
628 {-# INLINE unsafeUpd #-}
629 unsafeUpd v us = unsafeUpdate_stream v (Stream.fromList us)
630
631 -- | For each pair @(i,a)@ from the vector of index/value pairs, replace the
632 -- vector element at position @i@ by @a@. The indices are not checked.
633 --
634 -- > unsafeUpdate <5,9,2,7> <(2,1),(0,3),(2,8)> = <3,9,8,7>
635 --
636 unsafeUpdate :: (Vector v a, Vector v (Int, a)) => v a -> v (Int, a) -> v a
637 {-# INLINE unsafeUpdate #-}
638 unsafeUpdate v w = unsafeUpdate_stream v (stream w)
639
640 -- | For each index @i@ from the index vector and the corresponding value @a@
641 -- from the update vector, replace the element of the value vector at position
642 -- @i@ by @a@. The indices are not checked.
643 --
644 -- > unsafeUpdate_ <5,9,2,7> <2,0,2> <1,3,8> = <3,9,8,7>
645 --
646 -- This function is useful for instances of 'Vector' that cannot store pairs.
647 -- Otherwise, 'unsafeUpdate' is probably more convenient.
648 --
649 -- > unsafeUpdate_ xs is ys = unsafeUpdate xs (zip is ys)
650 --
651 unsafeUpdate_ :: (Vector v a, Vector v Int) => v a -> v Int -> v a -> v a
652 {-# INLINE unsafeUpdate_ #-}
653 unsafeUpdate_ v is w
654 = unsafeUpdate_stream v (Stream.zipWith (,) (stream is) (stream w))
655
656 unsafeUpdate_stream :: Vector v a => v a -> Stream (Int,a) -> v a
657 {-# INLINE unsafeUpdate_stream #-}
658 unsafeUpdate_stream = modifyWithStream M.unsafeUpdate
659
660 -- Accumulations
661 -- -------------
662
663 -- | For each pair @(i,b)@ from the list, replace the vector element @a@ at
664 -- position @i@ by @f a b@ where @f@ is the given accumulating function.
665 --
666 -- > accum (+) <5,9,2> [(2,4),(1,6),(0,3),(1,7)] = <5+3, 9+6+7, 2+4>
667 accum :: Vector v a => (a -> b -> a) -> v a -> [(Int,b)] -> v a
668 {-# INLINE accum #-}
669 accum f v us = accum_stream f v (Stream.fromList us)
670
671 -- | For each pair @(i,b)@ from the vector of pairs, replace the vector element
672 -- @a@ at position @i@ by @f a b@ where @f@ is the given accumulating function.
673 --
674 -- > accumulate (+) <5,9,2> <(2,4),(1,6),(0,3),(1,7)> = <5+3, 9+6+7, 2+4>
675 accumulate :: (Vector v a, Vector v (Int, b))
676 => (a -> b -> a) -> v a -> v (Int,b) -> v a
677 {-# INLINE accumulate #-}
678 accumulate f v us = accum_stream f v (stream us)
679
680 -- | For each index @i@ from the vector of indices (@v Int@) and the
681 -- corresponding value @b@ from the the vector of values to accumulate (@v b@),
682 -- replace the element of the initial vector (@v a@) at
683 -- position @i@ by @f a b@ where @f@ is the given accumulating function.
684 --
685 -- > accumulate_ (+) <5,9,2> <2,1,0,1> <4,6,3,7> = <5+3, 9+6+7, 2+4>
686 --
687 -- This function is useful for instances of 'Vector' that cannot store pairs.
688 -- Otherwise, 'accumulate' is probably more convenient:
689 --
690 -- > accumulate_ f as is bs = accumulate f as (zip is bs)
691 --
692 accumulate_ :: (Vector v a, Vector v Int, Vector v b)
693 => (a -> b -> a) -- ^ accumulating function
694 -> v a -- ^ initial values
695 -> v Int -- ^ indices
696 -> v b -- ^ values to accumulate
697 -> v a
698 {-# INLINE accumulate_ #-}
699 accumulate_ f v is xs = accum_stream f v (Stream.zipWith (,) (stream is)
700 (stream xs))
701
702
703 accum_stream :: Vector v a => (a -> b -> a) -> v a -> Stream (Int,b) -> v a
704 {-# INLINE accum_stream #-}
705 accum_stream f = modifyWithStream (M.accum f)
706
707 -- | For each pair @(i,b)@ from the list, replace the vector element @a@ at
708 -- position @i@ by @f a b@ where @f@ is the given accumulating function. The
709 -- indices are not checked.
710 unsafeAccum :: Vector v a => (a -> b -> a) -> v a -> [(Int,b)] -> v a
711 {-# INLINE unsafeAccum #-}
712 unsafeAccum f v us = unsafeAccum_stream f v (Stream.fromList us)
713
714 -- | For each pair @(i,b)@ from the vector of pairs, replace the vector element
715 -- @a@ at position @i@ by @f a b@ where @f@ is the given accumulating function.
716 -- The indices are not checked.
717 unsafeAccumulate :: (Vector v a, Vector v (Int, b))
718 => (a -> b -> a) -> v a -> v (Int,b) -> v a
719 {-# INLINE unsafeAccumulate #-}
720 unsafeAccumulate f v us = unsafeAccum_stream f v (stream us)
721
722 -- | For each index @i@ from the vector of indices (@v Int@) and the
723 -- corresponding value @b@ from the the vector of values to accumulate (@v b@),
724 -- replace the element of the initial vector (@v a@) at
725 -- position @i@ by @f a b@ where @f@ is the given accumulating function. The
726 -- indices are not checked.
727 --
728 -- This function is useful for instances of 'Vector' that cannot store pairs.
729 -- In other cases, 'unsafeAccumulate' is probably more convenient.
730 unsafeAccumulate_ :: (Vector v a, Vector v Int, Vector v b)
731 => (a -> b -> a) -> v a -> v Int -> v b -> v a
732 {-# INLINE unsafeAccumulate_ #-}
733 unsafeAccumulate_ f v is xs
734 = unsafeAccum_stream f v (Stream.zipWith (,) (stream is) (stream xs))
735
736 unsafeAccum_stream
737 :: Vector v a => (a -> b -> a) -> v a -> Stream (Int,b) -> v a
738 {-# INLINE unsafeAccum_stream #-}
739 unsafeAccum_stream f = modifyWithStream (M.unsafeAccum f)
740
741 -- Permutations
742 -- ------------
743
744 -- | Reverse a vector
745 reverse :: (Vector v a) => v a -> v a
746 {-# INLINE reverse #-}
747 -- FIXME: make this fuse better, add support for recycling
748 reverse = unstream . streamR
749
750 -- | Given a vector of values @xs@ and a vector of indices @is@, yield a vector
751 -- obtained by replacing each @i@ from @is@ by @xs!i@.
752 --
753 -- > backpermute xs is = map (v!) is
754 --
755 -- > backpermute <a,b,c,d> <0,3,2,3,1,0> = <a,d,c,d,b,a>
756 --
757 backpermute :: (Vector v a, Vector v Int) => v a -> v Int -> v a
758 {-# INLINE backpermute #-}
759 -- This somewhat non-intuitive definition ensures that the resulting vector
760 -- does not retain references to the original one even if it is lazy in its
761 -- elements. This would not be the case if we simply used map (v!)
762 backpermute v is = seq v
763 $ unstream
764 $ Stream.unbox
765 $ Stream.map (indexM v)
766 $ stream is
767
768 -- | Given a vector of values @xs@ and a vector of indices @is@, yield a vector
769 -- obtained by replacing each @i@ from @is@ by @xs!i@. The indices are not
770 -- checked.
771 --
772 -- > unsafeBackpermute xs is = map (unsafeIndex v) is
773 --
774 -- > unsafeBackpermute <a,b,c,d> <0,3,2,3,1,0> = <a,d,c,d,b,a>
775 --
776 unsafeBackpermute :: (Vector v a, Vector v Int) => v a -> v Int -> v a
777 {-# INLINE unsafeBackpermute #-}
778 unsafeBackpermute v is = seq v
779 $ unstream
780 $ Stream.unbox
781 $ Stream.map (unsafeIndexM v)
782 $ stream is
783
784 -- Safe destructive updates
785 -- ------------------------
786
787 -- | Apply a destructive operation to a vector. The operation will be
788 -- performed in place if it is safe to do so and will modify a copy of the
789 -- vector otherwise.
790 modify :: Vector v a => (forall s. Mutable v s a -> ST s ()) -> v a -> v a
791 {-# INLINE modify #-}
792 modify p = new . New.modify p . clone
793
794 -- We have to make sure that this is strict in the stream but we can't seq on
795 -- it while fusion is happening. Hence this ugliness.
796 modifyWithStream :: Vector v a
797 => (forall s. Mutable v s a -> Stream b -> ST s ())
798 -> v a -> Stream b -> v a
799 {-# INLINE modifyWithStream #-}
800 modifyWithStream p v s = new (New.modifyWithStream p (clone v) s)
801
802 -- Mapping
803 -- -------
804
805 -- | Map a function over a vector
806 map :: (Vector v a, Vector v b) => (a -> b) -> v a -> v b
807 {-# INLINE map #-}
808 map f = unstream . inplace (MStream.map f) . stream
809
810 -- | Apply a function to every element of a vector and its index
811 imap :: (Vector v a, Vector v b) => (Int -> a -> b) -> v a -> v b
812 {-# INLINE imap #-}
813 imap f = unstream . inplace (MStream.map (uncurry f) . MStream.indexed)
814 . stream
815
816 -- | Map a function over a vector and concatenate the results.
817 concatMap :: (Vector v a, Vector v b) => (a -> v b) -> v a -> v b
818 {-# INLINE concatMap #-}
819 concatMap f = unstream . Stream.concatMap (stream . f) . stream
820
821 -- Monadic mapping
822 -- ---------------
823
824 -- | Apply the monadic action to all elements of the vector, yielding a vector
825 -- of results
826 mapM :: (Monad m, Vector v a, Vector v b) => (a -> m b) -> v a -> m (v b)
827 -- FIXME: specialise for ST and IO?
828 {-# INLINE mapM #-}
829 mapM f = unstreamM . Stream.mapM f . stream
830
831 -- | Apply the monadic action to all elements of a vector and ignore the
832 -- results
833 mapM_ :: (Monad m, Vector v a) => (a -> m b) -> v a -> m ()
834 {-# INLINE mapM_ #-}
835 mapM_ f = Stream.mapM_ f . stream
836
837 -- | Apply the monadic action to all elements of the vector, yielding a vector
838 -- of results
839 forM :: (Monad m, Vector v a, Vector v b) => v a -> (a -> m b) -> m (v b)
840 {-# INLINE forM #-}
841 forM as f = mapM f as
842
843 -- | Apply the monadic action to all elements of a vector and ignore the
844 -- results
845 forM_ :: (Monad m, Vector v a) => v a -> (a -> m b) -> m ()
846 {-# INLINE forM_ #-}
847 forM_ as f = mapM_ f as
848
849 -- Zipping
850 -- -------
851
852 -- | Zip two vectors with the given function.
853 zipWith :: (Vector v a, Vector v b, Vector v c)
854 => (a -> b -> c) -> v a -> v b -> v c
855 {-# INLINE zipWith #-}
856 zipWith f xs ys = unstream (Stream.zipWith f (stream xs) (stream ys))
857
858 -- | Zip three vectors with the given function.
859 zipWith3 :: (Vector v a, Vector v b, Vector v c, Vector v d)
860 => (a -> b -> c -> d) -> v a -> v b -> v c -> v d
861 {-# INLINE zipWith3 #-}
862 zipWith3 f as bs cs = unstream (Stream.zipWith3 f (stream as)
863 (stream bs)
864 (stream cs))
865
866 zipWith4 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e)
867 => (a -> b -> c -> d -> e) -> v a -> v b -> v c -> v d -> v e
868 {-# INLINE zipWith4 #-}
869 zipWith4 f as bs cs ds
870 = unstream (Stream.zipWith4 f (stream as)
871 (stream bs)
872 (stream cs)
873 (stream ds))
874
875 zipWith5 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
876 Vector v f)
877 => (a -> b -> c -> d -> e -> f) -> v a -> v b -> v c -> v d -> v e
878 -> v f
879 {-# INLINE zipWith5 #-}
880 zipWith5 f as bs cs ds es
881 = unstream (Stream.zipWith5 f (stream as)
882 (stream bs)
883 (stream cs)
884 (stream ds)
885 (stream es))
886
887 zipWith6 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
888 Vector v f, Vector v g)
889 => (a -> b -> c -> d -> e -> f -> g)
890 -> v a -> v b -> v c -> v d -> v e -> v f -> v g
891 {-# INLINE zipWith6 #-}
892 zipWith6 f as bs cs ds es fs
893 = unstream (Stream.zipWith6 f (stream as)
894 (stream bs)
895 (stream cs)
896 (stream ds)
897 (stream es)
898 (stream fs))
899
900 -- | Zip two vectors with a function that also takes the elements' indices.
901 izipWith :: (Vector v a, Vector v b, Vector v c)
902 => (Int -> a -> b -> c) -> v a -> v b -> v c
903 {-# INLINE izipWith #-}
904 izipWith f xs ys = unstream
905 (Stream.zipWith (uncurry f) (Stream.indexed (stream xs))
906 (stream ys))
907
908 izipWith3 :: (Vector v a, Vector v b, Vector v c, Vector v d)
909 => (Int -> a -> b -> c -> d) -> v a -> v b -> v c -> v d
910 {-# INLINE izipWith3 #-}
911 izipWith3 f as bs cs
912 = unstream (Stream.zipWith3 (uncurry f) (Stream.indexed (stream as))
913 (stream bs)
914 (stream cs))
915
916 izipWith4 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e)
917 => (Int -> a -> b -> c -> d -> e) -> v a -> v b -> v c -> v d -> v e
918 {-# INLINE izipWith4 #-}
919 izipWith4 f as bs cs ds
920 = unstream (Stream.zipWith4 (uncurry f) (Stream.indexed (stream as))
921 (stream bs)
922 (stream cs)
923 (stream ds))
924
925 izipWith5 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
926 Vector v f)
927 => (Int -> a -> b -> c -> d -> e -> f) -> v a -> v b -> v c -> v d
928 -> v e -> v f
929 {-# INLINE izipWith5 #-}
930 izipWith5 f as bs cs ds es
931 = unstream (Stream.zipWith5 (uncurry f) (Stream.indexed (stream as))
932 (stream bs)
933 (stream cs)
934 (stream ds)
935 (stream es))
936
937 izipWith6 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
938 Vector v f, Vector v g)
939 => (Int -> a -> b -> c -> d -> e -> f -> g)
940 -> v a -> v b -> v c -> v d -> v e -> v f -> v g
941 {-# INLINE izipWith6 #-}
942 izipWith6 f as bs cs ds es fs
943 = unstream (Stream.zipWith6 (uncurry f) (Stream.indexed (stream as))
944 (stream bs)
945 (stream cs)
946 (stream ds)
947 (stream es)
948 (stream fs))
949
950 -- | Zip two vectors
951 zip :: (Vector v a, Vector v b, Vector v (a,b)) => v a -> v b -> v (a, b)
952 {-# INLINE zip #-}
953 zip = zipWith (,)
954
955 zip3 :: (Vector v a, Vector v b, Vector v c, Vector v (a, b, c))
956 => v a -> v b -> v c -> v (a, b, c)
957 {-# INLINE zip3 #-}
958 zip3 = zipWith3 (,,)
959
960 zip4 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v (a, b, c, d))
961 => v a -> v b -> v c -> v d -> v (a, b, c, d)
962 {-# INLINE zip4 #-}
963 zip4 = zipWith4 (,,,)
964
965 zip5 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
966 Vector v (a, b, c, d, e))
967 => v a -> v b -> v c -> v d -> v e -> v (a, b, c, d, e)
968 {-# INLINE zip5 #-}
969 zip5 = zipWith5 (,,,,)
970
971 zip6 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
972 Vector v f, Vector v (a, b, c, d, e, f))
973 => v a -> v b -> v c -> v d -> v e -> v f -> v (a, b, c, d, e, f)
974 {-# INLINE zip6 #-}
975 zip6 = zipWith6 (,,,,,)
976
977 -- Monadic zipping
978 -- ---------------
979
980 -- | Zip the two vectors with the monadic action and yield a vector of results
981 zipWithM :: (Monad m, Vector v a, Vector v b, Vector v c)
982 => (a -> b -> m c) -> v a -> v b -> m (v c)
983 -- FIXME: specialise for ST and IO?
984 {-# INLINE zipWithM #-}
985 zipWithM f as bs = unstreamM $ Stream.zipWithM f (stream as) (stream bs)
986
987 -- | Zip the two vectors with the monadic action and ignore the results
988 zipWithM_ :: (Monad m, Vector v a, Vector v b)
989 => (a -> b -> m c) -> v a -> v b -> m ()
990 {-# INLINE zipWithM_ #-}
991 zipWithM_ f as bs = Stream.zipWithM_ f (stream as) (stream bs)
992
993 -- Unzipping
994 -- ---------
995
996 unzip :: (Vector v a, Vector v b, Vector v (a,b)) => v (a, b) -> (v a, v b)
997 {-# INLINE unzip #-}
998 unzip xs = (map fst xs, map snd xs)
999
1000 unzip3 :: (Vector v a, Vector v b, Vector v c, Vector v (a, b, c))
1001 => v (a, b, c) -> (v a, v b, v c)
1002 {-# INLINE unzip3 #-}
1003 unzip3 xs = (map (\(a, b, c) -> a) xs,
1004 map (\(a, b, c) -> b) xs,
1005 map (\(a, b, c) -> c) xs)
1006
1007 unzip4 :: (Vector v a, Vector v b, Vector v c, Vector v d,
1008 Vector v (a, b, c, d))
1009 => v (a, b, c, d) -> (v a, v b, v c, v d)
1010 {-# INLINE unzip4 #-}
1011 unzip4 xs = (map (\(a, b, c, d) -> a) xs,
1012 map (\(a, b, c, d) -> b) xs,
1013 map (\(a, b, c, d) -> c) xs,
1014 map (\(a, b, c, d) -> d) xs)
1015
1016 unzip5 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
1017 Vector v (a, b, c, d, e))
1018 => v (a, b, c, d, e) -> (v a, v b, v c, v d, v e)
1019 {-# INLINE unzip5 #-}
1020 unzip5 xs = (map (\(a, b, c, d, e) -> a) xs,
1021 map (\(a, b, c, d, e) -> b) xs,
1022 map (\(a, b, c, d, e) -> c) xs,
1023 map (\(a, b, c, d, e) -> d) xs,
1024 map (\(a, b, c, d, e) -> e) xs)
1025
1026 unzip6 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
1027 Vector v f, Vector v (a, b, c, d, e, f))
1028 => v (a, b, c, d, e, f) -> (v a, v b, v c, v d, v e, v f)
1029 {-# INLINE unzip6 #-}
1030 unzip6 xs = (map (\(a, b, c, d, e, f) -> a) xs,
1031 map (\(a, b, c, d, e, f) -> b) xs,
1032 map (\(a, b, c, d, e, f) -> c) xs,
1033 map (\(a, b, c, d, e, f) -> d) xs,
1034 map (\(a, b, c, d, e, f) -> e) xs,
1035 map (\(a, b, c, d, e, f) -> f) xs)
1036
1037 -- Filtering
1038 -- ---------
1039
1040 -- | Drop elements that do not satisfy the predicate
1041 filter :: Vector v a => (a -> Bool) -> v a -> v a
1042 {-# INLINE filter #-}
1043 filter f = unstream . inplace (MStream.filter f) . stream
1044
1045 -- | Drop elements that do not satisfy the predicate which is applied to values
1046 -- and their indices
1047 ifilter :: Vector v a => (Int -> a -> Bool) -> v a -> v a
1048 {-# INLINE ifilter #-}
1049 ifilter f = unstream
1050 . inplace (MStream.map snd . MStream.filter (uncurry f)
1051 . MStream.indexed)
1052 . stream
1053
1054 -- | Drop elements that do not satisfy the monadic predicate
1055 filterM :: (Monad m, Vector v a) => (a -> m Bool) -> v a -> m (v a)
1056 {-# INLINE filterM #-}
1057 filterM f = unstreamM . Stream.filterM f . stream
1058
1059 -- | Yield the longest prefix of elements satisfying the predicate without
1060 -- copying.
1061 takeWhile :: Vector v a => (a -> Bool) -> v a -> v a
1062 {-# INLINE takeWhile #-}
1063 takeWhile f = unstream . Stream.takeWhile f . stream
1064
1065 -- | Drop the longest prefix of elements that satisfy the predicate without
1066 -- copying.
1067 dropWhile :: Vector v a => (a -> Bool) -> v a -> v a
1068 {-# INLINE dropWhile #-}
1069 dropWhile f = unstream . Stream.dropWhile f . stream
1070
1071 -- Parititioning
1072 -- -------------
1073
1074 -- | Split the vector in two parts, the first one containing those elements
1075 -- that satisfy the predicate and the second one those that don't. The
1076 -- relative order of the elements is preserved at the cost of a (sometimes)
1077 -- reduced performance compared to 'unstablePartition'.
1078 partition :: Vector v a => (a -> Bool) -> v a -> (v a, v a)
1079 {-# INLINE partition #-}
1080 partition f = partition_stream f . stream
1081
1082 -- FIXME: Make this inplace-fusible (look at how stable_partition is
1083 -- implemented in C++)
1084
1085 partition_stream :: Vector v a => (a -> Bool) -> Stream a -> (v a, v a)
1086 {-# INLINE_STREAM partition_stream #-}
1087 partition_stream f s = s `seq` runST (
1088 do
1089 (mv1,mv2) <- M.partitionStream f s
1090 v1 <- unsafeFreeze mv1
1091 v2 <- unsafeFreeze mv2
1092 return (v1,v2))
1093
1094 -- | Split the vector in two parts, the first one containing those elements
1095 -- that satisfy the predicate and the second one those that don't. The order
1096 -- of the elements is not preserved but the operation is often faster than
1097 -- 'partition'.
1098 unstablePartition :: Vector v a => (a -> Bool) -> v a -> (v a, v a)
1099 {-# INLINE unstablePartition #-}
1100 unstablePartition f = unstablePartition_stream f . stream
1101
1102 unstablePartition_stream
1103 :: Vector v a => (a -> Bool) -> Stream a -> (v a, v a)
1104 {-# INLINE_STREAM unstablePartition_stream #-}
1105 unstablePartition_stream f s = s `seq` runST (
1106 do
1107 (mv1,mv2) <- M.unstablePartitionStream f s
1108 v1 <- unsafeFreeze mv1
1109 v2 <- unsafeFreeze mv2
1110 return (v1,v2))
1111
1112 unstablePartition_new :: Vector v a => (a -> Bool) -> New v a -> (v a, v a)
1113 {-# INLINE_STREAM unstablePartition_new #-}
1114 unstablePartition_new f (New.New p) = runST (
1115 do
1116 mv <- p
1117 i <- M.unstablePartition f mv
1118 v <- unsafeFreeze mv
1119 return (unsafeTake i v, unsafeDrop i v))
1120
1121 {-# RULES
1122
1123 "unstablePartition" forall f p.
1124 unstablePartition_stream f (stream (new p))
1125 = unstablePartition_new f p
1126
1127 #-}
1128
1129
1130 -- FIXME: make span and break fusible
1131
1132 -- | Split the vector into the longest prefix of elements that satisfy the
1133 -- predicate and the rest without copying.
1134 span :: Vector v a => (a -> Bool) -> v a -> (v a, v a)
1135 {-# INLINE span #-}
1136 span f = break (not . f)
1137
1138 -- | Split the vector into the longest prefix of elements that do not satisfy
1139 -- the predicate and the rest without copying.
1140 break :: Vector v a => (a -> Bool) -> v a -> (v a, v a)
1141 {-# INLINE break #-}
1142 break f xs = case findIndex f xs of
1143 Just i -> (unsafeSlice 0 i xs, unsafeSlice i (length xs - i) xs)
1144 Nothing -> (xs, empty)
1145
1146
1147 -- Searching
1148 -- ---------
1149
1150 infix 4 `elem`
1151 -- | Check if the vector contains an element
1152 elem :: (Vector v a, Eq a) => a -> v a -> Bool
1153 {-# INLINE elem #-}
1154 elem x = Stream.elem x . stream
1155
1156 infix 4 `notElem`
1157 -- | Check if the vector does not contain an element (inverse of 'elem')
1158 notElem :: (Vector v a, Eq a) => a -> v a -> Bool
1159 {-# INLINE notElem #-}
1160 notElem x = Stream.notElem x . stream
1161
1162 -- | Yield 'Just' the first element matching the predicate or 'Nothing' if no
1163 -- such element exists.
1164 find :: Vector v a => (a -> Bool) -> v a -> Maybe a
1165 {-# INLINE find #-}
1166 find f = Stream.find f . stream
1167
1168 -- | Yield 'Just' the index of the first element matching the predicate or
1169 -- 'Nothing' if no such element exists.
1170 findIndex :: Vector v a => (a -> Bool) -> v a -> Maybe Int
1171 {-# INLINE findIndex #-}
1172 findIndex f = Stream.findIndex f . stream
1173
1174 -- | Yield the indices of elements satisfying the predicate in ascending
1175 -- order.
1176 findIndices :: (Vector v a, Vector v Int) => (a -> Bool) -> v a -> v Int
1177 {-# INLINE findIndices #-}
1178 findIndices f = unstream
1179 . inplace (MStream.map fst . MStream.filter (f . snd)
1180 . MStream.indexed)
1181 . stream
1182
1183 -- | Yield 'Just' the index of the first occurence of the given element or
1184 -- 'Nothing' if the vector does not contain the element
1185 elemIndex :: (Vector v a, Eq a) => a -> v a -> Maybe Int
1186 {-# INLINE elemIndex #-}
1187 elemIndex x = findIndex (x==)
1188
1189 -- | Yield the indices of all occurences of the given element in ascending
1190 -- order.
1191 elemIndices :: (Vector v a, Vector v Int, Eq a) => a -> v a -> v Int
1192 {-# INLINE elemIndices #-}
1193 elemIndices x = findIndices (x==)
1194
1195 -- Folding
1196 -- -------
1197
1198 -- | Left fold
1199 foldl :: Vector v b => (a -> b -> a) -> a -> v b -> a
1200 {-# INLINE foldl #-}
1201 foldl f z = Stream.foldl f z . stream
1202
1203 -- | Left fold on non-empty vectors
1204 foldl1 :: Vector v a => (a -> a -> a) -> v a -> a
1205 {-# INLINE foldl1 #-}
1206 foldl1 f = Stream.foldl1 f . stream
1207
1208 -- | Left fold with strict accumulator
1209 foldl' :: Vector v b => (a -> b -> a) -> a -> v b -> a
1210 {-# INLINE foldl' #-}
1211 foldl' f z = Stream.foldl' f z . stream
1212
1213 -- | Left fold on non-empty vectors with strict accumulator
1214 foldl1' :: Vector v a => (a -> a -> a) -> v a -> a
1215 {-# INLINE foldl1' #-}
1216 foldl1' f = Stream.foldl1' f . stream
1217
1218 -- | Right fold
1219 foldr :: Vector v a => (a -> b -> b) -> b -> v a -> b
1220 {-# INLINE foldr #-}
1221 foldr f z = Stream.foldr f z . stream
1222
1223 -- | Right fold on non-empty vectors
1224 foldr1 :: Vector v a => (a -> a -> a) -> v a -> a
1225 {-# INLINE foldr1 #-}
1226 foldr1 f = Stream.foldr1 f . stream
1227
1228 -- | Right fold with a strict accumulator
1229 foldr' :: Vector v a => (a -> b -> b) -> b -> v a -> b
1230 {-# INLINE foldr' #-}
1231 foldr' f z = Stream.foldl' (flip f) z . streamR
1232
1233 -- | Right fold on non-empty vectors with strict accumulator
1234 foldr1' :: Vector v a => (a -> a -> a) -> v a -> a
1235 {-# INLINE foldr1' #-}
1236 foldr1' f = Stream.foldl1' (flip f) . streamR
1237
1238 -- | Left fold (function applied to each element and its index)
1239 ifoldl :: Vector v b => (a -> Int -> b -> a) -> a -> v b -> a
1240 {-# INLINE ifoldl #-}
1241 ifoldl f z = Stream.foldl (uncurry . f) z . Stream.indexed . stream
1242
1243 -- | Left fold with strict accumulator (function applied to each element and
1244 -- its index)
1245 ifoldl' :: Vector v b => (a -> Int -> b -> a) -> a -> v b -> a
1246 {-# INLINE ifoldl' #-}
1247 ifoldl' f z = Stream.foldl' (uncurry . f) z . Stream.indexed . stream
1248
1249 -- | Right fold (function applied to each element and its index)
1250 ifoldr :: Vector v a => (Int -> a -> b -> b) -> b -> v a -> b
1251 {-# INLINE ifoldr #-}
1252 ifoldr f z = Stream.foldr (uncurry f) z . Stream.indexed . stream
1253
1254 -- | Right fold with strict accumulator (function applied to each element and
1255 -- its index)
1256 ifoldr' :: Vector v a => (Int -> a -> b -> b) -> b -> v a -> b
1257 {-# INLINE ifoldr' #-}
1258 ifoldr' f z xs = Stream.foldl' (flip (uncurry f)) z
1259 $ Stream.indexedR (length xs) $ streamR xs
1260
1261 -- Specialised folds
1262 -- -----------------
1263
1264 -- | Determine if all elements satisfy the predicate.
1265 all :: Vector v a => (a -> Bool) -> v a -> Bool
1266 {-# INLINE all #-}
1267 all f = Stream.and . Stream.map f . stream
1268
1269 -- | Determine if any element satisfies the predicate.
1270 any :: Vector v a => (a -> Bool) -> v a -> Bool
1271 {-# INLINE any #-}
1272 any f = Stream.or . Stream.map f . stream
1273
1274 -- | Check if all elements are 'True'
1275 and :: Vector v Bool => v Bool -> Bool
1276 {-# INLINE and #-}
1277 and = Stream.and . stream
1278
1279 -- | Check if any element is 'True'
1280 or :: Vector v Bool => v Bool -> Bool
1281 {-# INLINE or #-}
1282 or = Stream.or . stream
1283
1284 -- | Compute the sum of the elements
1285 sum :: (Vector v a, Num a) => v a -> a
1286 {-# INLINE sum #-}
1287 sum = Stream.foldl' (+) 0 . stream
1288
1289 -- | Compute the produce of the elements
1290 product :: (Vector v a, Num a) => v a -> a
1291 {-# INLINE product #-}
1292 product = Stream.foldl' (*) 1 . stream
1293
1294 -- | Yield the element of the vector. The vector may not be empty.
1295 maximum :: (Vector v a, Ord a) => v a -> a
1296 {-# INLINE maximum #-}
1297 maximum = Stream.foldl1' max . stream
1298
1299 -- | Yield the maximum element of the vector according to the given comparison
1300 -- function. The vector may not be empty.
1301 maximumBy :: Vector v a => (a -> a -> Ordering) -> v a -> a
1302 {-# INLINE maximumBy #-}
1303 maximumBy cmp = Stream.foldl1' maxBy . stream
1304 where
1305 {-# INLINE maxBy #-}
1306 maxBy x y = case cmp x y of
1307 LT -> y
1308 _ -> x
1309
1310 -- | Yield the minimum element of the vector. The vector may not be empty.
1311 minimum :: (Vector v a, Ord a) => v a -> a
1312 {-# INLINE minimum #-}
1313 minimum = Stream.foldl1' min . stream
1314
1315 -- | Yield the minimum element of the vector according to the given comparison
1316 -- function. The vector may not be empty.
1317 minimumBy :: Vector v a => (a -> a -> Ordering) -> v a -> a
1318 {-# INLINE minimumBy #-}
1319 minimumBy cmp = Stream.foldl1' minBy . stream
1320 where
1321 {-# INLINE minBy #-}
1322 minBy x y = case cmp x y of
1323 GT -> y
1324 _ -> x
1325
1326 -- | Yield the index of the maximum element of the vector. The vector may not
1327 -- be empty.
1328 maxIndex :: (Vector v a, Ord a) => v a -> Int
1329 {-# INLINE maxIndex #-}
1330 maxIndex = maxIndexBy compare
1331
1332 -- | Yield the index of the maximum element of the vector according to the
1333 -- given comparison function. The vector may not be empty.
1334 maxIndexBy :: Vector v a => (a -> a -> Ordering) -> v a -> Int
1335 {-# INLINE maxIndexBy #-}
1336 maxIndexBy cmp = fst . Stream.foldl1' imax . Stream.indexed . stream
1337 where
1338 imax (i,x) (j,y) = i `seq` j `seq`
1339 case cmp x y of
1340 LT -> (j,y)
1341 _ -> (i,x)
1342
1343 -- | Yield the index of the minimum element of the vector. The vector may not
1344 -- be empty.
1345 minIndex :: (Vector v a, Ord a) => v a -> Int
1346 {-# INLINE minIndex #-}
1347 minIndex = minIndexBy compare
1348
1349 -- | Yield the index of the minimum element of the vector according to the
1350 -- given comparison function. The vector may not be empty.
1351 minIndexBy :: Vector v a => (a -> a -> Ordering) -> v a -> Int
1352 {-# INLINE minIndexBy #-}
1353 minIndexBy cmp = fst . Stream.foldl1' imin . Stream.indexed . stream
1354 where
1355 imin (i,x) (j,y) = i `seq` j `seq`
1356 case cmp x y of
1357 GT -> (j,y)
1358 _ -> (i,x)
1359
1360 -- Monadic folds
1361 -- -------------
1362
1363 -- | Monadic fold
1364 foldM :: (Monad m, Vector v b) => (a -> b -> m a) -> a -> v b -> m a
1365 {-# INLINE foldM #-}
1366 foldM m z = Stream.foldM m z . stream
1367
1368 -- | Monadic fold over non-empty vectors
1369 fold1M :: (Monad m, Vector v a) => (a -> a -> m a) -> v a -> m a
1370 {-# INLINE fold1M #-}
1371 fold1M m = Stream.fold1M m . stream
1372
1373 -- | Monadic fold with strict accumulator
1374 foldM' :: (Monad m, Vector v b) => (a -> b -> m a) -> a -> v b -> m a
1375 {-# INLINE foldM' #-}
1376 foldM' m z = Stream.foldM' m z . stream
1377
1378 -- | Monad fold over non-empty vectors with strict accumulator
1379 fold1M' :: (Monad m, Vector v a) => (a -> a -> m a) -> v a -> m a
1380 {-# INLINE fold1M' #-}
1381 fold1M' m = Stream.fold1M' m . stream
1382
1383 -- Prefix sums (scans)
1384 -- -------------------
1385
1386 -- | Prescan
1387 prescanl :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1388 {-# INLINE prescanl #-}
1389 prescanl f z = unstream . inplace (MStream.prescanl f z) . stream
1390
1391 -- | Prescan with strict accumulator
1392 prescanl' :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1393 {-# INLINE prescanl' #-}
1394 prescanl' f z = unstream . inplace (MStream.prescanl' f z) . stream
1395
1396 -- | Scan
1397 postscanl :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1398 {-# INLINE postscanl #-}
1399 postscanl f z = unstream . inplace (MStream.postscanl f z) . stream
1400
1401 -- | Scan with strict accumulator
1402 postscanl' :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1403 {-# INLINE postscanl' #-}
1404 postscanl' f z = unstream . inplace (MStream.postscanl' f z) . stream
1405
1406 -- | Haskell-style scan
1407 scanl :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1408 {-# INLINE scanl #-}
1409 scanl f z = unstream . Stream.scanl f z . stream
1410
1411 -- | Haskell-style scan with strict accumulator
1412 scanl' :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1413 {-# INLINE scanl' #-}
1414 scanl' f z = unstream . Stream.scanl' f z . stream
1415
1416 -- | Scan over a non-empty vector
1417 scanl1 :: Vector v a => (a -> a -> a) -> v a -> v a
1418 {-# INLINE scanl1 #-}
1419 scanl1 f = unstream . inplace (MStream.scanl1 f) . stream
1420
1421 -- | Scan over a non-empty vector with a strict accumulator
1422 scanl1' :: Vector v a => (a -> a -> a) -> v a -> v a
1423 {-# INLINE scanl1' #-}
1424 scanl1' f = unstream . inplace (MStream.scanl1' f) . stream
1425
1426 -- | Right-to-left prescan
1427 prescanr :: (Vector v a, Vector v b) => (a -> b -> b) -> b -> v a -> v b
1428 {-# INLINE prescanr #-}
1429 prescanr f z = unstreamR . inplace (MStream.prescanl (flip f) z) . streamR
1430
1431 -- | Right-to-left prescan with strict accumulator
1432 prescanr' :: (Vector v a, Vector v b) => (a -> b -> b) -> b -> v a -> v b
1433 {-# INLINE prescanr' #-}
1434 prescanr' f z = unstreamR . inplace (MStream.prescanl' (flip f) z) . streamR
1435
1436 -- | Right-to-left scan
1437 postscanr :: (Vector v a, Vector v b) => (a -> b -> b) -> b -> v a -> v b
1438 {-# INLINE postscanr #-}
1439 postscanr f z = unstreamR . inplace (MStream.postscanl (flip f) z) . streamR
1440
1441 -- | Right-to-left scan with strict accumulator
1442 postscanr' :: (Vector v a, Vector v b) => (a -> b -> b) -> b -> v a -> v b
1443 {-# INLINE postscanr' #-}
1444 postscanr' f z = unstreamR . inplace (MStream.postscanl' (flip f) z) . streamR
1445
1446 -- | Haskell-style right-to-left scan
1447 scanr :: (Vector v a, Vector v b) => (a -> b -> b) -> b -> v a -> v b
1448 {-# INLINE scanr #-}
1449 scanr f z = unstreamR . Stream.scanl (flip f) z . streamR
1450
1451 -- | Haskell-style right-to-left scan with strict accumulator
1452 scanr' :: (Vector v a, Vector v b) => (a -> b -> b) -> b -> v a -> v b
1453 {-# INLINE scanr' #-}
1454 scanr' f z = unstreamR . Stream.scanl' (flip f) z . streamR
1455
1456 -- | Right-to-left scan over a non-empty vector
1457 scanr1 :: Vector v a => (a -> a -> a) -> v a -> v a
1458 {-# INLINE scanr1 #-}
1459 scanr1 f = unstreamR . inplace (MStream.scanl1 (flip f)) . streamR
1460
1461 -- | Right-to-left scan over a non-empty vector with a strict accumulator
1462 scanr1' :: Vector v a => (a -> a -> a) -> v a -> v a
1463 {-# INLINE scanr1' #-}
1464 scanr1' f = unstreamR . inplace (MStream.scanl1' (flip f)) . streamR
1465
1466 -- Conversions - Lists
1467 -- ------------------------
1468
1469 -- | Convert a vector to a list
1470 toList :: Vector v a => v a -> [a]
1471 {-# INLINE toList #-}
1472 toList = Stream.toList . stream
1473
1474 -- | Convert a list to a vector
1475 fromList :: Vector v a => [a] -> v a
1476 {-# INLINE fromList #-}
1477 fromList = unstream . Stream.fromList
1478
1479 -- | Convert the first @n@ elements of a list to a vector
1480 --
1481 -- > fromListN n xs = fromList (take n xs)
1482 fromListN :: Vector v a => Int -> [a] -> v a
1483 {-# INLINE fromListN #-}
1484 fromListN n = unstream . Stream.fromListN n
1485
1486 -- Conversions - Mutable vectors
1487 -- -----------------------------
1488
1489 -- | Copy an immutable vector into a mutable one. The two vectors must have the
1490 -- same length.
1491 copy
1492 :: (PrimMonad m, Vector v a) => Mutable v (PrimState m) a -> v a -> m ()
1493 {-# INLINE copy #-}
1494 copy dst src = BOUNDS_CHECK(check) "copy" "length mismatch"
1495 (M.length dst == length src)
1496 $ unsafeCopy dst src
1497
1498 -- | Copy an immutable vector into a mutable one. The two vectors must have
1499 -- the same length. This is not checked.
1500 unsafeCopy
1501 :: (PrimMonad m, Vector v a) => Mutable v (PrimState m) a -> v a -> m ()
1502 {-# INLINE unsafeCopy #-}
1503 unsafeCopy dst src = UNSAFE_CHECK(check) "unsafeCopy" "length mismatch"
1504 (M.length dst == length src)
1505 $ (dst `seq` src `seq` basicUnsafeCopy dst src)
1506
1507 -- Conversions to/from Streams
1508 -- ---------------------------
1509
1510 -- | Convert a vector to a 'Stream'
1511 stream :: Vector v a => v a -> Stream a
1512 {-# INLINE_STREAM stream #-}
1513 stream v = v `seq` (Stream.unfoldr get 0 `Stream.sized` Exact n)
1514 where
1515 n = length v
1516
1517 -- NOTE: the False case comes first in Core so making it the recursive one
1518 -- makes the code easier to read
1519 {-# INLINE get #-}
1520 get i | i >= n = Nothing
1521 | otherwise = case basicUnsafeIndexM v i of Box x -> Just (x, i+1)
1522
1523 -- | Construct a vector from a 'Stream'
1524 unstream :: Vector v a => Stream a -> v a
1525 {-# INLINE unstream #-}
1526 unstream s = new (New.unstream s)
1527
1528 {-# RULES
1529
1530 "stream/unstream [Vector]" forall s.
1531 stream (new (New.unstream s)) = s
1532
1533 "New.unstream/stream [Vector]" forall v.
1534 New.unstream (stream v) = clone v
1535
1536 "clone/new [Vector]" forall p.
1537 clone (new p) = p
1538
1539 "inplace [Vector]"
1540 forall (f :: forall m. Monad m => MStream m a -> MStream m a) m.
1541 New.unstream (inplace f (stream (new m))) = New.transform f m
1542
1543 "uninplace [Vector]"
1544 forall (f :: forall m. Monad m => MStream m a -> MStream m a) m.
1545 stream (new (New.transform f m)) = inplace f (stream (new m))
1546
1547 #-}
1548
1549 -- | Convert a vector to a 'Stream', proceeding from right to left
1550 streamR :: Vector v a => v a -> Stream a
1551 {-# INLINE_STREAM streamR #-}
1552 streamR v = v `seq` (Stream.unfoldr get n `Stream.sized` Exact n)
1553 where
1554 n = length v
1555
1556 {-# INLINE get #-}
1557 get 0 = Nothing
1558 get i = let i' = i-1
1559 in
1560 case basicUnsafeIndexM v i' of Box x -> Just (x, i')
1561
1562 -- | Construct a vector from a 'Stream', proceeding from right to left
1563 unstreamR :: Vector v a => Stream a -> v a
1564 {-# INLINE unstreamR #-}
1565 unstreamR s = new (New.unstreamR s)
1566
1567 {-# RULES
1568
1569 "streamR/unstreamR [Vector]" forall s.
1570 streamR (new (New.unstreamR s)) = s
1571
1572 "New.unstreamR/streamR/new [Vector]" forall p.
1573 New.unstreamR (streamR (new p)) = p
1574
1575 "inplace right [Vector]"
1576 forall (f :: forall m. Monad m => MStream m a -> MStream m a) m.
1577 New.unstreamR (inplace f (streamR (new m))) = New.transformR f m
1578
1579 "uninplace right [Vector]"
1580 forall (f :: forall m. Monad m => MStream m a -> MStream m a) m.
1581 streamR (new (New.transformR f m)) = inplace f (streamR (new m))
1582
1583 #-}
1584
1585 unstreamM :: (Vector v a, Monad m) => MStream m a -> m (v a)
1586 {-# INLINE_STREAM unstreamM #-}
1587 unstreamM s = do
1588 xs <- MStream.toList s
1589 return $ unstream $ Stream.unsafeFromList (MStream.size s) xs
1590
1591 -- Recycling support
1592 -- -----------------
1593
1594 -- | Construct a vector from a monadic initialiser
1595 new :: Vector v a => New v a -> v a
1596 {-# INLINE_STREAM new #-}
1597 new m = m `seq` runST (unsafeFreeze =<< New.run m)
1598
1599 -- | Convert a vector to an initialiser which, when run, produces a copy of
1600 -- the vector
1601 clone :: Vector v a => v a -> New v a
1602 {-# INLINE_STREAM clone #-}
1603 clone v = v `seq` New.create (
1604 do
1605 mv <- M.new (length v)
1606 unsafeCopy mv v
1607 return mv)
1608
1609 -- Comparisons
1610 -- -----------
1611
1612 -- | Check if two vectors are equal. All 'Vector' instances are also instances
1613 -- of 'Eq' and it is usually more appropriate to use those. This function is
1614 -- primarily intended for implementing 'Eq' instances for new vector types.
1615 eq :: (Vector v a, Eq a) => v a -> v a -> Bool
1616 {-# INLINE eq #-}
1617 xs `eq` ys = stream xs == stream ys
1618
1619 -- | Compare two vectors lexicographically. All 'Vector' instances are also
1620 -- instances of 'Ord' and it is usually more appropriate to use those. This
1621 -- function is primarily intended for implementing 'Ord' instances for new
1622 -- vector types.
1623 cmp :: (Vector v a, Ord a) => v a -> v a -> Ordering
1624 {-# INLINE cmp #-}
1625 cmp xs ys = compare (stream xs) (stream ys)
1626
1627 -- Data and Typeable
1628 -- -----------------
1629
1630 -- | Generic definion of 'Data.Data.gfoldl' that views a 'Vector' as a
1631 -- list.
1632 gfoldl :: (Vector v a, Data a)
1633 => (forall d b. Data d => c (d -> b) -> d -> c b)
1634 -> (forall g. g -> c g)
1635 -> v a
1636 -> c (v a)
1637 {-# INLINE gfoldl #-}
1638 gfoldl f z v = z fromList `f` toList v
1639
1640 mkType :: String -> DataType
1641 {-# INLINE mkType #-}
1642 mkType = mkNorepType
1643
1644 dataCast :: (Vector v a, Data a, Typeable1 v, Typeable1 t)
1645 => (forall d. Data d => c (t d)) -> Maybe (c (v a))
1646 {-# INLINE dataCast #-}
1647 dataCast f = gcast1 f
1648