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