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