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