Documentation
[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 -- | /O(1)/ 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 -- | /O(1)/ 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 -- | O(1) 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 -- | /O(1)/ First element
222 head :: Vector v a => v a -> a
223 {-# INLINE_STREAM head #-}
224 head v = v ! 0
225
226 -- | /O(1)/ Last element
227 last :: Vector v a => v a -> a
228 {-# INLINE_STREAM last #-}
229 last v = v ! (length v - 1)
230
231 -- | /O(1)/ 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 -- | /O(1)/ First element without checking if the vector is empty
238 unsafeHead :: Vector v a => v a -> a
239 {-# INLINE_STREAM unsafeHead #-}
240 unsafeHead v = unsafeIndex v 0
241
242 -- | /O(1)/ Last element without checking if the vector is 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 -- | /O(1)/ 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 -- | /O(1)/ 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 -- | /O(1)/ 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 -- | /O(1)/ 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 -- | /O(1)/ 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 -- | /O(1)/ 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 -- | /O(1)/ Yield a slice of the vector without copying it. The vector must
345 -- contain at least @i+n@ elements.
346 slice :: Vector v a => Int -- ^ @i@, starting index
347 -> Int -- ^ @n@, 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 -- | /O(1)/ Yield all but the last element without copying. The vector may not
355 -- be 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 -- | /O(1)/ Yield all but the first element without copying. The vector may not
361 -- be 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 -- | /O(1)/ Yield at the first @n@ elements without copying. The vector may
367 -- contain less than @n@ elements in which case it is returned unchanged.
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 -- | /O(1)/ 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 -- | /O(1)/ Yield a slice of the vector without copying. The vector must
383 -- contain at least @i+n@ elements but this is not checked.
384 unsafeSlice :: Vector v a => Int -- ^ @i@ starting index
385 -> Int -- ^ @n@ 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 -- | /O(1)/ Yield all but the last element without copying. The vector may not
393 -- be 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 -- | /O(1)/ Yield all but the first element without copying. The vector may not
399 -- be 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 -- | /O(1)/ Yield the first @n@ elements without copying. The vector must
405 -- contain at 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 -- | /O(1)/ Yield all but the first @n@ elements without copying. The vector
411 -- must 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 -- | /O(1)/ Empty vector
448 empty :: Vector v a => v a
449 {-# INLINE empty #-}
450 empty = unstream Stream.empty
451
452 -- | /O(1)/ 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 -- | /O(n)/ 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 -- | /O(n)/ Construct a vector of the given length by applying the function to
466 -- each 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 -- | /O(n)/ 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 -- | /O(n)/ Unfold 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 -- | /O(n)/ Yield a vector of the given length containing the values @x@, @x+1@
488 -- etc. 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 -- | /O(n)/ 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 'enumFromThenTo'.
495 enumFromStepN :: forall v a. (Vector v a, Num a) => a -> a -> Int -> v a
496 {-# INLINE enumFromStepN #-}
497 enumFromStepN x y n = elemseq (undefined :: v a) x
498 $ elemseq (undefined :: v a) y
499 $ unstream
500 $ Stream.enumFromStepN x y n
501
502 -- | /O(n)/ Enumerate values from @x@ to @y@.
503 --
504 -- /WARNING:/ This operation can be very inefficient. If at all possible, use
505 -- 'enumFromN' instead.
506 enumFromTo :: (Vector v a, Enum a) => a -> a -> v a
507 {-# INLINE enumFromTo #-}
508 enumFromTo x y = unstream (Stream.enumFromTo x y)
509
510 -- | /O(n)/ Enumerate values from @x@ to @y@ with a specific step @z@.
511 --
512 -- /WARNING:/ This operation can be very inefficient. If at all possible, use
513 -- 'enumFromStepN' instead.
514 enumFromThenTo :: (Vector v a, Enum a) => a -> a -> a -> v a
515 {-# INLINE enumFromThenTo #-}
516 enumFromThenTo x y z = unstream (Stream.enumFromThenTo x y z)
517
518 -- Concatenation
519 -- -------------
520
521 -- | /O(n)/ Prepend an element
522 cons :: forall v a. Vector v a => a -> v a -> v a
523 {-# INLINE cons #-}
524 cons x v = elemseq (undefined :: v a) x
525 $ unstream
526 $ Stream.cons x
527 $ stream v
528
529 -- | /O(n)/ Append an element
530 snoc :: forall v a. Vector v a => v a -> a -> v a
531 {-# INLINE snoc #-}
532 snoc v x = elemseq (undefined :: v a) x
533 $ unstream
534 $ Stream.snoc (stream v) x
535
536 infixr 5 ++
537 -- | /O(m+n)/ Concatenate two vectors
538 (++) :: Vector v a => v a -> v a -> v a
539 {-# INLINE (++) #-}
540 v ++ w = unstream (stream v Stream.++ stream w)
541
542 -- Monadic initialisation
543 -- ----------------------
544
545 -- | /O(n)/ Perform the monadic action the given number of times and store the
546 -- results in a vector.
547 replicateM :: (Monad m, Vector v a) => Int -> m a -> m (v a)
548 -- FIXME: specialise for ST and IO?
549 {-# INLINE replicateM #-}
550 replicateM n m = fromListN n `Monad.liftM` Monad.replicateM n m
551
552 -- | Execute the monadic action and freeze the resulting vector.
553 create :: Vector v a => (forall s. ST s (Mutable v s a)) -> v a
554 {-# INLINE create #-}
555 create p = new (New.create p)
556
557 -- Restricting memory usage
558 -- ------------------------
559
560 -- | /O(n)/ Yield the argument but force it not to retain any extra memory,
561 -- possibly by copying it.
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 -- | /O(m+n)/ For each pair @(i,a)@ from the list, replace the vector
580 -- element at 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 -- ^ initial vector (of length @m@)
585 -> [(Int, a)] -- ^ list of index/value pairs (of length @n@)
586 -> v a
587 {-# INLINE (//) #-}
588 v // us = update_stream v (Stream.fromList us)
589
590 -- | /O(m+n)/ For each pair @(i,a)@ from the vector of index/value pairs,
591 -- replace the 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))
596 => v a -- ^ initial vector (of length @m@)
597 -> v (Int, a) -- ^ vector of index/value pairs (of length @n@)
598 -> v a
599 {-# INLINE update #-}
600 update v w = update_stream v (stream w)
601
602 -- | /O(m+min(n1,n2))/ For each index @i@ from the index vector and the
603 -- corresponding value @a@ from the value vector, replace the element of the
604 -- initial vector at position @i@ by @a@.
605 --
606 -- > update_ <5,9,2,7> <2,0,2> <1,3,8> = <3,9,8,7>
607 --
608 -- This function is useful for instances of 'Vector' that cannot store pairs.
609 -- Otherwise, 'update' is probably more convenient.
610 --
611 -- > update_ xs is ys = update xs (zip is ys)
612 --
613 update_ :: (Vector v a, Vector v Int)
614 => v a -- ^ initial vector (of length @m@)
615 -> v Int -- ^ index vector (of length @n1@)
616 -> v a -- ^ value vector (of length @n2@)
617 -> v a
618 {-# INLINE update_ #-}
619 update_ v is w = update_stream v (Stream.zipWith (,) (stream is) (stream w))
620
621 update_stream :: Vector v a => v a -> Stream (Int,a) -> v a
622 {-# INLINE update_stream #-}
623 update_stream = modifyWithStream M.update
624
625 -- | Same as ('//') but without bounds checking.
626 unsafeUpd :: Vector v a => v a -> [(Int, a)] -> v a
627 {-# INLINE unsafeUpd #-}
628 unsafeUpd v us = unsafeUpdate_stream v (Stream.fromList us)
629
630 -- | Same as 'update' but without bounds checking.
631 unsafeUpdate :: (Vector v a, Vector v (Int, a)) => v a -> v (Int, a) -> v a
632 {-# INLINE unsafeUpdate #-}
633 unsafeUpdate v w = unsafeUpdate_stream v (stream w)
634
635 -- | Same as 'update_' but without bounds checking.
636 unsafeUpdate_ :: (Vector v a, Vector v Int) => v a -> v Int -> v a -> v a
637 {-# INLINE unsafeUpdate_ #-}
638 unsafeUpdate_ v is w
639 = unsafeUpdate_stream v (Stream.zipWith (,) (stream is) (stream w))
640
641 unsafeUpdate_stream :: Vector v a => v a -> Stream (Int,a) -> v a
642 {-# INLINE unsafeUpdate_stream #-}
643 unsafeUpdate_stream = modifyWithStream M.unsafeUpdate
644
645 -- Accumulations
646 -- -------------
647
648 -- | /O(m+n)/ For each pair @(i,b)@ from the list, replace the vector element
649 -- @a@ at position @i@ by @f a b@.
650 --
651 -- > accum (+) <5,9,2> [(2,4),(1,6),(0,3),(1,7)] = <5+3, 9+6+7, 2+4>
652 accum :: Vector v a
653 => (a -> b -> a) -- ^ @f@ accumulating function
654 -> v a -- ^ initial vector (of length @m@)
655 -> [(Int,b)] -- ^ list of index/value pairs (of length @n@)
656 -> v a
657 {-# INLINE accum #-}
658 accum f v us = accum_stream f v (Stream.fromList us)
659
660 -- | /O(m+n)/ For each pair @(i,b)@ from the vector of pairs, replace the vector
661 -- element @a@ at position @i@ by @f a b@.
662 --
663 -- > accumulate (+) <5,9,2> <(2,4),(1,6),(0,3),(1,7)> = <5+3, 9+6+7, 2+4>
664 accumulate :: (Vector v a, Vector v (Int, b))
665 => (a -> b -> a) -- ^ @f@ accumulating function
666 -> v a -- ^ initial vector (of length @m@)
667 -> v (Int,b) -- ^ vector of index/value pairs (of length @n@)
668 -> v a
669 {-# INLINE accumulate #-}
670 accumulate f v us = accum_stream f v (stream us)
671
672 -- | /O(m+min(n1,n2))/ For each index @i@ from the index vector and the
673 -- corresponding value @b@ from the the value vector,
674 -- replace the element of the initial vector at
675 -- position @i@ by @f a b@.
676 --
677 -- > accumulate_ (+) <5,9,2> <2,1,0,1> <4,6,3,7> = <5+3, 9+6+7, 2+4>
678 --
679 -- This function is useful for instances of 'Vector' that cannot store pairs.
680 -- Otherwise, 'accumulate' is probably more convenient:
681 --
682 -- > accumulate_ f as is bs = accumulate f as (zip is bs)
683 --
684 accumulate_ :: (Vector v a, Vector v Int, Vector v b)
685 => (a -> b -> a) -- ^ @f@ accumulating function
686 -> v a -- ^ initial vector (of length @m@)
687 -> v Int -- ^ index vector (of length @n1@)
688 -> v b -- ^ value vector (of length @n2@)
689 -> v a
690 {-# INLINE accumulate_ #-}
691 accumulate_ f v is xs = accum_stream f v (Stream.zipWith (,) (stream is)
692 (stream xs))
693
694
695 accum_stream :: Vector v a => (a -> b -> a) -> v a -> Stream (Int,b) -> v a
696 {-# INLINE accum_stream #-}
697 accum_stream f = modifyWithStream (M.accum f)
698
699 -- | Same as 'accum' but without bounds checking.
700 unsafeAccum :: Vector v a => (a -> b -> a) -> v a -> [(Int,b)] -> v a
701 {-# INLINE unsafeAccum #-}
702 unsafeAccum f v us = unsafeAccum_stream f v (Stream.fromList us)
703
704 -- | Same as 'accumulate' but without bounds checking.
705 unsafeAccumulate :: (Vector v a, Vector v (Int, b))
706 => (a -> b -> a) -> v a -> v (Int,b) -> v a
707 {-# INLINE unsafeAccumulate #-}
708 unsafeAccumulate f v us = unsafeAccum_stream f v (stream us)
709
710 -- | Same as 'accumulate_' but without bounds checking.
711 unsafeAccumulate_ :: (Vector v a, Vector v Int, Vector v b)
712 => (a -> b -> a) -> v a -> v Int -> v b -> v a
713 {-# INLINE unsafeAccumulate_ #-}
714 unsafeAccumulate_ f v is xs
715 = unsafeAccum_stream f v (Stream.zipWith (,) (stream is) (stream xs))
716
717 unsafeAccum_stream
718 :: Vector v a => (a -> b -> a) -> v a -> Stream (Int,b) -> v a
719 {-# INLINE unsafeAccum_stream #-}
720 unsafeAccum_stream f = modifyWithStream (M.unsafeAccum f)
721
722 -- Permutations
723 -- ------------
724
725 -- | /O(n)/ Reverse a vector
726 reverse :: (Vector v a) => v a -> v a
727 {-# INLINE reverse #-}
728 -- FIXME: make this fuse better, add support for recycling
729 reverse = unstream . streamR
730
731 -- | /O(n)/ Yield the vector obtained by replacing each element @i@ of the
732 -- index vector by @xs!i@. This is equivalent to @map (xs!) is@ but is often
733 -- much more efficient.
734 --
735 -- > backpermute <a,b,c,d> <0,3,2,3,1,0> = <a,d,c,d,b,a>
736 backpermute :: (Vector v a, Vector v Int)
737 => v a -- ^ @xs@ value vector
738 -> v Int -- ^ @is@ index vector (of length @n@)
739 -> v a
740 {-# INLINE backpermute #-}
741 -- This somewhat non-intuitive definition ensures that the resulting vector
742 -- does not retain references to the original one even if it is lazy in its
743 -- elements. This would not be the case if we simply used map (v!)
744 backpermute v is = seq v
745 $ unstream
746 $ Stream.unbox
747 $ Stream.map (indexM v)
748 $ stream is
749
750 -- | Same as 'backpermute' but without bounds checking.
751 unsafeBackpermute :: (Vector v a, Vector v Int) => v a -> v Int -> v a
752 {-# INLINE unsafeBackpermute #-}
753 unsafeBackpermute v is = seq v
754 $ unstream
755 $ Stream.unbox
756 $ Stream.map (unsafeIndexM v)
757 $ stream is
758
759 -- Safe destructive updates
760 -- ------------------------
761
762 -- | Apply a destructive operation to a vector. The operation will be
763 -- performed in place if it is safe to do so and will modify a copy of the
764 -- vector otherwise.
765 modify :: Vector v a => (forall s. Mutable v s a -> ST s ()) -> v a -> v a
766 {-# INLINE modify #-}
767 modify p = new . New.modify p . clone
768
769 -- We have to make sure that this is strict in the stream but we can't seq on
770 -- it while fusion is happening. Hence this ugliness.
771 modifyWithStream :: Vector v a
772 => (forall s. Mutable v s a -> Stream b -> ST s ())
773 -> v a -> Stream b -> v a
774 {-# INLINE modifyWithStream #-}
775 modifyWithStream p v s = new (New.modifyWithStream p (clone v) s)
776
777 -- Mapping
778 -- -------
779
780 -- | /O(n)/ Map a function over a vector
781 map :: (Vector v a, Vector v b) => (a -> b) -> v a -> v b
782 {-# INLINE map #-}
783 map f = unstream . inplace (MStream.map f) . stream
784
785 -- | /O(n)/ Apply a function to every element of a vector and its index
786 imap :: (Vector v a, Vector v b) => (Int -> a -> b) -> v a -> v b
787 {-# INLINE imap #-}
788 imap f = unstream . inplace (MStream.map (uncurry f) . MStream.indexed)
789 . stream
790
791 -- | Map a function over a vector and concatenate the results.
792 concatMap :: (Vector v a, Vector v b) => (a -> v b) -> v a -> v b
793 {-# INLINE concatMap #-}
794 concatMap f = unstream . Stream.concatMap (stream . f) . stream
795
796 -- Monadic mapping
797 -- ---------------
798
799 -- | /O(n)/ Apply the monadic action to all elements of the vector, yielding a
800 -- vector of results
801 mapM :: (Monad m, Vector v a, Vector v b) => (a -> m b) -> v a -> m (v b)
802 -- FIXME: specialise for ST and IO?
803 {-# INLINE mapM #-}
804 mapM f = unstreamM . Stream.mapM f . stream
805
806 -- | /O(n)/ Apply the monadic action to all elements of a vector and ignore the
807 -- results
808 mapM_ :: (Monad m, Vector v a) => (a -> m b) -> v a -> m ()
809 {-# INLINE mapM_ #-}
810 mapM_ f = Stream.mapM_ f . stream
811
812 -- | /O(n)/ Apply the monadic action to all elements of the vector, yielding a
813 -- vector of results. Equvalent to @flip mapM@.
814 forM :: (Monad m, Vector v a, Vector v b) => v a -> (a -> m b) -> m (v b)
815 {-# INLINE forM #-}
816 forM as f = mapM f as
817
818 -- | /O(n)/ Apply the monadic action to all elements of a vector and ignore the
819 -- results. Equivalent to @flip mapM_@.
820 forM_ :: (Monad m, Vector v a) => v a -> (a -> m b) -> m ()
821 {-# INLINE forM_ #-}
822 forM_ as f = mapM_ f as
823
824 -- Zipping
825 -- -------
826
827 -- | /O(min(m,n))/ Zip two vectors with the given function.
828 zipWith :: (Vector v a, Vector v b, Vector v c)
829 => (a -> b -> c) -> v a -> v b -> v c
830 {-# INLINE zipWith #-}
831 zipWith f xs ys = unstream (Stream.zipWith f (stream xs) (stream ys))
832
833 -- | Zip three vectors with the given function.
834 zipWith3 :: (Vector v a, Vector v b, Vector v c, Vector v d)
835 => (a -> b -> c -> d) -> v a -> v b -> v c -> v d
836 {-# INLINE zipWith3 #-}
837 zipWith3 f as bs cs = unstream (Stream.zipWith3 f (stream as)
838 (stream bs)
839 (stream cs))
840
841 zipWith4 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e)
842 => (a -> b -> c -> d -> e) -> v a -> v b -> v c -> v d -> v e
843 {-# INLINE zipWith4 #-}
844 zipWith4 f as bs cs ds
845 = unstream (Stream.zipWith4 f (stream as)
846 (stream bs)
847 (stream cs)
848 (stream ds))
849
850 zipWith5 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
851 Vector v f)
852 => (a -> b -> c -> d -> e -> f) -> v a -> v b -> v c -> v d -> v e
853 -> v f
854 {-# INLINE zipWith5 #-}
855 zipWith5 f as bs cs ds es
856 = unstream (Stream.zipWith5 f (stream as)
857 (stream bs)
858 (stream cs)
859 (stream ds)
860 (stream es))
861
862 zipWith6 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
863 Vector v f, Vector v g)
864 => (a -> b -> c -> d -> e -> f -> g)
865 -> v a -> v b -> v c -> v d -> v e -> v f -> v g
866 {-# INLINE zipWith6 #-}
867 zipWith6 f as bs cs ds es fs
868 = unstream (Stream.zipWith6 f (stream as)
869 (stream bs)
870 (stream cs)
871 (stream ds)
872 (stream es)
873 (stream fs))
874
875 -- | /O(min(m,n))/ Zip two vectors with a function that also takes the
876 -- elements' indices.
877 izipWith :: (Vector v a, Vector v b, Vector v c)
878 => (Int -> a -> b -> c) -> v a -> v b -> v c
879 {-# INLINE izipWith #-}
880 izipWith f xs ys = unstream
881 (Stream.zipWith (uncurry f) (Stream.indexed (stream xs))
882 (stream ys))
883
884 izipWith3 :: (Vector v a, Vector v b, Vector v c, Vector v d)
885 => (Int -> a -> b -> c -> d) -> v a -> v b -> v c -> v d
886 {-# INLINE izipWith3 #-}
887 izipWith3 f as bs cs
888 = unstream (Stream.zipWith3 (uncurry f) (Stream.indexed (stream as))
889 (stream bs)
890 (stream cs))
891
892 izipWith4 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e)
893 => (Int -> a -> b -> c -> d -> e) -> v a -> v b -> v c -> v d -> v e
894 {-# INLINE izipWith4 #-}
895 izipWith4 f as bs cs ds
896 = unstream (Stream.zipWith4 (uncurry f) (Stream.indexed (stream as))
897 (stream bs)
898 (stream cs)
899 (stream ds))
900
901 izipWith5 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
902 Vector v f)
903 => (Int -> a -> b -> c -> d -> e -> f) -> v a -> v b -> v c -> v d
904 -> v e -> v f
905 {-# INLINE izipWith5 #-}
906 izipWith5 f as bs cs ds es
907 = unstream (Stream.zipWith5 (uncurry f) (Stream.indexed (stream as))
908 (stream bs)
909 (stream cs)
910 (stream ds)
911 (stream es))
912
913 izipWith6 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
914 Vector v f, Vector v g)
915 => (Int -> a -> b -> c -> d -> e -> f -> g)
916 -> v a -> v b -> v c -> v d -> v e -> v f -> v g
917 {-# INLINE izipWith6 #-}
918 izipWith6 f as bs cs ds es fs
919 = unstream (Stream.zipWith6 (uncurry f) (Stream.indexed (stream as))
920 (stream bs)
921 (stream cs)
922 (stream ds)
923 (stream es)
924 (stream fs))
925
926 -- | /O(min(m,n))/ Zip two vectors
927 zip :: (Vector v a, Vector v b, Vector v (a,b)) => v a -> v b -> v (a, b)
928 {-# INLINE zip #-}
929 zip = zipWith (,)
930
931 zip3 :: (Vector v a, Vector v b, Vector v c, Vector v (a, b, c))
932 => v a -> v b -> v c -> v (a, b, c)
933 {-# INLINE zip3 #-}
934 zip3 = zipWith3 (,,)
935
936 zip4 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v (a, b, c, d))
937 => v a -> v b -> v c -> v d -> v (a, b, c, d)
938 {-# INLINE zip4 #-}
939 zip4 = zipWith4 (,,,)
940
941 zip5 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
942 Vector v (a, b, c, d, e))
943 => v a -> v b -> v c -> v d -> v e -> v (a, b, c, d, e)
944 {-# INLINE zip5 #-}
945 zip5 = zipWith5 (,,,,)
946
947 zip6 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
948 Vector v f, Vector v (a, b, c, d, e, f))
949 => v a -> v b -> v c -> v d -> v e -> v f -> v (a, b, c, d, e, f)
950 {-# INLINE zip6 #-}
951 zip6 = zipWith6 (,,,,,)
952
953 -- Monadic zipping
954 -- ---------------
955
956 -- | /O(min(m,n))/ Zip the two vectors with the monadic action and yield a
957 -- vector of results
958 zipWithM :: (Monad m, Vector v a, Vector v b, Vector v c)
959 => (a -> b -> m c) -> v a -> v b -> m (v c)
960 -- FIXME: specialise for ST and IO?
961 {-# INLINE zipWithM #-}
962 zipWithM f as bs = unstreamM $ Stream.zipWithM f (stream as) (stream bs)
963
964 -- | /O(min(m,n))/ Zip the two vectors with the monadic action and ignore the
965 -- results
966 zipWithM_ :: (Monad m, Vector v a, Vector v b)
967 => (a -> b -> m c) -> v a -> v b -> m ()
968 {-# INLINE zipWithM_ #-}
969 zipWithM_ f as bs = Stream.zipWithM_ f (stream as) (stream bs)
970
971 -- Unzipping
972 -- ---------
973
974 -- | /O(min(m,n))/ Unzip a vector of pairs.
975 unzip :: (Vector v a, Vector v b, Vector v (a,b)) => v (a, b) -> (v a, v b)
976 {-# INLINE unzip #-}
977 unzip xs = (map fst xs, map snd xs)
978
979 unzip3 :: (Vector v a, Vector v b, Vector v c, Vector v (a, b, c))
980 => v (a, b, c) -> (v a, v b, v c)
981 {-# INLINE unzip3 #-}
982 unzip3 xs = (map (\(a, b, c) -> a) xs,
983 map (\(a, b, c) -> b) xs,
984 map (\(a, b, c) -> c) xs)
985
986 unzip4 :: (Vector v a, Vector v b, Vector v c, Vector v d,
987 Vector v (a, b, c, d))
988 => v (a, b, c, d) -> (v a, v b, v c, v d)
989 {-# INLINE unzip4 #-}
990 unzip4 xs = (map (\(a, b, c, d) -> a) xs,
991 map (\(a, b, c, d) -> b) xs,
992 map (\(a, b, c, d) -> c) xs,
993 map (\(a, b, c, d) -> d) xs)
994
995 unzip5 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
996 Vector v (a, b, c, d, e))
997 => v (a, b, c, d, e) -> (v a, v b, v c, v d, v e)
998 {-# INLINE unzip5 #-}
999 unzip5 xs = (map (\(a, b, c, d, e) -> a) xs,
1000 map (\(a, b, c, d, e) -> b) xs,
1001 map (\(a, b, c, d, e) -> c) xs,
1002 map (\(a, b, c, d, e) -> d) xs,
1003 map (\(a, b, c, d, e) -> e) xs)
1004
1005 unzip6 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
1006 Vector v f, Vector v (a, b, c, d, e, f))
1007 => v (a, b, c, d, e, f) -> (v a, v b, v c, v d, v e, v f)
1008 {-# INLINE unzip6 #-}
1009 unzip6 xs = (map (\(a, b, c, d, e, f) -> a) xs,
1010 map (\(a, b, c, d, e, f) -> b) xs,
1011 map (\(a, b, c, d, e, f) -> c) xs,
1012 map (\(a, b, c, d, e, f) -> d) xs,
1013 map (\(a, b, c, d, e, f) -> e) xs,
1014 map (\(a, b, c, d, e, f) -> f) xs)
1015
1016 -- Filtering
1017 -- ---------
1018
1019 -- | /O(n)/ Drop elements that do not satisfy the predicate
1020 filter :: Vector v a => (a -> Bool) -> v a -> v a
1021 {-# INLINE filter #-}
1022 filter f = unstream . inplace (MStream.filter f) . stream
1023
1024 -- | /O(n)/ Drop elements that do not satisfy the predicate which is applied to
1025 -- values and their indices
1026 ifilter :: Vector v a => (Int -> a -> Bool) -> v a -> v a
1027 {-# INLINE ifilter #-}
1028 ifilter f = unstream
1029 . inplace (MStream.map snd . MStream.filter (uncurry f)
1030 . MStream.indexed)
1031 . stream
1032
1033 -- | /O(n)/ Drop elements that do not satisfy the monadic predicate
1034 filterM :: (Monad m, Vector v a) => (a -> m Bool) -> v a -> m (v a)
1035 {-# INLINE filterM #-}
1036 filterM f = unstreamM . Stream.filterM f . stream
1037
1038 -- | /O(n)/ Yield the longest prefix of elements satisfying the predicate
1039 -- without copying.
1040 takeWhile :: Vector v a => (a -> Bool) -> v a -> v a
1041 {-# INLINE takeWhile #-}
1042 takeWhile f = unstream . Stream.takeWhile f . stream
1043
1044 -- | /O(n)/ Drop the longest prefix of elements that satisfy the predicate
1045 -- without copying.
1046 dropWhile :: Vector v a => (a -> Bool) -> v a -> v a
1047 {-# INLINE dropWhile #-}
1048 dropWhile f = unstream . Stream.dropWhile f . stream
1049
1050 -- Parititioning
1051 -- -------------
1052
1053 -- | /O(n)/ Split the vector in two parts, the first one containing those
1054 -- elements that satisfy the predicate and the second one those that don't. The
1055 -- relative order of the elements is preserved at the cost of a sometimes
1056 -- reduced performance compared to 'unstablePartition'.
1057 partition :: Vector v a => (a -> Bool) -> v a -> (v a, v a)
1058 {-# INLINE partition #-}
1059 partition f = partition_stream f . stream
1060
1061 -- FIXME: Make this inplace-fusible (look at how stable_partition is
1062 -- implemented in C++)
1063
1064 partition_stream :: Vector v a => (a -> Bool) -> Stream a -> (v a, v a)
1065 {-# INLINE_STREAM partition_stream #-}
1066 partition_stream f s = s `seq` runST (
1067 do
1068 (mv1,mv2) <- M.partitionStream f s
1069 v1 <- unsafeFreeze mv1
1070 v2 <- unsafeFreeze mv2
1071 return (v1,v2))
1072
1073 -- | /O(n)/ Split the vector in two parts, the first one containing those
1074 -- elements that satisfy the predicate and the second one those that don't.
1075 -- The order of the elements is not preserved but the operation is often
1076 -- faster than 'partition'.
1077 unstablePartition :: Vector v a => (a -> Bool) -> v a -> (v a, v a)
1078 {-# INLINE unstablePartition #-}
1079 unstablePartition f = unstablePartition_stream f . stream
1080
1081 unstablePartition_stream
1082 :: Vector v a => (a -> Bool) -> Stream a -> (v a, v a)
1083 {-# INLINE_STREAM unstablePartition_stream #-}
1084 unstablePartition_stream f s = s `seq` runST (
1085 do
1086 (mv1,mv2) <- M.unstablePartitionStream f s
1087 v1 <- unsafeFreeze mv1
1088 v2 <- unsafeFreeze mv2
1089 return (v1,v2))
1090
1091 unstablePartition_new :: Vector v a => (a -> Bool) -> New v a -> (v a, v a)
1092 {-# INLINE_STREAM unstablePartition_new #-}
1093 unstablePartition_new f (New.New p) = runST (
1094 do
1095 mv <- p
1096 i <- M.unstablePartition f mv
1097 v <- unsafeFreeze mv
1098 return (unsafeTake i v, unsafeDrop i v))
1099
1100 {-# RULES
1101
1102 "unstablePartition" forall f p.
1103 unstablePartition_stream f (stream (new p))
1104 = unstablePartition_new f p
1105
1106 #-}
1107
1108
1109 -- FIXME: make span and break fusible
1110
1111 -- | /O(n)/ Split the vector into the longest prefix of elements that satisfy
1112 -- the predicate and the rest without copying.
1113 span :: Vector v a => (a -> Bool) -> v a -> (v a, v a)
1114 {-# INLINE span #-}
1115 span f = break (not . f)
1116
1117 -- | /O(n)/ Split the vector into the longest prefix of elements that do not
1118 -- satisfy the predicate and the rest without copying.
1119 break :: Vector v a => (a -> Bool) -> v a -> (v a, v a)
1120 {-# INLINE break #-}
1121 break f xs = case findIndex f xs of
1122 Just i -> (unsafeSlice 0 i xs, unsafeSlice i (length xs - i) xs)
1123 Nothing -> (xs, empty)
1124
1125
1126 -- Searching
1127 -- ---------
1128
1129 infix 4 `elem`
1130 -- | /O(n)/ Check if the vector contains an element
1131 elem :: (Vector v a, Eq a) => a -> v a -> Bool
1132 {-# INLINE elem #-}
1133 elem x = Stream.elem x . stream
1134
1135 infix 4 `notElem`
1136 -- | /O(n)/ Check if the vector does not contain an element (inverse of 'elem')
1137 notElem :: (Vector v a, Eq a) => a -> v a -> Bool
1138 {-# INLINE notElem #-}
1139 notElem x = Stream.notElem x . stream
1140
1141 -- | /O(n)/ Yield 'Just' the first element matching the predicate or 'Nothing'
1142 -- if no such element exists.
1143 find :: Vector v a => (a -> Bool) -> v a -> Maybe a
1144 {-# INLINE find #-}
1145 find f = Stream.find f . stream
1146
1147 -- | /O(n)/ Yield 'Just' the index of the first element matching the predicate
1148 -- or 'Nothing' if no such element exists.
1149 findIndex :: Vector v a => (a -> Bool) -> v a -> Maybe Int
1150 {-# INLINE findIndex #-}
1151 findIndex f = Stream.findIndex f . stream
1152
1153 -- | /O(n)/ Yield the indices of elements satisfying the predicate in ascending
1154 -- order.
1155 findIndices :: (Vector v a, Vector v Int) => (a -> Bool) -> v a -> v Int
1156 {-# INLINE findIndices #-}
1157 findIndices f = unstream
1158 . inplace (MStream.map fst . MStream.filter (f . snd)
1159 . MStream.indexed)
1160 . stream
1161
1162 -- | /O(n)/ Yield 'Just' the index of the first occurence of the given element or
1163 -- 'Nothing' if the vector does not contain the element. This is a specialised
1164 -- version of 'findIndex'.
1165 elemIndex :: (Vector v a, Eq a) => a -> v a -> Maybe Int
1166 {-# INLINE elemIndex #-}
1167 elemIndex x = findIndex (x==)
1168
1169 -- | /O(n)/ Yield the indices of all occurences of the given element in
1170 -- ascending order. This is a specialised version of 'findIndices'.
1171 elemIndices :: (Vector v a, Vector v Int, Eq a) => a -> v a -> v Int
1172 {-# INLINE elemIndices #-}
1173 elemIndices x = findIndices (x==)
1174
1175 -- Folding
1176 -- -------
1177
1178 -- | /O(n)/ Left fold
1179 foldl :: Vector v b => (a -> b -> a) -> a -> v b -> a
1180 {-# INLINE foldl #-}
1181 foldl f z = Stream.foldl f z . stream
1182
1183 -- | /O(n)/ Left fold on non-empty vectors
1184 foldl1 :: Vector v a => (a -> a -> a) -> v a -> a
1185 {-# INLINE foldl1 #-}
1186 foldl1 f = Stream.foldl1 f . stream
1187
1188 -- | /O(n)/ Left fold with strict accumulator
1189 foldl' :: Vector v b => (a -> b -> a) -> a -> v b -> a
1190 {-# INLINE foldl' #-}
1191 foldl' f z = Stream.foldl' f z . stream
1192
1193 -- | /O(n)/ Left fold on non-empty vectors with strict accumulator
1194 foldl1' :: Vector v a => (a -> a -> a) -> v a -> a
1195 {-# INLINE foldl1' #-}
1196 foldl1' f = Stream.foldl1' f . stream
1197
1198 -- | /O(n)/ Right fold
1199 foldr :: Vector v a => (a -> b -> b) -> b -> v a -> b
1200 {-# INLINE foldr #-}
1201 foldr f z = Stream.foldr f z . stream
1202
1203 -- | /O(n)/ Right fold on non-empty vectors
1204 foldr1 :: Vector v a => (a -> a -> a) -> v a -> a
1205 {-# INLINE foldr1 #-}
1206 foldr1 f = Stream.foldr1 f . stream
1207
1208 -- | /O(n)/ Right fold with a strict accumulator
1209 foldr' :: Vector v a => (a -> b -> b) -> b -> v a -> b
1210 {-# INLINE foldr' #-}
1211 foldr' f z = Stream.foldl' (flip f) z . streamR
1212
1213 -- | /O(n)/ Right fold on non-empty vectors with strict accumulator
1214 foldr1' :: Vector v a => (a -> a -> a) -> v a -> a
1215 {-# INLINE foldr1' #-}
1216 foldr1' f = Stream.foldl1' (flip f) . streamR
1217
1218 -- | /O(n)/ Left fold (function applied to each element and its index)
1219 ifoldl :: Vector v b => (a -> Int -> b -> a) -> a -> v b -> a
1220 {-# INLINE ifoldl #-}
1221 ifoldl f z = Stream.foldl (uncurry . f) z . Stream.indexed . stream
1222
1223 -- | /O(n)/ Left fold with strict accumulator (function applied to each element
1224 -- and its index)
1225 ifoldl' :: Vector v b => (a -> Int -> b -> a) -> a -> v b -> a
1226 {-# INLINE ifoldl' #-}
1227 ifoldl' f z = Stream.foldl' (uncurry . f) z . Stream.indexed . stream
1228
1229 -- | /O(n)/ Right fold (function applied to each element and its index)
1230 ifoldr :: Vector v a => (Int -> a -> b -> b) -> b -> v a -> b
1231 {-# INLINE ifoldr #-}
1232 ifoldr f z = Stream.foldr (uncurry f) z . Stream.indexed . stream
1233
1234 -- | /O(n)/ Right fold with strict accumulator (function applied to each
1235 -- element and its index)
1236 ifoldr' :: Vector v a => (Int -> a -> b -> b) -> b -> v a -> b
1237 {-# INLINE ifoldr' #-}
1238 ifoldr' f z xs = Stream.foldl' (flip (uncurry f)) z
1239 $ Stream.indexedR (length xs) $ streamR xs
1240
1241 -- Specialised folds
1242 -- -----------------
1243
1244 -- | /O(n)/ Check if all elements satisfy the predicate.
1245 all :: Vector v a => (a -> Bool) -> v a -> Bool
1246 {-# INLINE all #-}
1247 all f = Stream.and . Stream.map f . stream
1248
1249 -- | /O(n)/ Check if any element satisfies the predicate.
1250 any :: Vector v a => (a -> Bool) -> v a -> Bool
1251 {-# INLINE any #-}
1252 any f = Stream.or . Stream.map f . stream
1253
1254 -- | /O(n)/ Check if all elements are 'True'
1255 and :: Vector v Bool => v Bool -> Bool
1256 {-# INLINE and #-}
1257 and = Stream.and . stream
1258
1259 -- | /O(n)/ Check if any element is 'True'
1260 or :: Vector v Bool => v Bool -> Bool
1261 {-# INLINE or #-}
1262 or = Stream.or . stream
1263
1264 -- | /O(n)/ Compute the sum of the elements
1265 sum :: (Vector v a, Num a) => v a -> a
1266 {-# INLINE sum #-}
1267 sum = Stream.foldl' (+) 0 . stream
1268
1269 -- | /O(n)/ Compute the produce of the elements
1270 product :: (Vector v a, Num a) => v a -> a
1271 {-# INLINE product #-}
1272 product = Stream.foldl' (*) 1 . stream
1273
1274 -- | /O(n)/ Yield the maximum element of the vector. The vector may not be
1275 -- empty.
1276 maximum :: (Vector v a, Ord a) => v a -> a
1277 {-# INLINE maximum #-}
1278 maximum = Stream.foldl1' max . stream
1279
1280 -- | /O(n)/ Yield the maximum element of the vector according to the given
1281 -- comparison function. The vector may not be empty.
1282 maximumBy :: Vector v a => (a -> a -> Ordering) -> v a -> a
1283 {-# INLINE maximumBy #-}
1284 maximumBy cmp = Stream.foldl1' maxBy . stream
1285 where
1286 {-# INLINE maxBy #-}
1287 maxBy x y = case cmp x y of
1288 LT -> y
1289 _ -> x
1290
1291 -- | /O(n)/ Yield the minimum element of the vector. The vector may not be
1292 -- empty.
1293 minimum :: (Vector v a, Ord a) => v a -> a
1294 {-# INLINE minimum #-}
1295 minimum = Stream.foldl1' min . stream
1296
1297 -- | /O(n)/ Yield the minimum element of the vector according to the given
1298 -- comparison function. The vector may not be empty.
1299 minimumBy :: Vector v a => (a -> a -> Ordering) -> v a -> a
1300 {-# INLINE minimumBy #-}
1301 minimumBy cmp = Stream.foldl1' minBy . stream
1302 where
1303 {-# INLINE minBy #-}
1304 minBy x y = case cmp x y of
1305 GT -> y
1306 _ -> x
1307
1308 -- | /O(n)/ Yield the index of the maximum element of the vector. The vector
1309 -- may not be empty.
1310 maxIndex :: (Vector v a, Ord a) => v a -> Int
1311 {-# INLINE maxIndex #-}
1312 maxIndex = maxIndexBy compare
1313
1314 -- | /O(n)/ Yield the index of the maximum element of the vector according to
1315 -- the given comparison function. The vector may not be empty.
1316 maxIndexBy :: Vector v a => (a -> a -> Ordering) -> v a -> Int
1317 {-# INLINE maxIndexBy #-}
1318 maxIndexBy cmp = fst . Stream.foldl1' imax . Stream.indexed . stream
1319 where
1320 imax (i,x) (j,y) = i `seq` j `seq`
1321 case cmp x y of
1322 LT -> (j,y)
1323 _ -> (i,x)
1324
1325 -- | /O(n)/ Yield the index of the minimum element of the vector. The vector
1326 -- may not be empty.
1327 minIndex :: (Vector v a, Ord a) => v a -> Int
1328 {-# INLINE minIndex #-}
1329 minIndex = minIndexBy compare
1330
1331 -- | /O(n)/ Yield the index of the minimum element of the vector according to
1332 -- the given comparison function. The vector may not be empty.
1333 minIndexBy :: Vector v a => (a -> a -> Ordering) -> v a -> Int
1334 {-# INLINE minIndexBy #-}
1335 minIndexBy cmp = fst . Stream.foldl1' imin . Stream.indexed . stream
1336 where
1337 imin (i,x) (j,y) = i `seq` j `seq`
1338 case cmp x y of
1339 GT -> (j,y)
1340 _ -> (i,x)
1341
1342 -- Monadic folds
1343 -- -------------
1344
1345 -- | /O(n)/ Monadic fold
1346 foldM :: (Monad m, Vector v b) => (a -> b -> m a) -> a -> v b -> m a
1347 {-# INLINE foldM #-}
1348 foldM m z = Stream.foldM m z . stream
1349
1350 -- | /O(n)/ Monadic fold over non-empty vectors
1351 fold1M :: (Monad m, Vector v a) => (a -> a -> m a) -> v a -> m a
1352 {-# INLINE fold1M #-}
1353 fold1M m = Stream.fold1M m . stream
1354
1355 -- | /O(n)/ Monadic fold with strict accumulator
1356 foldM' :: (Monad m, Vector v b) => (a -> b -> m a) -> a -> v b -> m a
1357 {-# INLINE foldM' #-}
1358 foldM' m z = Stream.foldM' m z . stream
1359
1360 -- | /O(n)/ Monad fold over non-empty vectors with strict accumulator
1361 fold1M' :: (Monad m, Vector v a) => (a -> a -> m a) -> v a -> m a
1362 {-# INLINE fold1M' #-}
1363 fold1M' m = Stream.fold1M' m . stream
1364
1365 -- Prefix sums (scans)
1366 -- -------------------
1367
1368 -- | /O(n)/ Prescan
1369 --
1370 -- @
1371 -- prescanl f z = 'init' . 'scanl' f z
1372 -- @
1373 --
1374 -- Example: @prescanl (+) 0 \<1,2,3,4\> = \<0,1,3,6\>@
1375 --
1376 prescanl :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1377 {-# INLINE prescanl #-}
1378 prescanl f z = unstream . inplace (MStream.prescanl f z) . stream
1379
1380 -- | /O(n)/ Prescan with strict accumulator
1381 prescanl' :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1382 {-# INLINE prescanl' #-}
1383 prescanl' f z = unstream . inplace (MStream.prescanl' f z) . stream
1384
1385 -- | /O(n)/ Scan
1386 --
1387 -- @
1388 -- postscanl f z = 'tail' . 'scanl' f z
1389 -- @
1390 --
1391 -- Example: @postscanl (+) 0 \<1,2,3,4\> = \<1,3,6,10\>@
1392 --
1393 postscanl :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1394 {-# INLINE postscanl #-}
1395 postscanl f z = unstream . inplace (MStream.postscanl f z) . stream
1396
1397 -- | /O(n)/ Scan with strict accumulator
1398 postscanl' :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1399 {-# INLINE postscanl' #-}
1400 postscanl' f z = unstream . inplace (MStream.postscanl' f z) . stream
1401
1402 -- | /O(n)/ Haskell-style scan
1403 --
1404 -- > scanl f z <x1,...,xn> = <y1,...,y(n+1)>
1405 -- > where y1 = z
1406 -- > yi = f y(i-1) x(i-1)
1407 --
1408 -- Example: @scanl (+) 0 \<1,2,3,4\> = \<0,1,3,6,10\>@
1409 --
1410 scanl :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1411 {-# INLINE scanl #-}
1412 scanl f z = unstream . Stream.scanl f z . stream
1413
1414 -- | /O(n)/ Haskell-style scan with strict accumulator
1415 scanl' :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1416 {-# INLINE scanl' #-}
1417 scanl' f z = unstream . Stream.scanl' f z . stream
1418
1419 -- | /O(n)/ Scan over a non-empty vector
1420 --
1421 -- > scanl f <x1,...,xn> = <y1,...,yn>
1422 -- > where y1 = x1
1423 -- > yi = f y(i-1) xi
1424 --
1425 scanl1 :: Vector v a => (a -> a -> a) -> v a -> v a
1426 {-# INLINE scanl1 #-}
1427 scanl1 f = unstream . inplace (MStream.scanl1 f) . stream
1428
1429 -- | /O(n)/ Scan over a non-empty vector with a strict accumulator
1430 scanl1' :: Vector v a => (a -> a -> a) -> v a -> v a
1431 {-# INLINE scanl1' #-}
1432 scanl1' f = unstream . inplace (MStream.scanl1' f) . stream
1433
1434 -- | /O(n)/ Right-to-left prescan
1435 --
1436 -- @
1437 -- prescanr f z = 'reverse' . 'prescanl' (flip f) z . 'reverse'
1438 -- @
1439 --
1440 prescanr :: (Vector v a, Vector v b) => (a -> b -> b) -> b -> v a -> v b
1441 {-# INLINE prescanr #-}
1442 prescanr f z = unstreamR . inplace (MStream.prescanl (flip f) z) . streamR
1443
1444 -- | /O(n)/ Right-to-left prescan with strict accumulator
1445 prescanr' :: (Vector v a, Vector v b) => (a -> b -> b) -> b -> v a -> v b
1446 {-# INLINE prescanr' #-}
1447 prescanr' f z = unstreamR . inplace (MStream.prescanl' (flip f) z) . streamR
1448
1449 -- | /O(n)/ Right-to-left scan
1450 postscanr :: (Vector v a, Vector v b) => (a -> b -> b) -> b -> v a -> v b
1451 {-# INLINE postscanr #-}
1452 postscanr f z = unstreamR . inplace (MStream.postscanl (flip f) z) . streamR
1453
1454 -- | /O(n)/ Right-to-left scan with strict accumulator
1455 postscanr' :: (Vector v a, Vector v b) => (a -> b -> b) -> b -> v a -> v b
1456 {-# INLINE postscanr' #-}
1457 postscanr' f z = unstreamR . inplace (MStream.postscanl' (flip f) z) . streamR
1458
1459 -- | /O(n)/ Right-to-left Haskell-style scan
1460 scanr :: (Vector v a, Vector v b) => (a -> b -> b) -> b -> v a -> v b
1461 {-# INLINE scanr #-}
1462 scanr f z = unstreamR . Stream.scanl (flip f) z . streamR
1463
1464 -- | /O(n)/ Right-to-left Haskell-style scan with strict accumulator
1465 scanr' :: (Vector v a, Vector v b) => (a -> b -> b) -> b -> v a -> v b
1466 {-# INLINE scanr' #-}
1467 scanr' f z = unstreamR . Stream.scanl' (flip f) z . streamR
1468
1469 -- | /O(n)/ Right-to-left scan over a non-empty vector
1470 scanr1 :: Vector v a => (a -> a -> a) -> v a -> v a
1471 {-# INLINE scanr1 #-}
1472 scanr1 f = unstreamR . inplace (MStream.scanl1 (flip f)) . streamR
1473
1474 -- | /O(n)/ Right-to-left scan over a non-empty vector with a strict
1475 -- accumulator
1476 scanr1' :: Vector v a => (a -> a -> a) -> v a -> v a
1477 {-# INLINE scanr1' #-}
1478 scanr1' f = unstreamR . inplace (MStream.scanl1' (flip f)) . streamR
1479
1480 -- Conversions - Lists
1481 -- ------------------------
1482
1483 -- | /O(n)/ Convert a vector to a list
1484 toList :: Vector v a => v a -> [a]
1485 {-# INLINE toList #-}
1486 toList = Stream.toList . stream
1487
1488 -- | /O(n)/ Convert a list to a vector
1489 fromList :: Vector v a => [a] -> v a
1490 {-# INLINE fromList #-}
1491 fromList = unstream . Stream.fromList
1492
1493 -- | /O(n)/ Convert the first @n@ elements of a list to a vector
1494 --
1495 -- > fromListN n xs = fromList (take n xs)
1496 fromListN :: Vector v a => Int -> [a] -> v a
1497 {-# INLINE fromListN #-}
1498 fromListN n = unstream . Stream.fromListN n
1499
1500 -- Conversions - Mutable vectors
1501 -- -----------------------------
1502
1503 -- | /O(n)/ Copy an immutable vector into a mutable one. The two vectors must
1504 -- have the same length.
1505 copy
1506 :: (PrimMonad m, Vector v a) => Mutable v (PrimState m) a -> v a -> m ()
1507 {-# INLINE copy #-}
1508 copy dst src = BOUNDS_CHECK(check) "copy" "length mismatch"
1509 (M.length dst == length src)
1510 $ unsafeCopy dst src
1511
1512 -- | /O(n)/ Copy an immutable vector into a mutable one. The two vectors must
1513 -- have the same length. This is not checked.
1514 unsafeCopy
1515 :: (PrimMonad m, Vector v a) => Mutable v (PrimState m) a -> v a -> m ()
1516 {-# INLINE unsafeCopy #-}
1517 unsafeCopy dst src = UNSAFE_CHECK(check) "unsafeCopy" "length mismatch"
1518 (M.length dst == length src)
1519 $ (dst `seq` src `seq` basicUnsafeCopy dst src)
1520
1521 -- Conversions to/from Streams
1522 -- ---------------------------
1523
1524 -- | /O(1)/ Convert a vector to a 'Stream'
1525 stream :: Vector v a => v a -> Stream a
1526 {-# INLINE_STREAM stream #-}
1527 stream v = v `seq` (Stream.unfoldr get 0 `Stream.sized` Exact n)
1528 where
1529 n = length v
1530
1531 -- NOTE: the False case comes first in Core so making it the recursive one
1532 -- makes the code easier to read
1533 {-# INLINE get #-}
1534 get i | i >= n = Nothing
1535 | otherwise = case basicUnsafeIndexM v i of Box x -> Just (x, i+1)
1536
1537 -- | /O(n)/ Construct a vector from a 'Stream'
1538 unstream :: Vector v a => Stream a -> v a
1539 {-# INLINE unstream #-}
1540 unstream s = new (New.unstream s)
1541
1542 {-# RULES
1543
1544 "stream/unstream [Vector]" forall s.
1545 stream (new (New.unstream s)) = s
1546
1547 "New.unstream/stream [Vector]" forall v.
1548 New.unstream (stream v) = clone v
1549
1550 "clone/new [Vector]" forall p.
1551 clone (new p) = p
1552
1553 "inplace [Vector]"
1554 forall (f :: forall m. Monad m => MStream m a -> MStream m a) m.
1555 New.unstream (inplace f (stream (new m))) = New.transform f m
1556
1557 "uninplace [Vector]"
1558 forall (f :: forall m. Monad m => MStream m a -> MStream m a) m.
1559 stream (new (New.transform f m)) = inplace f (stream (new m))
1560
1561 #-}
1562
1563 -- | /O(1)/ Convert a vector to a 'Stream', proceeding from right to left
1564 streamR :: Vector v a => v a -> Stream a
1565 {-# INLINE_STREAM streamR #-}
1566 streamR v = v `seq` (Stream.unfoldr get n `Stream.sized` Exact n)
1567 where
1568 n = length v
1569
1570 {-# INLINE get #-}
1571 get 0 = Nothing
1572 get i = let i' = i-1
1573 in
1574 case basicUnsafeIndexM v i' of Box x -> Just (x, i')
1575
1576 -- | /O(n)/ Construct a vector from a 'Stream', proceeding from right to left
1577 unstreamR :: Vector v a => Stream a -> v a
1578 {-# INLINE unstreamR #-}
1579 unstreamR s = new (New.unstreamR s)
1580
1581 {-# RULES
1582
1583 "streamR/unstreamR [Vector]" forall s.
1584 streamR (new (New.unstreamR s)) = s
1585
1586 "New.unstreamR/streamR/new [Vector]" forall p.
1587 New.unstreamR (streamR (new p)) = p
1588
1589 "inplace right [Vector]"
1590 forall (f :: forall m. Monad m => MStream m a -> MStream m a) m.
1591 New.unstreamR (inplace f (streamR (new m))) = New.transformR f m
1592
1593 "uninplace right [Vector]"
1594 forall (f :: forall m. Monad m => MStream m a -> MStream m a) m.
1595 streamR (new (New.transformR f m)) = inplace f (streamR (new m))
1596
1597 #-}
1598
1599 unstreamM :: (Vector v a, Monad m) => MStream m a -> m (v a)
1600 {-# INLINE_STREAM unstreamM #-}
1601 unstreamM s = do
1602 xs <- MStream.toList s
1603 return $ unstream $ Stream.unsafeFromList (MStream.size s) xs
1604
1605 -- Recycling support
1606 -- -----------------
1607
1608 -- | Construct a vector from a monadic initialiser.
1609 new :: Vector v a => New v a -> v a
1610 {-# INLINE_STREAM new #-}
1611 new m = m `seq` runST (unsafeFreeze =<< New.run m)
1612
1613 -- | Convert a vector to an initialiser which, when run, produces a copy of
1614 -- the vector.
1615 clone :: Vector v a => v a -> New v a
1616 {-# INLINE_STREAM clone #-}
1617 clone v = v `seq` New.create (
1618 do
1619 mv <- M.new (length v)
1620 unsafeCopy mv v
1621 return mv)
1622
1623 -- Comparisons
1624 -- -----------
1625
1626 -- | /O(n)/ Check if two vectors are equal. All 'Vector' instances are also
1627 -- instances of 'Eq' and it is usually more appropriate to use those. This
1628 -- function is primarily intended for implementing 'Eq' instances for new
1629 -- vector types.
1630 eq :: (Vector v a, Eq a) => v a -> v a -> Bool
1631 {-# INLINE eq #-}
1632 xs `eq` ys = stream xs == stream ys
1633
1634 -- | Compare two vectors lexicographically. All 'Vector' instances are also
1635 -- instances of 'Ord' and it is usually more appropriate to use those. This
1636 -- function is primarily intended for implementing 'Ord' instances for new
1637 -- vector types.
1638 cmp :: (Vector v a, Ord a) => v a -> v a -> Ordering
1639 {-# INLINE cmp #-}
1640 cmp xs ys = compare (stream xs) (stream ys)
1641
1642 -- Data and Typeable
1643 -- -----------------
1644
1645 -- | Generic definion of 'Data.Data.gfoldl' that views a 'Vector' as a
1646 -- list.
1647 gfoldl :: (Vector v a, Data a)
1648 => (forall d b. Data d => c (d -> b) -> d -> c b)
1649 -> (forall g. g -> c g)
1650 -> v a
1651 -> c (v a)
1652 {-# INLINE gfoldl #-}
1653 gfoldl f z v = z fromList `f` toList v
1654
1655 mkType :: String -> DataType
1656 {-# INLINE mkType #-}
1657 mkType = mkNorepType
1658
1659 dataCast :: (Vector v a, Data a, Typeable1 v, Typeable1 t)
1660 => (forall d. Data d => c (t d)) -> Maybe (c (v a))
1661 {-# INLINE dataCast #-}
1662 dataCast f = gcast1 f
1663