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