Rename iterate -> iterateN
[darcs-mirrors/vector.git] / Data / Vector / Generic.hs
1 {-# LANGUAGE Rank2Types, MultiParamTypeClasses, FlexibleContexts,
2 TypeFamilies, ScopedTypeVariables #-}
3 -- |
4 -- Module : Data.Vector.Generic
5 -- Copyright : (c) Roman Leshchinskiy 2008-2010
6 -- License : BSD-style
7 --
8 -- Maintainer : Roman Leshchinskiy <rl@cse.unsw.edu.au>
9 -- Stability : experimental
10 -- Portability : non-portable
11 --
12 -- Generic interface to pure vectors.
13 --
14
15 module Data.Vector.Generic (
16 -- * Immutable vectors
17 Vector(..), Mutable,
18
19 -- * Accessors
20
21 -- ** Length information
22 length, null,
23
24 -- ** Indexing
25 (!), (!?), head, last,
26 unsafeIndex, unsafeHead, unsafeLast,
27
28 -- ** Monadic indexing
29 indexM, headM, lastM,
30 unsafeIndexM, unsafeHeadM, unsafeLastM,
31
32 -- ** Extracting subvectors (slicing)
33 slice, init, tail, take, drop, splitAt,
34 unsafeSlice, unsafeInit, unsafeTail, unsafeTake, unsafeDrop,
35
36 -- * Construction
37
38 -- ** Initialisation
39 empty, singleton, replicate, generate, iterateN,
40
41 -- ** Monadic initialisation
42 replicateM, generateM, create,
43
44 -- ** Unfolding
45 unfoldr, unfoldrN,
46
47 -- ** Enumeration
48 enumFromN, enumFromStepN, enumFromTo, enumFromThenTo,
49
50 -- ** Concatenation
51 cons, snoc, (++), concat,
52
53 -- ** Restricting memory usage
54 force,
55
56 -- * Modifying vectors
57
58 -- ** Bulk updates
59 (//), update, update_,
60 unsafeUpd, unsafeUpdate, unsafeUpdate_,
61
62 -- ** Accumulations
63 accum, accumulate, accumulate_,
64 unsafeAccum, unsafeAccumulate, unsafeAccumulate_,
65
66 -- ** Permutations
67 reverse, backpermute, unsafeBackpermute,
68
69 -- ** Safe destructive updates
70 modify,
71
72 -- * Elementwise operations
73
74 -- ** Indexing
75 indexed,
76
77 -- ** Mapping
78 map, imap, concatMap,
79
80 -- ** Monadic mapping
81 mapM, mapM_, forM, forM_,
82
83 -- ** Zipping
84 zipWith, zipWith3, zipWith4, zipWith5, zipWith6,
85 izipWith, izipWith3, izipWith4, izipWith5, izipWith6,
86 zip, zip3, zip4, zip5, zip6,
87
88 -- ** Monadic zipping
89 zipWithM, zipWithM_,
90
91 -- ** Unzipping
92 unzip, unzip3, unzip4, unzip5, unzip6,
93
94 -- * Working with predicates
95
96 -- ** Filtering
97 filter, ifilter, filterM,
98 takeWhile, dropWhile,
99
100 -- ** Partitioning
101 partition, unstablePartition, span, break,
102
103 -- ** Searching
104 elem, notElem, find, findIndex, findIndices, elemIndex, elemIndices,
105
106 -- * Folding
107 foldl, foldl1, foldl', foldl1', foldr, foldr1, foldr', foldr1',
108 ifoldl, ifoldl', ifoldr, ifoldr',
109
110 -- ** Specialised folds
111 all, any, and, or,
112 sum, product,
113 maximum, maximumBy, minimum, minimumBy,
114 minIndex, minIndexBy, maxIndex, maxIndexBy,
115
116 -- ** Monadic folds
117 foldM, foldM', fold1M, fold1M',
118 foldM_, foldM'_, fold1M_, fold1M'_,
119
120 -- ** Monadic sequencing
121 sequence, sequence_,
122
123 -- * Prefix sums (scans)
124 prescanl, prescanl',
125 postscanl, postscanl',
126 scanl, scanl', scanl1, scanl1',
127 prescanr, prescanr',
128 postscanr, postscanr',
129 scanr, scanr', scanr1, scanr1',
130
131 -- * Conversions
132
133 -- ** Lists
134 toList, fromList, fromListN,
135
136 -- ** Different vector types
137 convert,
138
139 -- ** Mutable vectors
140 freeze, thaw, copy, unsafeFreeze, unsafeThaw, unsafeCopy,
141
142 -- * Fusion support
143
144 -- ** Conversion to/from Streams
145 stream, unstream, streamR, unstreamR,
146
147 -- ** Recycling support
148 new, clone,
149
150 -- * Utilities
151
152 -- ** Comparisons
153 eq, cmp,
154
155 -- ** @Data@ and @Typeable@
156 gfoldl, dataCast, mkType
157 ) where
158
159 import Data.Vector.Generic.Base
160
161 import Data.Vector.Generic.Mutable ( MVector )
162 import qualified Data.Vector.Generic.Mutable as M
163
164 import qualified Data.Vector.Generic.New as New
165 import Data.Vector.Generic.New ( New )
166
167 import qualified Data.Vector.Fusion.Stream as Stream
168 import Data.Vector.Fusion.Stream ( Stream, MStream, inplace, liftStream )
169 import qualified Data.Vector.Fusion.Stream.Monadic as MStream
170 import Data.Vector.Fusion.Stream.Size
171 import Data.Vector.Fusion.Util
172
173 import Control.Monad.ST ( ST, runST )
174 import Control.Monad.Primitive
175 import qualified Control.Monad as Monad
176 import qualified Data.List as List
177 import Prelude hiding ( length, null,
178 replicate, (++), concat,
179 head, last,
180 init, tail, take, drop, splitAt, reverse,
181 map, concat, concatMap,
182 zipWith, zipWith3, zip, zip3, unzip, unzip3,
183 filter, takeWhile, dropWhile, span, break,
184 elem, notElem,
185 foldl, foldl1, foldr, foldr1,
186 all, any, and, or, sum, product, maximum, minimum,
187 scanl, scanl1, scanr, scanr1,
188 enumFromTo, enumFromThenTo,
189 mapM, mapM_, sequence, sequence_ )
190
191 import Data.Typeable ( Typeable1, gcast1 )
192
193 #include "vector.h"
194
195 import Data.Data ( Data, DataType )
196 #if MIN_VERSION_base(4,2,0)
197 import Data.Data ( mkNoRepType )
198 #else
199 import Data.Data ( mkNorepType )
200 mkNoRepType :: String -> DataType
201 mkNoRepType = mkNorepType
202 #endif
203
204 -- Length information
205 -- ------------------
206
207 -- | /O(1)/ Yield the length of the vector.
208 length :: Vector v a => v a -> Int
209 {-# INLINE_STREAM length #-}
210 length v = basicLength v
211
212 {-# RULES
213
214 "length/unstream [Vector]" forall s.
215 length (new (New.unstream s)) = Stream.length s
216
217 #-}
218
219 -- | /O(1)/ Test whether a vector if empty
220 null :: Vector v a => v a -> Bool
221 {-# INLINE_STREAM null #-}
222 null v = basicLength v == 0
223
224 {-# RULES
225
226 "null/unstream [Vector]" forall s.
227 null (new (New.unstream s)) = Stream.null s
228
229 #-}
230
231 -- Indexing
232 -- --------
233
234 -- | O(1) Indexing
235 (!) :: Vector v a => v a -> Int -> a
236 {-# INLINE_STREAM (!) #-}
237 v ! i = BOUNDS_CHECK(checkIndex) "(!)" i (length v)
238 $ unId (basicUnsafeIndexM v i)
239
240 -- | O(1) Safe indexing
241 (!?) :: Vector v a => v a -> Int -> Maybe a
242 {-# INLINE_STREAM (!?) #-}
243 v !? i | i < 0 || i >= length v = Nothing
244 | otherwise = Just $ unsafeIndex v i
245
246 -- | /O(1)/ First element
247 head :: Vector v a => v a -> a
248 {-# INLINE_STREAM head #-}
249 head v = v ! 0
250
251 -- | /O(1)/ Last element
252 last :: Vector v a => v a -> a
253 {-# INLINE_STREAM last #-}
254 last v = v ! (length v - 1)
255
256 -- | /O(1)/ Unsafe indexing without bounds checking
257 unsafeIndex :: Vector v a => v a -> Int -> a
258 {-# INLINE_STREAM unsafeIndex #-}
259 unsafeIndex v i = UNSAFE_CHECK(checkIndex) "unsafeIndex" i (length v)
260 $ unId (basicUnsafeIndexM v i)
261
262 -- | /O(1)/ First element without checking if the vector is empty
263 unsafeHead :: Vector v a => v a -> a
264 {-# INLINE_STREAM unsafeHead #-}
265 unsafeHead v = unsafeIndex v 0
266
267 -- | /O(1)/ Last element without checking if the vector is empty
268 unsafeLast :: Vector v a => v a -> a
269 {-# INLINE_STREAM unsafeLast #-}
270 unsafeLast v = unsafeIndex v (length v - 1)
271
272 {-# RULES
273
274 "(!)/unstream [Vector]" forall i s.
275 new (New.unstream s) ! i = s Stream.!! i
276
277 "head/unstream [Vector]" forall s.
278 head (new (New.unstream s)) = Stream.head s
279
280 "last/unstream [Vector]" forall s.
281 last (new (New.unstream s)) = Stream.last s
282
283 "unsafeIndex/unstream [Vector]" forall i s.
284 unsafeIndex (new (New.unstream s)) i = s Stream.!! i
285
286 "unsafeHead/unstream [Vector]" forall s.
287 unsafeHead (new (New.unstream s)) = Stream.head s
288
289 "unsafeLast/unstream [Vector]" forall s.
290 unsafeLast (new (New.unstream s)) = Stream.last s
291
292 #-}
293
294 -- Monadic indexing
295 -- ----------------
296
297 -- | /O(1)/ Indexing in a monad.
298 --
299 -- The monad allows operations to be strict in the vector when necessary.
300 -- Suppose vector copying is implemented like this:
301 --
302 -- > copy mv v = ... write mv i (v ! i) ...
303 --
304 -- For lazy vectors, @v ! i@ would not be evaluated which means that @mv@
305 -- would unnecessarily retain a reference to @v@ in each element written.
306 --
307 -- With 'indexM', copying can be implemented like this instead:
308 --
309 -- > copy mv v = ... do
310 -- > x <- indexM v i
311 -- > write mv i x
312 --
313 -- Here, no references to @v@ are retained because indexing (but /not/ the
314 -- elements) is evaluated eagerly.
315 --
316 indexM :: (Vector v a, Monad m) => v a -> Int -> m a
317 {-# INLINE_STREAM indexM #-}
318 indexM v i = BOUNDS_CHECK(checkIndex) "indexM" i (length v)
319 $ basicUnsafeIndexM v i
320
321 -- | /O(1)/ First element of a vector in a monad. See 'indexM' for an
322 -- explanation of why this is useful.
323 headM :: (Vector v a, Monad m) => v a -> m a
324 {-# INLINE_STREAM headM #-}
325 headM v = indexM v 0
326
327 -- | /O(1)/ Last element of a vector in a monad. See 'indexM' for an
328 -- explanation of why this is useful.
329 lastM :: (Vector v a, Monad m) => v a -> m a
330 {-# INLINE_STREAM lastM #-}
331 lastM v = indexM v (length v - 1)
332
333 -- | /O(1)/ Indexing in a monad without bounds checks. See 'indexM' for an
334 -- explanation of why this is useful.
335 unsafeIndexM :: (Vector v a, Monad m) => v a -> Int -> m a
336 {-# INLINE_STREAM unsafeIndexM #-}
337 unsafeIndexM v i = UNSAFE_CHECK(checkIndex) "unsafeIndexM" i (length v)
338 $ basicUnsafeIndexM v i
339
340 -- | /O(1)/ First element in a monad without checking for empty vectors.
341 -- See 'indexM' for an explanation of why this is useful.
342 unsafeHeadM :: (Vector v a, Monad m) => v a -> m a
343 {-# INLINE_STREAM unsafeHeadM #-}
344 unsafeHeadM v = unsafeIndexM v 0
345
346 -- | /O(1)/ Last element in a monad without checking for empty vectors.
347 -- See 'indexM' for an explanation of why this is useful.
348 unsafeLastM :: (Vector v a, Monad m) => v a -> m a
349 {-# INLINE_STREAM unsafeLastM #-}
350 unsafeLastM v = unsafeIndexM v (length v - 1)
351
352 {-# RULES
353
354 "indexM/unstream [Vector]" forall s i.
355 indexM (new (New.unstream s)) i = liftStream s MStream.!! i
356
357 "headM/unstream [Vector]" forall s.
358 headM (new (New.unstream s)) = MStream.head (liftStream s)
359
360 "lastM/unstream [Vector]" forall s.
361 lastM (new (New.unstream s)) = MStream.last (liftStream s)
362
363 "unsafeIndexM/unstream [Vector]" forall s i.
364 unsafeIndexM (new (New.unstream s)) i = liftStream s MStream.!! i
365
366 "unsafeHeadM/unstream [Vector]" forall s.
367 unsafeHeadM (new (New.unstream s)) = MStream.head (liftStream s)
368
369 "unsafeLastM/unstream [Vector]" forall s.
370 unsafeLastM (new (New.unstream s)) = MStream.last (liftStream s)
371
372 #-}
373
374 -- Extracting subvectors (slicing)
375 -- -------------------------------
376
377 -- | /O(1)/ Yield a slice of the vector without copying it. The vector must
378 -- contain at least @i+n@ elements.
379 slice :: Vector v a => Int -- ^ @i@ starting index
380 -> Int -- ^ @n@ length
381 -> v a
382 -> v a
383 {-# INLINE_STREAM slice #-}
384 slice i n v = BOUNDS_CHECK(checkSlice) "slice" i n (length v)
385 $ basicUnsafeSlice i n v
386
387 -- | /O(1)/ Yield all but the last element without copying. The vector may not
388 -- be empty.
389 init :: Vector v a => v a -> v a
390 {-# INLINE_STREAM init #-}
391 init v = slice 0 (length v - 1) v
392
393 -- | /O(1)/ Yield all but the first element without copying. The vector may not
394 -- be empty.
395 tail :: Vector v a => v a -> v a
396 {-# INLINE_STREAM tail #-}
397 tail v = slice 1 (length v - 1) v
398
399 -- | /O(1)/ Yield the first @n@ elements without copying. The vector may
400 -- contain less than @n@ elements in which case it is returned unchanged.
401 take :: Vector v a => Int -> v a -> v a
402 {-# INLINE_STREAM take #-}
403 take n v = unsafeSlice 0 (delay_inline min n' (length v)) v
404 where n' = max n 0
405
406 -- | /O(1)/ Yield all but the first @n@ elements without copying. The vector may
407 -- contain less than @n@ elements in which case an empty vector is returned.
408 drop :: Vector v a => Int -> v a -> v a
409 {-# INLINE_STREAM drop #-}
410 drop n v = unsafeSlice (delay_inline min n' len)
411 (delay_inline max 0 (len - n')) v
412 where n' = max n 0
413 len = length v
414
415 -- | /O(1)/ Yield the first @n@ elements paired with the remainder without copying.
416 --
417 -- Note that @'splitAt' n v@ is equivalent to @('take' n v, 'drop' n v)@
418 -- but slightly more efficient.
419 {-# INLINE_STREAM splitAt #-}
420 splitAt :: Vector v a => Int -> v a -> (v a, v a)
421 splitAt n v = ( unsafeSlice 0 m v
422 , unsafeSlice m (delay_inline max 0 (len - n')) v
423 )
424 where
425 m = delay_inline min n' len
426 n' = max n 0
427 len = length v
428
429 -- | /O(1)/ Yield a slice of the vector without copying. The vector must
430 -- contain at least @i+n@ elements but this is not checked.
431 unsafeSlice :: Vector v a => Int -- ^ @i@ starting index
432 -> Int -- ^ @n@ length
433 -> v a
434 -> v a
435 {-# INLINE_STREAM unsafeSlice #-}
436 unsafeSlice i n v = UNSAFE_CHECK(checkSlice) "unsafeSlice" i n (length v)
437 $ basicUnsafeSlice i n v
438
439 -- | /O(1)/ Yield all but the last element without copying. The vector may not
440 -- be empty but this is not checked.
441 unsafeInit :: Vector v a => v a -> v a
442 {-# INLINE_STREAM unsafeInit #-}
443 unsafeInit v = unsafeSlice 0 (length v - 1) v
444
445 -- | /O(1)/ Yield all but the first element without copying. The vector may not
446 -- be empty but this is not checked.
447 unsafeTail :: Vector v a => v a -> v a
448 {-# INLINE_STREAM unsafeTail #-}
449 unsafeTail v = unsafeSlice 1 (length v - 1) v
450
451 -- | /O(1)/ Yield the first @n@ elements without copying. The vector must
452 -- contain at least @n@ elements but this is not checked.
453 unsafeTake :: Vector v a => Int -> v a -> v a
454 {-# INLINE unsafeTake #-}
455 unsafeTake n v = unsafeSlice 0 n v
456
457 -- | /O(1)/ Yield all but the first @n@ elements without copying. The vector
458 -- must contain at least @n@ elements but this is not checked.
459 unsafeDrop :: Vector v a => Int -> v a -> v a
460 {-# INLINE unsafeDrop #-}
461 unsafeDrop n v = unsafeSlice n (length v - n) v
462
463 {-# RULES
464
465 "slice/new [Vector]" forall i n p.
466 slice i n (new p) = new (New.slice i n p)
467
468 "init/new [Vector]" forall p.
469 init (new p) = new (New.init p)
470
471 "tail/new [Vector]" forall p.
472 tail (new p) = new (New.tail p)
473
474 "take/new [Vector]" forall n p.
475 take n (new p) = new (New.take n p)
476
477 "drop/new [Vector]" forall n p.
478 drop n (new p) = new (New.drop n p)
479
480 "unsafeSlice/new [Vector]" forall i n p.
481 unsafeSlice i n (new p) = new (New.unsafeSlice i n p)
482
483 "unsafeInit/new [Vector]" forall p.
484 unsafeInit (new p) = new (New.unsafeInit p)
485
486 "unsafeTail/new [Vector]" forall p.
487 unsafeTail (new p) = new (New.unsafeTail p)
488
489 #-}
490
491 -- Initialisation
492 -- --------------
493
494 -- | /O(1)/ Empty vector
495 empty :: Vector v a => v a
496 {-# INLINE empty #-}
497 empty = unstream Stream.empty
498
499 -- | /O(1)/ Vector with exactly one element
500 singleton :: forall v a. Vector v a => a -> v a
501 {-# INLINE singleton #-}
502 singleton x = elemseq (undefined :: v a) x
503 $ unstream (Stream.singleton x)
504
505 -- | /O(n)/ Vector of the given length with the same value in each position
506 replicate :: forall v a. Vector v a => Int -> a -> v a
507 {-# INLINE replicate #-}
508 replicate n x = elemseq (undefined :: v a) x
509 $ unstream
510 $ Stream.replicate n x
511
512 -- | /O(n)/ Construct a vector of the given length by applying the function to
513 -- each index
514 generate :: Vector v a => Int -> (Int -> a) -> v a
515 {-# INLINE generate #-}
516 generate n f = unstream (Stream.generate n f)
517
518 -- | /O(n)/ Apply function n times to value. Zeroth element is original value.
519 iterateN :: Vector v a => Int -> (a -> a) -> a -> v a
520 {-# INLINE iterateN #-}
521 iterateN n f x = unstream (Stream.iterateN n f x)
522
523 -- Unfolding
524 -- ---------
525
526 -- | /O(n)/ Construct a vector by repeatedly applying the generator function
527 -- to a seed. The generator function yields 'Just' the next element and the
528 -- new seed or 'Nothing' if there are no more elements.
529 --
530 -- > unfoldr (\n -> if n == 0 then Nothing else Just (n,n-1)) 10
531 -- > = <10,9,8,7,6,5,4,3,2,1>
532 unfoldr :: Vector v a => (b -> Maybe (a, b)) -> b -> v a
533 {-# INLINE unfoldr #-}
534 unfoldr f = unstream . Stream.unfoldr f
535
536 -- | /O(n)/ Construct a vector with at most @n@ by repeatedly applying the
537 -- generator function to the a seed. The generator function yields 'Just' the
538 -- next element and the new seed or 'Nothing' if there are no more elements.
539 --
540 -- > unfoldrN 3 (\n -> Just (n,n-1)) 10 = <10,9,8>
541 unfoldrN :: Vector v a => Int -> (b -> Maybe (a, b)) -> b -> v a
542 {-# INLINE unfoldrN #-}
543 unfoldrN n f = unstream . Stream.unfoldrN n f
544
545 -- Enumeration
546 -- -----------
547
548 -- | /O(n)/ Yield a vector of the given length containing the values @x@, @x+1@
549 -- etc. This operation is usually more efficient than 'enumFromTo'.
550 --
551 -- > enumFromN 5 3 = <5,6,7>
552 enumFromN :: (Vector v a, Num a) => a -> Int -> v a
553 {-# INLINE enumFromN #-}
554 enumFromN x n = enumFromStepN x 1 n
555
556 -- | /O(n)/ Yield a vector of the given length containing the values @x@, @x+y@,
557 -- @x+y+y@ etc. This operations is usually more efficient than 'enumFromThenTo'.
558 --
559 -- > enumFromStepN 1 0.1 5 = <1,1.1,1.2,1.3,1.4>
560 enumFromStepN :: forall v a. (Vector v a, Num a) => a -> a -> Int -> v a
561 {-# INLINE enumFromStepN #-}
562 enumFromStepN x y n = elemseq (undefined :: v a) x
563 $ elemseq (undefined :: v a) y
564 $ unstream
565 $ Stream.enumFromStepN x y n
566
567 -- | /O(n)/ Enumerate values from @x@ to @y@.
568 --
569 -- /WARNING:/ This operation can be very inefficient. If at all possible, use
570 -- 'enumFromN' instead.
571 enumFromTo :: (Vector v a, Enum a) => a -> a -> v a
572 {-# INLINE enumFromTo #-}
573 enumFromTo x y = unstream (Stream.enumFromTo x y)
574
575 -- | /O(n)/ Enumerate values from @x@ to @y@ with a specific step @z@.
576 --
577 -- /WARNING:/ This operation can be very inefficient. If at all possible, use
578 -- 'enumFromStepN' instead.
579 enumFromThenTo :: (Vector v a, Enum a) => a -> a -> a -> v a
580 {-# INLINE enumFromThenTo #-}
581 enumFromThenTo x y z = unstream (Stream.enumFromThenTo x y z)
582
583 -- Concatenation
584 -- -------------
585
586 -- | /O(n)/ Prepend an element
587 cons :: forall v a. Vector v a => a -> v a -> v a
588 {-# INLINE cons #-}
589 cons x v = elemseq (undefined :: v a) x
590 $ unstream
591 $ Stream.cons x
592 $ stream v
593
594 -- | /O(n)/ Append an element
595 snoc :: forall v a. Vector v a => v a -> a -> v a
596 {-# INLINE snoc #-}
597 snoc v x = elemseq (undefined :: v a) x
598 $ unstream
599 $ Stream.snoc (stream v) x
600
601 infixr 5 ++
602 -- | /O(m+n)/ Concatenate two vectors
603 (++) :: Vector v a => v a -> v a -> v a
604 {-# INLINE (++) #-}
605 v ++ w = unstream (stream v Stream.++ stream w)
606
607 -- | /O(n)/ Concatenate all vectors in the list
608 concat :: Vector v a => [v a] -> v a
609 {-# INLINE concat #-}
610 -- concat vs = create (thawMany vs)
611 concat vs = unstream (Stream.flatten mk step (Exact n) (Stream.fromList vs))
612 where
613 n = List.foldl' (\k v -> k + length v) 0 vs
614
615 {-# INLINE_INNER step #-}
616 step (v,i,k)
617 | i < k = case unsafeIndexM v i of
618 Box x -> Stream.Yield x (v,i+1,k)
619 | otherwise = Stream.Done
620
621 {-# INLINE mk #-}
622 mk v = let k = length v
623 in
624 k `seq` (v,0,k)
625
626 -- Monadic initialisation
627 -- ----------------------
628
629 -- | /O(n)/ Execute the monadic action the given number of times and store the
630 -- results in a vector.
631 replicateM :: (Monad m, Vector v a) => Int -> m a -> m (v a)
632 {-# INLINE replicateM #-}
633 replicateM n m = unstreamM (MStream.replicateM n m)
634
635 -- | /O(n)/ Construct a vector of the given length by applying the monadic
636 -- action to each index
637 generateM :: (Monad m, Vector v a) => Int -> (Int -> m a) -> m (v a)
638 {-# INLINE generateM #-}
639 generateM n f = unstreamM (MStream.generateM n f)
640
641 -- | Execute the monadic action and freeze the resulting vector.
642 --
643 -- @
644 -- create (do { v \<- 'M.new' 2; 'M.write' v 0 \'a\'; 'M.write' v 1 \'b\' }) = \<'a','b'\>
645 -- @
646 create :: Vector v a => (forall s. ST s (Mutable v s a)) -> v a
647 {-# INLINE create #-}
648 create p = new (New.create p)
649
650 -- Restricting memory usage
651 -- ------------------------
652
653 -- | /O(n)/ Yield the argument but force it not to retain any extra memory,
654 -- possibly by copying it.
655 --
656 -- This is especially useful when dealing with slices. For example:
657 --
658 -- > force (slice 0 2 <huge vector>)
659 --
660 -- Here, the slice retains a reference to the huge vector. Forcing it creates
661 -- a copy of just the elements that belong to the slice and allows the huge
662 -- vector to be garbage collected.
663 force :: Vector v a => v a -> v a
664 -- FIXME: we probably ought to inline this later as the rules still might fire
665 -- otherwise
666 {-# INLINE_STREAM force #-}
667 force v = new (clone v)
668
669 -- Bulk updates
670 -- ------------
671
672 -- | /O(m+n)/ For each pair @(i,a)@ from the list, replace the vector
673 -- element at position @i@ by @a@.
674 --
675 -- > <5,9,2,7> // [(2,1),(0,3),(2,8)] = <3,9,8,7>
676 --
677 (//) :: Vector v a => v a -- ^ initial vector (of length @m@)
678 -> [(Int, a)] -- ^ list of index/value pairs (of length @n@)
679 -> v a
680 {-# INLINE (//) #-}
681 v // us = update_stream v (Stream.fromList us)
682
683 -- | /O(m+n)/ For each pair @(i,a)@ from the vector of index/value pairs,
684 -- replace the vector element at position @i@ by @a@.
685 --
686 -- > update <5,9,2,7> <(2,1),(0,3),(2,8)> = <3,9,8,7>
687 --
688 update :: (Vector v a, Vector v (Int, a))
689 => v a -- ^ initial vector (of length @m@)
690 -> v (Int, a) -- ^ vector of index/value pairs (of length @n@)
691 -> v a
692 {-# INLINE update #-}
693 update v w = update_stream v (stream w)
694
695 -- | /O(m+min(n1,n2))/ For each index @i@ from the index vector and the
696 -- corresponding value @a@ from the value vector, replace the element of the
697 -- initial vector at position @i@ by @a@.
698 --
699 -- > update_ <5,9,2,7> <2,0,2> <1,3,8> = <3,9,8,7>
700 --
701 -- This function is useful for instances of 'Vector' that cannot store pairs.
702 -- Otherwise, 'update' is probably more convenient.
703 --
704 -- @
705 -- update_ xs is ys = 'update' xs ('zip' is ys)
706 -- @
707 update_ :: (Vector v a, Vector v Int)
708 => v a -- ^ initial vector (of length @m@)
709 -> v Int -- ^ index vector (of length @n1@)
710 -> v a -- ^ value vector (of length @n2@)
711 -> v a
712 {-# INLINE update_ #-}
713 update_ v is w = update_stream v (Stream.zipWith (,) (stream is) (stream w))
714
715 update_stream :: Vector v a => v a -> Stream (Int,a) -> v a
716 {-# INLINE update_stream #-}
717 update_stream = modifyWithStream M.update
718
719 -- | Same as ('//') but without bounds checking.
720 unsafeUpd :: Vector v a => v a -> [(Int, a)] -> v a
721 {-# INLINE unsafeUpd #-}
722 unsafeUpd v us = unsafeUpdate_stream v (Stream.fromList us)
723
724 -- | Same as 'update' but without bounds checking.
725 unsafeUpdate :: (Vector v a, Vector v (Int, a)) => v a -> v (Int, a) -> v a
726 {-# INLINE unsafeUpdate #-}
727 unsafeUpdate v w = unsafeUpdate_stream v (stream w)
728
729 -- | Same as 'update_' but without bounds checking.
730 unsafeUpdate_ :: (Vector v a, Vector v Int) => v a -> v Int -> v a -> v a
731 {-# INLINE unsafeUpdate_ #-}
732 unsafeUpdate_ v is w
733 = unsafeUpdate_stream v (Stream.zipWith (,) (stream is) (stream w))
734
735 unsafeUpdate_stream :: Vector v a => v a -> Stream (Int,a) -> v a
736 {-# INLINE unsafeUpdate_stream #-}
737 unsafeUpdate_stream = modifyWithStream M.unsafeUpdate
738
739 -- Accumulations
740 -- -------------
741
742 -- | /O(m+n)/ For each pair @(i,b)@ from the list, replace the vector element
743 -- @a@ at position @i@ by @f a b@.
744 --
745 -- > accum (+) <5,9,2> [(2,4),(1,6),(0,3),(1,7)] = <5+3, 9+6+7, 2+4>
746 accum :: Vector v a
747 => (a -> b -> a) -- ^ accumulating function @f@
748 -> v a -- ^ initial vector (of length @m@)
749 -> [(Int,b)] -- ^ list of index/value pairs (of length @n@)
750 -> v a
751 {-# INLINE accum #-}
752 accum f v us = accum_stream f v (Stream.fromList us)
753
754 -- | /O(m+n)/ For each pair @(i,b)@ from the vector of pairs, replace the vector
755 -- element @a@ at position @i@ by @f a b@.
756 --
757 -- > accumulate (+) <5,9,2> <(2,4),(1,6),(0,3),(1,7)> = <5+3, 9+6+7, 2+4>
758 accumulate :: (Vector v a, Vector v (Int, b))
759 => (a -> b -> a) -- ^ accumulating function @f@
760 -> v a -- ^ initial vector (of length @m@)
761 -> v (Int,b) -- ^ vector of index/value pairs (of length @n@)
762 -> v a
763 {-# INLINE accumulate #-}
764 accumulate f v us = accum_stream f v (stream us)
765
766 -- | /O(m+min(n1,n2))/ For each index @i@ from the index vector and the
767 -- corresponding value @b@ from the the value vector,
768 -- replace the element of the initial vector at
769 -- position @i@ by @f a b@.
770 --
771 -- > accumulate_ (+) <5,9,2> <2,1,0,1> <4,6,3,7> = <5+3, 9+6+7, 2+4>
772 --
773 -- This function is useful for instances of 'Vector' that cannot store pairs.
774 -- Otherwise, 'accumulate' is probably more convenient:
775 --
776 -- @
777 -- accumulate_ f as is bs = 'accumulate' f as ('zip' is bs)
778 -- @
779 accumulate_ :: (Vector v a, Vector v Int, Vector v b)
780 => (a -> b -> a) -- ^ accumulating function @f@
781 -> v a -- ^ initial vector (of length @m@)
782 -> v Int -- ^ index vector (of length @n1@)
783 -> v b -- ^ value vector (of length @n2@)
784 -> v a
785 {-# INLINE accumulate_ #-}
786 accumulate_ f v is xs = accum_stream f v (Stream.zipWith (,) (stream is)
787 (stream xs))
788
789
790 accum_stream :: Vector v a => (a -> b -> a) -> v a -> Stream (Int,b) -> v a
791 {-# INLINE accum_stream #-}
792 accum_stream f = modifyWithStream (M.accum f)
793
794 -- | Same as 'accum' but without bounds checking.
795 unsafeAccum :: Vector v a => (a -> b -> a) -> v a -> [(Int,b)] -> v a
796 {-# INLINE unsafeAccum #-}
797 unsafeAccum f v us = unsafeAccum_stream f v (Stream.fromList us)
798
799 -- | Same as 'accumulate' but without bounds checking.
800 unsafeAccumulate :: (Vector v a, Vector v (Int, b))
801 => (a -> b -> a) -> v a -> v (Int,b) -> v a
802 {-# INLINE unsafeAccumulate #-}
803 unsafeAccumulate f v us = unsafeAccum_stream f v (stream us)
804
805 -- | Same as 'accumulate_' but without bounds checking.
806 unsafeAccumulate_ :: (Vector v a, Vector v Int, Vector v b)
807 => (a -> b -> a) -> v a -> v Int -> v b -> v a
808 {-# INLINE unsafeAccumulate_ #-}
809 unsafeAccumulate_ f v is xs
810 = unsafeAccum_stream f v (Stream.zipWith (,) (stream is) (stream xs))
811
812 unsafeAccum_stream
813 :: Vector v a => (a -> b -> a) -> v a -> Stream (Int,b) -> v a
814 {-# INLINE unsafeAccum_stream #-}
815 unsafeAccum_stream f = modifyWithStream (M.unsafeAccum f)
816
817 -- Permutations
818 -- ------------
819
820 -- | /O(n)/ Reverse a vector
821 reverse :: (Vector v a) => v a -> v a
822 {-# INLINE reverse #-}
823 -- FIXME: make this fuse better, add support for recycling
824 reverse = unstream . streamR
825
826 -- | /O(n)/ Yield the vector obtained by replacing each element @i@ of the
827 -- index vector by @xs'!'i@. This is equivalent to @'map' (xs'!') is@ but is
828 -- often much more efficient.
829 --
830 -- > backpermute <a,b,c,d> <0,3,2,3,1,0> = <a,d,c,d,b,a>
831 backpermute :: (Vector v a, Vector v Int)
832 => v a -- ^ @xs@ value vector
833 -> v Int -- ^ @is@ index vector (of length @n@)
834 -> v a
835 {-# INLINE backpermute #-}
836 -- This somewhat non-intuitive definition ensures that the resulting vector
837 -- does not retain references to the original one even if it is lazy in its
838 -- elements. This would not be the case if we simply used map (v!)
839 backpermute v is = seq v
840 $ seq n
841 $ unstream
842 $ Stream.unbox
843 $ Stream.map index
844 $ stream is
845 where
846 n = length v
847
848 {-# INLINE index #-}
849 -- NOTE: we do it this way to avoid triggering LiberateCase on n in
850 -- polymorphic code
851 index i = BOUNDS_CHECK(checkIndex) "backpermute" i n
852 $ basicUnsafeIndexM v i
853
854 -- | Same as 'backpermute' but without bounds checking.
855 unsafeBackpermute :: (Vector v a, Vector v Int) => v a -> v Int -> v a
856 {-# INLINE unsafeBackpermute #-}
857 unsafeBackpermute v is = seq v
858 $ seq n
859 $ unstream
860 $ Stream.unbox
861 $ Stream.map index
862 $ stream is
863 where
864 n = length v
865
866 {-# INLINE index #-}
867 -- NOTE: we do it this way to avoid triggering LiberateCase on n in
868 -- polymorphic code
869 index i = UNSAFE_CHECK(checkIndex) "unsafeBackpermute" i n
870 $ basicUnsafeIndexM v i
871
872 -- Safe destructive updates
873 -- ------------------------
874
875 -- | Apply a destructive operation to a vector. The operation will be
876 -- performed in place if it is safe to do so and will modify a copy of the
877 -- vector otherwise.
878 --
879 -- @
880 -- modify (\\v -> 'M.write' v 0 \'x\') ('replicate' 3 \'a\') = \<\'x\',\'a\',\'a\'\>
881 -- @
882 modify :: Vector v a => (forall s. Mutable v s a -> ST s ()) -> v a -> v a
883 {-# INLINE modify #-}
884 modify p = new . New.modify p . clone
885
886 -- We have to make sure that this is strict in the stream but we can't seq on
887 -- it while fusion is happening. Hence this ugliness.
888 modifyWithStream :: Vector v a
889 => (forall s. Mutable v s a -> Stream b -> ST s ())
890 -> v a -> Stream b -> v a
891 {-# INLINE modifyWithStream #-}
892 modifyWithStream p v s = new (New.modifyWithStream p (clone v) s)
893
894 -- Indexing
895 -- --------
896
897 -- | /O(n)/ Pair each element in a vector with its index
898 indexed :: (Vector v a, Vector v (Int,a)) => v a -> v (Int,a)
899 {-# INLINE indexed #-}
900 indexed = unstream . Stream.indexed . stream
901
902 -- Mapping
903 -- -------
904
905 -- | /O(n)/ Map a function over a vector
906 map :: (Vector v a, Vector v b) => (a -> b) -> v a -> v b
907 {-# INLINE map #-}
908 map f = unstream . inplace (MStream.map f) . stream
909
910 -- | /O(n)/ Apply a function to every element of a vector and its index
911 imap :: (Vector v a, Vector v b) => (Int -> a -> b) -> v a -> v b
912 {-# INLINE imap #-}
913 imap f = unstream . inplace (MStream.map (uncurry f) . MStream.indexed)
914 . stream
915
916 -- | Map a function over a vector and concatenate the results.
917 concatMap :: (Vector v a, Vector v b) => (a -> v b) -> v a -> v b
918 {-# INLINE concatMap #-}
919 -- NOTE: We can't fuse concatMap anyway so don't pretend we do.
920 -- concatMap f = unstream . Stream.concatMap (stream . f) . stream
921 concatMap f = concat . Stream.toList . Stream.map f . stream
922
923 -- Monadic mapping
924 -- ---------------
925
926 -- | /O(n)/ Apply the monadic action to all elements of the vector, yielding a
927 -- vector of results
928 mapM :: (Monad m, Vector v a, Vector v b) => (a -> m b) -> v a -> m (v b)
929 {-# INLINE mapM #-}
930 mapM f = unstreamM . Stream.mapM f . stream
931
932 -- | /O(n)/ Apply the monadic action to all elements of a vector and ignore the
933 -- results
934 mapM_ :: (Monad m, Vector v a) => (a -> m b) -> v a -> m ()
935 {-# INLINE mapM_ #-}
936 mapM_ f = Stream.mapM_ f . stream
937
938 -- | /O(n)/ Apply the monadic action to all elements of the vector, yielding a
939 -- vector of results. Equvalent to @flip 'mapM'@.
940 forM :: (Monad m, Vector v a, Vector v b) => v a -> (a -> m b) -> m (v b)
941 {-# INLINE forM #-}
942 forM as f = mapM f as
943
944 -- | /O(n)/ Apply the monadic action to all elements of a vector and ignore the
945 -- results. Equivalent to @flip 'mapM_'@.
946 forM_ :: (Monad m, Vector v a) => v a -> (a -> m b) -> m ()
947 {-# INLINE forM_ #-}
948 forM_ as f = mapM_ f as
949
950 -- Zipping
951 -- -------
952
953 -- | /O(min(m,n))/ Zip two vectors with the given function.
954 zipWith :: (Vector v a, Vector v b, Vector v c)
955 => (a -> b -> c) -> v a -> v b -> v c
956 {-# INLINE zipWith #-}
957 zipWith f xs ys = unstream (Stream.zipWith f (stream xs) (stream ys))
958
959 -- | Zip three vectors with the given function.
960 zipWith3 :: (Vector v a, Vector v b, Vector v c, Vector v d)
961 => (a -> b -> c -> d) -> v a -> v b -> v c -> v d
962 {-# INLINE zipWith3 #-}
963 zipWith3 f as bs cs = unstream (Stream.zipWith3 f (stream as)
964 (stream bs)
965 (stream cs))
966
967 zipWith4 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e)
968 => (a -> b -> c -> d -> e) -> v a -> v b -> v c -> v d -> v e
969 {-# INLINE zipWith4 #-}
970 zipWith4 f as bs cs ds
971 = unstream (Stream.zipWith4 f (stream as)
972 (stream bs)
973 (stream cs)
974 (stream ds))
975
976 zipWith5 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
977 Vector v f)
978 => (a -> b -> c -> d -> e -> f) -> v a -> v b -> v c -> v d -> v e
979 -> v f
980 {-# INLINE zipWith5 #-}
981 zipWith5 f as bs cs ds es
982 = unstream (Stream.zipWith5 f (stream as)
983 (stream bs)
984 (stream cs)
985 (stream ds)
986 (stream es))
987
988 zipWith6 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
989 Vector v f, Vector v g)
990 => (a -> b -> c -> d -> e -> f -> g)
991 -> v a -> v b -> v c -> v d -> v e -> v f -> v g
992 {-# INLINE zipWith6 #-}
993 zipWith6 f as bs cs ds es fs
994 = unstream (Stream.zipWith6 f (stream as)
995 (stream bs)
996 (stream cs)
997 (stream ds)
998 (stream es)
999 (stream fs))
1000
1001 -- | /O(min(m,n))/ Zip two vectors with a function that also takes the
1002 -- elements' indices.
1003 izipWith :: (Vector v a, Vector v b, Vector v c)
1004 => (Int -> a -> b -> c) -> v a -> v b -> v c
1005 {-# INLINE izipWith #-}
1006 izipWith f xs ys = unstream
1007 (Stream.zipWith (uncurry f) (Stream.indexed (stream xs))
1008 (stream ys))
1009
1010 izipWith3 :: (Vector v a, Vector v b, Vector v c, Vector v d)
1011 => (Int -> a -> b -> c -> d) -> v a -> v b -> v c -> v d
1012 {-# INLINE izipWith3 #-}
1013 izipWith3 f as bs cs
1014 = unstream (Stream.zipWith3 (uncurry f) (Stream.indexed (stream as))
1015 (stream bs)
1016 (stream cs))
1017
1018 izipWith4 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e)
1019 => (Int -> a -> b -> c -> d -> e) -> v a -> v b -> v c -> v d -> v e
1020 {-# INLINE izipWith4 #-}
1021 izipWith4 f as bs cs ds
1022 = unstream (Stream.zipWith4 (uncurry f) (Stream.indexed (stream as))
1023 (stream bs)
1024 (stream cs)
1025 (stream ds))
1026
1027 izipWith5 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
1028 Vector v f)
1029 => (Int -> a -> b -> c -> d -> e -> f) -> v a -> v b -> v c -> v d
1030 -> v e -> v f
1031 {-# INLINE izipWith5 #-}
1032 izipWith5 f as bs cs ds es
1033 = unstream (Stream.zipWith5 (uncurry f) (Stream.indexed (stream as))
1034 (stream bs)
1035 (stream cs)
1036 (stream ds)
1037 (stream es))
1038
1039 izipWith6 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
1040 Vector v f, Vector v g)
1041 => (Int -> a -> b -> c -> d -> e -> f -> g)
1042 -> v a -> v b -> v c -> v d -> v e -> v f -> v g
1043 {-# INLINE izipWith6 #-}
1044 izipWith6 f as bs cs ds es fs
1045 = unstream (Stream.zipWith6 (uncurry f) (Stream.indexed (stream as))
1046 (stream bs)
1047 (stream cs)
1048 (stream ds)
1049 (stream es)
1050 (stream fs))
1051
1052 -- | /O(min(m,n))/ Zip two vectors
1053 zip :: (Vector v a, Vector v b, Vector v (a,b)) => v a -> v b -> v (a, b)
1054 {-# INLINE zip #-}
1055 zip = zipWith (,)
1056
1057 zip3 :: (Vector v a, Vector v b, Vector v c, Vector v (a, b, c))
1058 => v a -> v b -> v c -> v (a, b, c)
1059 {-# INLINE zip3 #-}
1060 zip3 = zipWith3 (,,)
1061
1062 zip4 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v (a, b, c, d))
1063 => v a -> v b -> v c -> v d -> v (a, b, c, d)
1064 {-# INLINE zip4 #-}
1065 zip4 = zipWith4 (,,,)
1066
1067 zip5 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
1068 Vector v (a, b, c, d, e))
1069 => v a -> v b -> v c -> v d -> v e -> v (a, b, c, d, e)
1070 {-# INLINE zip5 #-}
1071 zip5 = zipWith5 (,,,,)
1072
1073 zip6 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
1074 Vector v f, Vector v (a, b, c, d, e, f))
1075 => v a -> v b -> v c -> v d -> v e -> v f -> v (a, b, c, d, e, f)
1076 {-# INLINE zip6 #-}
1077 zip6 = zipWith6 (,,,,,)
1078
1079 -- Monadic zipping
1080 -- ---------------
1081
1082 -- | /O(min(m,n))/ Zip the two vectors with the monadic action and yield a
1083 -- vector of results
1084 zipWithM :: (Monad m, Vector v a, Vector v b, Vector v c)
1085 => (a -> b -> m c) -> v a -> v b -> m (v c)
1086 -- FIXME: specialise for ST and IO?
1087 {-# INLINE zipWithM #-}
1088 zipWithM f as bs = unstreamM $ Stream.zipWithM f (stream as) (stream bs)
1089
1090 -- | /O(min(m,n))/ Zip the two vectors with the monadic action and ignore the
1091 -- results
1092 zipWithM_ :: (Monad m, Vector v a, Vector v b)
1093 => (a -> b -> m c) -> v a -> v b -> m ()
1094 {-# INLINE zipWithM_ #-}
1095 zipWithM_ f as bs = Stream.zipWithM_ f (stream as) (stream bs)
1096
1097 -- Unzipping
1098 -- ---------
1099
1100 -- | /O(min(m,n))/ Unzip a vector of pairs.
1101 unzip :: (Vector v a, Vector v b, Vector v (a,b)) => v (a, b) -> (v a, v b)
1102 {-# INLINE unzip #-}
1103 unzip xs = (map fst xs, map snd xs)
1104
1105 unzip3 :: (Vector v a, Vector v b, Vector v c, Vector v (a, b, c))
1106 => v (a, b, c) -> (v a, v b, v c)
1107 {-# INLINE unzip3 #-}
1108 unzip3 xs = (map (\(a, b, c) -> a) xs,
1109 map (\(a, b, c) -> b) xs,
1110 map (\(a, b, c) -> c) xs)
1111
1112 unzip4 :: (Vector v a, Vector v b, Vector v c, Vector v d,
1113 Vector v (a, b, c, d))
1114 => v (a, b, c, d) -> (v a, v b, v c, v d)
1115 {-# INLINE unzip4 #-}
1116 unzip4 xs = (map (\(a, b, c, d) -> a) xs,
1117 map (\(a, b, c, d) -> b) xs,
1118 map (\(a, b, c, d) -> c) xs,
1119 map (\(a, b, c, d) -> d) xs)
1120
1121 unzip5 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
1122 Vector v (a, b, c, d, e))
1123 => v (a, b, c, d, e) -> (v a, v b, v c, v d, v e)
1124 {-# INLINE unzip5 #-}
1125 unzip5 xs = (map (\(a, b, c, d, e) -> a) xs,
1126 map (\(a, b, c, d, e) -> b) xs,
1127 map (\(a, b, c, d, e) -> c) xs,
1128 map (\(a, b, c, d, e) -> d) xs,
1129 map (\(a, b, c, d, e) -> e) xs)
1130
1131 unzip6 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
1132 Vector v f, Vector v (a, b, c, d, e, f))
1133 => v (a, b, c, d, e, f) -> (v a, v b, v c, v d, v e, v f)
1134 {-# INLINE unzip6 #-}
1135 unzip6 xs = (map (\(a, b, c, d, e, f) -> a) xs,
1136 map (\(a, b, c, d, e, f) -> b) xs,
1137 map (\(a, b, c, d, e, f) -> c) xs,
1138 map (\(a, b, c, d, e, f) -> d) xs,
1139 map (\(a, b, c, d, e, f) -> e) xs,
1140 map (\(a, b, c, d, e, f) -> f) xs)
1141
1142 -- Filtering
1143 -- ---------
1144
1145 -- | /O(n)/ Drop elements that do not satisfy the predicate
1146 filter :: Vector v a => (a -> Bool) -> v a -> v a
1147 {-# INLINE filter #-}
1148 filter f = unstream . inplace (MStream.filter f) . stream
1149
1150 -- | /O(n)/ Drop elements that do not satisfy the predicate which is applied to
1151 -- values and their indices
1152 ifilter :: Vector v a => (Int -> a -> Bool) -> v a -> v a
1153 {-# INLINE ifilter #-}
1154 ifilter f = unstream
1155 . inplace (MStream.map snd . MStream.filter (uncurry f)
1156 . MStream.indexed)
1157 . stream
1158
1159 -- | /O(n)/ Drop elements that do not satisfy the monadic predicate
1160 filterM :: (Monad m, Vector v a) => (a -> m Bool) -> v a -> m (v a)
1161 {-# INLINE filterM #-}
1162 filterM f = unstreamM . Stream.filterM f . stream
1163
1164 -- | /O(n)/ Yield the longest prefix of elements satisfying the predicate
1165 -- without copying.
1166 takeWhile :: Vector v a => (a -> Bool) -> v a -> v a
1167 {-# INLINE takeWhile #-}
1168 takeWhile f = unstream . Stream.takeWhile f . stream
1169
1170 -- | /O(n)/ Drop the longest prefix of elements that satisfy the predicate
1171 -- without copying.
1172 dropWhile :: Vector v a => (a -> Bool) -> v a -> v a
1173 {-# INLINE dropWhile #-}
1174 dropWhile f = unstream . Stream.dropWhile f . stream
1175
1176 -- Parititioning
1177 -- -------------
1178
1179 -- | /O(n)/ Split the vector in two parts, the first one containing those
1180 -- elements that satisfy the predicate and the second one those that don't. The
1181 -- relative order of the elements is preserved at the cost of a sometimes
1182 -- reduced performance compared to 'unstablePartition'.
1183 partition :: Vector v a => (a -> Bool) -> v a -> (v a, v a)
1184 {-# INLINE partition #-}
1185 partition f = partition_stream f . stream
1186
1187 -- FIXME: Make this inplace-fusible (look at how stable_partition is
1188 -- implemented in C++)
1189
1190 partition_stream :: Vector v a => (a -> Bool) -> Stream a -> (v a, v a)
1191 {-# INLINE_STREAM partition_stream #-}
1192 partition_stream f s = s `seq` runST (
1193 do
1194 (mv1,mv2) <- M.partitionStream f s
1195 v1 <- unsafeFreeze mv1
1196 v2 <- unsafeFreeze mv2
1197 return (v1,v2))
1198
1199 -- | /O(n)/ Split the vector in two parts, the first one containing those
1200 -- elements that satisfy the predicate and the second one those that don't.
1201 -- The order of the elements is not preserved but the operation is often
1202 -- faster than 'partition'.
1203 unstablePartition :: Vector v a => (a -> Bool) -> v a -> (v a, v a)
1204 {-# INLINE unstablePartition #-}
1205 unstablePartition f = unstablePartition_stream f . stream
1206
1207 unstablePartition_stream
1208 :: Vector v a => (a -> Bool) -> Stream a -> (v a, v a)
1209 {-# INLINE_STREAM unstablePartition_stream #-}
1210 unstablePartition_stream f s = s `seq` runST (
1211 do
1212 (mv1,mv2) <- M.unstablePartitionStream f s
1213 v1 <- unsafeFreeze mv1
1214 v2 <- unsafeFreeze mv2
1215 return (v1,v2))
1216
1217 unstablePartition_new :: Vector v a => (a -> Bool) -> New v a -> (v a, v a)
1218 {-# INLINE_STREAM unstablePartition_new #-}
1219 unstablePartition_new f (New.New p) = runST (
1220 do
1221 mv <- p
1222 i <- M.unstablePartition f mv
1223 v <- unsafeFreeze mv
1224 return (unsafeTake i v, unsafeDrop i v))
1225
1226 {-# RULES
1227
1228 "unstablePartition" forall f p.
1229 unstablePartition_stream f (stream (new p))
1230 = unstablePartition_new f p
1231
1232 #-}
1233
1234
1235 -- FIXME: make span and break fusible
1236
1237 -- | /O(n)/ Split the vector into the longest prefix of elements that satisfy
1238 -- the predicate and the rest without copying.
1239 span :: Vector v a => (a -> Bool) -> v a -> (v a, v a)
1240 {-# INLINE span #-}
1241 span f = break (not . f)
1242
1243 -- | /O(n)/ Split the vector into the longest prefix of elements that do not
1244 -- satisfy the predicate and the rest without copying.
1245 break :: Vector v a => (a -> Bool) -> v a -> (v a, v a)
1246 {-# INLINE break #-}
1247 break f xs = case findIndex f xs of
1248 Just i -> (unsafeSlice 0 i xs, unsafeSlice i (length xs - i) xs)
1249 Nothing -> (xs, empty)
1250
1251
1252 -- Searching
1253 -- ---------
1254
1255 infix 4 `elem`
1256 -- | /O(n)/ Check if the vector contains an element
1257 elem :: (Vector v a, Eq a) => a -> v a -> Bool
1258 {-# INLINE elem #-}
1259 elem x = Stream.elem x . stream
1260
1261 infix 4 `notElem`
1262 -- | /O(n)/ Check if the vector does not contain an element (inverse of 'elem')
1263 notElem :: (Vector v a, Eq a) => a -> v a -> Bool
1264 {-# INLINE notElem #-}
1265 notElem x = Stream.notElem x . stream
1266
1267 -- | /O(n)/ Yield 'Just' the first element matching the predicate or 'Nothing'
1268 -- if no such element exists.
1269 find :: Vector v a => (a -> Bool) -> v a -> Maybe a
1270 {-# INLINE find #-}
1271 find f = Stream.find f . stream
1272
1273 -- | /O(n)/ Yield 'Just' the index of the first element matching the predicate
1274 -- or 'Nothing' if no such element exists.
1275 findIndex :: Vector v a => (a -> Bool) -> v a -> Maybe Int
1276 {-# INLINE findIndex #-}
1277 findIndex f = Stream.findIndex f . stream
1278
1279 -- | /O(n)/ Yield the indices of elements satisfying the predicate in ascending
1280 -- order.
1281 findIndices :: (Vector v a, Vector v Int) => (a -> Bool) -> v a -> v Int
1282 {-# INLINE findIndices #-}
1283 findIndices f = unstream
1284 . inplace (MStream.map fst . MStream.filter (f . snd)
1285 . MStream.indexed)
1286 . stream
1287
1288 -- | /O(n)/ Yield 'Just' the index of the first occurence of the given element or
1289 -- 'Nothing' if the vector does not contain the element. This is a specialised
1290 -- version of 'findIndex'.
1291 elemIndex :: (Vector v a, Eq a) => a -> v a -> Maybe Int
1292 {-# INLINE elemIndex #-}
1293 elemIndex x = findIndex (x==)
1294
1295 -- | /O(n)/ Yield the indices of all occurences of the given element in
1296 -- ascending order. This is a specialised version of 'findIndices'.
1297 elemIndices :: (Vector v a, Vector v Int, Eq a) => a -> v a -> v Int
1298 {-# INLINE elemIndices #-}
1299 elemIndices x = findIndices (x==)
1300
1301 -- Folding
1302 -- -------
1303
1304 -- | /O(n)/ Left fold
1305 foldl :: Vector v b => (a -> b -> a) -> a -> v b -> a
1306 {-# INLINE foldl #-}
1307 foldl f z = Stream.foldl f z . stream
1308
1309 -- | /O(n)/ Left fold on non-empty vectors
1310 foldl1 :: Vector v a => (a -> a -> a) -> v a -> a
1311 {-# INLINE foldl1 #-}
1312 foldl1 f = Stream.foldl1 f . stream
1313
1314 -- | /O(n)/ Left fold with strict accumulator
1315 foldl' :: Vector v b => (a -> b -> a) -> a -> v b -> a
1316 {-# INLINE foldl' #-}
1317 foldl' f z = Stream.foldl' f z . stream
1318
1319 -- | /O(n)/ Left fold on non-empty vectors with strict accumulator
1320 foldl1' :: Vector v a => (a -> a -> a) -> v a -> a
1321 {-# INLINE foldl1' #-}
1322 foldl1' f = Stream.foldl1' f . stream
1323
1324 -- | /O(n)/ Right fold
1325 foldr :: Vector v a => (a -> b -> b) -> b -> v a -> b
1326 {-# INLINE foldr #-}
1327 foldr f z = Stream.foldr f z . stream
1328
1329 -- | /O(n)/ Right fold on non-empty vectors
1330 foldr1 :: Vector v a => (a -> a -> a) -> v a -> a
1331 {-# INLINE foldr1 #-}
1332 foldr1 f = Stream.foldr1 f . stream
1333
1334 -- | /O(n)/ Right fold with a strict accumulator
1335 foldr' :: Vector v a => (a -> b -> b) -> b -> v a -> b
1336 {-# INLINE foldr' #-}
1337 foldr' f z = Stream.foldl' (flip f) z . streamR
1338
1339 -- | /O(n)/ Right fold on non-empty vectors with strict accumulator
1340 foldr1' :: Vector v a => (a -> a -> a) -> v a -> a
1341 {-# INLINE foldr1' #-}
1342 foldr1' f = Stream.foldl1' (flip f) . streamR
1343
1344 -- | /O(n)/ Left fold (function applied to each element and its index)
1345 ifoldl :: Vector v b => (a -> Int -> b -> a) -> a -> v b -> a
1346 {-# INLINE ifoldl #-}
1347 ifoldl f z = Stream.foldl (uncurry . f) z . Stream.indexed . stream
1348
1349 -- | /O(n)/ Left fold with strict accumulator (function applied to each element
1350 -- and its index)
1351 ifoldl' :: Vector v b => (a -> Int -> b -> a) -> a -> v b -> a
1352 {-# INLINE ifoldl' #-}
1353 ifoldl' f z = Stream.foldl' (uncurry . f) z . Stream.indexed . stream
1354
1355 -- | /O(n)/ Right fold (function applied to each element and its index)
1356 ifoldr :: Vector v a => (Int -> a -> b -> b) -> b -> v a -> b
1357 {-# INLINE ifoldr #-}
1358 ifoldr f z = Stream.foldr (uncurry f) z . Stream.indexed . stream
1359
1360 -- | /O(n)/ Right fold with strict accumulator (function applied to each
1361 -- element and its index)
1362 ifoldr' :: Vector v a => (Int -> a -> b -> b) -> b -> v a -> b
1363 {-# INLINE ifoldr' #-}
1364 ifoldr' f z xs = Stream.foldl' (flip (uncurry f)) z
1365 $ Stream.indexedR (length xs) $ streamR xs
1366
1367 -- Specialised folds
1368 -- -----------------
1369
1370 -- | /O(n)/ Check if all elements satisfy the predicate.
1371 all :: Vector v a => (a -> Bool) -> v a -> Bool
1372 {-# INLINE all #-}
1373 all f = Stream.and . Stream.map f . stream
1374
1375 -- | /O(n)/ Check if any element satisfies the predicate.
1376 any :: Vector v a => (a -> Bool) -> v a -> Bool
1377 {-# INLINE any #-}
1378 any f = Stream.or . Stream.map f . stream
1379
1380 -- | /O(n)/ Check if all elements are 'True'
1381 and :: Vector v Bool => v Bool -> Bool
1382 {-# INLINE and #-}
1383 and = Stream.and . stream
1384
1385 -- | /O(n)/ Check if any element is 'True'
1386 or :: Vector v Bool => v Bool -> Bool
1387 {-# INLINE or #-}
1388 or = Stream.or . stream
1389
1390 -- | /O(n)/ Compute the sum of the elements
1391 sum :: (Vector v a, Num a) => v a -> a
1392 {-# INLINE sum #-}
1393 sum = Stream.foldl' (+) 0 . stream
1394
1395 -- | /O(n)/ Compute the produce of the elements
1396 product :: (Vector v a, Num a) => v a -> a
1397 {-# INLINE product #-}
1398 product = Stream.foldl' (*) 1 . stream
1399
1400 -- | /O(n)/ Yield the maximum element of the vector. The vector may not be
1401 -- empty.
1402 maximum :: (Vector v a, Ord a) => v a -> a
1403 {-# INLINE maximum #-}
1404 maximum = Stream.foldl1' max . stream
1405
1406 -- | /O(n)/ Yield the maximum element of the vector according to the given
1407 -- comparison function. The vector may not be empty.
1408 maximumBy :: Vector v a => (a -> a -> Ordering) -> v a -> a
1409 {-# INLINE maximumBy #-}
1410 maximumBy cmp = Stream.foldl1' maxBy . stream
1411 where
1412 {-# INLINE maxBy #-}
1413 maxBy x y = case cmp x y of
1414 LT -> y
1415 _ -> x
1416
1417 -- | /O(n)/ Yield the minimum element of the vector. The vector may not be
1418 -- empty.
1419 minimum :: (Vector v a, Ord a) => v a -> a
1420 {-# INLINE minimum #-}
1421 minimum = Stream.foldl1' min . stream
1422
1423 -- | /O(n)/ Yield the minimum element of the vector according to the given
1424 -- comparison function. The vector may not be empty.
1425 minimumBy :: Vector v a => (a -> a -> Ordering) -> v a -> a
1426 {-# INLINE minimumBy #-}
1427 minimumBy cmp = Stream.foldl1' minBy . stream
1428 where
1429 {-# INLINE minBy #-}
1430 minBy x y = case cmp x y of
1431 GT -> y
1432 _ -> x
1433
1434 -- | /O(n)/ Yield the index of the maximum element of the vector. The vector
1435 -- may not be empty.
1436 maxIndex :: (Vector v a, Ord a) => v a -> Int
1437 {-# INLINE maxIndex #-}
1438 maxIndex = maxIndexBy compare
1439
1440 -- | /O(n)/ Yield the index of the maximum element of the vector according to
1441 -- the given comparison function. The vector may not be empty.
1442 maxIndexBy :: Vector v a => (a -> a -> Ordering) -> v a -> Int
1443 {-# INLINE maxIndexBy #-}
1444 maxIndexBy cmp = fst . Stream.foldl1' imax . Stream.indexed . stream
1445 where
1446 imax (i,x) (j,y) = i `seq` j `seq`
1447 case cmp x y of
1448 LT -> (j,y)
1449 _ -> (i,x)
1450
1451 -- | /O(n)/ Yield the index of the minimum element of the vector. The vector
1452 -- may not be empty.
1453 minIndex :: (Vector v a, Ord a) => v a -> Int
1454 {-# INLINE minIndex #-}
1455 minIndex = minIndexBy compare
1456
1457 -- | /O(n)/ Yield the index of the minimum element of the vector according to
1458 -- the given comparison function. The vector may not be empty.
1459 minIndexBy :: Vector v a => (a -> a -> Ordering) -> v a -> Int
1460 {-# INLINE minIndexBy #-}
1461 minIndexBy cmp = fst . Stream.foldl1' imin . Stream.indexed . stream
1462 where
1463 imin (i,x) (j,y) = i `seq` j `seq`
1464 case cmp x y of
1465 GT -> (j,y)
1466 _ -> (i,x)
1467
1468 -- Monadic folds
1469 -- -------------
1470
1471 -- | /O(n)/ Monadic fold
1472 foldM :: (Monad m, Vector v b) => (a -> b -> m a) -> a -> v b -> m a
1473 {-# INLINE foldM #-}
1474 foldM m z = Stream.foldM m z . stream
1475
1476 -- | /O(n)/ Monadic fold over non-empty vectors
1477 fold1M :: (Monad m, Vector v a) => (a -> a -> m a) -> v a -> m a
1478 {-# INLINE fold1M #-}
1479 fold1M m = Stream.fold1M m . stream
1480
1481 -- | /O(n)/ Monadic fold with strict accumulator
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 -- | /O(n)/ Monadic fold over non-empty vectors with strict accumulator
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 discard :: Monad m => m a -> m ()
1492 {-# INLINE discard #-}
1493 discard m = m >> return ()
1494
1495 -- | /O(n)/ Monadic fold that discards the result
1496 foldM_ :: (Monad m, Vector v b) => (a -> b -> m a) -> a -> v b -> m ()
1497 {-# INLINE foldM_ #-}
1498 foldM_ m z = discard . Stream.foldM m z . stream
1499
1500 -- | /O(n)/ Monadic fold over non-empty vectors that discards the result
1501 fold1M_ :: (Monad m, Vector v a) => (a -> a -> m a) -> v a -> m ()
1502 {-# INLINE fold1M_ #-}
1503 fold1M_ m = discard . Stream.fold1M m . stream
1504
1505 -- | /O(n)/ Monadic fold with strict accumulator that discards the result
1506 foldM'_ :: (Monad m, Vector v b) => (a -> b -> m a) -> a -> v b -> m ()
1507 {-# INLINE foldM'_ #-}
1508 foldM'_ m z = discard . Stream.foldM' m z . stream
1509
1510 -- | /O(n)/ Monad fold over non-empty vectors with strict accumulator
1511 -- that discards the result
1512 fold1M'_ :: (Monad m, Vector v a) => (a -> a -> m a) -> v a -> m ()
1513 {-# INLINE fold1M'_ #-}
1514 fold1M'_ m = discard . Stream.fold1M' m . stream
1515
1516 -- Monadic sequencing
1517 -- ------------------
1518
1519 -- | Evaluate each action and collect the results
1520 sequence :: (Monad m, Vector v a, Vector v (m a)) => v (m a) -> m (v a)
1521 {-# INLINE sequence #-}
1522 sequence = mapM id
1523
1524 -- | Evaluate each action and discard the results
1525 sequence_ :: (Monad m, Vector v (m a)) => v (m a) -> m ()
1526 {-# INLINE sequence_ #-}
1527 sequence_ = mapM_ id
1528
1529 -- Prefix sums (scans)
1530 -- -------------------
1531
1532 -- | /O(n)/ Prescan
1533 --
1534 -- @
1535 -- prescanl f z = 'init' . 'scanl' f z
1536 -- @
1537 --
1538 -- Example: @prescanl (+) 0 \<1,2,3,4\> = \<0,1,3,6\>@
1539 --
1540 prescanl :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1541 {-# INLINE prescanl #-}
1542 prescanl f z = unstream . inplace (MStream.prescanl f z) . stream
1543
1544 -- | /O(n)/ Prescan with strict accumulator
1545 prescanl' :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1546 {-# INLINE prescanl' #-}
1547 prescanl' f z = unstream . inplace (MStream.prescanl' f z) . stream
1548
1549 -- | /O(n)/ Scan
1550 --
1551 -- @
1552 -- postscanl f z = 'tail' . 'scanl' f z
1553 -- @
1554 --
1555 -- Example: @postscanl (+) 0 \<1,2,3,4\> = \<1,3,6,10\>@
1556 --
1557 postscanl :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1558 {-# INLINE postscanl #-}
1559 postscanl f z = unstream . inplace (MStream.postscanl f z) . stream
1560
1561 -- | /O(n)/ Scan with strict accumulator
1562 postscanl' :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1563 {-# INLINE postscanl' #-}
1564 postscanl' f z = unstream . inplace (MStream.postscanl' f z) . stream
1565
1566 -- | /O(n)/ Haskell-style scan
1567 --
1568 -- > scanl f z <x1,...,xn> = <y1,...,y(n+1)>
1569 -- > where y1 = z
1570 -- > yi = f y(i-1) x(i-1)
1571 --
1572 -- Example: @scanl (+) 0 \<1,2,3,4\> = \<0,1,3,6,10\>@
1573 --
1574 scanl :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1575 {-# INLINE scanl #-}
1576 scanl f z = unstream . Stream.scanl f z . stream
1577
1578 -- | /O(n)/ Haskell-style scan with strict accumulator
1579 scanl' :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1580 {-# INLINE scanl' #-}
1581 scanl' f z = unstream . Stream.scanl' f z . stream
1582
1583 -- | /O(n)/ Scan over a non-empty vector
1584 --
1585 -- > scanl f <x1,...,xn> = <y1,...,yn>
1586 -- > where y1 = x1
1587 -- > yi = f y(i-1) xi
1588 --
1589 scanl1 :: Vector v a => (a -> a -> a) -> v a -> v a
1590 {-# INLINE scanl1 #-}
1591 scanl1 f = unstream . inplace (MStream.scanl1 f) . stream
1592
1593 -- | /O(n)/ Scan over a non-empty vector with a strict accumulator
1594 scanl1' :: Vector v a => (a -> a -> a) -> v a -> v a
1595 {-# INLINE scanl1' #-}
1596 scanl1' f = unstream . inplace (MStream.scanl1' f) . stream
1597
1598 -- | /O(n)/ Right-to-left prescan
1599 --
1600 -- @
1601 -- prescanr f z = 'reverse' . 'prescanl' (flip f) z . 'reverse'
1602 -- @
1603 --
1604 prescanr :: (Vector v a, Vector v b) => (a -> b -> b) -> b -> v a -> v b
1605 {-# INLINE prescanr #-}
1606 prescanr f z = unstreamR . inplace (MStream.prescanl (flip f) z) . streamR
1607
1608 -- | /O(n)/ Right-to-left prescan with strict accumulator
1609 prescanr' :: (Vector v a, Vector v b) => (a -> b -> b) -> b -> v a -> v b
1610 {-# INLINE prescanr' #-}
1611 prescanr' f z = unstreamR . inplace (MStream.prescanl' (flip f) z) . streamR
1612
1613 -- | /O(n)/ Right-to-left scan
1614 postscanr :: (Vector v a, Vector v b) => (a -> b -> b) -> b -> v a -> v b
1615 {-# INLINE postscanr #-}
1616 postscanr f z = unstreamR . inplace (MStream.postscanl (flip f) z) . streamR
1617
1618 -- | /O(n)/ Right-to-left scan with strict accumulator
1619 postscanr' :: (Vector v a, Vector v b) => (a -> b -> b) -> b -> v a -> v b
1620 {-# INLINE postscanr' #-}
1621 postscanr' f z = unstreamR . inplace (MStream.postscanl' (flip f) z) . streamR
1622
1623 -- | /O(n)/ Right-to-left Haskell-style scan
1624 scanr :: (Vector v a, Vector v b) => (a -> b -> b) -> b -> v a -> v b
1625 {-# INLINE scanr #-}
1626 scanr f z = unstreamR . Stream.scanl (flip f) z . streamR
1627
1628 -- | /O(n)/ Right-to-left Haskell-style scan with strict accumulator
1629 scanr' :: (Vector v a, Vector v b) => (a -> b -> b) -> b -> v a -> v b
1630 {-# INLINE scanr' #-}
1631 scanr' f z = unstreamR . Stream.scanl' (flip f) z . streamR
1632
1633 -- | /O(n)/ Right-to-left scan over a non-empty vector
1634 scanr1 :: Vector v a => (a -> a -> a) -> v a -> v a
1635 {-# INLINE scanr1 #-}
1636 scanr1 f = unstreamR . inplace (MStream.scanl1 (flip f)) . streamR
1637
1638 -- | /O(n)/ Right-to-left scan over a non-empty vector with a strict
1639 -- accumulator
1640 scanr1' :: Vector v a => (a -> a -> a) -> v a -> v a
1641 {-# INLINE scanr1' #-}
1642 scanr1' f = unstreamR . inplace (MStream.scanl1' (flip f)) . streamR
1643
1644 -- Conversions - Lists
1645 -- ------------------------
1646
1647 -- | /O(n)/ Convert a vector to a list
1648 toList :: Vector v a => v a -> [a]
1649 {-# INLINE toList #-}
1650 toList = Stream.toList . stream
1651
1652 -- | /O(n)/ Convert a list to a vector
1653 fromList :: Vector v a => [a] -> v a
1654 {-# INLINE fromList #-}
1655 fromList = unstream . Stream.fromList
1656
1657 -- | /O(n)/ Convert the first @n@ elements of a list to a vector
1658 --
1659 -- @
1660 -- fromListN n xs = 'fromList' ('take' n xs)
1661 -- @
1662 fromListN :: Vector v a => Int -> [a] -> v a
1663 {-# INLINE fromListN #-}
1664 fromListN n = unstream . Stream.fromListN n
1665
1666 -- Conversions - Immutable vectors
1667 -- -------------------------------
1668
1669 -- | /O(n)/ Convert different vector types
1670 convert :: (Vector v a, Vector w a) => v a -> w a
1671 {-# INLINE convert #-}
1672 convert = unstream . stream
1673
1674 -- Conversions - Mutable vectors
1675 -- -----------------------------
1676
1677 -- | /O(1)/ Unsafe convert a mutable vector to an immutable one without
1678 -- copying. The mutable vector may not be used after this operation.
1679 unsafeFreeze
1680 :: (PrimMonad m, Vector v a) => Mutable v (PrimState m) a -> m (v a)
1681 {-# INLINE unsafeFreeze #-}
1682 unsafeFreeze = basicUnsafeFreeze
1683
1684
1685 -- | /O(1)/ Unsafely convert an immutable vector to a mutable one without
1686 -- copying. The immutable vector may not be used after this operation.
1687 unsafeThaw :: (PrimMonad m, Vector v a) => v a -> m (Mutable v (PrimState m) a)
1688 {-# INLINE unsafeThaw #-}
1689 unsafeThaw = basicUnsafeThaw
1690
1691 -- | /O(n)/ Yield a mutable copy of the immutable vector.
1692 thaw :: (PrimMonad m, Vector v a) => v a -> m (Mutable v (PrimState m) a)
1693 {-# INLINE_STREAM thaw #-}
1694 thaw v = do
1695 mv <- M.unsafeNew (length v)
1696 unsafeCopy mv v
1697 return mv
1698
1699 -- | /O(n)/ Yield an immutable copy of the mutable vector.
1700 freeze :: (PrimMonad m, Vector v a) => Mutable v (PrimState m) a -> m (v a)
1701 {-# INLINE freeze #-}
1702 freeze mv = unsafeFreeze =<< M.clone mv
1703
1704 {-
1705 -- | /O(n)/ Yield a mutable vector containing copies of each vector in the
1706 -- list.
1707 thawMany :: (PrimMonad m, Vector v a) => [v a] -> m (Mutable v (PrimState m) a)
1708 {-# INLINE_STREAM thawMany #-}
1709 -- FIXME: add rule for (stream (new (New.create (thawMany vs))))
1710 -- NOTE: We don't try to consume the list lazily as this wouldn't significantly
1711 -- change the space requirements anyway.
1712 thawMany vs = do
1713 mv <- M.new n
1714 thaw_loop mv vs
1715 return mv
1716 where
1717 n = List.foldl' (\k v -> k + length v) 0 vs
1718
1719 thaw_loop mv [] = mv `seq` return ()
1720 thaw_loop mv (v:vs)
1721 = do
1722 let n = length v
1723 unsafeCopy (M.unsafeTake n mv) v
1724 thaw_loop (M.unsafeDrop n mv) vs
1725 -}
1726
1727 -- | /O(n)/ Copy an immutable vector into a mutable one. The two vectors must
1728 -- have the same length.
1729 copy
1730 :: (PrimMonad m, Vector v a) => Mutable v (PrimState m) a -> v a -> m ()
1731 {-# INLINE copy #-}
1732 copy dst src = BOUNDS_CHECK(check) "copy" "length mismatch"
1733 (M.length dst == length src)
1734 $ unsafeCopy dst src
1735
1736 -- | /O(n)/ Copy an immutable vector into a mutable one. The two vectors must
1737 -- have the same length. This is not checked.
1738 unsafeCopy
1739 :: (PrimMonad m, Vector v a) => Mutable v (PrimState m) a -> v a -> m ()
1740 {-# INLINE unsafeCopy #-}
1741 unsafeCopy dst src = UNSAFE_CHECK(check) "unsafeCopy" "length mismatch"
1742 (M.length dst == length src)
1743 $ (dst `seq` src `seq` basicUnsafeCopy dst src)
1744
1745 -- Conversions to/from Streams
1746 -- ---------------------------
1747
1748 -- | /O(1)/ Convert a vector to a 'Stream'
1749 stream :: Vector v a => v a -> Stream a
1750 {-# INLINE_STREAM stream #-}
1751 stream v = v `seq` n `seq` (Stream.unfoldr get 0 `Stream.sized` Exact n)
1752 where
1753 n = length v
1754
1755 -- NOTE: the False case comes first in Core so making it the recursive one
1756 -- makes the code easier to read
1757 {-# INLINE get #-}
1758 get i | i >= n = Nothing
1759 | otherwise = case basicUnsafeIndexM v i of Box x -> Just (x, i+1)
1760
1761 -- | /O(n)/ Construct a vector from a 'Stream'
1762 unstream :: Vector v a => Stream a -> v a
1763 {-# INLINE unstream #-}
1764 unstream s = new (New.unstream s)
1765
1766 {-# RULES
1767
1768 "stream/unstream [Vector]" forall s.
1769 stream (new (New.unstream s)) = s
1770
1771 "New.unstream/stream [Vector]" forall v.
1772 New.unstream (stream v) = clone v
1773
1774 "clone/new [Vector]" forall p.
1775 clone (new p) = p
1776
1777 "inplace [Vector]"
1778 forall (f :: forall m. Monad m => MStream m a -> MStream m a) m.
1779 New.unstream (inplace f (stream (new m))) = New.transform f m
1780
1781 "uninplace [Vector]"
1782 forall (f :: forall m. Monad m => MStream m a -> MStream m a) m.
1783 stream (new (New.transform f m)) = inplace f (stream (new m))
1784
1785 #-}
1786
1787 -- | /O(1)/ Convert a vector to a 'Stream', proceeding from right to left
1788 streamR :: Vector v a => v a -> Stream a
1789 {-# INLINE_STREAM streamR #-}
1790 streamR v = v `seq` n `seq` (Stream.unfoldr get n `Stream.sized` Exact n)
1791 where
1792 n = length v
1793
1794 {-# INLINE get #-}
1795 get 0 = Nothing
1796 get i = let i' = i-1
1797 in
1798 case basicUnsafeIndexM v i' of Box x -> Just (x, i')
1799
1800 -- | /O(n)/ Construct a vector from a 'Stream', proceeding from right to left
1801 unstreamR :: Vector v a => Stream a -> v a
1802 {-# INLINE unstreamR #-}
1803 unstreamR s = new (New.unstreamR s)
1804
1805 {-# RULES
1806
1807 "streamR/unstreamR [Vector]" forall s.
1808 streamR (new (New.unstreamR s)) = s
1809
1810 "New.unstreamR/streamR/new [Vector]" forall p.
1811 New.unstreamR (streamR (new p)) = p
1812
1813 "inplace right [Vector]"
1814 forall (f :: forall m. Monad m => MStream m a -> MStream m a) m.
1815 New.unstreamR (inplace f (streamR (new m))) = New.transformR f m
1816
1817 "uninplace right [Vector]"
1818 forall (f :: forall m. Monad m => MStream m a -> MStream m a) m.
1819 streamR (new (New.transformR f m)) = inplace f (streamR (new m))
1820
1821 #-}
1822
1823 unstreamM :: (Monad m, Vector v a) => MStream m a -> m (v a)
1824 {-# INLINE_STREAM unstreamM #-}
1825 unstreamM s = do
1826 xs <- MStream.toList s
1827 return $ unstream $ Stream.unsafeFromList (MStream.size s) xs
1828
1829 unstreamPrimM :: (PrimMonad m, Vector v a) => MStream m a -> m (v a)
1830 {-# INLINE_STREAM unstreamPrimM #-}
1831 unstreamPrimM s = M.munstream s >>= unsafeFreeze
1832
1833 -- FIXME: the next two functions are only necessary for the specialisations
1834 unstreamPrimM_IO :: Vector v a => MStream IO a -> IO (v a)
1835 {-# INLINE unstreamPrimM_IO #-}
1836 unstreamPrimM_IO = unstreamPrimM
1837
1838 unstreamPrimM_ST :: Vector v a => MStream (ST s) a -> ST s (v a)
1839 {-# INLINE unstreamPrimM_ST #-}
1840 unstreamPrimM_ST = unstreamPrimM
1841
1842 {-# RULES
1843
1844 "unstreamM[IO]" unstreamM = unstreamPrimM_IO
1845 "unstreamM[ST]" unstreamM = unstreamPrimM_ST
1846
1847 #-}
1848
1849
1850 -- Recycling support
1851 -- -----------------
1852
1853 -- | Construct a vector from a monadic initialiser.
1854 new :: Vector v a => New v a -> v a
1855 {-# INLINE_STREAM new #-}
1856 new m = m `seq` runST (unsafeFreeze =<< New.run m)
1857
1858 -- | Convert a vector to an initialiser which, when run, produces a copy of
1859 -- the vector.
1860 clone :: Vector v a => v a -> New v a
1861 {-# INLINE_STREAM clone #-}
1862 clone v = v `seq` New.create (
1863 do
1864 mv <- M.new (length v)
1865 unsafeCopy mv v
1866 return mv)
1867
1868 -- Comparisons
1869 -- -----------
1870
1871 -- | /O(n)/ Check if two vectors are equal. All 'Vector' instances are also
1872 -- instances of 'Eq' and it is usually more appropriate to use those. This
1873 -- function is primarily intended for implementing 'Eq' instances for new
1874 -- vector types.
1875 eq :: (Vector v a, Eq a) => v a -> v a -> Bool
1876 {-# INLINE eq #-}
1877 xs `eq` ys = stream xs == stream ys
1878
1879 -- | /O(n)/ Compare two vectors lexicographically. All 'Vector' instances are
1880 -- also instances of 'Ord' and it is usually more appropriate to use those. This
1881 -- function is primarily intended for implementing 'Ord' instances for new
1882 -- vector types.
1883 cmp :: (Vector v a, Ord a) => v a -> v a -> Ordering
1884 {-# INLINE cmp #-}
1885 cmp xs ys = compare (stream xs) (stream ys)
1886
1887 -- Data and Typeable
1888 -- -----------------
1889
1890 -- | Generic definion of 'Data.Data.gfoldl' that views a 'Vector' as a
1891 -- list.
1892 gfoldl :: (Vector v a, Data a)
1893 => (forall d b. Data d => c (d -> b) -> d -> c b)
1894 -> (forall g. g -> c g)
1895 -> v a
1896 -> c (v a)
1897 {-# INLINE gfoldl #-}
1898 gfoldl f z v = z fromList `f` toList v
1899
1900 mkType :: String -> DataType
1901 {-# INLINE mkType #-}
1902 mkType = mkNoRepType
1903
1904 dataCast :: (Vector v a, Data a, Typeable1 v, Typeable1 t)
1905 => (forall d. Data d => c (t d)) -> Maybe (c (v a))
1906 {-# INLINE dataCast #-}
1907 dataCast f = gcast1 f
1908