dbccb1642315035ed5fc24c81bddaba8816e8f21
[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 where 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 exaclty 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 given 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 index/value pair
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 concatMap :: (Vector v a, Vector v b) => (a -> v b) -> v a -> v b
767 {-# INLINE concatMap #-}
768 concatMap f = unstream . Stream.concatMap (stream . f) . stream
769
770 -- Zipping/unzipping
771 -- -----------------
772
773 -- | Zip two vectors with the given function.
774 zipWith :: (Vector v a, Vector v b, Vector v c)
775 => (a -> b -> c) -> v a -> v b -> v c
776 {-# INLINE zipWith #-}
777 zipWith f xs ys = unstream (Stream.zipWith f (stream xs) (stream ys))
778
779 -- | Zip three vectors with the given function.
780 zipWith3 :: (Vector v a, Vector v b, Vector v c, Vector v d)
781 => (a -> b -> c -> d) -> v a -> v b -> v c -> v d
782 {-# INLINE zipWith3 #-}
783 zipWith3 f as bs cs = unstream (Stream.zipWith3 f (stream as)
784 (stream bs)
785 (stream cs))
786
787 zipWith4 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e)
788 => (a -> b -> c -> d -> e) -> v a -> v b -> v c -> v d -> v e
789 {-# INLINE zipWith4 #-}
790 zipWith4 f as bs cs ds
791 = unstream (Stream.zipWith4 f (stream as)
792 (stream bs)
793 (stream cs)
794 (stream ds))
795
796 zipWith5 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
797 Vector v f)
798 => (a -> b -> c -> d -> e -> f) -> v a -> v b -> v c -> v d -> v e
799 -> v f
800 {-# INLINE zipWith5 #-}
801 zipWith5 f as bs cs ds es
802 = unstream (Stream.zipWith5 f (stream as)
803 (stream bs)
804 (stream cs)
805 (stream ds)
806 (stream es))
807
808 zipWith6 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
809 Vector v f, Vector v g)
810 => (a -> b -> c -> d -> e -> f -> g)
811 -> v a -> v b -> v c -> v d -> v e -> v f -> v g
812 {-# INLINE zipWith6 #-}
813 zipWith6 f as bs cs ds es fs
814 = unstream (Stream.zipWith6 f (stream as)
815 (stream bs)
816 (stream cs)
817 (stream ds)
818 (stream es)
819 (stream fs))
820
821 -- | Zip two vectors and their indices with the given function.
822 izipWith :: (Vector v a, Vector v b, Vector v c)
823 => (Int -> a -> b -> c) -> v a -> v b -> v c
824 {-# INLINE izipWith #-}
825 izipWith f xs ys = unstream
826 (Stream.zipWith (uncurry f) (Stream.indexed (stream xs))
827 (stream ys))
828
829 -- | Zip three vectors and their indices with the given function.
830 izipWith3 :: (Vector v a, Vector v b, Vector v c, Vector v d)
831 => (Int -> a -> b -> c -> d) -> v a -> v b -> v c -> v d
832 {-# INLINE izipWith3 #-}
833 izipWith3 f as bs cs
834 = unstream (Stream.zipWith3 (uncurry f) (Stream.indexed (stream as))
835 (stream bs)
836 (stream cs))
837
838 izipWith4 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e)
839 => (Int -> a -> b -> c -> d -> e) -> v a -> v b -> v c -> v d -> v e
840 {-# INLINE izipWith4 #-}
841 izipWith4 f as bs cs ds
842 = unstream (Stream.zipWith4 (uncurry f) (Stream.indexed (stream as))
843 (stream bs)
844 (stream cs)
845 (stream ds))
846
847 izipWith5 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
848 Vector v f)
849 => (Int -> a -> b -> c -> d -> e -> f) -> v a -> v b -> v c -> v d
850 -> v e -> v f
851 {-# INLINE izipWith5 #-}
852 izipWith5 f as bs cs ds es
853 = unstream (Stream.zipWith5 (uncurry f) (Stream.indexed (stream as))
854 (stream bs)
855 (stream cs)
856 (stream ds)
857 (stream es))
858
859 izipWith6 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
860 Vector v f, Vector v g)
861 => (Int -> a -> b -> c -> d -> e -> f -> g)
862 -> v a -> v b -> v c -> v d -> v e -> v f -> v g
863 {-# INLINE izipWith6 #-}
864 izipWith6 f as bs cs ds es fs
865 = unstream (Stream.zipWith6 (uncurry f) (Stream.indexed (stream as))
866 (stream bs)
867 (stream cs)
868 (stream ds)
869 (stream es)
870 (stream fs))
871
872 zip :: (Vector v a, Vector v b, Vector v (a,b)) => v a -> v b -> v (a, b)
873 {-# INLINE zip #-}
874 zip = zipWith (,)
875
876 zip3 :: (Vector v a, Vector v b, Vector v c, Vector v (a, b, c))
877 => v a -> v b -> v c -> v (a, b, c)
878 {-# INLINE zip3 #-}
879 zip3 = zipWith3 (,,)
880
881 zip4 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v (a, b, c, d))
882 => v a -> v b -> v c -> v d -> v (a, b, c, d)
883 {-# INLINE zip4 #-}
884 zip4 = zipWith4 (,,,)
885
886 zip5 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
887 Vector v (a, b, c, d, e))
888 => v a -> v b -> v c -> v d -> v e -> v (a, b, c, d, e)
889 {-# INLINE zip5 #-}
890 zip5 = zipWith5 (,,,,)
891
892 zip6 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
893 Vector v f, Vector v (a, b, c, d, e, f))
894 => v a -> v b -> v c -> v d -> v e -> v f -> v (a, b, c, d, e, f)
895 {-# INLINE zip6 #-}
896 zip6 = zipWith6 (,,,,,)
897
898 unzip :: (Vector v a, Vector v b, Vector v (a,b)) => v (a, b) -> (v a, v b)
899 {-# INLINE unzip #-}
900 unzip xs = (map fst xs, map snd xs)
901
902 unzip3 :: (Vector v a, Vector v b, Vector v c, Vector v (a, b, c))
903 => v (a, b, c) -> (v a, v b, v c)
904 {-# INLINE unzip3 #-}
905 unzip3 xs = (map (\(a, b, c) -> a) xs,
906 map (\(a, b, c) -> b) xs,
907 map (\(a, b, c) -> c) xs)
908
909 unzip4 :: (Vector v a, Vector v b, Vector v c, Vector v d,
910 Vector v (a, b, c, d))
911 => v (a, b, c, d) -> (v a, v b, v c, v d)
912 {-# INLINE unzip4 #-}
913 unzip4 xs = (map (\(a, b, c, d) -> a) xs,
914 map (\(a, b, c, d) -> b) xs,
915 map (\(a, b, c, d) -> c) xs,
916 map (\(a, b, c, d) -> d) xs)
917
918 unzip5 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
919 Vector v (a, b, c, d, e))
920 => v (a, b, c, d, e) -> (v a, v b, v c, v d, v e)
921 {-# INLINE unzip5 #-}
922 unzip5 xs = (map (\(a, b, c, d, e) -> a) xs,
923 map (\(a, b, c, d, e) -> b) xs,
924 map (\(a, b, c, d, e) -> c) xs,
925 map (\(a, b, c, d, e) -> d) xs,
926 map (\(a, b, c, d, e) -> e) xs)
927
928 unzip6 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
929 Vector v f, Vector v (a, b, c, d, e, f))
930 => v (a, b, c, d, e, f) -> (v a, v b, v c, v d, v e, v f)
931 {-# INLINE unzip6 #-}
932 unzip6 xs = (map (\(a, b, c, d, e, f) -> a) xs,
933 map (\(a, b, c, d, e, f) -> b) xs,
934 map (\(a, b, c, d, e, f) -> c) xs,
935 map (\(a, b, c, d, e, f) -> d) xs,
936 map (\(a, b, c, d, e, f) -> e) xs,
937 map (\(a, b, c, d, e, f) -> f) xs)
938
939 -- Comparisons
940 -- -----------
941
942 eq :: (Vector v a, Eq a) => v a -> v a -> Bool
943 {-# INLINE eq #-}
944 xs `eq` ys = stream xs == stream ys
945
946 cmp :: (Vector v a, Ord a) => v a -> v a -> Ordering
947 {-# INLINE cmp #-}
948 cmp xs ys = compare (stream xs) (stream ys)
949
950 -- Filtering
951 -- ---------
952
953 -- | Drop elements that do not satisfy the predicate
954 filter :: Vector v a => (a -> Bool) -> v a -> v a
955 {-# INLINE filter #-}
956 filter f = unstream . inplace (MStream.filter f) . stream
957
958 -- | Drop elements that do not satisfy the predicate (applied to values and
959 -- their indices)
960 ifilter :: Vector v a => (Int -> a -> Bool) -> v a -> v a
961 {-# INLINE ifilter #-}
962 ifilter f = unstream
963 . inplace (MStream.map snd . MStream.filter (uncurry f)
964 . MStream.indexed)
965 . stream
966
967 -- | Yield the longest prefix of elements satisfying the predicate.
968 takeWhile :: Vector v a => (a -> Bool) -> v a -> v a
969 {-# INLINE takeWhile #-}
970 takeWhile f = unstream . Stream.takeWhile f . stream
971
972 -- | Drop the longest prefix of elements that satisfy the predicate.
973 dropWhile :: Vector v a => (a -> Bool) -> v a -> v a
974 {-# INLINE dropWhile #-}
975 dropWhile f = unstream . Stream.dropWhile f . stream
976
977 -- | Split the vector in two parts, the first one containing those elements
978 -- that satisfy the predicate and the second one those that don't. The
979 -- relative order of the elements is preserved at the cost of a (sometimes)
980 -- reduced performance compared to 'unstablePartition'.
981 partition :: Vector v a => (a -> Bool) -> v a -> (v a, v a)
982 {-# INLINE partition #-}
983 partition f = partition_stream f . stream
984
985 -- FIXME: Make this inplace-fusible (look at how stable_partition is
986 -- implemented in C++)
987
988 partition_stream :: Vector v a => (a -> Bool) -> Stream a -> (v a, v a)
989 {-# INLINE_STREAM partition_stream #-}
990 partition_stream f s = s `seq` runST (
991 do
992 (mv1,mv2) <- M.partitionStream f s
993 v1 <- unsafeFreeze mv1
994 v2 <- unsafeFreeze mv2
995 return (v1,v2))
996
997 -- | Split the vector in two parts, the first one containing those elements
998 -- that satisfy the predicate and the second one those that don't. The order
999 -- of the elements is not preserved but the operation is often faster than
1000 -- 'partition'.
1001 unstablePartition :: Vector v a => (a -> Bool) -> v a -> (v a, v a)
1002 {-# INLINE unstablePartition #-}
1003 unstablePartition f = unstablePartition_stream f . stream
1004
1005 unstablePartition_stream
1006 :: Vector v a => (a -> Bool) -> Stream a -> (v a, v a)
1007 {-# INLINE_STREAM unstablePartition_stream #-}
1008 unstablePartition_stream f s = s `seq` runST (
1009 do
1010 (mv1,mv2) <- M.unstablePartitionStream f s
1011 v1 <- unsafeFreeze mv1
1012 v2 <- unsafeFreeze mv2
1013 return (v1,v2))
1014
1015 unstablePartition_new :: Vector v a => (a -> Bool) -> New v a -> (v a, v a)
1016 {-# INLINE_STREAM unstablePartition_new #-}
1017 unstablePartition_new f (New.New p) = runST (
1018 do
1019 mv <- p
1020 i <- M.unstablePartition f mv
1021 v <- unsafeFreeze mv
1022 return (unsafeTake i v, unsafeDrop i v))
1023
1024 {-# RULES
1025
1026 "unstablePartition" forall f p.
1027 unstablePartition_stream f (stream (new p))
1028 = unstablePartition_new f p
1029
1030 #-}
1031
1032
1033 -- FIXME: make span and break fusible
1034
1035 -- | Split the vector into the longest prefix of elements that satisfy the
1036 -- predicate and the rest.
1037 span :: Vector v a => (a -> Bool) -> v a -> (v a, v a)
1038 {-# INLINE span #-}
1039 span f = break (not . f)
1040
1041 -- | Split the vector into the longest prefix of elements that do not satisfy
1042 -- the predicate and the rest.
1043 break :: Vector v a => (a -> Bool) -> v a -> (v a, v a)
1044 {-# INLINE break #-}
1045 break f xs = case findIndex f xs of
1046 Just i -> (unsafeSlice 0 i xs, unsafeSlice i (length xs - i) xs)
1047 Nothing -> (xs, empty)
1048
1049
1050 -- Searching
1051 -- ---------
1052
1053 infix 4 `elem`
1054 -- | Check whether the vector contains an element
1055 elem :: (Vector v a, Eq a) => a -> v a -> Bool
1056 {-# INLINE elem #-}
1057 elem x = Stream.elem x . stream
1058
1059 infix 4 `notElem`
1060 -- | Inverse of `elem`
1061 notElem :: (Vector v a, Eq a) => a -> v a -> Bool
1062 {-# INLINE notElem #-}
1063 notElem x = Stream.notElem x . stream
1064
1065 -- | Yield 'Just' the first element matching the predicate or 'Nothing' if no
1066 -- such element exists.
1067 find :: Vector v a => (a -> Bool) -> v a -> Maybe a
1068 {-# INLINE find #-}
1069 find f = Stream.find f . stream
1070
1071 -- | Yield 'Just' the index of the first element matching the predicate or
1072 -- 'Nothing' if no such element exists.
1073 findIndex :: Vector v a => (a -> Bool) -> v a -> Maybe Int
1074 {-# INLINE findIndex #-}
1075 findIndex f = Stream.findIndex f . stream
1076
1077 -- | Yield the indices of elements satisfying the predicate
1078 findIndices :: (Vector v a, Vector v Int) => (a -> Bool) -> v a -> v Int
1079 {-# INLINE findIndices #-}
1080 findIndices f = unstream
1081 . inplace (MStream.map fst . MStream.filter (f . snd)
1082 . MStream.indexed)
1083 . stream
1084
1085 -- | Yield 'Just' the index of the first occurence of the given element or
1086 -- 'Nothing' if the vector does not contain the element
1087 elemIndex :: (Vector v a, Eq a) => a -> v a -> Maybe Int
1088 {-# INLINE elemIndex #-}
1089 elemIndex x = findIndex (x==)
1090
1091 -- | Yield the indices of all occurences of the given element
1092 elemIndices :: (Vector v a, Vector v Int, Eq a) => a -> v a -> v Int
1093 {-# INLINE elemIndices #-}
1094 elemIndices x = findIndices (x==)
1095
1096 -- Folding
1097 -- -------
1098
1099 -- | Left fold
1100 foldl :: Vector v b => (a -> b -> a) -> a -> v b -> a
1101 {-# INLINE foldl #-}
1102 foldl f z = Stream.foldl f z . stream
1103
1104 -- | Left fold on non-empty vectors
1105 foldl1 :: Vector v a => (a -> a -> a) -> v a -> a
1106 {-# INLINE foldl1 #-}
1107 foldl1 f = Stream.foldl1 f . stream
1108
1109 -- | Left fold with strict accumulator
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 with strict accumulator
1115 foldl1' :: Vector v a => (a -> a -> a) -> v a -> a
1116 {-# INLINE foldl1' #-}
1117 foldl1' f = Stream.foldl1' f . stream
1118
1119 -- | Right fold
1120 foldr :: Vector v a => (a -> b -> b) -> b -> v a -> b
1121 {-# INLINE foldr #-}
1122 foldr f z = Stream.foldr f z . stream
1123
1124 -- | Right fold on non-empty vectors
1125 foldr1 :: Vector v a => (a -> a -> a) -> v a -> a
1126 {-# INLINE foldr1 #-}
1127 foldr1 f = Stream.foldr1 f . stream
1128
1129 -- | Right fold with a strict accumulator
1130 foldr' :: Vector v a => (a -> b -> b) -> b -> v a -> b
1131 {-# INLINE foldr' #-}
1132 foldr' f z = Stream.foldl' (flip f) z . streamR
1133
1134 -- | Right fold on non-empty vectors with strict accumulator
1135 foldr1' :: Vector v a => (a -> a -> a) -> v a -> a
1136 {-# INLINE foldr1' #-}
1137 foldr1' f = Stream.foldl1' (flip f) . streamR
1138
1139 -- | Left fold (function applied to each element and its index)
1140 ifoldl :: Vector v b => (a -> Int -> b -> a) -> a -> v b -> a
1141 {-# INLINE ifoldl #-}
1142 ifoldl f z = Stream.foldl (uncurry . f) z . Stream.indexed . stream
1143
1144 -- | Left fold with strict accumulator (function applied to each element and
1145 -- its index)
1146 ifoldl' :: Vector v b => (a -> Int -> b -> a) -> a -> v b -> a
1147 {-# INLINE ifoldl' #-}
1148 ifoldl' f z = Stream.foldl' (uncurry . f) z . Stream.indexed . stream
1149
1150 -- | Right fold (function applied to each element and its index)
1151 ifoldr :: Vector v a => (Int -> a -> b -> b) -> b -> v a -> b
1152 {-# INLINE ifoldr #-}
1153 ifoldr f z = Stream.foldr (uncurry f) z . Stream.indexed . stream
1154
1155 -- | Right fold with strict accumulator (function applied to each element and
1156 -- its index)
1157 ifoldr' :: Vector v a => (Int -> a -> b -> b) -> b -> v a -> b
1158 {-# INLINE ifoldr' #-}
1159 ifoldr' f z xs = Stream.foldl' (flip (uncurry f)) z
1160 $ Stream.indexedR (length xs) $ streamR xs
1161
1162 -- Specialised folds
1163 -- -----------------
1164
1165 all :: Vector v a => (a -> Bool) -> v a -> Bool
1166 {-# INLINE all #-}
1167 all f = Stream.and . Stream.map f . stream
1168
1169 any :: Vector v a => (a -> Bool) -> v a -> Bool
1170 {-# INLINE any #-}
1171 any f = Stream.or . Stream.map f . stream
1172
1173 and :: Vector v Bool => v Bool -> Bool
1174 {-# INLINE and #-}
1175 and = Stream.and . stream
1176
1177 or :: Vector v Bool => v Bool -> Bool
1178 {-# INLINE or #-}
1179 or = Stream.or . stream
1180
1181 sum :: (Vector v a, Num a) => v a -> a
1182 {-# INLINE sum #-}
1183 sum = Stream.foldl' (+) 0 . stream
1184
1185 product :: (Vector v a, Num a) => v a -> a
1186 {-# INLINE product #-}
1187 product = Stream.foldl' (*) 1 . stream
1188
1189 maximum :: (Vector v a, Ord a) => v a -> a
1190 {-# INLINE maximum #-}
1191 maximum = Stream.foldl1' max . stream
1192
1193 maximumBy :: Vector v a => (a -> a -> Ordering) -> v a -> a
1194 {-# INLINE maximumBy #-}
1195 maximumBy cmp = Stream.foldl1' maxBy . stream
1196 where
1197 {-# INLINE maxBy #-}
1198 maxBy x y = case cmp x y of
1199 LT -> y
1200 _ -> x
1201
1202 minimum :: (Vector v a, Ord a) => v a -> a
1203 {-# INLINE minimum #-}
1204 minimum = Stream.foldl1' min . stream
1205
1206 minimumBy :: Vector v a => (a -> a -> Ordering) -> v a -> a
1207 {-# INLINE minimumBy #-}
1208 minimumBy cmp = Stream.foldl1' minBy . stream
1209 where
1210 {-# INLINE minBy #-}
1211 minBy x y = case cmp x y of
1212 GT -> y
1213 _ -> x
1214
1215 maxIndex :: (Vector v a, Ord a) => v a -> Int
1216 {-# INLINE maxIndex #-}
1217 maxIndex = maxIndexBy compare
1218
1219 maxIndexBy :: Vector v a => (a -> a -> Ordering) -> v a -> Int
1220 {-# INLINE maxIndexBy #-}
1221 maxIndexBy cmp = fst . Stream.foldl1' imax . Stream.indexed . stream
1222 where
1223 imax (i,x) (j,y) = i `seq` j `seq`
1224 case cmp x y of
1225 LT -> (j,y)
1226 _ -> (i,x)
1227
1228 minIndex :: (Vector v a, Ord a) => v a -> Int
1229 {-# INLINE minIndex #-}
1230 minIndex = minIndexBy compare
1231
1232 minIndexBy :: Vector v a => (a -> a -> Ordering) -> v a -> Int
1233 {-# INLINE minIndexBy #-}
1234 minIndexBy cmp = fst . Stream.foldl1' imin . Stream.indexed . stream
1235 where
1236 imin (i,x) (j,y) = i `seq` j `seq`
1237 case cmp x y of
1238 GT -> (j,y)
1239 _ -> (i,x)
1240
1241
1242 -- Unfolding
1243 -- ---------
1244
1245 -- | Unfold
1246 unfoldr :: Vector v a => (b -> Maybe (a, b)) -> b -> v a
1247 {-# INLINE unfoldr #-}
1248 unfoldr f = unstream . Stream.unfoldr f
1249
1250 -- | Unfoldr at most @n@ elements.
1251 unfoldrN :: Vector v a => Int -> (b -> Maybe (a, b)) -> b -> v a
1252 {-# INLINE unfoldrN #-}
1253 unfoldrN n f = unstream . Stream.unfoldrN n f
1254
1255 -- Scans
1256 -- -----
1257
1258 -- | Prefix scan
1259 prescanl :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1260 {-# INLINE prescanl #-}
1261 prescanl f z = unstream . inplace (MStream.prescanl f z) . stream
1262
1263 -- | Prefix scan with strict accumulator
1264 prescanl' :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1265 {-# INLINE prescanl' #-}
1266 prescanl' f z = unstream . inplace (MStream.prescanl' f z) . stream
1267
1268 -- | Suffix scan
1269 postscanl :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1270 {-# INLINE postscanl #-}
1271 postscanl f z = unstream . inplace (MStream.postscanl f z) . stream
1272
1273 -- | Suffix scan with strict accumulator
1274 postscanl' :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1275 {-# INLINE postscanl' #-}
1276 postscanl' f z = unstream . inplace (MStream.postscanl' f z) . stream
1277
1278 -- | Haskell-style scan
1279 scanl :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1280 {-# INLINE scanl #-}
1281 scanl f z = unstream . Stream.scanl f z . stream
1282
1283 -- | Haskell-style scan with strict accumulator
1284 scanl' :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1285 {-# INLINE scanl' #-}
1286 scanl' f z = unstream . Stream.scanl' f z . stream
1287
1288 -- | Scan over a non-empty vector
1289 scanl1 :: Vector v a => (a -> a -> a) -> v a -> v a
1290 {-# INLINE scanl1 #-}
1291 scanl1 f = unstream . inplace (MStream.scanl1 f) . stream
1292
1293 -- | Scan over a non-empty vector with a strict accumulator
1294 scanl1' :: Vector v a => (a -> a -> a) -> v a -> v a
1295 {-# INLINE scanl1' #-}
1296 scanl1' f = unstream . inplace (MStream.scanl1' f) . stream
1297
1298
1299 -- | Prefix right-to-left scan
1300 prescanr :: (Vector v a, Vector v b) => (a -> b -> b) -> b -> v a -> v b
1301 {-# INLINE prescanr #-}
1302 prescanr f z = unstreamR . inplace (MStream.prescanl (flip f) z) . streamR
1303
1304 -- | Prefix right-to-left scan with strict accumulator
1305 prescanr' :: (Vector v a, Vector v b) => (a -> b -> b) -> b -> v a -> v b
1306 {-# INLINE prescanr' #-}
1307 prescanr' f z = unstreamR . inplace (MStream.prescanl' (flip f) z) . streamR
1308
1309 -- | Suffix right-to-left scan
1310 postscanr :: (Vector v a, Vector v b) => (a -> b -> b) -> b -> v a -> v b
1311 {-# INLINE postscanr #-}
1312 postscanr f z = unstreamR . inplace (MStream.postscanl (flip f) z) . streamR
1313
1314 -- | Suffix right-to-left scan with strict accumulator
1315 postscanr' :: (Vector v a, Vector v b) => (a -> b -> b) -> b -> v a -> v b
1316 {-# INLINE postscanr' #-}
1317 postscanr' f z = unstreamR . inplace (MStream.postscanl' (flip f) z) . streamR
1318
1319 -- | Haskell-style right-to-left scan
1320 scanr :: (Vector v a, Vector v b) => (a -> b -> b) -> b -> v a -> v b
1321 {-# INLINE scanr #-}
1322 scanr f z = unstreamR . Stream.scanl (flip f) z . streamR
1323
1324 -- | Haskell-style right-to-left scan with strict accumulator
1325 scanr' :: (Vector v a, Vector v b) => (a -> b -> b) -> b -> v a -> v b
1326 {-# INLINE scanr' #-}
1327 scanr' f z = unstreamR . Stream.scanl' (flip f) z . streamR
1328
1329 -- | Right-to-left scan over a non-empty vector
1330 scanr1 :: Vector v a => (a -> a -> a) -> v a -> v a
1331 {-# INLINE scanr1 #-}
1332 scanr1 f = unstreamR . inplace (MStream.scanl1 (flip f)) . streamR
1333
1334 -- | Right-to-left scan over a non-empty vector with a strict accumulator
1335 scanr1' :: Vector v a => (a -> a -> a) -> v a -> v a
1336 {-# INLINE scanr1' #-}
1337 scanr1' f = unstreamR . inplace (MStream.scanl1' (flip f)) . streamR
1338
1339 -- Enumeration
1340 -- -----------
1341
1342 -- | Yield a vector of the given length containing the values @x@, @x+1@ etc.
1343 -- This operation is usually more efficient than 'enumFromTo'.
1344 enumFromN :: (Vector v a, Num a) => a -> Int -> v a
1345 {-# INLINE enumFromN #-}
1346 enumFromN x n = enumFromStepN x 1 n
1347
1348 -- | Yield a vector of the given length containing the values @x@, @x+y@,
1349 -- @x+y+y@ etc. This operations is usually more efficient than
1350 -- 'enumFromThenTo'.
1351 enumFromStepN :: forall v a. (Vector v a, Num a) => a -> a -> Int -> v a
1352 {-# INLINE enumFromStepN #-}
1353 enumFromStepN x y n = elemseq (undefined :: v a) x
1354 $ elemseq (undefined :: v a) y
1355 $ unstream
1356 $ Stream.enumFromStepN x y n
1357
1358 -- | Enumerate values from @x@ to @y@.
1359 --
1360 -- /WARNING:/ This operation can be very inefficient. If at all possible, use
1361 -- 'enumFromN' instead.
1362 enumFromTo :: (Vector v a, Enum a) => a -> a -> v a
1363 {-# INLINE enumFromTo #-}
1364 enumFromTo x y = unstream (Stream.enumFromTo x y)
1365
1366 -- | Enumerate values from @x@ to @y@ with a specific step @z@.
1367 --
1368 -- /WARNING:/ This operation can be very inefficient. If at all possible, use
1369 -- 'enumFromStepN' instead.
1370 enumFromThenTo :: (Vector v a, Enum a) => a -> a -> a -> v a
1371 {-# INLINE enumFromThenTo #-}
1372 enumFromThenTo x y z = unstream (Stream.enumFromThenTo x y z)
1373
1374 -- Conversion to/from lists
1375 -- ------------------------
1376
1377 -- | Convert a vector to a list
1378 toList :: Vector v a => v a -> [a]
1379 {-# INLINE toList #-}
1380 toList = Stream.toList . stream
1381
1382 -- | Convert a list to a vector
1383 fromList :: Vector v a => [a] -> v a
1384 {-# INLINE fromList #-}
1385 fromList = unstream . Stream.fromList
1386
1387 -- | Convert the first @n@ elements of a list to a vector
1388 --
1389 -- > fromListN n xs = fromList (take n xs)
1390 fromListN :: Vector v a => Int -> [a] -> v a
1391 {-# INLINE fromListN #-}
1392 fromListN n = unstream . Stream.fromListN n
1393
1394 unstreamM :: (Vector v a, Monad m) => MStream m a -> m (v a)
1395 {-# INLINE_STREAM unstreamM #-}
1396 unstreamM s = do
1397 xs <- MStream.toList s
1398 return $ unstream $ Stream.unsafeFromList (MStream.size s) xs
1399
1400 -- Monadic operations
1401 -- ------------------
1402
1403 -- FIXME: specialise various combinators for ST and IO?
1404
1405 -- | Perform the monadic action the given number of times and store the
1406 -- results in a vector.
1407 replicateM :: (Monad m, Vector v a) => Int -> m a -> m (v a)
1408 {-# INLINE replicateM #-}
1409 replicateM n m = fromListN n `Monad.liftM` Monad.replicateM n m
1410
1411 -- | Apply the monadic action to all elements of the vector, yielding a vector
1412 -- of results
1413 mapM :: (Monad m, Vector v a, Vector v b) => (a -> m b) -> v a -> m (v b)
1414 {-# INLINE mapM #-}
1415 mapM f = unstreamM . Stream.mapM f . stream
1416
1417 -- | Apply the monadic action to all elements of a vector and ignore the
1418 -- results
1419 mapM_ :: (Monad m, Vector v a) => (a -> m b) -> v a -> m ()
1420 {-# INLINE mapM_ #-}
1421 mapM_ f = Stream.mapM_ f . stream
1422
1423 -- | Apply the monadic action to all elements of the vector, yielding a vector
1424 -- of results
1425 forM :: (Monad m, Vector v a, Vector v b) => v a -> (a -> m b) -> m (v b)
1426 {-# INLINE forM #-}
1427 forM as f = mapM f as
1428
1429 -- | Apply the monadic action to all elements of a vector and ignore the
1430 -- results
1431 forM_ :: (Monad m, Vector v a) => v a -> (a -> m b) -> m ()
1432 {-# INLINE forM_ #-}
1433 forM_ as f = mapM_ f as
1434
1435 -- | Zip the two vectors with the monadic action and yield a vector of results
1436 zipWithM :: (Monad m, Vector v a, Vector v b, Vector v c)
1437 => (a -> b -> m c) -> v a -> v b -> m (v c)
1438 {-# INLINE zipWithM #-}
1439 zipWithM f as bs = unstreamM $ Stream.zipWithM f (stream as) (stream bs)
1440
1441 -- | Zip the two vectors with the monadic action and ignore the results
1442 zipWithM_ :: (Monad m, Vector v a, Vector v b)
1443 => (a -> b -> m c) -> v a -> v b -> m ()
1444 {-# INLINE zipWithM_ #-}
1445 zipWithM_ f as bs = Stream.zipWithM_ f (stream as) (stream bs)
1446
1447 -- | Drop elements that do not satisfy the monadic predicate
1448 filterM :: (Monad m, Vector v a) => (a -> m Bool) -> v a -> m (v a)
1449 {-# INLINE filterM #-}
1450 filterM f = unstreamM . Stream.filterM f . stream
1451
1452 -- | Monadic fold
1453 foldM :: (Monad m, Vector v b) => (a -> b -> m a) -> a -> v b -> m a
1454 {-# INLINE foldM #-}
1455 foldM m z = Stream.foldM m z . stream
1456
1457 -- | Monadic fold over non-empty vectors
1458 fold1M :: (Monad m, Vector v a) => (a -> a -> m a) -> v a -> m a
1459 {-# INLINE fold1M #-}
1460 fold1M m = Stream.fold1M m . stream
1461
1462 -- | Monadic fold with strict accumulator
1463 foldM' :: (Monad m, Vector v b) => (a -> b -> m a) -> a -> v b -> m a
1464 {-# INLINE foldM' #-}
1465 foldM' m z = Stream.foldM' m z . stream
1466
1467 -- | Monad fold over non-empty vectors with strict accumulator
1468 fold1M' :: (Monad m, Vector v a) => (a -> a -> m a) -> v a -> m a
1469 {-# INLINE fold1M' #-}
1470 fold1M' m = Stream.fold1M' m . stream
1471
1472 -- Destructive operations
1473 -- ----------------------
1474
1475 -- | Destructively initialise a vector.
1476 create :: Vector v a => (forall s. ST s (Mutable v s a)) -> v a
1477 {-# INLINE create #-}
1478 create p = new (New.create p)
1479
1480 -- | Apply a destructive operation to a vector. The operation modifies a
1481 -- copy of the vector unless it can be safely performed in place.
1482 modify :: Vector v a => (forall s. Mutable v s a -> ST s ()) -> v a -> v a
1483 {-# INLINE modify #-}
1484 modify p = new . New.modify p . clone
1485
1486 -- We have to make sure that this is strict in the stream but we can't seq on
1487 -- it while fusion is happening. Hence this ugliness.
1488 modifyWithStream :: Vector v a
1489 => (forall s. Mutable v s a -> Stream b -> ST s ())
1490 -> v a -> Stream b -> v a
1491 {-# INLINE modifyWithStream #-}
1492 modifyWithStream p v s = new (New.modifyWithStream p (clone v) s)
1493
1494 -- | Copy an immutable vector into a mutable one. The two vectors must have
1495 -- the same length. This is not checked.
1496 unsafeCopy
1497 :: (PrimMonad m, Vector v a) => Mutable v (PrimState m) a -> v a -> m ()
1498 {-# INLINE unsafeCopy #-}
1499 unsafeCopy dst src = UNSAFE_CHECK(check) "unsafeCopy" "length mismatch"
1500 (M.length dst == length src)
1501 $ (dst `seq` src `seq` basicUnsafeCopy dst src)
1502
1503 -- | Copy an immutable vector into a mutable one. The two vectors must have the
1504 -- same length.
1505 copy
1506 :: (PrimMonad m, Vector v a) => Mutable v (PrimState m) a -> v a -> m ()
1507 {-# INLINE copy #-}
1508 copy dst src = BOUNDS_CHECK(check) "copy" "length mismatch"
1509 (M.length dst == length src)
1510 $ unsafeCopy dst src
1511
1512 -- Utilities for defining Data instances
1513 -- -------------------------------------
1514
1515 -- | Generic definion of 'Data.Data.gfoldl' that views a 'Vector' as a
1516 -- list.
1517 gfoldl :: (Vector v a, Data a)
1518 => (forall d b. Data d => c (d -> b) -> d -> c b)
1519 -> (forall g. g -> c g)
1520 -> v a
1521 -> c (v a)
1522 {-# INLINE gfoldl #-}
1523 gfoldl f z v = z fromList `f` toList v
1524
1525 mkType :: String -> DataType
1526 {-# INLINE mkType #-}
1527 mkType = mkNorepType
1528
1529 dataCast :: (Vector v a, Data a, Typeable1 v, Typeable1 t)
1530 => (forall d. Data d => c (t d)) -> Maybe (c (v a))
1531 {-# INLINE dataCast #-}
1532 dataCast f = gcast1 f
1533