Faster concatMap
[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 -- Used to be fastest
998 {-
999 concatMap f = unstream
1000 . Stream.flatten mk step Unknown
1001 . stream
1002 where
1003 {-# INLINE_INNER step #-}
1004 step (v,i,k)
1005 | i < k = case unsafeIndexM v i of
1006 Box x -> Stream.Yield x (v,i+1,k)
1007 | otherwise = Stream.Done
1008
1009 {-# INLINE mk #-}
1010 mk x = let v = f x
1011 k = length v
1012 in
1013 k `seq` (v,0,k)
1014 -}
1015
1016 -- This seems to be fastest now
1017 concatMap f = unstream
1018 . Stream.fromVectorStream
1019 . Stream.map f
1020 . stream
1021
1022 -- Monadic mapping
1023 -- ---------------
1024
1025 -- | /O(n)/ Apply the monadic action to all elements of the vector, yielding a
1026 -- vector of results
1027 mapM :: (Monad m, Vector v a, Vector v b) => (a -> m b) -> v a -> m (v b)
1028 {-# INLINE mapM #-}
1029 mapM f = unstreamM . Stream.mapM f . stream
1030
1031 -- | /O(n)/ Apply the monadic action to all elements of a vector and ignore the
1032 -- results
1033 mapM_ :: (Monad m, Vector v a) => (a -> m b) -> v a -> m ()
1034 {-# INLINE mapM_ #-}
1035 mapM_ f = Stream.mapM_ f . stream
1036
1037 -- | /O(n)/ Apply the monadic action to all elements of the vector, yielding a
1038 -- vector of results. Equvalent to @flip 'mapM'@.
1039 forM :: (Monad m, Vector v a, Vector v b) => v a -> (a -> m b) -> m (v b)
1040 {-# INLINE forM #-}
1041 forM as f = mapM f as
1042
1043 -- | /O(n)/ Apply the monadic action to all elements of a vector and ignore the
1044 -- results. Equivalent to @flip 'mapM_'@.
1045 forM_ :: (Monad m, Vector v a) => v a -> (a -> m b) -> m ()
1046 {-# INLINE forM_ #-}
1047 forM_ as f = mapM_ f as
1048
1049 -- Zipping
1050 -- -------
1051
1052 -- | /O(min(m,n))/ Zip two vectors with the given function.
1053 zipWith :: (Vector v a, Vector v b, Vector v c)
1054 => (a -> b -> c) -> v a -> v b -> v c
1055 {-# INLINE zipWith #-}
1056 zipWith f xs ys = unstream (Stream.zipWith f (stream xs) (stream ys))
1057
1058 -- | Zip three vectors with the given function.
1059 zipWith3 :: (Vector v a, Vector v b, Vector v c, Vector v d)
1060 => (a -> b -> c -> d) -> v a -> v b -> v c -> v d
1061 {-# INLINE zipWith3 #-}
1062 zipWith3 f as bs cs = unstream (Stream.zipWith3 f (stream as)
1063 (stream bs)
1064 (stream cs))
1065
1066 zipWith4 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e)
1067 => (a -> b -> c -> d -> e) -> v a -> v b -> v c -> v d -> v e
1068 {-# INLINE zipWith4 #-}
1069 zipWith4 f as bs cs ds
1070 = unstream (Stream.zipWith4 f (stream as)
1071 (stream bs)
1072 (stream cs)
1073 (stream ds))
1074
1075 zipWith5 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
1076 Vector v f)
1077 => (a -> b -> c -> d -> e -> f) -> v a -> v b -> v c -> v d -> v e
1078 -> v f
1079 {-# INLINE zipWith5 #-}
1080 zipWith5 f as bs cs ds es
1081 = unstream (Stream.zipWith5 f (stream as)
1082 (stream bs)
1083 (stream cs)
1084 (stream ds)
1085 (stream es))
1086
1087 zipWith6 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
1088 Vector v f, Vector v g)
1089 => (a -> b -> c -> d -> e -> f -> g)
1090 -> v a -> v b -> v c -> v d -> v e -> v f -> v g
1091 {-# INLINE zipWith6 #-}
1092 zipWith6 f as bs cs ds es fs
1093 = unstream (Stream.zipWith6 f (stream as)
1094 (stream bs)
1095 (stream cs)
1096 (stream ds)
1097 (stream es)
1098 (stream fs))
1099
1100 -- | /O(min(m,n))/ Zip two vectors with a function that also takes the
1101 -- elements' indices.
1102 izipWith :: (Vector v a, Vector v b, Vector v c)
1103 => (Int -> a -> b -> c) -> v a -> v b -> v c
1104 {-# INLINE izipWith #-}
1105 izipWith f xs ys = unstream
1106 (Stream.zipWith (uncurry f) (Stream.indexed (stream xs))
1107 (stream ys))
1108
1109 izipWith3 :: (Vector v a, Vector v b, Vector v c, Vector v d)
1110 => (Int -> a -> b -> c -> d) -> v a -> v b -> v c -> v d
1111 {-# INLINE izipWith3 #-}
1112 izipWith3 f as bs cs
1113 = unstream (Stream.zipWith3 (uncurry f) (Stream.indexed (stream as))
1114 (stream bs)
1115 (stream cs))
1116
1117 izipWith4 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e)
1118 => (Int -> a -> b -> c -> d -> e) -> v a -> v b -> v c -> v d -> v e
1119 {-# INLINE izipWith4 #-}
1120 izipWith4 f as bs cs ds
1121 = unstream (Stream.zipWith4 (uncurry f) (Stream.indexed (stream as))
1122 (stream bs)
1123 (stream cs)
1124 (stream ds))
1125
1126 izipWith5 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
1127 Vector v f)
1128 => (Int -> a -> b -> c -> d -> e -> f) -> v a -> v b -> v c -> v d
1129 -> v e -> v f
1130 {-# INLINE izipWith5 #-}
1131 izipWith5 f as bs cs ds es
1132 = unstream (Stream.zipWith5 (uncurry f) (Stream.indexed (stream as))
1133 (stream bs)
1134 (stream cs)
1135 (stream ds)
1136 (stream es))
1137
1138 izipWith6 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
1139 Vector v f, Vector v g)
1140 => (Int -> a -> b -> c -> d -> e -> f -> g)
1141 -> v a -> v b -> v c -> v d -> v e -> v f -> v g
1142 {-# INLINE izipWith6 #-}
1143 izipWith6 f as bs cs ds es fs
1144 = unstream (Stream.zipWith6 (uncurry f) (Stream.indexed (stream as))
1145 (stream bs)
1146 (stream cs)
1147 (stream ds)
1148 (stream es)
1149 (stream fs))
1150
1151 -- | /O(min(m,n))/ Zip two vectors
1152 zip :: (Vector v a, Vector v b, Vector v (a,b)) => v a -> v b -> v (a, b)
1153 {-# INLINE zip #-}
1154 zip = zipWith (,)
1155
1156 zip3 :: (Vector v a, Vector v b, Vector v c, Vector v (a, b, c))
1157 => v a -> v b -> v c -> v (a, b, c)
1158 {-# INLINE zip3 #-}
1159 zip3 = zipWith3 (,,)
1160
1161 zip4 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v (a, b, c, d))
1162 => v a -> v b -> v c -> v d -> v (a, b, c, d)
1163 {-# INLINE zip4 #-}
1164 zip4 = zipWith4 (,,,)
1165
1166 zip5 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
1167 Vector v (a, b, c, d, e))
1168 => v a -> v b -> v c -> v d -> v e -> v (a, b, c, d, e)
1169 {-# INLINE zip5 #-}
1170 zip5 = zipWith5 (,,,,)
1171
1172 zip6 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
1173 Vector v f, Vector v (a, b, c, d, e, f))
1174 => v a -> v b -> v c -> v d -> v e -> v f -> v (a, b, c, d, e, f)
1175 {-# INLINE zip6 #-}
1176 zip6 = zipWith6 (,,,,,)
1177
1178 -- Monadic zipping
1179 -- ---------------
1180
1181 -- | /O(min(m,n))/ Zip the two vectors with the monadic action and yield a
1182 -- vector of results
1183 zipWithM :: (Monad m, Vector v a, Vector v b, Vector v c)
1184 => (a -> b -> m c) -> v a -> v b -> m (v c)
1185 -- FIXME: specialise for ST and IO?
1186 {-# INLINE zipWithM #-}
1187 zipWithM f as bs = unstreamM $ Stream.zipWithM f (stream as) (stream bs)
1188
1189 -- | /O(min(m,n))/ Zip the two vectors with the monadic action and ignore the
1190 -- results
1191 zipWithM_ :: (Monad m, Vector v a, Vector v b)
1192 => (a -> b -> m c) -> v a -> v b -> m ()
1193 {-# INLINE zipWithM_ #-}
1194 zipWithM_ f as bs = Stream.zipWithM_ f (stream as) (stream bs)
1195
1196 -- Unzipping
1197 -- ---------
1198
1199 -- | /O(min(m,n))/ Unzip a vector of pairs.
1200 unzip :: (Vector v a, Vector v b, Vector v (a,b)) => v (a, b) -> (v a, v b)
1201 {-# INLINE unzip #-}
1202 unzip xs = (map fst xs, map snd xs)
1203
1204 unzip3 :: (Vector v a, Vector v b, Vector v c, Vector v (a, b, c))
1205 => v (a, b, c) -> (v a, v b, v c)
1206 {-# INLINE unzip3 #-}
1207 unzip3 xs = (map (\(a, b, c) -> a) xs,
1208 map (\(a, b, c) -> b) xs,
1209 map (\(a, b, c) -> c) xs)
1210
1211 unzip4 :: (Vector v a, Vector v b, Vector v c, Vector v d,
1212 Vector v (a, b, c, d))
1213 => v (a, b, c, d) -> (v a, v b, v c, v d)
1214 {-# INLINE unzip4 #-}
1215 unzip4 xs = (map (\(a, b, c, d) -> a) xs,
1216 map (\(a, b, c, d) -> b) xs,
1217 map (\(a, b, c, d) -> c) xs,
1218 map (\(a, b, c, d) -> d) xs)
1219
1220 unzip5 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
1221 Vector v (a, b, c, d, e))
1222 => v (a, b, c, d, e) -> (v a, v b, v c, v d, v e)
1223 {-# INLINE unzip5 #-}
1224 unzip5 xs = (map (\(a, b, c, d, e) -> a) xs,
1225 map (\(a, b, c, d, e) -> b) xs,
1226 map (\(a, b, c, d, e) -> c) xs,
1227 map (\(a, b, c, d, e) -> d) xs,
1228 map (\(a, b, c, d, e) -> e) xs)
1229
1230 unzip6 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
1231 Vector v f, Vector v (a, b, c, d, e, f))
1232 => v (a, b, c, d, e, f) -> (v a, v b, v c, v d, v e, v f)
1233 {-# INLINE unzip6 #-}
1234 unzip6 xs = (map (\(a, b, c, d, e, f) -> a) xs,
1235 map (\(a, b, c, d, e, f) -> b) xs,
1236 map (\(a, b, c, d, e, f) -> c) xs,
1237 map (\(a, b, c, d, e, f) -> d) xs,
1238 map (\(a, b, c, d, e, f) -> e) xs,
1239 map (\(a, b, c, d, e, f) -> f) xs)
1240
1241 -- Filtering
1242 -- ---------
1243
1244 -- | /O(n)/ Drop elements that do not satisfy the predicate
1245 filter :: Vector v a => (a -> Bool) -> v a -> v a
1246 {-# INLINE filter #-}
1247 filter f = unstream . inplace (MStream.filter f) . stream
1248
1249 -- | /O(n)/ Drop elements that do not satisfy the predicate which is applied to
1250 -- values and their indices
1251 ifilter :: Vector v a => (Int -> a -> Bool) -> v a -> v a
1252 {-# INLINE ifilter #-}
1253 ifilter f = unstream
1254 . inplace (MStream.map snd . MStream.filter (uncurry f)
1255 . MStream.indexed)
1256 . stream
1257
1258 -- | /O(n)/ Drop elements that do not satisfy the monadic predicate
1259 filterM :: (Monad m, Vector v a) => (a -> m Bool) -> v a -> m (v a)
1260 {-# INLINE filterM #-}
1261 filterM f = unstreamM . Stream.filterM f . stream
1262
1263 -- | /O(n)/ Yield the longest prefix of elements satisfying the predicate
1264 -- without copying.
1265 takeWhile :: Vector v a => (a -> Bool) -> v a -> v a
1266 {-# INLINE takeWhile #-}
1267 takeWhile f = unstream . Stream.takeWhile f . stream
1268
1269 -- | /O(n)/ Drop the longest prefix of elements that satisfy the predicate
1270 -- without copying.
1271 dropWhile :: Vector v a => (a -> Bool) -> v a -> v a
1272 {-# INLINE dropWhile #-}
1273 dropWhile f = unstream . Stream.dropWhile f . stream
1274
1275 -- Parititioning
1276 -- -------------
1277
1278 -- | /O(n)/ Split the vector in two parts, the first one containing those
1279 -- elements that satisfy the predicate and the second one those that don't. The
1280 -- relative order of the elements is preserved at the cost of a sometimes
1281 -- reduced performance compared to 'unstablePartition'.
1282 partition :: Vector v a => (a -> Bool) -> v a -> (v a, v a)
1283 {-# INLINE partition #-}
1284 partition f = partition_stream f . stream
1285
1286 -- FIXME: Make this inplace-fusible (look at how stable_partition is
1287 -- implemented in C++)
1288
1289 partition_stream :: Vector v a => (a -> Bool) -> Stream u a -> (v a, v a)
1290 {-# INLINE_STREAM partition_stream #-}
1291 partition_stream f s = s `seq` runST (
1292 do
1293 (mv1,mv2) <- M.partitionStream f s
1294 v1 <- unsafeFreeze mv1
1295 v2 <- unsafeFreeze mv2
1296 return (v1,v2))
1297
1298 -- | /O(n)/ Split the vector in two parts, the first one containing those
1299 -- elements that satisfy the predicate and the second one those that don't.
1300 -- The order of the elements is not preserved but the operation is often
1301 -- faster than 'partition'.
1302 unstablePartition :: Vector v a => (a -> Bool) -> v a -> (v a, v a)
1303 {-# INLINE unstablePartition #-}
1304 unstablePartition f = unstablePartition_stream f . stream
1305
1306 unstablePartition_stream
1307 :: Vector v a => (a -> Bool) -> Stream u a -> (v a, v a)
1308 {-# INLINE_STREAM unstablePartition_stream #-}
1309 unstablePartition_stream f s = s `seq` runST (
1310 do
1311 (mv1,mv2) <- M.unstablePartitionStream f s
1312 v1 <- unsafeFreeze mv1
1313 v2 <- unsafeFreeze mv2
1314 return (v1,v2))
1315
1316 unstablePartition_new :: Vector v a => (a -> Bool) -> New v a -> (v a, v a)
1317 {-# INLINE_STREAM unstablePartition_new #-}
1318 unstablePartition_new f (New.New p) = runST (
1319 do
1320 mv <- p
1321 i <- M.unstablePartition f mv
1322 v <- unsafeFreeze mv
1323 return (unsafeTake i v, unsafeDrop i v))
1324
1325 {-# RULES
1326
1327 "unstablePartition" forall f p.
1328 unstablePartition_stream f (stream (new p))
1329 = unstablePartition_new f p
1330
1331 #-}
1332
1333
1334 -- FIXME: make span and break fusible
1335
1336 -- | /O(n)/ Split the vector into the longest prefix of elements that satisfy
1337 -- the predicate and the rest without copying.
1338 span :: Vector v a => (a -> Bool) -> v a -> (v a, v a)
1339 {-# INLINE span #-}
1340 span f = break (not . f)
1341
1342 -- | /O(n)/ Split the vector into the longest prefix of elements that do not
1343 -- satisfy the predicate and the rest without copying.
1344 break :: Vector v a => (a -> Bool) -> v a -> (v a, v a)
1345 {-# INLINE break #-}
1346 break f xs = case findIndex f xs of
1347 Just i -> (unsafeSlice 0 i xs, unsafeSlice i (length xs - i) xs)
1348 Nothing -> (xs, empty)
1349
1350
1351 -- Searching
1352 -- ---------
1353
1354 infix 4 `elem`
1355 -- | /O(n)/ Check if the vector contains an element
1356 elem :: (Vector v a, Eq a) => a -> v a -> Bool
1357 {-# INLINE elem #-}
1358 elem x = Stream.elem x . stream
1359
1360 infix 4 `notElem`
1361 -- | /O(n)/ Check if the vector does not contain an element (inverse of 'elem')
1362 notElem :: (Vector v a, Eq a) => a -> v a -> Bool
1363 {-# INLINE notElem #-}
1364 notElem x = Stream.notElem x . stream
1365
1366 -- | /O(n)/ Yield 'Just' the first element matching the predicate or 'Nothing'
1367 -- if no such element exists.
1368 find :: Vector v a => (a -> Bool) -> v a -> Maybe a
1369 {-# INLINE find #-}
1370 find f = Stream.find f . stream
1371
1372 -- | /O(n)/ Yield 'Just' the index of the first element matching the predicate
1373 -- or 'Nothing' if no such element exists.
1374 findIndex :: Vector v a => (a -> Bool) -> v a -> Maybe Int
1375 {-# INLINE findIndex #-}
1376 findIndex f = Stream.findIndex f . stream
1377
1378 -- | /O(n)/ Yield the indices of elements satisfying the predicate in ascending
1379 -- order.
1380 findIndices :: (Vector v a, Vector v Int) => (a -> Bool) -> v a -> v Int
1381 {-# INLINE findIndices #-}
1382 findIndices f = unstream
1383 . inplace (MStream.map fst . MStream.filter (f . snd)
1384 . MStream.indexed)
1385 . stream
1386
1387 -- | /O(n)/ Yield 'Just' the index of the first occurence of the given element or
1388 -- 'Nothing' if the vector does not contain the element. This is a specialised
1389 -- version of 'findIndex'.
1390 elemIndex :: (Vector v a, Eq a) => a -> v a -> Maybe Int
1391 {-# INLINE elemIndex #-}
1392 elemIndex x = findIndex (x==)
1393
1394 -- | /O(n)/ Yield the indices of all occurences of the given element in
1395 -- ascending order. This is a specialised version of 'findIndices'.
1396 elemIndices :: (Vector v a, Vector v Int, Eq a) => a -> v a -> v Int
1397 {-# INLINE elemIndices #-}
1398 elemIndices x = findIndices (x==)
1399
1400 -- Folding
1401 -- -------
1402
1403 -- | /O(n)/ Left fold
1404 foldl :: Vector v b => (a -> b -> a) -> a -> v b -> a
1405 {-# INLINE foldl #-}
1406 foldl f z = Stream.foldl f z . stream
1407
1408 -- | /O(n)/ Left fold on non-empty vectors
1409 foldl1 :: Vector v a => (a -> a -> a) -> v a -> a
1410 {-# INLINE foldl1 #-}
1411 foldl1 f = Stream.foldl1 f . stream
1412
1413 -- | /O(n)/ Left fold with strict accumulator
1414 foldl' :: Vector v b => (a -> b -> a) -> a -> v b -> a
1415 {-# INLINE foldl' #-}
1416 foldl' f z = Stream.foldl' f z . stream
1417
1418 -- | /O(n)/ Left fold on non-empty vectors with strict accumulator
1419 foldl1' :: Vector v a => (a -> a -> a) -> v a -> a
1420 {-# INLINE foldl1' #-}
1421 foldl1' f = Stream.foldl1' f . stream
1422
1423 -- | /O(n)/ Right fold
1424 foldr :: Vector v a => (a -> b -> b) -> b -> v a -> b
1425 {-# INLINE foldr #-}
1426 foldr f z = Stream.foldr f z . stream
1427
1428 -- | /O(n)/ Right fold on non-empty vectors
1429 foldr1 :: Vector v a => (a -> a -> a) -> v a -> a
1430 {-# INLINE foldr1 #-}
1431 foldr1 f = Stream.foldr1 f . stream
1432
1433 -- | /O(n)/ Right fold with a strict accumulator
1434 foldr' :: Vector v a => (a -> b -> b) -> b -> v a -> b
1435 {-# INLINE foldr' #-}
1436 foldr' f z = Stream.foldl' (flip f) z . streamR
1437
1438 -- | /O(n)/ Right fold on non-empty vectors with strict accumulator
1439 foldr1' :: Vector v a => (a -> a -> a) -> v a -> a
1440 {-# INLINE foldr1' #-}
1441 foldr1' f = Stream.foldl1' (flip f) . streamR
1442
1443 -- | /O(n)/ Left fold (function applied to each element and its index)
1444 ifoldl :: Vector v b => (a -> Int -> b -> a) -> a -> v b -> a
1445 {-# INLINE ifoldl #-}
1446 ifoldl f z = Stream.foldl (uncurry . f) z . Stream.indexed . stream
1447
1448 -- | /O(n)/ Left fold with strict accumulator (function applied to each element
1449 -- and its index)
1450 ifoldl' :: Vector v b => (a -> Int -> b -> a) -> a -> v b -> a
1451 {-# INLINE ifoldl' #-}
1452 ifoldl' f z = Stream.foldl' (uncurry . f) z . Stream.indexed . stream
1453
1454 -- | /O(n)/ Right fold (function applied to each element and its index)
1455 ifoldr :: Vector v a => (Int -> a -> b -> b) -> b -> v a -> b
1456 {-# INLINE ifoldr #-}
1457 ifoldr f z = Stream.foldr (uncurry f) z . Stream.indexed . stream
1458
1459 -- | /O(n)/ Right fold with strict accumulator (function applied to each
1460 -- element and its index)
1461 ifoldr' :: Vector v a => (Int -> a -> b -> b) -> b -> v a -> b
1462 {-# INLINE ifoldr' #-}
1463 ifoldr' f z xs = Stream.foldl' (flip (uncurry f)) z
1464 $ Stream.indexedR (length xs) $ streamR xs
1465
1466 -- Specialised folds
1467 -- -----------------
1468
1469 -- | /O(n)/ Check if all elements satisfy the predicate.
1470 all :: Vector v a => (a -> Bool) -> v a -> Bool
1471 {-# INLINE all #-}
1472 all f = Stream.and . Stream.map f . stream
1473
1474 -- | /O(n)/ Check if any element satisfies the predicate.
1475 any :: Vector v a => (a -> Bool) -> v a -> Bool
1476 {-# INLINE any #-}
1477 any f = Stream.or . Stream.map f . stream
1478
1479 -- | /O(n)/ Check if all elements are 'True'
1480 and :: Vector v Bool => v Bool -> Bool
1481 {-# INLINE and #-}
1482 and = Stream.and . stream
1483
1484 -- | /O(n)/ Check if any element is 'True'
1485 or :: Vector v Bool => v Bool -> Bool
1486 {-# INLINE or #-}
1487 or = Stream.or . stream
1488
1489 -- | /O(n)/ Compute the sum of the elements
1490 sum :: (Vector v a, Num a) => v a -> a
1491 {-# INLINE sum #-}
1492 sum = Stream.foldl' (+) 0 . stream
1493
1494 -- | /O(n)/ Compute the produce of the elements
1495 product :: (Vector v a, Num a) => v a -> a
1496 {-# INLINE product #-}
1497 product = Stream.foldl' (*) 1 . stream
1498
1499 -- | /O(n)/ Yield the maximum element of the vector. The vector may not be
1500 -- empty.
1501 maximum :: (Vector v a, Ord a) => v a -> a
1502 {-# INLINE maximum #-}
1503 maximum = Stream.foldl1' max . stream
1504
1505 -- | /O(n)/ Yield the maximum element of the vector according to the given
1506 -- comparison function. The vector may not be empty.
1507 maximumBy :: Vector v a => (a -> a -> Ordering) -> v a -> a
1508 {-# INLINE maximumBy #-}
1509 maximumBy cmp = Stream.foldl1' maxBy . stream
1510 where
1511 {-# INLINE maxBy #-}
1512 maxBy x y = case cmp x y of
1513 LT -> y
1514 _ -> x
1515
1516 -- | /O(n)/ Yield the minimum element of the vector. The vector may not be
1517 -- empty.
1518 minimum :: (Vector v a, Ord a) => v a -> a
1519 {-# INLINE minimum #-}
1520 minimum = Stream.foldl1' min . stream
1521
1522 -- | /O(n)/ Yield the minimum element of the vector according to the given
1523 -- comparison function. The vector may not be empty.
1524 minimumBy :: Vector v a => (a -> a -> Ordering) -> v a -> a
1525 {-# INLINE minimumBy #-}
1526 minimumBy cmp = Stream.foldl1' minBy . stream
1527 where
1528 {-# INLINE minBy #-}
1529 minBy x y = case cmp x y of
1530 GT -> y
1531 _ -> x
1532
1533 -- | /O(n)/ Yield the index of the maximum element of the vector. The vector
1534 -- may not be empty.
1535 maxIndex :: (Vector v a, Ord a) => v a -> Int
1536 {-# INLINE maxIndex #-}
1537 maxIndex = maxIndexBy compare
1538
1539 -- | /O(n)/ Yield the index of the maximum element of the vector according to
1540 -- the given comparison function. The vector may not be empty.
1541 maxIndexBy :: Vector v a => (a -> a -> Ordering) -> v a -> Int
1542 {-# INLINE maxIndexBy #-}
1543 maxIndexBy cmp = fst . Stream.foldl1' imax . Stream.indexed . stream
1544 where
1545 imax (i,x) (j,y) = i `seq` j `seq`
1546 case cmp x y of
1547 LT -> (j,y)
1548 _ -> (i,x)
1549
1550 -- | /O(n)/ Yield the index of the minimum element of the vector. The vector
1551 -- may not be empty.
1552 minIndex :: (Vector v a, Ord a) => v a -> Int
1553 {-# INLINE minIndex #-}
1554 minIndex = minIndexBy compare
1555
1556 -- | /O(n)/ Yield the index of the minimum element of the vector according to
1557 -- the given comparison function. The vector may not be empty.
1558 minIndexBy :: Vector v a => (a -> a -> Ordering) -> v a -> Int
1559 {-# INLINE minIndexBy #-}
1560 minIndexBy cmp = fst . Stream.foldl1' imin . Stream.indexed . stream
1561 where
1562 imin (i,x) (j,y) = i `seq` j `seq`
1563 case cmp x y of
1564 GT -> (j,y)
1565 _ -> (i,x)
1566
1567 -- Monadic folds
1568 -- -------------
1569
1570 -- | /O(n)/ Monadic fold
1571 foldM :: (Monad m, Vector v b) => (a -> b -> m a) -> a -> v b -> m a
1572 {-# INLINE foldM #-}
1573 foldM m z = Stream.foldM m z . stream
1574
1575 -- | /O(n)/ Monadic fold over non-empty vectors
1576 fold1M :: (Monad m, Vector v a) => (a -> a -> m a) -> v a -> m a
1577 {-# INLINE fold1M #-}
1578 fold1M m = Stream.fold1M m . stream
1579
1580 -- | /O(n)/ Monadic fold with strict accumulator
1581 foldM' :: (Monad m, Vector v b) => (a -> b -> m a) -> a -> v b -> m a
1582 {-# INLINE foldM' #-}
1583 foldM' m z = Stream.foldM' m z . stream
1584
1585 -- | /O(n)/ Monadic fold over non-empty vectors with strict accumulator
1586 fold1M' :: (Monad m, Vector v a) => (a -> a -> m a) -> v a -> m a
1587 {-# INLINE fold1M' #-}
1588 fold1M' m = Stream.fold1M' m . stream
1589
1590 discard :: Monad m => m a -> m ()
1591 {-# INLINE discard #-}
1592 discard m = m >> return ()
1593
1594 -- | /O(n)/ Monadic fold that discards the result
1595 foldM_ :: (Monad m, Vector v b) => (a -> b -> m a) -> a -> v b -> m ()
1596 {-# INLINE foldM_ #-}
1597 foldM_ m z = discard . Stream.foldM m z . stream
1598
1599 -- | /O(n)/ Monadic fold over non-empty vectors that discards the result
1600 fold1M_ :: (Monad m, Vector v a) => (a -> a -> m a) -> v a -> m ()
1601 {-# INLINE fold1M_ #-}
1602 fold1M_ m = discard . Stream.fold1M m . stream
1603
1604 -- | /O(n)/ Monadic fold with strict accumulator that discards the result
1605 foldM'_ :: (Monad m, Vector v b) => (a -> b -> m a) -> a -> v b -> m ()
1606 {-# INLINE foldM'_ #-}
1607 foldM'_ m z = discard . Stream.foldM' m z . stream
1608
1609 -- | /O(n)/ Monad fold over non-empty vectors with strict accumulator
1610 -- that discards the result
1611 fold1M'_ :: (Monad m, Vector v a) => (a -> a -> m a) -> v a -> m ()
1612 {-# INLINE fold1M'_ #-}
1613 fold1M'_ m = discard . Stream.fold1M' m . stream
1614
1615 -- Monadic sequencing
1616 -- ------------------
1617
1618 -- | Evaluate each action and collect the results
1619 sequence :: (Monad m, Vector v a, Vector v (m a)) => v (m a) -> m (v a)
1620 {-# INLINE sequence #-}
1621 sequence = mapM id
1622
1623 -- | Evaluate each action and discard the results
1624 sequence_ :: (Monad m, Vector v (m a)) => v (m a) -> m ()
1625 {-# INLINE sequence_ #-}
1626 sequence_ = mapM_ id
1627
1628 -- Prefix sums (scans)
1629 -- -------------------
1630
1631 -- | /O(n)/ Prescan
1632 --
1633 -- @
1634 -- prescanl f z = 'init' . 'scanl' f z
1635 -- @
1636 --
1637 -- Example: @prescanl (+) 0 \<1,2,3,4\> = \<0,1,3,6\>@
1638 --
1639 prescanl :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1640 {-# INLINE prescanl #-}
1641 prescanl f z = unstream . inplace (MStream.prescanl f z) . stream
1642
1643 -- | /O(n)/ Prescan with strict accumulator
1644 prescanl' :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1645 {-# INLINE prescanl' #-}
1646 prescanl' f z = unstream . inplace (MStream.prescanl' f z) . stream
1647
1648 -- | /O(n)/ Scan
1649 --
1650 -- @
1651 -- postscanl f z = 'tail' . 'scanl' f z
1652 -- @
1653 --
1654 -- Example: @postscanl (+) 0 \<1,2,3,4\> = \<1,3,6,10\>@
1655 --
1656 postscanl :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1657 {-# INLINE postscanl #-}
1658 postscanl f z = unstream . inplace (MStream.postscanl f z) . stream
1659
1660 -- | /O(n)/ Scan with strict accumulator
1661 postscanl' :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1662 {-# INLINE postscanl' #-}
1663 postscanl' f z = unstream . inplace (MStream.postscanl' f z) . stream
1664
1665 -- | /O(n)/ Haskell-style scan
1666 --
1667 -- > scanl f z <x1,...,xn> = <y1,...,y(n+1)>
1668 -- > where y1 = z
1669 -- > yi = f y(i-1) x(i-1)
1670 --
1671 -- Example: @scanl (+) 0 \<1,2,3,4\> = \<0,1,3,6,10\>@
1672 --
1673 scanl :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1674 {-# INLINE scanl #-}
1675 scanl f z = unstream . Stream.scanl f z . stream
1676
1677 -- | /O(n)/ Haskell-style scan with strict accumulator
1678 scanl' :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1679 {-# INLINE scanl' #-}
1680 scanl' f z = unstream . Stream.scanl' f z . stream
1681
1682 -- | /O(n)/ Scan over a non-empty vector
1683 --
1684 -- > scanl f <x1,...,xn> = <y1,...,yn>
1685 -- > where y1 = x1
1686 -- > yi = f y(i-1) xi
1687 --
1688 scanl1 :: Vector v a => (a -> a -> a) -> v a -> v a
1689 {-# INLINE scanl1 #-}
1690 scanl1 f = unstream . inplace (MStream.scanl1 f) . stream
1691
1692 -- | /O(n)/ Scan over a non-empty vector with a strict accumulator
1693 scanl1' :: Vector v a => (a -> a -> a) -> v a -> v a
1694 {-# INLINE scanl1' #-}
1695 scanl1' f = unstream . inplace (MStream.scanl1' f) . stream
1696
1697 -- | /O(n)/ Right-to-left prescan
1698 --
1699 -- @
1700 -- prescanr f z = 'reverse' . 'prescanl' (flip f) z . 'reverse'
1701 -- @
1702 --
1703 prescanr :: (Vector v a, Vector v b) => (a -> b -> b) -> b -> v a -> v b
1704 {-# INLINE prescanr #-}
1705 prescanr f z = unstreamR . inplace (MStream.prescanl (flip f) z) . streamR
1706
1707 -- | /O(n)/ Right-to-left prescan with strict accumulator
1708 prescanr' :: (Vector v a, Vector v b) => (a -> b -> b) -> b -> v a -> v b
1709 {-# INLINE prescanr' #-}
1710 prescanr' f z = unstreamR . inplace (MStream.prescanl' (flip f) z) . streamR
1711
1712 -- | /O(n)/ Right-to-left scan
1713 postscanr :: (Vector v a, Vector v b) => (a -> b -> b) -> b -> v a -> v b
1714 {-# INLINE postscanr #-}
1715 postscanr f z = unstreamR . inplace (MStream.postscanl (flip f) z) . streamR
1716
1717 -- | /O(n)/ Right-to-left scan with strict accumulator
1718 postscanr' :: (Vector v a, Vector v b) => (a -> b -> b) -> b -> v a -> v b
1719 {-# INLINE postscanr' #-}
1720 postscanr' f z = unstreamR . inplace (MStream.postscanl' (flip f) z) . streamR
1721
1722 -- | /O(n)/ Right-to-left Haskell-style scan
1723 scanr :: (Vector v a, Vector v b) => (a -> b -> b) -> b -> v a -> v b
1724 {-# INLINE scanr #-}
1725 scanr f z = unstreamR . Stream.scanl (flip f) z . streamR
1726
1727 -- | /O(n)/ Right-to-left Haskell-style scan with strict accumulator
1728 scanr' :: (Vector v a, Vector v b) => (a -> b -> b) -> b -> v a -> v b
1729 {-# INLINE scanr' #-}
1730 scanr' f z = unstreamR . Stream.scanl' (flip f) z . streamR
1731
1732 -- | /O(n)/ Right-to-left scan over a non-empty vector
1733 scanr1 :: Vector v a => (a -> a -> a) -> v a -> v a
1734 {-# INLINE scanr1 #-}
1735 scanr1 f = unstreamR . inplace (MStream.scanl1 (flip f)) . streamR
1736
1737 -- | /O(n)/ Right-to-left scan over a non-empty vector with a strict
1738 -- accumulator
1739 scanr1' :: Vector v a => (a -> a -> a) -> v a -> v a
1740 {-# INLINE scanr1' #-}
1741 scanr1' f = unstreamR . inplace (MStream.scanl1' (flip f)) . streamR
1742
1743 -- Conversions - Lists
1744 -- ------------------------
1745
1746 -- | /O(n)/ Convert a vector to a list
1747 toList :: Vector v a => v a -> [a]
1748 {-# INLINE toList #-}
1749 toList = Stream.toList . stream
1750
1751 -- | /O(n)/ Convert a list to a vector
1752 fromList :: Vector v a => [a] -> v a
1753 {-# INLINE fromList #-}
1754 fromList = unstream . Stream.fromList
1755
1756 -- | /O(n)/ Convert the first @n@ elements of a list to a vector
1757 --
1758 -- @
1759 -- fromListN n xs = 'fromList' ('take' n xs)
1760 -- @
1761 fromListN :: Vector v a => Int -> [a] -> v a
1762 {-# INLINE fromListN #-}
1763 fromListN n = unstream . Stream.fromListN n
1764
1765 -- Conversions - Immutable vectors
1766 -- -------------------------------
1767
1768 -- | /O(n)/ Convert different vector types
1769 convert :: (Vector v a, Vector w a) => v a -> w a
1770 {-# INLINE convert #-}
1771 convert = unstream . Stream.reVector . stream
1772
1773 -- Conversions - Mutable vectors
1774 -- -----------------------------
1775
1776 -- | /O(1)/ Unsafe convert a mutable vector to an immutable one without
1777 -- copying. The mutable vector may not be used after this operation.
1778 unsafeFreeze
1779 :: (PrimMonad m, Vector v a) => Mutable v (PrimState m) a -> m (v a)
1780 {-# INLINE unsafeFreeze #-}
1781 unsafeFreeze = basicUnsafeFreeze
1782
1783 -- | /O(n)/ Yield an immutable copy of the mutable vector.
1784 freeze :: (PrimMonad m, Vector v a) => Mutable v (PrimState m) a -> m (v a)
1785 {-# INLINE freeze #-}
1786 freeze mv = unsafeFreeze =<< M.clone mv
1787
1788 -- | /O(1)/ Unsafely convert an immutable vector to a mutable one without
1789 -- copying. The immutable vector may not be used after this operation.
1790 unsafeThaw :: (PrimMonad m, Vector v a) => v a -> m (Mutable v (PrimState m) a)
1791 {-# INLINE_STREAM unsafeThaw #-}
1792 unsafeThaw = basicUnsafeThaw
1793
1794 -- | /O(n)/ Yield a mutable copy of the immutable vector.
1795 thaw :: (PrimMonad m, Vector v a) => v a -> m (Mutable v (PrimState m) a)
1796 {-# INLINE_STREAM thaw #-}
1797 thaw v = do
1798 mv <- M.unsafeNew (length v)
1799 unsafeCopy mv v
1800 return mv
1801
1802 {-# RULES
1803
1804 "unsafeThaw/new [Vector]" forall p.
1805 unsafeThaw (new p) = New.runPrim p
1806
1807 "thaw/new [Vector]" forall p.
1808 thaw (new p) = New.runPrim p
1809
1810 #-}
1811
1812 {-
1813 -- | /O(n)/ Yield a mutable vector containing copies of each vector in the
1814 -- list.
1815 thawMany :: (PrimMonad m, Vector v a) => [v a] -> m (Mutable v (PrimState m) a)
1816 {-# INLINE_STREAM thawMany #-}
1817 -- FIXME: add rule for (stream (new (New.create (thawMany vs))))
1818 -- NOTE: We don't try to consume the list lazily as this wouldn't significantly
1819 -- change the space requirements anyway.
1820 thawMany vs = do
1821 mv <- M.new n
1822 thaw_loop mv vs
1823 return mv
1824 where
1825 n = List.foldl' (\k v -> k + length v) 0 vs
1826
1827 thaw_loop mv [] = mv `seq` return ()
1828 thaw_loop mv (v:vs)
1829 = do
1830 let n = length v
1831 unsafeCopy (M.unsafeTake n mv) v
1832 thaw_loop (M.unsafeDrop n mv) vs
1833 -}
1834
1835 -- | /O(n)/ Copy an immutable vector into a mutable one. The two vectors must
1836 -- have the same length.
1837 copy
1838 :: (PrimMonad m, Vector v a) => Mutable v (PrimState m) a -> v a -> m ()
1839 {-# INLINE copy #-}
1840 copy dst src = BOUNDS_CHECK(check) "copy" "length mismatch"
1841 (M.length dst == length src)
1842 $ unsafeCopy dst src
1843
1844 -- | /O(n)/ Copy an immutable vector into a mutable one. The two vectors must
1845 -- have the same length. This is not checked.
1846 unsafeCopy
1847 :: (PrimMonad m, Vector v a) => Mutable v (PrimState m) a -> v a -> m ()
1848 {-# INLINE unsafeCopy #-}
1849 unsafeCopy dst src = UNSAFE_CHECK(check) "unsafeCopy" "length mismatch"
1850 (M.length dst == length src)
1851 $ (dst `seq` src `seq` basicUnsafeCopy dst src)
1852
1853 -- Conversions to/from Streams
1854 -- ---------------------------
1855
1856 -- | /O(1)/ Convert a vector to a 'Stream'
1857 stream :: Vector v a => v a -> Stream v a
1858 {-# INLINE_STREAM stream #-}
1859 stream v = Stream.fromVector v
1860
1861 {-
1862 stream v = v `seq` n `seq` (Stream.unfoldr get 0 `Stream.sized` Exact n)
1863 where
1864 n = length v
1865
1866 -- NOTE: the False case comes first in Core so making it the recursive one
1867 -- makes the code easier to read
1868 {-# INLINE get #-}
1869 get i | i >= n = Nothing
1870 | otherwise = case basicUnsafeIndexM v i of Box x -> Just (x, i+1)
1871 -}
1872
1873 -- | /O(n)/ Construct a vector from a 'Stream'
1874 unstream :: Vector v a => Stream v a -> v a
1875 {-# INLINE unstream #-}
1876 unstream s = new (New.unstream s)
1877
1878 {-# RULES
1879
1880 "stream/unstream [Vector]" forall s.
1881 stream (new (New.unstream s)) = s
1882
1883 "New.unstream/stream [Vector]" forall v.
1884 New.unstream (stream v) = clone v
1885
1886 "clone/new [Vector]" forall p.
1887 clone (new p) = p
1888
1889 "inplace [Vector]"
1890 forall (f :: forall m. Monad m => MStream m u a -> MStream m u a) m.
1891 New.unstream (inplace f (stream (new m))) = New.transform f m
1892
1893 "uninplace [Vector]"
1894 forall (f :: forall m. Monad m => MStream m u a -> MStream m u a) m.
1895 stream (new (New.transform f m)) = inplace f (stream (new m))
1896
1897 #-}
1898
1899 -- | /O(1)/ Convert a vector to a 'Stream', proceeding from right to left
1900 streamR :: Vector v a => v a -> Stream u a
1901 {-# INLINE_STREAM streamR #-}
1902 streamR v = v `seq` n `seq` (Stream.unfoldr get n `Stream.sized` Exact n)
1903 where
1904 n = length v
1905
1906 {-# INLINE get #-}
1907 get 0 = Nothing
1908 get i = let i' = i-1
1909 in
1910 case basicUnsafeIndexM v i' of Box x -> Just (x, i')
1911
1912 -- | /O(n)/ Construct a vector from a 'Stream', proceeding from right to left
1913 unstreamR :: Vector v a => Stream v a -> v a
1914 {-# INLINE unstreamR #-}
1915 unstreamR s = new (New.unstreamR s)
1916
1917 {-# RULES
1918
1919 "streamR/unstreamR [Vector]" forall s.
1920 streamR (new (New.unstreamR s)) = s
1921
1922 "New.unstreamR/streamR/new [Vector]" forall p.
1923 New.unstreamR (streamR (new p)) = p
1924
1925 "New.unstream/streamR/new [Vector]" forall p.
1926 New.unstream (streamR (new p)) = New.modify M.reverse p
1927
1928 "New.unstreamR/stream/new [Vector]" forall p.
1929 New.unstreamR (stream (new p)) = New.modify M.reverse p
1930
1931 "inplace right [Vector]"
1932 forall (f :: forall m. Monad m => MStream m u a -> MStream m u a) m.
1933 New.unstreamR (inplace f (streamR (new m))) = New.transformR f m
1934
1935 "uninplace right [Vector]"
1936 forall (f :: forall m. Monad m => MStream m u a -> MStream m u a) m.
1937 streamR (new (New.transformR f m)) = inplace f (streamR (new m))
1938
1939 #-}
1940
1941 unstreamM :: (Monad m, Vector v a) => MStream m u a -> m (v a)
1942 {-# INLINE_STREAM unstreamM #-}
1943 unstreamM s = do
1944 xs <- MStream.toList s
1945 return $ unstream $ Stream.unsafeFromList (MStream.size s) xs
1946
1947 unstreamPrimM :: (PrimMonad m, Vector v a) => MStream m u a -> m (v a)
1948 {-# INLINE_STREAM unstreamPrimM #-}
1949 unstreamPrimM s = M.munstream s >>= unsafeFreeze
1950
1951 -- FIXME: the next two functions are only necessary for the specialisations
1952 unstreamPrimM_IO :: Vector v a => MStream IO u a -> IO (v a)
1953 {-# INLINE unstreamPrimM_IO #-}
1954 unstreamPrimM_IO = unstreamPrimM
1955
1956 unstreamPrimM_ST :: Vector v a => MStream (ST s) u a -> ST s (v a)
1957 {-# INLINE unstreamPrimM_ST #-}
1958 unstreamPrimM_ST = unstreamPrimM
1959
1960 {-# RULES
1961
1962 "unstreamM[IO]" unstreamM = unstreamPrimM_IO
1963 "unstreamM[ST]" unstreamM = unstreamPrimM_ST
1964
1965 #-}
1966
1967
1968 -- Recycling support
1969 -- -----------------
1970
1971 -- | Construct a vector from a monadic initialiser.
1972 new :: Vector v a => New v a -> v a
1973 {-# INLINE_STREAM new #-}
1974 new m = m `seq` runST (unsafeFreeze =<< New.run m)
1975
1976 -- | Convert a vector to an initialiser which, when run, produces a copy of
1977 -- the vector.
1978 clone :: Vector v a => v a -> New v a
1979 {-# INLINE_STREAM clone #-}
1980 clone v = v `seq` New.create (
1981 do
1982 mv <- M.new (length v)
1983 unsafeCopy mv v
1984 return mv)
1985
1986 -- Comparisons
1987 -- -----------
1988
1989 -- | /O(n)/ Check if two vectors are equal. All 'Vector' instances are also
1990 -- instances of 'Eq' and it is usually more appropriate to use those. This
1991 -- function is primarily intended for implementing 'Eq' instances for new
1992 -- vector types.
1993 eq :: (Vector v a, Eq a) => v a -> v a -> Bool
1994 {-# INLINE eq #-}
1995 xs `eq` ys = stream xs == stream ys
1996
1997 -- | /O(n)/ Compare two vectors lexicographically. All 'Vector' instances are
1998 -- also instances of 'Ord' and it is usually more appropriate to use those. This
1999 -- function is primarily intended for implementing 'Ord' instances for new
2000 -- vector types.
2001 cmp :: (Vector v a, Ord a) => v a -> v a -> Ordering
2002 {-# INLINE cmp #-}
2003 cmp xs ys = compare (stream xs) (stream ys)
2004
2005 -- Show
2006 -- ----
2007
2008 -- | Generic definition of 'Prelude.showsPrec'
2009 showsPrec :: (Vector v a, Show a) => Int -> v a -> ShowS
2010 {-# INLINE showsPrec #-}
2011 showsPrec p v = showParen (p > 10) $ showString "fromList " . shows (toList v)
2012
2013 -- | Generic definition of 'Text.Read.readPrec'
2014 readPrec :: (Vector v a, Read a) => Read.ReadPrec (v a)
2015 {-# INLINE readPrec #-}
2016 readPrec = Read.parens $ Read.prec 10 $ do
2017 Read.Ident "fromList" <- Read.lexP
2018 xs <- Read.readPrec
2019 return (fromList xs)
2020
2021 -- Data and Typeable
2022 -- -----------------
2023
2024 -- | Generic definion of 'Data.Data.gfoldl' that views a 'Vector' as a
2025 -- list.
2026 gfoldl :: (Vector v a, Data a)
2027 => (forall d b. Data d => c (d -> b) -> d -> c b)
2028 -> (forall g. g -> c g)
2029 -> v a
2030 -> c (v a)
2031 {-# INLINE gfoldl #-}
2032 gfoldl f z v = z fromList `f` toList v
2033
2034 mkType :: String -> DataType
2035 {-# INLINE mkType #-}
2036 mkType = mkNoRepType
2037
2038 dataCast :: (Vector v a, Data a, Typeable1 v, Typeable1 t)
2039 => (forall d. Data d => c (t d)) -> Maybe (c (v a))
2040 {-# INLINE dataCast #-}
2041 dataCast f = gcast1 f
2042