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