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