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