Add constructN and constructrN
[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, 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 = create (thawMany vs)
675 concat vs = unstream (Stream.flatten mk step (Exact n) (Stream.fromList vs))
676 where
677 n = List.foldl' (\k v -> k + length v) 0 vs
678
679 {-# INLINE_INNER step #-}
680 step (v,i,k)
681 | i < k = case unsafeIndexM v i of
682 Box x -> Stream.Yield x (v,i+1,k)
683 | otherwise = Stream.Done
684
685 {-# INLINE mk #-}
686 mk v = let k = length v
687 in
688 k `seq` (v,0,k)
689
690 -- Monadic initialisation
691 -- ----------------------
692
693 -- | /O(n)/ Execute the monadic action the given number of times and store the
694 -- results in a vector.
695 replicateM :: (Monad m, Vector v a) => Int -> m a -> m (v a)
696 {-# INLINE replicateM #-}
697 replicateM n m = unstreamM (MStream.replicateM n m)
698
699 -- | /O(n)/ Construct a vector of the given length by applying the monadic
700 -- action to each index
701 generateM :: (Monad m, Vector v a) => Int -> (Int -> m a) -> m (v a)
702 {-# INLINE generateM #-}
703 generateM n f = unstreamM (MStream.generateM n f)
704
705 -- | Execute the monadic action and freeze the resulting vector.
706 --
707 -- @
708 -- create (do { v \<- 'M.new' 2; 'M.write' v 0 \'a\'; 'M.write' v 1 \'b\' }) = \<'a','b'\>
709 -- @
710 create :: Vector v a => (forall s. ST s (Mutable v s a)) -> v a
711 {-# INLINE create #-}
712 create p = new (New.create p)
713
714 -- Restricting memory usage
715 -- ------------------------
716
717 -- | /O(n)/ Yield the argument but force it not to retain any extra memory,
718 -- possibly by copying it.
719 --
720 -- This is especially useful when dealing with slices. For example:
721 --
722 -- > force (slice 0 2 <huge vector>)
723 --
724 -- Here, the slice retains a reference to the huge vector. Forcing it creates
725 -- a copy of just the elements that belong to the slice and allows the huge
726 -- vector to be garbage collected.
727 force :: Vector v a => v a -> v a
728 -- FIXME: we probably ought to inline this later as the rules still might fire
729 -- otherwise
730 {-# INLINE_STREAM force #-}
731 force v = new (clone v)
732
733 -- Bulk updates
734 -- ------------
735
736 -- | /O(m+n)/ For each pair @(i,a)@ from the list, replace the vector
737 -- element at position @i@ by @a@.
738 --
739 -- > <5,9,2,7> // [(2,1),(0,3),(2,8)] = <3,9,8,7>
740 --
741 (//) :: Vector v a => v a -- ^ initial vector (of length @m@)
742 -> [(Int, a)] -- ^ list of index/value pairs (of length @n@)
743 -> v a
744 {-# INLINE (//) #-}
745 v // us = update_stream v (Stream.fromList us)
746
747 -- | /O(m+n)/ For each pair @(i,a)@ from the vector of index/value pairs,
748 -- replace the vector element at position @i@ by @a@.
749 --
750 -- > update <5,9,2,7> <(2,1),(0,3),(2,8)> = <3,9,8,7>
751 --
752 update :: (Vector v a, Vector v (Int, a))
753 => v a -- ^ initial vector (of length @m@)
754 -> v (Int, a) -- ^ vector of index/value pairs (of length @n@)
755 -> v a
756 {-# INLINE update #-}
757 update v w = update_stream v (stream w)
758
759 -- | /O(m+min(n1,n2))/ For each index @i@ from the index vector and the
760 -- corresponding value @a@ from the value vector, replace the element of the
761 -- initial vector at position @i@ by @a@.
762 --
763 -- > update_ <5,9,2,7> <2,0,2> <1,3,8> = <3,9,8,7>
764 --
765 -- This function is useful for instances of 'Vector' that cannot store pairs.
766 -- Otherwise, 'update' is probably more convenient.
767 --
768 -- @
769 -- update_ xs is ys = 'update' xs ('zip' is ys)
770 -- @
771 update_ :: (Vector v a, Vector v Int)
772 => v a -- ^ initial vector (of length @m@)
773 -> v Int -- ^ index vector (of length @n1@)
774 -> v a -- ^ value vector (of length @n2@)
775 -> v a
776 {-# INLINE update_ #-}
777 update_ v is w = update_stream v (Stream.zipWith (,) (stream is) (stream w))
778
779 update_stream :: Vector v a => v a -> Stream (Int,a) -> v a
780 {-# INLINE update_stream #-}
781 update_stream = modifyWithStream M.update
782
783 -- | Same as ('//') but without bounds checking.
784 unsafeUpd :: Vector v a => v a -> [(Int, a)] -> v a
785 {-# INLINE unsafeUpd #-}
786 unsafeUpd v us = unsafeUpdate_stream v (Stream.fromList us)
787
788 -- | Same as 'update' but without bounds checking.
789 unsafeUpdate :: (Vector v a, Vector v (Int, a)) => v a -> v (Int, a) -> v a
790 {-# INLINE unsafeUpdate #-}
791 unsafeUpdate v w = unsafeUpdate_stream v (stream w)
792
793 -- | Same as 'update_' but without bounds checking.
794 unsafeUpdate_ :: (Vector v a, Vector v Int) => v a -> v Int -> v a -> v a
795 {-# INLINE unsafeUpdate_ #-}
796 unsafeUpdate_ v is w
797 = unsafeUpdate_stream v (Stream.zipWith (,) (stream is) (stream w))
798
799 unsafeUpdate_stream :: Vector v a => v a -> Stream (Int,a) -> v a
800 {-# INLINE unsafeUpdate_stream #-}
801 unsafeUpdate_stream = modifyWithStream M.unsafeUpdate
802
803 -- Accumulations
804 -- -------------
805
806 -- | /O(m+n)/ For each pair @(i,b)@ from the list, replace the vector element
807 -- @a@ at position @i@ by @f a b@.
808 --
809 -- > accum (+) <5,9,2> [(2,4),(1,6),(0,3),(1,7)] = <5+3, 9+6+7, 2+4>
810 accum :: Vector v a
811 => (a -> b -> a) -- ^ accumulating function @f@
812 -> v a -- ^ initial vector (of length @m@)
813 -> [(Int,b)] -- ^ list of index/value pairs (of length @n@)
814 -> v a
815 {-# INLINE accum #-}
816 accum f v us = accum_stream f v (Stream.fromList us)
817
818 -- | /O(m+n)/ For each pair @(i,b)@ from the vector of pairs, replace the vector
819 -- element @a@ at position @i@ by @f a b@.
820 --
821 -- > accumulate (+) <5,9,2> <(2,4),(1,6),(0,3),(1,7)> = <5+3, 9+6+7, 2+4>
822 accumulate :: (Vector v a, Vector v (Int, b))
823 => (a -> b -> a) -- ^ accumulating function @f@
824 -> v a -- ^ initial vector (of length @m@)
825 -> v (Int,b) -- ^ vector of index/value pairs (of length @n@)
826 -> v a
827 {-# INLINE accumulate #-}
828 accumulate f v us = accum_stream f v (stream us)
829
830 -- | /O(m+min(n1,n2))/ For each index @i@ from the index vector and the
831 -- corresponding value @b@ from the the value vector,
832 -- replace the element of the initial vector at
833 -- position @i@ by @f a b@.
834 --
835 -- > accumulate_ (+) <5,9,2> <2,1,0,1> <4,6,3,7> = <5+3, 9+6+7, 2+4>
836 --
837 -- This function is useful for instances of 'Vector' that cannot store pairs.
838 -- Otherwise, 'accumulate' is probably more convenient:
839 --
840 -- @
841 -- accumulate_ f as is bs = 'accumulate' f as ('zip' is bs)
842 -- @
843 accumulate_ :: (Vector v a, Vector v Int, Vector v b)
844 => (a -> b -> a) -- ^ accumulating function @f@
845 -> v a -- ^ initial vector (of length @m@)
846 -> v Int -- ^ index vector (of length @n1@)
847 -> v b -- ^ value vector (of length @n2@)
848 -> v a
849 {-# INLINE accumulate_ #-}
850 accumulate_ f v is xs = accum_stream f v (Stream.zipWith (,) (stream is)
851 (stream xs))
852
853
854 accum_stream :: Vector v a => (a -> b -> a) -> v a -> Stream (Int,b) -> v a
855 {-# INLINE accum_stream #-}
856 accum_stream f = modifyWithStream (M.accum f)
857
858 -- | Same as 'accum' but without bounds checking.
859 unsafeAccum :: Vector v a => (a -> b -> a) -> v a -> [(Int,b)] -> v a
860 {-# INLINE unsafeAccum #-}
861 unsafeAccum f v us = unsafeAccum_stream f v (Stream.fromList us)
862
863 -- | Same as 'accumulate' but without bounds checking.
864 unsafeAccumulate :: (Vector v a, Vector v (Int, b))
865 => (a -> b -> a) -> v a -> v (Int,b) -> v a
866 {-# INLINE unsafeAccumulate #-}
867 unsafeAccumulate f v us = unsafeAccum_stream f v (stream us)
868
869 -- | Same as 'accumulate_' but without bounds checking.
870 unsafeAccumulate_ :: (Vector v a, Vector v Int, Vector v b)
871 => (a -> b -> a) -> v a -> v Int -> v b -> v a
872 {-# INLINE unsafeAccumulate_ #-}
873 unsafeAccumulate_ f v is xs
874 = unsafeAccum_stream f v (Stream.zipWith (,) (stream is) (stream xs))
875
876 unsafeAccum_stream
877 :: Vector v a => (a -> b -> a) -> v a -> Stream (Int,b) -> v a
878 {-# INLINE unsafeAccum_stream #-}
879 unsafeAccum_stream f = modifyWithStream (M.unsafeAccum f)
880
881 -- Permutations
882 -- ------------
883
884 -- | /O(n)/ Reverse a vector
885 reverse :: (Vector v a) => v a -> v a
886 {-# INLINE reverse #-}
887 -- FIXME: make this fuse better, add support for recycling
888 reverse = unstream . streamR
889
890 -- | /O(n)/ Yield the vector obtained by replacing each element @i@ of the
891 -- index vector by @xs'!'i@. This is equivalent to @'map' (xs'!') is@ but is
892 -- often much more efficient.
893 --
894 -- > backpermute <a,b,c,d> <0,3,2,3,1,0> = <a,d,c,d,b,a>
895 backpermute :: (Vector v a, Vector v Int)
896 => v a -- ^ @xs@ value vector
897 -> v Int -- ^ @is@ index vector (of length @n@)
898 -> v a
899 {-# INLINE backpermute #-}
900 -- This somewhat non-intuitive definition ensures that the resulting vector
901 -- does not retain references to the original one even if it is lazy in its
902 -- elements. This would not be the case if we simply used map (v!)
903 backpermute v is = seq v
904 $ seq n
905 $ unstream
906 $ Stream.unbox
907 $ Stream.map index
908 $ stream is
909 where
910 n = length v
911
912 {-# INLINE index #-}
913 -- NOTE: we do it this way to avoid triggering LiberateCase on n in
914 -- polymorphic code
915 index i = BOUNDS_CHECK(checkIndex) "backpermute" i n
916 $ basicUnsafeIndexM v i
917
918 -- | Same as 'backpermute' but without bounds checking.
919 unsafeBackpermute :: (Vector v a, Vector v Int) => v a -> v Int -> v a
920 {-# INLINE unsafeBackpermute #-}
921 unsafeBackpermute v is = seq v
922 $ seq n
923 $ unstream
924 $ Stream.unbox
925 $ Stream.map index
926 $ stream is
927 where
928 n = length v
929
930 {-# INLINE index #-}
931 -- NOTE: we do it this way to avoid triggering LiberateCase on n in
932 -- polymorphic code
933 index i = UNSAFE_CHECK(checkIndex) "unsafeBackpermute" i n
934 $ basicUnsafeIndexM v i
935
936 -- Safe destructive updates
937 -- ------------------------
938
939 -- | Apply a destructive operation to a vector. The operation will be
940 -- performed in place if it is safe to do so and will modify a copy of the
941 -- vector otherwise.
942 --
943 -- @
944 -- modify (\\v -> 'M.write' v 0 \'x\') ('replicate' 3 \'a\') = \<\'x\',\'a\',\'a\'\>
945 -- @
946 modify :: Vector v a => (forall s. Mutable v s a -> ST s ()) -> v a -> v a
947 {-# INLINE modify #-}
948 modify p = new . New.modify p . clone
949
950 -- We have to make sure that this is strict in the stream but we can't seq on
951 -- it while fusion is happening. Hence this ugliness.
952 modifyWithStream :: Vector v a
953 => (forall s. Mutable v s a -> Stream b -> ST s ())
954 -> v a -> Stream b -> v a
955 {-# INLINE modifyWithStream #-}
956 modifyWithStream p v s = new (New.modifyWithStream p (clone v) s)
957
958 -- Indexing
959 -- --------
960
961 -- | /O(n)/ Pair each element in a vector with its index
962 indexed :: (Vector v a, Vector v (Int,a)) => v a -> v (Int,a)
963 {-# INLINE indexed #-}
964 indexed = unstream . Stream.indexed . stream
965
966 -- Mapping
967 -- -------
968
969 -- | /O(n)/ Map a function over a vector
970 map :: (Vector v a, Vector v b) => (a -> b) -> v a -> v b
971 {-# INLINE map #-}
972 map f = unstream . inplace (MStream.map f) . stream
973
974 -- | /O(n)/ Apply a function to every element of a vector and its index
975 imap :: (Vector v a, Vector v b) => (Int -> a -> b) -> v a -> v b
976 {-# INLINE imap #-}
977 imap f = unstream . inplace (MStream.map (uncurry f) . MStream.indexed)
978 . stream
979
980 -- | Map a function over a vector and concatenate the results.
981 concatMap :: (Vector v a, Vector v b) => (a -> v b) -> v a -> v b
982 {-# INLINE concatMap #-}
983 -- NOTE: We can't fuse concatMap anyway so don't pretend we do.
984 -- concatMap f = unstream . Stream.concatMap (stream . f) . stream
985 concatMap f = concat . Stream.toList . Stream.map f . stream
986
987 -- Monadic mapping
988 -- ---------------
989
990 -- | /O(n)/ Apply the monadic action to all elements of the vector, yielding a
991 -- vector of results
992 mapM :: (Monad m, Vector v a, Vector v b) => (a -> m b) -> v a -> m (v b)
993 {-# INLINE mapM #-}
994 mapM f = unstreamM . Stream.mapM f . stream
995
996 -- | /O(n)/ Apply the monadic action to all elements of a vector and ignore the
997 -- results
998 mapM_ :: (Monad m, Vector v a) => (a -> m b) -> v a -> m ()
999 {-# INLINE mapM_ #-}
1000 mapM_ f = Stream.mapM_ f . stream
1001
1002 -- | /O(n)/ Apply the monadic action to all elements of the vector, yielding a
1003 -- vector of results. Equvalent to @flip 'mapM'@.
1004 forM :: (Monad m, Vector v a, Vector v b) => v a -> (a -> m b) -> m (v b)
1005 {-# INLINE forM #-}
1006 forM as f = mapM f as
1007
1008 -- | /O(n)/ Apply the monadic action to all elements of a vector and ignore the
1009 -- results. Equivalent to @flip 'mapM_'@.
1010 forM_ :: (Monad m, Vector v a) => v a -> (a -> m b) -> m ()
1011 {-# INLINE forM_ #-}
1012 forM_ as f = mapM_ f as
1013
1014 -- Zipping
1015 -- -------
1016
1017 -- | /O(min(m,n))/ Zip two vectors with the given function.
1018 zipWith :: (Vector v a, Vector v b, Vector v c)
1019 => (a -> b -> c) -> v a -> v b -> v c
1020 {-# INLINE zipWith #-}
1021 zipWith f xs ys = unstream (Stream.zipWith f (stream xs) (stream ys))
1022
1023 -- | Zip three vectors with the given function.
1024 zipWith3 :: (Vector v a, Vector v b, Vector v c, Vector v d)
1025 => (a -> b -> c -> d) -> v a -> v b -> v c -> v d
1026 {-# INLINE zipWith3 #-}
1027 zipWith3 f as bs cs = unstream (Stream.zipWith3 f (stream as)
1028 (stream bs)
1029 (stream cs))
1030
1031 zipWith4 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e)
1032 => (a -> b -> c -> d -> e) -> v a -> v b -> v c -> v d -> v e
1033 {-# INLINE zipWith4 #-}
1034 zipWith4 f as bs cs ds
1035 = unstream (Stream.zipWith4 f (stream as)
1036 (stream bs)
1037 (stream cs)
1038 (stream ds))
1039
1040 zipWith5 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
1041 Vector v f)
1042 => (a -> b -> c -> d -> e -> f) -> v a -> v b -> v c -> v d -> v e
1043 -> v f
1044 {-# INLINE zipWith5 #-}
1045 zipWith5 f as bs cs ds es
1046 = unstream (Stream.zipWith5 f (stream as)
1047 (stream bs)
1048 (stream cs)
1049 (stream ds)
1050 (stream es))
1051
1052 zipWith6 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
1053 Vector v f, Vector v g)
1054 => (a -> b -> c -> d -> e -> f -> g)
1055 -> v a -> v b -> v c -> v d -> v e -> v f -> v g
1056 {-# INLINE zipWith6 #-}
1057 zipWith6 f as bs cs ds es fs
1058 = unstream (Stream.zipWith6 f (stream as)
1059 (stream bs)
1060 (stream cs)
1061 (stream ds)
1062 (stream es)
1063 (stream fs))
1064
1065 -- | /O(min(m,n))/ Zip two vectors with a function that also takes the
1066 -- elements' indices.
1067 izipWith :: (Vector v a, Vector v b, Vector v c)
1068 => (Int -> a -> b -> c) -> v a -> v b -> v c
1069 {-# INLINE izipWith #-}
1070 izipWith f xs ys = unstream
1071 (Stream.zipWith (uncurry f) (Stream.indexed (stream xs))
1072 (stream ys))
1073
1074 izipWith3 :: (Vector v a, Vector v b, Vector v c, Vector v d)
1075 => (Int -> a -> b -> c -> d) -> v a -> v b -> v c -> v d
1076 {-# INLINE izipWith3 #-}
1077 izipWith3 f as bs cs
1078 = unstream (Stream.zipWith3 (uncurry f) (Stream.indexed (stream as))
1079 (stream bs)
1080 (stream cs))
1081
1082 izipWith4 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e)
1083 => (Int -> a -> b -> c -> d -> e) -> v a -> v b -> v c -> v d -> v e
1084 {-# INLINE izipWith4 #-}
1085 izipWith4 f as bs cs ds
1086 = unstream (Stream.zipWith4 (uncurry f) (Stream.indexed (stream as))
1087 (stream bs)
1088 (stream cs)
1089 (stream ds))
1090
1091 izipWith5 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
1092 Vector v f)
1093 => (Int -> a -> b -> c -> d -> e -> f) -> v a -> v b -> v c -> v d
1094 -> v e -> v f
1095 {-# INLINE izipWith5 #-}
1096 izipWith5 f as bs cs ds es
1097 = unstream (Stream.zipWith5 (uncurry f) (Stream.indexed (stream as))
1098 (stream bs)
1099 (stream cs)
1100 (stream ds)
1101 (stream es))
1102
1103 izipWith6 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
1104 Vector v f, Vector v g)
1105 => (Int -> a -> b -> c -> d -> e -> f -> g)
1106 -> v a -> v b -> v c -> v d -> v e -> v f -> v g
1107 {-# INLINE izipWith6 #-}
1108 izipWith6 f as bs cs ds es fs
1109 = unstream (Stream.zipWith6 (uncurry f) (Stream.indexed (stream as))
1110 (stream bs)
1111 (stream cs)
1112 (stream ds)
1113 (stream es)
1114 (stream fs))
1115
1116 -- | /O(min(m,n))/ Zip two vectors
1117 zip :: (Vector v a, Vector v b, Vector v (a,b)) => v a -> v b -> v (a, b)
1118 {-# INLINE zip #-}
1119 zip = zipWith (,)
1120
1121 zip3 :: (Vector v a, Vector v b, Vector v c, Vector v (a, b, c))
1122 => v a -> v b -> v c -> v (a, b, c)
1123 {-# INLINE zip3 #-}
1124 zip3 = zipWith3 (,,)
1125
1126 zip4 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v (a, b, c, d))
1127 => v a -> v b -> v c -> v d -> v (a, b, c, d)
1128 {-# INLINE zip4 #-}
1129 zip4 = zipWith4 (,,,)
1130
1131 zip5 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
1132 Vector v (a, b, c, d, e))
1133 => v a -> v b -> v c -> v d -> v e -> v (a, b, c, d, e)
1134 {-# INLINE zip5 #-}
1135 zip5 = zipWith5 (,,,,)
1136
1137 zip6 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
1138 Vector v f, Vector v (a, b, c, d, e, f))
1139 => v a -> v b -> v c -> v d -> v e -> v f -> v (a, b, c, d, e, f)
1140 {-# INLINE zip6 #-}
1141 zip6 = zipWith6 (,,,,,)
1142
1143 -- Monadic zipping
1144 -- ---------------
1145
1146 -- | /O(min(m,n))/ Zip the two vectors with the monadic action and yield a
1147 -- vector of results
1148 zipWithM :: (Monad m, Vector v a, Vector v b, Vector v c)
1149 => (a -> b -> m c) -> v a -> v b -> m (v c)
1150 -- FIXME: specialise for ST and IO?
1151 {-# INLINE zipWithM #-}
1152 zipWithM f as bs = unstreamM $ Stream.zipWithM f (stream as) (stream bs)
1153
1154 -- | /O(min(m,n))/ Zip the two vectors with the monadic action and ignore the
1155 -- results
1156 zipWithM_ :: (Monad m, Vector v a, Vector v b)
1157 => (a -> b -> m c) -> v a -> v b -> m ()
1158 {-# INLINE zipWithM_ #-}
1159 zipWithM_ f as bs = Stream.zipWithM_ f (stream as) (stream bs)
1160
1161 -- Unzipping
1162 -- ---------
1163
1164 -- | /O(min(m,n))/ Unzip a vector of pairs.
1165 unzip :: (Vector v a, Vector v b, Vector v (a,b)) => v (a, b) -> (v a, v b)
1166 {-# INLINE unzip #-}
1167 unzip xs = (map fst xs, map snd xs)
1168
1169 unzip3 :: (Vector v a, Vector v b, Vector v c, Vector v (a, b, c))
1170 => v (a, b, c) -> (v a, v b, v c)
1171 {-# INLINE unzip3 #-}
1172 unzip3 xs = (map (\(a, b, c) -> a) xs,
1173 map (\(a, b, c) -> b) xs,
1174 map (\(a, b, c) -> c) xs)
1175
1176 unzip4 :: (Vector v a, Vector v b, Vector v c, Vector v d,
1177 Vector v (a, b, c, d))
1178 => v (a, b, c, d) -> (v a, v b, v c, v d)
1179 {-# INLINE unzip4 #-}
1180 unzip4 xs = (map (\(a, b, c, d) -> a) xs,
1181 map (\(a, b, c, d) -> b) xs,
1182 map (\(a, b, c, d) -> c) xs,
1183 map (\(a, b, c, d) -> d) xs)
1184
1185 unzip5 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
1186 Vector v (a, b, c, d, e))
1187 => v (a, b, c, d, e) -> (v a, v b, v c, v d, v e)
1188 {-# INLINE unzip5 #-}
1189 unzip5 xs = (map (\(a, b, c, d, e) -> a) xs,
1190 map (\(a, b, c, d, e) -> b) xs,
1191 map (\(a, b, c, d, e) -> c) xs,
1192 map (\(a, b, c, d, e) -> d) xs,
1193 map (\(a, b, c, d, e) -> e) xs)
1194
1195 unzip6 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
1196 Vector v f, Vector v (a, b, c, d, e, f))
1197 => v (a, b, c, d, e, f) -> (v a, v b, v c, v d, v e, v f)
1198 {-# INLINE unzip6 #-}
1199 unzip6 xs = (map (\(a, b, c, d, e, f) -> a) xs,
1200 map (\(a, b, c, d, e, f) -> b) xs,
1201 map (\(a, b, c, d, e, f) -> c) xs,
1202 map (\(a, b, c, d, e, f) -> d) xs,
1203 map (\(a, b, c, d, e, f) -> e) xs,
1204 map (\(a, b, c, d, e, f) -> f) xs)
1205
1206 -- Filtering
1207 -- ---------
1208
1209 -- | /O(n)/ Drop elements that do not satisfy the predicate
1210 filter :: Vector v a => (a -> Bool) -> v a -> v a
1211 {-# INLINE filter #-}
1212 filter f = unstream . inplace (MStream.filter f) . stream
1213
1214 -- | /O(n)/ Drop elements that do not satisfy the predicate which is applied to
1215 -- values and their indices
1216 ifilter :: Vector v a => (Int -> a -> Bool) -> v a -> v a
1217 {-# INLINE ifilter #-}
1218 ifilter f = unstream
1219 . inplace (MStream.map snd . MStream.filter (uncurry f)
1220 . MStream.indexed)
1221 . stream
1222
1223 -- | /O(n)/ Drop elements that do not satisfy the monadic predicate
1224 filterM :: (Monad m, Vector v a) => (a -> m Bool) -> v a -> m (v a)
1225 {-# INLINE filterM #-}
1226 filterM f = unstreamM . Stream.filterM f . stream
1227
1228 -- | /O(n)/ Yield the longest prefix of elements satisfying the predicate
1229 -- without copying.
1230 takeWhile :: Vector v a => (a -> Bool) -> v a -> v a
1231 {-# INLINE takeWhile #-}
1232 takeWhile f = unstream . Stream.takeWhile f . stream
1233
1234 -- | /O(n)/ Drop the longest prefix of elements that satisfy the predicate
1235 -- without copying.
1236 dropWhile :: Vector v a => (a -> Bool) -> v a -> v a
1237 {-# INLINE dropWhile #-}
1238 dropWhile f = unstream . Stream.dropWhile f . stream
1239
1240 -- Parititioning
1241 -- -------------
1242
1243 -- | /O(n)/ Split the vector in two parts, the first one containing those
1244 -- elements that satisfy the predicate and the second one those that don't. The
1245 -- relative order of the elements is preserved at the cost of a sometimes
1246 -- reduced performance compared to 'unstablePartition'.
1247 partition :: Vector v a => (a -> Bool) -> v a -> (v a, v a)
1248 {-# INLINE partition #-}
1249 partition f = partition_stream f . stream
1250
1251 -- FIXME: Make this inplace-fusible (look at how stable_partition is
1252 -- implemented in C++)
1253
1254 partition_stream :: Vector v a => (a -> Bool) -> Stream a -> (v a, v a)
1255 {-# INLINE_STREAM partition_stream #-}
1256 partition_stream f s = s `seq` runST (
1257 do
1258 (mv1,mv2) <- M.partitionStream f s
1259 v1 <- unsafeFreeze mv1
1260 v2 <- unsafeFreeze mv2
1261 return (v1,v2))
1262
1263 -- | /O(n)/ Split the vector in two parts, the first one containing those
1264 -- elements that satisfy the predicate and the second one those that don't.
1265 -- The order of the elements is not preserved but the operation is often
1266 -- faster than 'partition'.
1267 unstablePartition :: Vector v a => (a -> Bool) -> v a -> (v a, v a)
1268 {-# INLINE unstablePartition #-}
1269 unstablePartition f = unstablePartition_stream f . stream
1270
1271 unstablePartition_stream
1272 :: Vector v a => (a -> Bool) -> Stream a -> (v a, v a)
1273 {-# INLINE_STREAM unstablePartition_stream #-}
1274 unstablePartition_stream f s = s `seq` runST (
1275 do
1276 (mv1,mv2) <- M.unstablePartitionStream f s
1277 v1 <- unsafeFreeze mv1
1278 v2 <- unsafeFreeze mv2
1279 return (v1,v2))
1280
1281 unstablePartition_new :: Vector v a => (a -> Bool) -> New v a -> (v a, v a)
1282 {-# INLINE_STREAM unstablePartition_new #-}
1283 unstablePartition_new f (New.New p) = runST (
1284 do
1285 mv <- p
1286 i <- M.unstablePartition f mv
1287 v <- unsafeFreeze mv
1288 return (unsafeTake i v, unsafeDrop i v))
1289
1290 {-# RULES
1291
1292 "unstablePartition" forall f p.
1293 unstablePartition_stream f (stream (new p))
1294 = unstablePartition_new f p
1295
1296 #-}
1297
1298
1299 -- FIXME: make span and break fusible
1300
1301 -- | /O(n)/ Split the vector into the longest prefix of elements that satisfy
1302 -- the predicate and the rest without copying.
1303 span :: Vector v a => (a -> Bool) -> v a -> (v a, v a)
1304 {-# INLINE span #-}
1305 span f = break (not . f)
1306
1307 -- | /O(n)/ Split the vector into the longest prefix of elements that do not
1308 -- satisfy the predicate and the rest without copying.
1309 break :: Vector v a => (a -> Bool) -> v a -> (v a, v a)
1310 {-# INLINE break #-}
1311 break f xs = case findIndex f xs of
1312 Just i -> (unsafeSlice 0 i xs, unsafeSlice i (length xs - i) xs)
1313 Nothing -> (xs, empty)
1314
1315
1316 -- Searching
1317 -- ---------
1318
1319 infix 4 `elem`
1320 -- | /O(n)/ Check if the vector contains an element
1321 elem :: (Vector v a, Eq a) => a -> v a -> Bool
1322 {-# INLINE elem #-}
1323 elem x = Stream.elem x . stream
1324
1325 infix 4 `notElem`
1326 -- | /O(n)/ Check if the vector does not contain an element (inverse of 'elem')
1327 notElem :: (Vector v a, Eq a) => a -> v a -> Bool
1328 {-# INLINE notElem #-}
1329 notElem x = Stream.notElem x . stream
1330
1331 -- | /O(n)/ Yield 'Just' the first element matching the predicate or 'Nothing'
1332 -- if no such element exists.
1333 find :: Vector v a => (a -> Bool) -> v a -> Maybe a
1334 {-# INLINE find #-}
1335 find f = Stream.find f . stream
1336
1337 -- | /O(n)/ Yield 'Just' the index of the first element matching the predicate
1338 -- or 'Nothing' if no such element exists.
1339 findIndex :: Vector v a => (a -> Bool) -> v a -> Maybe Int
1340 {-# INLINE findIndex #-}
1341 findIndex f = Stream.findIndex f . stream
1342
1343 -- | /O(n)/ Yield the indices of elements satisfying the predicate in ascending
1344 -- order.
1345 findIndices :: (Vector v a, Vector v Int) => (a -> Bool) -> v a -> v Int
1346 {-# INLINE findIndices #-}
1347 findIndices f = unstream
1348 . inplace (MStream.map fst . MStream.filter (f . snd)
1349 . MStream.indexed)
1350 . stream
1351
1352 -- | /O(n)/ Yield 'Just' the index of the first occurence of the given element or
1353 -- 'Nothing' if the vector does not contain the element. This is a specialised
1354 -- version of 'findIndex'.
1355 elemIndex :: (Vector v a, Eq a) => a -> v a -> Maybe Int
1356 {-# INLINE elemIndex #-}
1357 elemIndex x = findIndex (x==)
1358
1359 -- | /O(n)/ Yield the indices of all occurences of the given element in
1360 -- ascending order. This is a specialised version of 'findIndices'.
1361 elemIndices :: (Vector v a, Vector v Int, Eq a) => a -> v a -> v Int
1362 {-# INLINE elemIndices #-}
1363 elemIndices x = findIndices (x==)
1364
1365 -- Folding
1366 -- -------
1367
1368 -- | /O(n)/ Left fold
1369 foldl :: Vector v b => (a -> b -> a) -> a -> v b -> a
1370 {-# INLINE foldl #-}
1371 foldl f z = Stream.foldl f z . stream
1372
1373 -- | /O(n)/ Left fold on non-empty vectors
1374 foldl1 :: Vector v a => (a -> a -> a) -> v a -> a
1375 {-# INLINE foldl1 #-}
1376 foldl1 f = Stream.foldl1 f . stream
1377
1378 -- | /O(n)/ Left fold with strict accumulator
1379 foldl' :: Vector v b => (a -> b -> a) -> a -> v b -> a
1380 {-# INLINE foldl' #-}
1381 foldl' f z = Stream.foldl' f z . stream
1382
1383 -- | /O(n)/ Left fold on non-empty vectors with strict accumulator
1384 foldl1' :: Vector v a => (a -> a -> a) -> v a -> a
1385 {-# INLINE foldl1' #-}
1386 foldl1' f = Stream.foldl1' f . stream
1387
1388 -- | /O(n)/ Right fold
1389 foldr :: Vector v a => (a -> b -> b) -> b -> v a -> b
1390 {-# INLINE foldr #-}
1391 foldr f z = Stream.foldr f z . stream
1392
1393 -- | /O(n)/ Right fold on non-empty vectors
1394 foldr1 :: Vector v a => (a -> a -> a) -> v a -> a
1395 {-# INLINE foldr1 #-}
1396 foldr1 f = Stream.foldr1 f . stream
1397
1398 -- | /O(n)/ Right fold with a strict accumulator
1399 foldr' :: Vector v a => (a -> b -> b) -> b -> v a -> b
1400 {-# INLINE foldr' #-}
1401 foldr' f z = Stream.foldl' (flip f) z . streamR
1402
1403 -- | /O(n)/ Right fold on non-empty vectors with strict accumulator
1404 foldr1' :: Vector v a => (a -> a -> a) -> v a -> a
1405 {-# INLINE foldr1' #-}
1406 foldr1' f = Stream.foldl1' (flip f) . streamR
1407
1408 -- | /O(n)/ Left fold (function applied to each element and its index)
1409 ifoldl :: Vector v b => (a -> Int -> b -> a) -> a -> v b -> a
1410 {-# INLINE ifoldl #-}
1411 ifoldl f z = Stream.foldl (uncurry . f) z . Stream.indexed . stream
1412
1413 -- | /O(n)/ Left fold with strict accumulator (function applied to each element
1414 -- and its index)
1415 ifoldl' :: Vector v b => (a -> Int -> b -> a) -> a -> v b -> a
1416 {-# INLINE ifoldl' #-}
1417 ifoldl' f z = Stream.foldl' (uncurry . f) z . Stream.indexed . stream
1418
1419 -- | /O(n)/ Right fold (function applied to each element and its index)
1420 ifoldr :: Vector v a => (Int -> a -> b -> b) -> b -> v a -> b
1421 {-# INLINE ifoldr #-}
1422 ifoldr f z = Stream.foldr (uncurry f) z . Stream.indexed . stream
1423
1424 -- | /O(n)/ Right fold with strict accumulator (function applied to each
1425 -- element and its index)
1426 ifoldr' :: Vector v a => (Int -> a -> b -> b) -> b -> v a -> b
1427 {-# INLINE ifoldr' #-}
1428 ifoldr' f z xs = Stream.foldl' (flip (uncurry f)) z
1429 $ Stream.indexedR (length xs) $ streamR xs
1430
1431 -- Specialised folds
1432 -- -----------------
1433
1434 -- | /O(n)/ Check if all elements satisfy the predicate.
1435 all :: Vector v a => (a -> Bool) -> v a -> Bool
1436 {-# INLINE all #-}
1437 all f = Stream.and . Stream.map f . stream
1438
1439 -- | /O(n)/ Check if any element satisfies the predicate.
1440 any :: Vector v a => (a -> Bool) -> v a -> Bool
1441 {-# INLINE any #-}
1442 any f = Stream.or . Stream.map f . stream
1443
1444 -- | /O(n)/ Check if all elements are 'True'
1445 and :: Vector v Bool => v Bool -> Bool
1446 {-# INLINE and #-}
1447 and = Stream.and . stream
1448
1449 -- | /O(n)/ Check if any element is 'True'
1450 or :: Vector v Bool => v Bool -> Bool
1451 {-# INLINE or #-}
1452 or = Stream.or . stream
1453
1454 -- | /O(n)/ Compute the sum of the elements
1455 sum :: (Vector v a, Num a) => v a -> a
1456 {-# INLINE sum #-}
1457 sum = Stream.foldl' (+) 0 . stream
1458
1459 -- | /O(n)/ Compute the produce of the elements
1460 product :: (Vector v a, Num a) => v a -> a
1461 {-# INLINE product #-}
1462 product = Stream.foldl' (*) 1 . stream
1463
1464 -- | /O(n)/ Yield the maximum element of the vector. The vector may not be
1465 -- empty.
1466 maximum :: (Vector v a, Ord a) => v a -> a
1467 {-# INLINE maximum #-}
1468 maximum = Stream.foldl1' max . stream
1469
1470 -- | /O(n)/ Yield the maximum element of the vector according to the given
1471 -- comparison function. The vector may not be empty.
1472 maximumBy :: Vector v a => (a -> a -> Ordering) -> v a -> a
1473 {-# INLINE maximumBy #-}
1474 maximumBy cmp = Stream.foldl1' maxBy . stream
1475 where
1476 {-# INLINE maxBy #-}
1477 maxBy x y = case cmp x y of
1478 LT -> y
1479 _ -> x
1480
1481 -- | /O(n)/ Yield the minimum element of the vector. The vector may not be
1482 -- empty.
1483 minimum :: (Vector v a, Ord a) => v a -> a
1484 {-# INLINE minimum #-}
1485 minimum = Stream.foldl1' min . stream
1486
1487 -- | /O(n)/ Yield the minimum element of the vector according to the given
1488 -- comparison function. The vector may not be empty.
1489 minimumBy :: Vector v a => (a -> a -> Ordering) -> v a -> a
1490 {-# INLINE minimumBy #-}
1491 minimumBy cmp = Stream.foldl1' minBy . stream
1492 where
1493 {-# INLINE minBy #-}
1494 minBy x y = case cmp x y of
1495 GT -> y
1496 _ -> x
1497
1498 -- | /O(n)/ Yield the index of the maximum element of the vector. The vector
1499 -- may not be empty.
1500 maxIndex :: (Vector v a, Ord a) => v a -> Int
1501 {-# INLINE maxIndex #-}
1502 maxIndex = maxIndexBy compare
1503
1504 -- | /O(n)/ Yield the index of the maximum element of the vector according to
1505 -- the given comparison function. The vector may not be empty.
1506 maxIndexBy :: Vector v a => (a -> a -> Ordering) -> v a -> Int
1507 {-# INLINE maxIndexBy #-}
1508 maxIndexBy cmp = fst . Stream.foldl1' imax . Stream.indexed . stream
1509 where
1510 imax (i,x) (j,y) = i `seq` j `seq`
1511 case cmp x y of
1512 LT -> (j,y)
1513 _ -> (i,x)
1514
1515 -- | /O(n)/ Yield the index of the minimum element of the vector. The vector
1516 -- may not be empty.
1517 minIndex :: (Vector v a, Ord a) => v a -> Int
1518 {-# INLINE minIndex #-}
1519 minIndex = minIndexBy compare
1520
1521 -- | /O(n)/ Yield the index of the minimum element of the vector according to
1522 -- the given comparison function. The vector may not be empty.
1523 minIndexBy :: Vector v a => (a -> a -> Ordering) -> v a -> Int
1524 {-# INLINE minIndexBy #-}
1525 minIndexBy cmp = fst . Stream.foldl1' imin . Stream.indexed . stream
1526 where
1527 imin (i,x) (j,y) = i `seq` j `seq`
1528 case cmp x y of
1529 GT -> (j,y)
1530 _ -> (i,x)
1531
1532 -- Monadic folds
1533 -- -------------
1534
1535 -- | /O(n)/ Monadic fold
1536 foldM :: (Monad m, Vector v b) => (a -> b -> m a) -> a -> v b -> m a
1537 {-# INLINE foldM #-}
1538 foldM m z = Stream.foldM m z . stream
1539
1540 -- | /O(n)/ Monadic fold over non-empty vectors
1541 fold1M :: (Monad m, Vector v a) => (a -> a -> m a) -> v a -> m a
1542 {-# INLINE fold1M #-}
1543 fold1M m = Stream.fold1M m . stream
1544
1545 -- | /O(n)/ Monadic fold with strict accumulator
1546 foldM' :: (Monad m, Vector v b) => (a -> b -> m a) -> a -> v b -> m a
1547 {-# INLINE foldM' #-}
1548 foldM' m z = Stream.foldM' m z . stream
1549
1550 -- | /O(n)/ Monadic fold over non-empty vectors with strict accumulator
1551 fold1M' :: (Monad m, Vector v a) => (a -> a -> m a) -> v a -> m a
1552 {-# INLINE fold1M' #-}
1553 fold1M' m = Stream.fold1M' m . stream
1554
1555 discard :: Monad m => m a -> m ()
1556 {-# INLINE discard #-}
1557 discard m = m >> return ()
1558
1559 -- | /O(n)/ Monadic fold that discards the result
1560 foldM_ :: (Monad m, Vector v b) => (a -> b -> m a) -> a -> v b -> m ()
1561 {-# INLINE foldM_ #-}
1562 foldM_ m z = discard . Stream.foldM m z . stream
1563
1564 -- | /O(n)/ Monadic fold over non-empty vectors that discards the result
1565 fold1M_ :: (Monad m, Vector v a) => (a -> a -> m a) -> v a -> m ()
1566 {-# INLINE fold1M_ #-}
1567 fold1M_ m = discard . Stream.fold1M m . stream
1568
1569 -- | /O(n)/ Monadic fold with strict accumulator that discards the result
1570 foldM'_ :: (Monad m, Vector v b) => (a -> b -> m a) -> a -> v b -> m ()
1571 {-# INLINE foldM'_ #-}
1572 foldM'_ m z = discard . Stream.foldM' m z . stream
1573
1574 -- | /O(n)/ Monad fold over non-empty vectors with strict accumulator
1575 -- that discards the result
1576 fold1M'_ :: (Monad m, Vector v a) => (a -> a -> m a) -> v a -> m ()
1577 {-# INLINE fold1M'_ #-}
1578 fold1M'_ m = discard . Stream.fold1M' m . stream
1579
1580 -- Monadic sequencing
1581 -- ------------------
1582
1583 -- | Evaluate each action and collect the results
1584 sequence :: (Monad m, Vector v a, Vector v (m a)) => v (m a) -> m (v a)
1585 {-# INLINE sequence #-}
1586 sequence = mapM id
1587
1588 -- | Evaluate each action and discard the results
1589 sequence_ :: (Monad m, Vector v (m a)) => v (m a) -> m ()
1590 {-# INLINE sequence_ #-}
1591 sequence_ = mapM_ id
1592
1593 -- Prefix sums (scans)
1594 -- -------------------
1595
1596 -- | /O(n)/ Prescan
1597 --
1598 -- @
1599 -- prescanl f z = 'init' . 'scanl' f z
1600 -- @
1601 --
1602 -- Example: @prescanl (+) 0 \<1,2,3,4\> = \<0,1,3,6\>@
1603 --
1604 prescanl :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1605 {-# INLINE prescanl #-}
1606 prescanl f z = unstream . inplace (MStream.prescanl f z) . stream
1607
1608 -- | /O(n)/ Prescan with strict accumulator
1609 prescanl' :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1610 {-# INLINE prescanl' #-}
1611 prescanl' f z = unstream . inplace (MStream.prescanl' f z) . stream
1612
1613 -- | /O(n)/ Scan
1614 --
1615 -- @
1616 -- postscanl f z = 'tail' . 'scanl' f z
1617 -- @
1618 --
1619 -- Example: @postscanl (+) 0 \<1,2,3,4\> = \<1,3,6,10\>@
1620 --
1621 postscanl :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1622 {-# INLINE postscanl #-}
1623 postscanl f z = unstream . inplace (MStream.postscanl f z) . stream
1624
1625 -- | /O(n)/ Scan with strict accumulator
1626 postscanl' :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1627 {-# INLINE postscanl' #-}
1628 postscanl' f z = unstream . inplace (MStream.postscanl' f z) . stream
1629
1630 -- | /O(n)/ Haskell-style scan
1631 --
1632 -- > scanl f z <x1,...,xn> = <y1,...,y(n+1)>
1633 -- > where y1 = z
1634 -- > yi = f y(i-1) x(i-1)
1635 --
1636 -- Example: @scanl (+) 0 \<1,2,3,4\> = \<0,1,3,6,10\>@
1637 --
1638 scanl :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1639 {-# INLINE scanl #-}
1640 scanl f z = unstream . Stream.scanl f z . stream
1641
1642 -- | /O(n)/ Haskell-style scan with strict accumulator
1643 scanl' :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1644 {-# INLINE scanl' #-}
1645 scanl' f z = unstream . Stream.scanl' f z . stream
1646
1647 -- | /O(n)/ Scan over a non-empty vector
1648 --
1649 -- > scanl f <x1,...,xn> = <y1,...,yn>
1650 -- > where y1 = x1
1651 -- > yi = f y(i-1) xi
1652 --
1653 scanl1 :: Vector v a => (a -> a -> a) -> v a -> v a
1654 {-# INLINE scanl1 #-}
1655 scanl1 f = unstream . inplace (MStream.scanl1 f) . stream
1656
1657 -- | /O(n)/ Scan over a non-empty vector with a strict accumulator
1658 scanl1' :: Vector v a => (a -> a -> a) -> v a -> v a
1659 {-# INLINE scanl1' #-}
1660 scanl1' f = unstream . inplace (MStream.scanl1' f) . stream
1661
1662 -- | /O(n)/ Right-to-left prescan
1663 --
1664 -- @
1665 -- prescanr f z = 'reverse' . 'prescanl' (flip f) z . 'reverse'
1666 -- @
1667 --
1668 prescanr :: (Vector v a, Vector v b) => (a -> b -> b) -> b -> v a -> v b
1669 {-# INLINE prescanr #-}
1670 prescanr f z = unstreamR . inplace (MStream.prescanl (flip f) z) . streamR
1671
1672 -- | /O(n)/ Right-to-left prescan with strict accumulator
1673 prescanr' :: (Vector v a, Vector v b) => (a -> b -> b) -> b -> v a -> v b
1674 {-# INLINE prescanr' #-}
1675 prescanr' f z = unstreamR . inplace (MStream.prescanl' (flip f) z) . streamR
1676
1677 -- | /O(n)/ Right-to-left scan
1678 postscanr :: (Vector v a, Vector v b) => (a -> b -> b) -> b -> v a -> v b
1679 {-# INLINE postscanr #-}
1680 postscanr f z = unstreamR . inplace (MStream.postscanl (flip f) z) . streamR
1681
1682 -- | /O(n)/ Right-to-left scan with strict accumulator
1683 postscanr' :: (Vector v a, Vector v b) => (a -> b -> b) -> b -> v a -> v b
1684 {-# INLINE postscanr' #-}
1685 postscanr' f z = unstreamR . inplace (MStream.postscanl' (flip f) z) . streamR
1686
1687 -- | /O(n)/ Right-to-left Haskell-style scan
1688 scanr :: (Vector v a, Vector v b) => (a -> b -> b) -> b -> v a -> v b
1689 {-# INLINE scanr #-}
1690 scanr f z = unstreamR . Stream.scanl (flip f) z . streamR
1691
1692 -- | /O(n)/ Right-to-left Haskell-style scan with strict accumulator
1693 scanr' :: (Vector v a, Vector v b) => (a -> b -> b) -> b -> v a -> v b
1694 {-# INLINE scanr' #-}
1695 scanr' f z = unstreamR . Stream.scanl' (flip f) z . streamR
1696
1697 -- | /O(n)/ Right-to-left scan over a non-empty vector
1698 scanr1 :: Vector v a => (a -> a -> a) -> v a -> v a
1699 {-# INLINE scanr1 #-}
1700 scanr1 f = unstreamR . inplace (MStream.scanl1 (flip f)) . streamR
1701
1702 -- | /O(n)/ Right-to-left scan over a non-empty vector with a strict
1703 -- accumulator
1704 scanr1' :: Vector v a => (a -> a -> a) -> v a -> v a
1705 {-# INLINE scanr1' #-}
1706 scanr1' f = unstreamR . inplace (MStream.scanl1' (flip f)) . streamR
1707
1708 -- Conversions - Lists
1709 -- ------------------------
1710
1711 -- | /O(n)/ Convert a vector to a list
1712 toList :: Vector v a => v a -> [a]
1713 {-# INLINE toList #-}
1714 toList = Stream.toList . stream
1715
1716 -- | /O(n)/ Convert a list to a vector
1717 fromList :: Vector v a => [a] -> v a
1718 {-# INLINE fromList #-}
1719 fromList = unstream . Stream.fromList
1720
1721 -- | /O(n)/ Convert the first @n@ elements of a list to a vector
1722 --
1723 -- @
1724 -- fromListN n xs = 'fromList' ('take' n xs)
1725 -- @
1726 fromListN :: Vector v a => Int -> [a] -> v a
1727 {-# INLINE fromListN #-}
1728 fromListN n = unstream . Stream.fromListN n
1729
1730 -- Conversions - Immutable vectors
1731 -- -------------------------------
1732
1733 -- | /O(n)/ Convert different vector types
1734 convert :: (Vector v a, Vector w a) => v a -> w a
1735 {-# INLINE convert #-}
1736 convert = unstream . stream
1737
1738 -- Conversions - Mutable vectors
1739 -- -----------------------------
1740
1741 -- | /O(1)/ Unsafe convert a mutable vector to an immutable one without
1742 -- copying. The mutable vector may not be used after this operation.
1743 unsafeFreeze
1744 :: (PrimMonad m, Vector v a) => Mutable v (PrimState m) a -> m (v a)
1745 {-# INLINE unsafeFreeze #-}
1746 unsafeFreeze = basicUnsafeFreeze
1747
1748 -- | /O(n)/ Yield an immutable copy of the mutable vector.
1749 freeze :: (PrimMonad m, Vector v a) => Mutable v (PrimState m) a -> m (v a)
1750 {-# INLINE freeze #-}
1751 freeze mv = unsafeFreeze =<< M.clone mv
1752
1753 -- | /O(1)/ Unsafely convert an immutable vector to a mutable one without
1754 -- copying. The immutable vector may not be used after this operation.
1755 unsafeThaw :: (PrimMonad m, Vector v a) => v a -> m (Mutable v (PrimState m) a)
1756 {-# INLINE_STREAM unsafeThaw #-}
1757 unsafeThaw = basicUnsafeThaw
1758
1759 -- | /O(n)/ Yield a mutable copy of the immutable vector.
1760 thaw :: (PrimMonad m, Vector v a) => v a -> m (Mutable v (PrimState m) a)
1761 {-# INLINE_STREAM thaw #-}
1762 thaw v = do
1763 mv <- M.unsafeNew (length v)
1764 unsafeCopy mv v
1765 return mv
1766
1767 {-# RULES
1768
1769 "unsafeThaw/new [Vector]" forall p.
1770 unsafeThaw (new p) = New.runPrim p
1771
1772 "thaw/new [Vector]" forall p.
1773 thaw (new p) = New.runPrim p
1774
1775 #-}
1776
1777 {-
1778 -- | /O(n)/ Yield a mutable vector containing copies of each vector in the
1779 -- list.
1780 thawMany :: (PrimMonad m, Vector v a) => [v a] -> m (Mutable v (PrimState m) a)
1781 {-# INLINE_STREAM thawMany #-}
1782 -- FIXME: add rule for (stream (new (New.create (thawMany vs))))
1783 -- NOTE: We don't try to consume the list lazily as this wouldn't significantly
1784 -- change the space requirements anyway.
1785 thawMany vs = do
1786 mv <- M.new n
1787 thaw_loop mv vs
1788 return mv
1789 where
1790 n = List.foldl' (\k v -> k + length v) 0 vs
1791
1792 thaw_loop mv [] = mv `seq` return ()
1793 thaw_loop mv (v:vs)
1794 = do
1795 let n = length v
1796 unsafeCopy (M.unsafeTake n mv) v
1797 thaw_loop (M.unsafeDrop n mv) vs
1798 -}
1799
1800 -- | /O(n)/ Copy an immutable vector into a mutable one. The two vectors must
1801 -- have the same length.
1802 copy
1803 :: (PrimMonad m, Vector v a) => Mutable v (PrimState m) a -> v a -> m ()
1804 {-# INLINE copy #-}
1805 copy dst src = BOUNDS_CHECK(check) "copy" "length mismatch"
1806 (M.length dst == length src)
1807 $ unsafeCopy dst src
1808
1809 -- | /O(n)/ Copy an immutable vector into a mutable one. The two vectors must
1810 -- have the same length. This is not checked.
1811 unsafeCopy
1812 :: (PrimMonad m, Vector v a) => Mutable v (PrimState m) a -> v a -> m ()
1813 {-# INLINE unsafeCopy #-}
1814 unsafeCopy dst src = UNSAFE_CHECK(check) "unsafeCopy" "length mismatch"
1815 (M.length dst == length src)
1816 $ (dst `seq` src `seq` basicUnsafeCopy dst src)
1817
1818 -- Conversions to/from Streams
1819 -- ---------------------------
1820
1821 -- | /O(1)/ Convert a vector to a 'Stream'
1822 stream :: Vector v a => v a -> Stream a
1823 {-# INLINE_STREAM stream #-}
1824 stream v = v `seq` n `seq` (Stream.unfoldr get 0 `Stream.sized` Exact n)
1825 where
1826 n = length v
1827
1828 -- NOTE: the False case comes first in Core so making it the recursive one
1829 -- makes the code easier to read
1830 {-# INLINE get #-}
1831 get i | i >= n = Nothing
1832 | otherwise = case basicUnsafeIndexM v i of Box x -> Just (x, i+1)
1833
1834 -- | /O(n)/ Construct a vector from a 'Stream'
1835 unstream :: Vector v a => Stream a -> v a
1836 {-# INLINE unstream #-}
1837 unstream s = new (New.unstream s)
1838
1839 {-# RULES
1840
1841 "stream/unstream [Vector]" forall s.
1842 stream (new (New.unstream s)) = s
1843
1844 "New.unstream/stream [Vector]" forall v.
1845 New.unstream (stream v) = clone v
1846
1847 "clone/new [Vector]" forall p.
1848 clone (new p) = p
1849
1850 "inplace [Vector]"
1851 forall (f :: forall m. Monad m => MStream m a -> MStream m a) m.
1852 New.unstream (inplace f (stream (new m))) = New.transform f m
1853
1854 "uninplace [Vector]"
1855 forall (f :: forall m. Monad m => MStream m a -> MStream m a) m.
1856 stream (new (New.transform f m)) = inplace f (stream (new m))
1857
1858 #-}
1859
1860 -- | /O(1)/ Convert a vector to a 'Stream', proceeding from right to left
1861 streamR :: Vector v a => v a -> Stream a
1862 {-# INLINE_STREAM streamR #-}
1863 streamR v = v `seq` n `seq` (Stream.unfoldr get n `Stream.sized` Exact n)
1864 where
1865 n = length v
1866
1867 {-# INLINE get #-}
1868 get 0 = Nothing
1869 get i = let i' = i-1
1870 in
1871 case basicUnsafeIndexM v i' of Box x -> Just (x, i')
1872
1873 -- | /O(n)/ Construct a vector from a 'Stream', proceeding from right to left
1874 unstreamR :: Vector v a => Stream a -> v a
1875 {-# INLINE unstreamR #-}
1876 unstreamR s = new (New.unstreamR s)
1877
1878 {-# RULES
1879
1880 "streamR/unstreamR [Vector]" forall s.
1881 streamR (new (New.unstreamR s)) = s
1882
1883 "New.unstreamR/streamR/new [Vector]" forall p.
1884 New.unstreamR (streamR (new p)) = p
1885
1886 "New.unstream/streamR/new [Vector]" forall p.
1887 New.unstream (streamR (new p)) = New.modify M.reverse p
1888
1889 "New.unstreamR/stream/new [Vector]" forall p.
1890 New.unstreamR (stream (new p)) = New.modify M.reverse p
1891
1892 "inplace right [Vector]"
1893 forall (f :: forall m. Monad m => MStream m a -> MStream m a) m.
1894 New.unstreamR (inplace f (streamR (new m))) = New.transformR f m
1895
1896 "uninplace right [Vector]"
1897 forall (f :: forall m. Monad m => MStream m a -> MStream m a) m.
1898 streamR (new (New.transformR f m)) = inplace f (streamR (new m))
1899
1900 #-}
1901
1902 unstreamM :: (Monad m, Vector v a) => MStream m a -> m (v a)
1903 {-# INLINE_STREAM unstreamM #-}
1904 unstreamM s = do
1905 xs <- MStream.toList s
1906 return $ unstream $ Stream.unsafeFromList (MStream.size s) xs
1907
1908 unstreamPrimM :: (PrimMonad m, Vector v a) => MStream m a -> m (v a)
1909 {-# INLINE_STREAM unstreamPrimM #-}
1910 unstreamPrimM s = M.munstream s >>= unsafeFreeze
1911
1912 -- FIXME: the next two functions are only necessary for the specialisations
1913 unstreamPrimM_IO :: Vector v a => MStream IO a -> IO (v a)
1914 {-# INLINE unstreamPrimM_IO #-}
1915 unstreamPrimM_IO = unstreamPrimM
1916
1917 unstreamPrimM_ST :: Vector v a => MStream (ST s) a -> ST s (v a)
1918 {-# INLINE unstreamPrimM_ST #-}
1919 unstreamPrimM_ST = unstreamPrimM
1920
1921 {-# RULES
1922
1923 "unstreamM[IO]" unstreamM = unstreamPrimM_IO
1924 "unstreamM[ST]" unstreamM = unstreamPrimM_ST
1925
1926 #-}
1927
1928
1929 -- Recycling support
1930 -- -----------------
1931
1932 -- | Construct a vector from a monadic initialiser.
1933 new :: Vector v a => New v a -> v a
1934 {-# INLINE_STREAM new #-}
1935 new m = m `seq` runST (unsafeFreeze =<< New.run m)
1936
1937 -- | Convert a vector to an initialiser which, when run, produces a copy of
1938 -- the vector.
1939 clone :: Vector v a => v a -> New v a
1940 {-# INLINE_STREAM clone #-}
1941 clone v = v `seq` New.create (
1942 do
1943 mv <- M.new (length v)
1944 unsafeCopy mv v
1945 return mv)
1946
1947 -- Comparisons
1948 -- -----------
1949
1950 -- | /O(n)/ Check if two vectors are equal. All 'Vector' instances are also
1951 -- instances of 'Eq' and it is usually more appropriate to use those. This
1952 -- function is primarily intended for implementing 'Eq' instances for new
1953 -- vector types.
1954 eq :: (Vector v a, Eq a) => v a -> v a -> Bool
1955 {-# INLINE eq #-}
1956 xs `eq` ys = stream xs == stream ys
1957
1958 -- | /O(n)/ Compare two vectors lexicographically. All 'Vector' instances are
1959 -- also instances of 'Ord' and it is usually more appropriate to use those. This
1960 -- function is primarily intended for implementing 'Ord' instances for new
1961 -- vector types.
1962 cmp :: (Vector v a, Ord a) => v a -> v a -> Ordering
1963 {-# INLINE cmp #-}
1964 cmp xs ys = compare (stream xs) (stream ys)
1965
1966 -- Data and Typeable
1967 -- -----------------
1968
1969 -- | Generic definion of 'Data.Data.gfoldl' that views a 'Vector' as a
1970 -- list.
1971 gfoldl :: (Vector v a, Data a)
1972 => (forall d b. Data d => c (d -> b) -> d -> c b)
1973 -> (forall g. g -> c g)
1974 -> v a
1975 -> c (v a)
1976 {-# INLINE gfoldl #-}
1977 gfoldl f z v = z fromList `f` toList v
1978
1979 mkType :: String -> DataType
1980 {-# INLINE mkType #-}
1981 mkType = mkNoRepType
1982
1983 dataCast :: (Vector v a, Data a, Typeable1 v, Typeable1 t)
1984 => (forall d. Data d => c (t d)) -> Maybe (c (v a))
1985 {-# INLINE dataCast #-}
1986 dataCast f = gcast1 f
1987