Switch the order of argument to slice and unsafeSlice
[darcs-mirrors/vector.git] / Data / Vector / Generic.hs
1 {-# LANGUAGE Rank2Types, MultiParamTypeClasses, FlexibleContexts,
2 TypeFamilies, ScopedTypeVariables #-}
3 -- |
4 -- Module : Data.Vector.Generic
5 -- Copyright : (c) Roman Leshchinskiy 2008-2009
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 -- * Length information
20 length, null,
21
22 -- * Construction
23 empty, singleton, cons, snoc, replicate, generate, (++), copy,
24
25 -- * Accessing individual elements
26 (!), head, last, indexM, headM, lastM,
27 unsafeIndex, unsafeHead, unsafeLast,
28 unsafeIndexM, unsafeHeadM, unsafeLastM,
29
30 -- * Subvectors
31 slice, init, tail, take, drop,
32 unsafeSlice,
33
34 -- * Permutations
35 accum, accumulate, accumulate_,
36 (//), update, update_,
37 backpermute, reverse,
38
39 -- * Mapping
40 map, imap, concatMap,
41
42 -- * Zipping and unzipping
43 zipWith, zipWith3, zipWith4, zipWith5, zipWith6,
44 izipWith, izipWith3, izipWith4, izipWith5, izipWith6,
45 zip, zip3, zip4, zip5, zip6,
46 unzip, unzip3, unzip4, unzip5, unzip6,
47
48 -- * Comparisons
49 eq, cmp,
50
51 -- * Filtering
52 filter, ifilter, takeWhile, dropWhile,
53 unstablePartition, span, break,
54
55 -- * Searching
56 elem, notElem, find, findIndex, findIndices, elemIndex, elemIndices,
57
58 -- * Folding
59 foldl, foldl1, foldl', foldl1', foldr, foldr1,
60 ifoldl, ifoldl', ifoldr,
61
62 -- * Specialised folds
63 all, any, and, or,
64 sum, product,
65 maximum, maximumBy, minimum, minimumBy,
66 minIndex, minIndexBy, maxIndex, maxIndexBy,
67
68 -- * Unfolding
69 unfoldr,
70
71 -- * Scans
72 prescanl, prescanl',
73 postscanl, postscanl',
74 scanl, scanl', scanl1, scanl1',
75
76 -- * Enumeration
77 enumFromTo, enumFromThenTo,
78
79 -- * Conversion to/from lists
80 toList, fromList,
81
82 -- * Conversion to/from Streams
83 stream, unstream,
84
85 -- * MVector-based initialisation
86 new,
87
88 -- * Unsafe operations
89 unsafeAccum, unsafeAccumulate, unsafeAccumulate_,
90 unsafeUpd, unsafeUpdate, unsafeUpdate_
91 ) where
92
93 import Data.Vector.Generic.Mutable ( MVector )
94 import qualified Data.Vector.Generic.Mutable as M
95
96 import qualified Data.Vector.Generic.New as New
97 import Data.Vector.Generic.New ( New )
98
99 import qualified Data.Vector.Fusion.Stream as Stream
100 import Data.Vector.Fusion.Stream ( Stream, MStream, inplace )
101 import qualified Data.Vector.Fusion.Stream.Monadic as MStream
102 import Data.Vector.Fusion.Stream.Size
103 import Data.Vector.Fusion.Util
104
105 import Control.Monad.ST ( runST )
106 import Control.Monad.Primitive
107 import Prelude hiding ( length, null,
108 replicate, (++),
109 head, last,
110 init, tail, take, drop, reverse,
111 map, concatMap,
112 zipWith, zipWith3, zip, zip3, unzip, unzip3,
113 filter, takeWhile, dropWhile, span, break,
114 elem, notElem,
115 foldl, foldl1, foldr, foldr1,
116 all, any, and, or, sum, product, maximum, minimum,
117 scanl, scanl1,
118 enumFromTo, enumFromThenTo )
119
120 #include "vector.h"
121
122 type family Mutable (v :: * -> *) :: * -> * -> *
123
124 -- | Class of immutable vectors.
125 --
126 class MVector (Mutable v) a => Vector v a where
127 -- | Unsafely convert a mutable vector to its immutable version
128 -- without copying. The mutable vector may not be used after
129 -- this operation.
130 unsafeFreeze :: PrimMonad m => Mutable v (PrimState m) a -> m (v a)
131
132 -- | Length of the vector (not fusible!)
133 basicLength :: v a -> Int
134
135 -- | Yield a part of the vector without copying it. No range checks!
136 basicUnsafeSlice :: Int -> Int -> v a -> v a
137
138 -- | Yield the element at the given position in a monad. The monad allows us
139 -- to be strict in the vector if we want. Suppose we had
140 --
141 -- > unsafeIndex :: v a -> Int -> a
142 --
143 -- instead. Now, if we wanted to copy a vector, we'd do something like
144 --
145 -- > copy mv v ... = ... unsafeWrite mv i (unsafeIndex v i) ...
146 --
147 -- For lazy vectors, the indexing would not be evaluated which means that we
148 -- would retain a reference to the original vector in each element we write.
149 -- This is not what we want!
150 --
151 -- With 'basicUnsafeIndexM', we can do
152 --
153 -- > copy mv v ... = ... case basicUnsafeIndexM v i of
154 -- > Box x -> unsafeWrite mv i x ...
155 --
156 -- which does not have this problem because indexing (but not the returned
157 -- element!) is evaluated immediately.
158 --
159 basicUnsafeIndexM :: Monad m => v a -> Int -> m a
160
161 elemseq :: v a -> a -> b -> b
162
163 {-# INLINE elemseq #-}
164 elemseq _ = \_ x -> x
165
166 -- Fusion
167 -- ------
168
169 -- | Construct a pure vector from a monadic initialiser
170 new :: Vector v a => New a -> v a
171 {-# INLINE new #-}
172 new m = new' undefined m
173
174 -- | Same as 'new' but with a dummy argument necessary for correctly typing
175 -- the rule @uninplace@.
176 --
177 -- See http://hackage.haskell.org/trac/ghc/ticket/2600
178 new' :: Vector v a => v a -> New a -> v a
179 {-# INLINE_STREAM new' #-}
180 new' _ m = m `seq` runST (do
181 mv <- New.run m
182 unsafeFreeze mv)
183
184 -- | Convert a vector to a 'Stream'
185 stream :: Vector v a => v a -> Stream a
186 {-# INLINE_STREAM stream #-}
187 stream v = v `seq` (Stream.unfoldr get 0 `Stream.sized` Exact n)
188 where
189 n = length v
190
191 -- NOTE: the False case comes first in Core so making it the recursive one
192 -- makes the code easier to read
193 {-# INLINE get #-}
194 get i | i >= n = Nothing
195 | otherwise = case basicUnsafeIndexM v i of Box x -> Just (x, i+1)
196
197 -- | Create a vector from a 'Stream'
198 unstream :: Vector v a => Stream a -> v a
199 {-# INLINE unstream #-}
200 unstream s = new (New.unstream s)
201
202 {-# RULES
203
204 "stream/unstream [Vector]" forall v s.
205 stream (new' v (New.unstream s)) = s
206
207 "New.unstream/stream/new [Vector]" forall v p.
208 New.unstream (stream (new' v p)) = p
209
210 #-}
211
212 {-# RULES
213
214 "inplace [Vector]"
215 forall (f :: forall m. Monad m => MStream m a -> MStream m a) v m.
216 New.unstream (inplace f (stream (new' v m))) = New.transform f m
217
218 "uninplace [Vector]"
219 forall (f :: forall m. Monad m => MStream m a -> MStream m a) v m.
220 stream (new' v (New.transform f m)) = inplace f (stream (new' v m))
221
222 #-}
223
224 -- Length
225 -- ------
226
227 length :: Vector v a => v a -> Int
228 {-# INLINE_STREAM length #-}
229 length v = basicLength v
230
231 {-# RULES
232
233 "length/unstream [Vector]" forall v s.
234 length (new' v (New.unstream s)) = Stream.length s
235
236 #-}
237
238 null :: Vector v a => v a -> Bool
239 {-# INLINE null #-}
240 null v = length v == 0
241
242 -- Construction
243 -- ------------
244
245 -- | Empty vector
246 empty :: Vector v a => v a
247 {-# INLINE empty #-}
248 empty = unstream Stream.empty
249
250 -- | Vector with exaclty one element
251 singleton :: forall v a. Vector v a => a -> v a
252 {-# INLINE singleton #-}
253 singleton x = elemseq (undefined :: v a) x
254 $ unstream (Stream.singleton x)
255
256 -- | Vector of the given length with the given value in each position
257 replicate :: forall v a. Vector v a => Int -> a -> v a
258 {-# INLINE replicate #-}
259 replicate n x = elemseq (undefined :: v a) x
260 $ unstream
261 $ Stream.replicate n x
262
263 -- | Generate a vector of the given length by applying the function to each
264 -- index
265 generate :: Vector v a => Int -> (Int -> a) -> v a
266 {-# INLINE generate #-}
267 generate n f = unstream (Stream.generate n f)
268
269 -- | Prepend an element
270 cons :: forall v a. Vector v a => a -> v a -> v a
271 {-# INLINE cons #-}
272 cons x v = elemseq (undefined :: v a) x
273 $ unstream
274 $ Stream.cons x
275 $ stream v
276
277 -- | Append an element
278 snoc :: forall v a. Vector v a => v a -> a -> v a
279 {-# INLINE snoc #-}
280 snoc v x = elemseq (undefined :: v a) x
281 $ unstream
282 $ Stream.snoc (stream v) x
283
284 infixr 5 ++
285 -- | Concatenate two vectors
286 (++) :: Vector v a => v a -> v a -> v a
287 {-# INLINE (++) #-}
288 v ++ w = unstream (stream v Stream.++ stream w)
289
290 -- | Create a copy of a vector. Useful when dealing with slices.
291 copy :: Vector v a => v a -> v a
292 {-# INLINE_STREAM copy #-}
293 copy = unstream . stream
294
295 {-# RULES
296
297 "copy/unstream [Vector]" forall v s.
298 copy (new' v (New.unstream s)) = new' v (New.unstream s)
299
300 #-}
301
302 -- Accessing individual elements
303 -- -----------------------------
304
305 -- | Indexing
306 (!) :: Vector v a => v a -> Int -> a
307 {-# INLINE_STREAM (!) #-}
308 v ! i = BOUNDS_CHECK(checkIndex) "(!)" i (length v)
309 $ unId (basicUnsafeIndexM v i)
310
311 -- | First element
312 head :: Vector v a => v a -> a
313 {-# INLINE_STREAM head #-}
314 head v = v ! 0
315
316 -- | Last element
317 last :: Vector v a => v a -> a
318 {-# INLINE_STREAM last #-}
319 last v = v ! (length v - 1)
320
321 -- | Unsafe indexing without bounds checking
322 unsafeIndex :: Vector v a => v a -> Int -> a
323 {-# INLINE_STREAM unsafeIndex #-}
324 unsafeIndex v i = UNSAFE_CHECK(checkIndex) "unsafeIndex" i (length v)
325 $ unId (basicUnsafeIndexM v i)
326
327 -- | Yield the first element of a vector without checking if the vector is
328 -- empty
329 unsafeHead :: Vector v a => v a -> a
330 {-# INLINE_STREAM unsafeHead #-}
331 unsafeHead v = unsafeIndex v 0
332
333 -- | Yield the last element of a vector without checking if the vector is
334 -- empty
335 unsafeLast :: Vector v a => v a -> a
336 {-# INLINE_STREAM unsafeLast #-}
337 unsafeLast v = unsafeIndex v (length v - 1)
338
339 {-# RULES
340
341 "(!)/unstream [Vector]" forall v i s.
342 new' v (New.unstream s) ! i = s Stream.!! i
343
344 "head/unstream [Vector]" forall v s.
345 head (new' v (New.unstream s)) = Stream.head s
346
347 "last/unstream [Vector]" forall v s.
348 last (new' v (New.unstream s)) = Stream.last s
349
350 "unsafeIndex/unstream [Vector]" forall v i s.
351 unsafeIndex (new' v (New.unstream s)) i = s Stream.!! i
352
353 "unsafeHead/unstream [Vector]" forall v s.
354 unsafeHead (new' v (New.unstream s)) = Stream.head s
355
356 "unsafeLast/unstream [Vector]" forall v s.
357 unsafeLast (new' v (New.unstream s)) = Stream.last s
358
359 #-}
360
361 -- | Monadic indexing which can be strict in the vector while remaining lazy in
362 -- the element.
363 indexM :: (Vector v a, Monad m) => v a -> Int -> m a
364 {-# INLINE_STREAM indexM #-}
365 indexM v i = BOUNDS_CHECK(checkIndex) "indexM" i (length v)
366 $ basicUnsafeIndexM v i
367
368 headM :: (Vector v a, Monad m) => v a -> m a
369 {-# INLINE_STREAM headM #-}
370 headM v = indexM v 0
371
372 lastM :: (Vector v a, Monad m) => v a -> m a
373 {-# INLINE_STREAM lastM #-}
374 lastM v = indexM v (length v - 1)
375
376 -- | Unsafe monadic indexing without bounds checks
377 unsafeIndexM :: (Vector v a, Monad m) => v a -> Int -> m a
378 {-# INLINE_STREAM unsafeIndexM #-}
379 unsafeIndexM v i = UNSAFE_CHECK(checkIndex) "unsafeIndexM" i (length v)
380 $ basicUnsafeIndexM v i
381
382 unsafeHeadM :: (Vector v a, Monad m) => v a -> m a
383 {-# INLINE_STREAM unsafeHeadM #-}
384 unsafeHeadM v = unsafeIndexM v 0
385
386 unsafeLastM :: (Vector v a, Monad m) => v a -> m a
387 {-# INLINE_STREAM unsafeLastM #-}
388 unsafeLastM v = unsafeIndexM v (length v - 1)
389
390 -- FIXME: the rhs of these rules are lazy in the stream which is WRONG
391 {- RULES
392
393 "indexM/unstream [Vector]" forall v i s.
394 indexM (new' v (New.unstream s)) i = return (s Stream.!! i)
395
396 "headM/unstream [Vector]" forall v s.
397 headM (new' v (New.unstream s)) = return (Stream.head s)
398
399 "lastM/unstream [Vector]" forall v s.
400 lastM (new' v (New.unstream s)) = return (Stream.last s)
401
402 -}
403
404 -- Subarrays
405 -- ---------
406
407 -- FIXME: slicing doesn't work with the inplace stuff at the moment
408
409 -- | Yield a part of the vector without copying it.
410 slice :: Vector v a => Int -- ^ starting index
411 -> Int -- ^ length
412 -> v a
413 -> v a
414 {-# INLINE_STREAM slice #-}
415 slice i n v = BOUNDS_CHECK(checkSlice) "slice" i n (length v)
416 $ basicUnsafeSlice i n v
417
418 -- | Unsafely yield a part of the vector without copying it and without
419 -- performing bounds checks.
420 unsafeSlice :: Vector v a => Int -- ^ starting index
421 -> Int -- ^ length
422 -> v a
423 -> v a
424 {-# INLINE_STREAM unsafeSlice #-}
425 unsafeSlice i n v = UNSAFE_CHECK(checkSlice) "unsafeSlice" i n (length v)
426 $ basicUnsafeSlice i n v
427
428 -- | Yield all but the last element without copying.
429 init :: Vector v a => v a -> v a
430 {-# INLINE_STREAM init #-}
431 init v = slice 0 (length v - 1) v
432
433 -- | All but the first element (without copying).
434 tail :: Vector v a => v a -> v a
435 {-# INLINE_STREAM tail #-}
436 tail v = slice 1 (length v - 1) v
437
438 -- | Yield the first @n@ elements without copying.
439 take :: Vector v a => Int -> v a -> v a
440 {-# INLINE_STREAM take #-}
441 take n v = slice 0 (min n' (length v)) v
442 where n' = max n 0
443
444 -- | Yield all but the first @n@ elements without copying.
445 drop :: Vector v a => Int -> v a -> v a
446 {-# INLINE_STREAM drop #-}
447 drop n v = slice (min n' len) (max 0 (len - n')) v
448 where n' = max n 0
449 len = length v
450
451 {-# RULES
452
453 "slice/new [Vector]" forall i n v p.
454 slice i n (new' v p) = new' v (New.slice i n p)
455
456 "unsafeSlice/new [Vector]" forall i n v p.
457 unsafeSlice i n (new' v p) = new' v (New.unsafeSlice i n p)
458
459 "init/new [Vector]" forall v p.
460 init (new' v p) = new' v (New.init p)
461
462 "tail/new [Vector]" forall v p.
463 tail (new' v p) = new' v (New.tail p)
464
465 "take/new [Vector]" forall n v p.
466 take n (new' v p) = new' v (New.take n p)
467
468 "drop/new [Vector]" forall n v p.
469 drop n (new' v p) = new' v (New.drop n p)
470
471 #-}
472
473 -- Permutations
474 -- ------------
475
476 unsafeAccum_stream
477 :: Vector v a => (a -> b -> a) -> v a -> Stream (Int,b) -> v a
478 {-# INLINE unsafeAccum_stream #-}
479 unsafeAccum_stream f v s = new (New.accum f (New.unstream (stream v)) s)
480
481 unsafeAccum :: Vector v a => (a -> b -> a) -> v a -> [(Int,b)] -> v a
482 {-# INLINE unsafeAccum #-}
483 unsafeAccum f v us = unsafeAccum_stream f v (Stream.fromList us)
484
485 unsafeAccumulate :: (Vector v a, Vector v (Int, b))
486 => (a -> b -> a) -> v a -> v (Int,b) -> v a
487 {-# INLINE unsafeAccumulate #-}
488 unsafeAccumulate f v us = unsafeAccum_stream f v (stream us)
489
490 unsafeAccumulate_ :: (Vector v a, Vector v Int, Vector v b)
491 => (a -> b -> a) -> v a -> v Int -> v b -> v a
492 {-# INLINE unsafeAccumulate_ #-}
493 unsafeAccumulate_ f v is xs
494 = unsafeAccum_stream f v (Stream.zipWith (,) (stream is) (stream xs))
495
496 accum_stream :: Vector v a => (a -> b -> a) -> v a -> Stream (Int,b) -> v a
497 {-# INLINE accum_stream #-}
498 accum_stream f v s = new (New.accum f (New.unstream (stream v)) s)
499
500 accum :: Vector v a => (a -> b -> a) -> v a -> [(Int,b)] -> v a
501 {-# INLINE accum #-}
502 accum f v us = accum_stream f v (Stream.fromList us)
503
504 accumulate :: (Vector v a, Vector v (Int, b))
505 => (a -> b -> a) -> v a -> v (Int,b) -> v a
506 {-# INLINE accumulate #-}
507 accumulate f v us = accum_stream f v (stream us)
508
509 accumulate_ :: (Vector v a, Vector v Int, Vector v b)
510 => (a -> b -> a) -> v a -> v Int -> v b -> v a
511 {-# INLINE accumulate_ #-}
512 accumulate_ f v is xs = accum_stream f v (Stream.zipWith (,) (stream is)
513 (stream xs))
514
515
516 unsafeUpdate_stream :: Vector v a => v a -> Stream (Int,a) -> v a
517 {-# INLINE unsafeUpdate_stream #-}
518 unsafeUpdate_stream v s = new (New.unsafeUpdate (New.unstream (stream v)) s)
519
520 unsafeUpd :: Vector v a => v a -> [(Int, a)] -> v a
521 {-# INLINE unsafeUpd #-}
522 unsafeUpd v us = unsafeUpdate_stream v (Stream.fromList us)
523
524 unsafeUpdate :: (Vector v a, Vector v (Int, a)) => v a -> v (Int, a) -> v a
525 {-# INLINE unsafeUpdate #-}
526 unsafeUpdate v w = unsafeUpdate_stream v (stream w)
527
528 unsafeUpdate_ :: (Vector v a, Vector v Int) => v a -> v Int -> v a -> v a
529 {-# INLINE unsafeUpdate_ #-}
530 unsafeUpdate_ v is w
531 = unsafeUpdate_stream v (Stream.zipWith (,) (stream is) (stream w))
532
533 update_stream :: Vector v a => v a -> Stream (Int,a) -> v a
534 {-# INLINE update_stream #-}
535 update_stream v s = new (New.update (New.unstream (stream v)) s)
536
537 (//) :: Vector v a => v a -> [(Int, a)] -> v a
538 {-# INLINE (//) #-}
539 v // us = update_stream v (Stream.fromList us)
540
541 update :: (Vector v a, Vector v (Int, a)) => v a -> v (Int, a) -> v a
542 {-# INLINE update #-}
543 update v w = update_stream v (stream w)
544
545 update_ :: (Vector v a, Vector v Int) => v a -> v Int -> v a -> v a
546 {-# INLINE update_ #-}
547 update_ v is w = update_stream v (Stream.zipWith (,) (stream is) (stream w))
548
549 -- This somewhat non-intuitive definition ensures that the resulting vector
550 -- does not retain references to the original one even if it is lazy in its
551 -- elements. This would not be the case if we simply used
552 --
553 -- backpermute v is = map (v!) is
554 backpermute :: (Vector v a, Vector v Int) => v a -> v Int -> v a
555 {-# INLINE backpermute #-}
556 backpermute v is = seq v
557 $ unstream
558 $ Stream.unbox
559 $ Stream.map (indexM v)
560 $ stream is
561
562 reverse :: (Vector v a) => v a -> v a
563 {-# INLINE reverse #-}
564 reverse = new . New.reverse . New.unstream . stream
565
566 -- Mapping
567 -- -------
568
569 -- | Map a function over a vector
570 map :: (Vector v a, Vector v b) => (a -> b) -> v a -> v b
571 {-# INLINE map #-}
572 map f = unstream . inplace (MStream.map f) . stream
573
574 -- | Apply a function to every index/value pair
575 imap :: (Vector v a, Vector v b) => (Int -> a -> b) -> v a -> v b
576 {-# INLINE imap #-}
577 imap f = unstream . inplace (MStream.map (uncurry f) . MStream.indexed)
578 . stream
579
580 concatMap :: (Vector v a, Vector v b) => (a -> v b) -> v a -> v b
581 {-# INLINE concatMap #-}
582 concatMap f = unstream . Stream.concatMap (stream . f) . stream
583
584 -- Zipping/unzipping
585 -- -----------------
586
587 -- | Zip two vectors with the given function.
588 zipWith :: (Vector v a, Vector v b, Vector v c)
589 => (a -> b -> c) -> v a -> v b -> v c
590 {-# INLINE zipWith #-}
591 zipWith f xs ys = unstream (Stream.zipWith f (stream xs) (stream ys))
592
593 -- | Zip three vectors with the given function.
594 zipWith3 :: (Vector v a, Vector v b, Vector v c, Vector v d)
595 => (a -> b -> c -> d) -> v a -> v b -> v c -> v d
596 {-# INLINE zipWith3 #-}
597 zipWith3 f as bs cs = unstream (Stream.zipWith3 f (stream as)
598 (stream bs)
599 (stream cs))
600
601 zipWith4 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e)
602 => (a -> b -> c -> d -> e) -> v a -> v b -> v c -> v d -> v e
603 {-# INLINE zipWith4 #-}
604 zipWith4 f as bs cs ds
605 = unstream (Stream.zipWith4 f (stream as)
606 (stream bs)
607 (stream cs)
608 (stream ds))
609
610 zipWith5 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
611 Vector v f)
612 => (a -> b -> c -> d -> e -> f) -> v a -> v b -> v c -> v d -> v e
613 -> v f
614 {-# INLINE zipWith5 #-}
615 zipWith5 f as bs cs ds es
616 = unstream (Stream.zipWith5 f (stream as)
617 (stream bs)
618 (stream cs)
619 (stream ds)
620 (stream es))
621
622 zipWith6 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
623 Vector v f, Vector v g)
624 => (a -> b -> c -> d -> e -> f -> g)
625 -> v a -> v b -> v c -> v d -> v e -> v f -> v g
626 {-# INLINE zipWith6 #-}
627 zipWith6 f as bs cs ds es fs
628 = unstream (Stream.zipWith6 f (stream as)
629 (stream bs)
630 (stream cs)
631 (stream ds)
632 (stream es)
633 (stream fs))
634
635 -- | Zip two vectors and their indices with the given function.
636 izipWith :: (Vector v a, Vector v b, Vector v c)
637 => (Int -> a -> b -> c) -> v a -> v b -> v c
638 {-# INLINE izipWith #-}
639 izipWith f xs ys = unstream
640 (Stream.zipWith (uncurry f) (Stream.indexed (stream xs))
641 (stream ys))
642
643 -- | Zip three vectors and their indices with the given function.
644 izipWith3 :: (Vector v a, Vector v b, Vector v c, Vector v d)
645 => (Int -> a -> b -> c -> d) -> v a -> v b -> v c -> v d
646 {-# INLINE izipWith3 #-}
647 izipWith3 f as bs cs
648 = unstream (Stream.zipWith3 (uncurry f) (Stream.indexed (stream as))
649 (stream bs)
650 (stream cs))
651
652 izipWith4 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e)
653 => (Int -> a -> b -> c -> d -> e) -> v a -> v b -> v c -> v d -> v e
654 {-# INLINE izipWith4 #-}
655 izipWith4 f as bs cs ds
656 = unstream (Stream.zipWith4 (uncurry f) (Stream.indexed (stream as))
657 (stream bs)
658 (stream cs)
659 (stream ds))
660
661 izipWith5 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
662 Vector v f)
663 => (Int -> a -> b -> c -> d -> e -> f) -> v a -> v b -> v c -> v d
664 -> v e -> v f
665 {-# INLINE izipWith5 #-}
666 izipWith5 f as bs cs ds es
667 = unstream (Stream.zipWith5 (uncurry f) (Stream.indexed (stream as))
668 (stream bs)
669 (stream cs)
670 (stream ds)
671 (stream es))
672
673 izipWith6 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
674 Vector v f, Vector v g)
675 => (Int -> a -> b -> c -> d -> e -> f -> g)
676 -> v a -> v b -> v c -> v d -> v e -> v f -> v g
677 {-# INLINE izipWith6 #-}
678 izipWith6 f as bs cs ds es fs
679 = unstream (Stream.zipWith6 (uncurry f) (Stream.indexed (stream as))
680 (stream bs)
681 (stream cs)
682 (stream ds)
683 (stream es)
684 (stream fs))
685
686 zip :: (Vector v a, Vector v b, Vector v (a,b)) => v a -> v b -> v (a, b)
687 {-# INLINE zip #-}
688 zip = zipWith (,)
689
690 zip3 :: (Vector v a, Vector v b, Vector v c, Vector v (a, b, c))
691 => v a -> v b -> v c -> v (a, b, c)
692 {-# INLINE zip3 #-}
693 zip3 = zipWith3 (,,)
694
695 zip4 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v (a, b, c, d))
696 => v a -> v b -> v c -> v d -> v (a, b, c, d)
697 {-# INLINE zip4 #-}
698 zip4 = zipWith4 (,,,)
699
700 zip5 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
701 Vector v (a, b, c, d, e))
702 => v a -> v b -> v c -> v d -> v e -> v (a, b, c, d, e)
703 {-# INLINE zip5 #-}
704 zip5 = zipWith5 (,,,,)
705
706 zip6 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
707 Vector v f, Vector v (a, b, c, d, e, f))
708 => v a -> v b -> v c -> v d -> v e -> v f -> v (a, b, c, d, e, f)
709 {-# INLINE zip6 #-}
710 zip6 = zipWith6 (,,,,,)
711
712 unzip :: (Vector v a, Vector v b, Vector v (a,b)) => v (a, b) -> (v a, v b)
713 {-# INLINE unzip #-}
714 unzip xs = (map fst xs, map snd xs)
715
716 unzip3 :: (Vector v a, Vector v b, Vector v c, Vector v (a, b, c))
717 => v (a, b, c) -> (v a, v b, v c)
718 {-# INLINE unzip3 #-}
719 unzip3 xs = (map (\(a, b, c) -> a) xs,
720 map (\(a, b, c) -> b) xs,
721 map (\(a, b, c) -> c) xs)
722
723 unzip4 :: (Vector v a, Vector v b, Vector v c, Vector v d,
724 Vector v (a, b, c, d))
725 => v (a, b, c, d) -> (v a, v b, v c, v d)
726 {-# INLINE unzip4 #-}
727 unzip4 xs = (map (\(a, b, c, d) -> a) xs,
728 map (\(a, b, c, d) -> b) xs,
729 map (\(a, b, c, d) -> c) xs,
730 map (\(a, b, c, d) -> d) xs)
731
732 unzip5 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
733 Vector v (a, b, c, d, e))
734 => v (a, b, c, d, e) -> (v a, v b, v c, v d, v e)
735 {-# INLINE unzip5 #-}
736 unzip5 xs = (map (\(a, b, c, d, e) -> a) xs,
737 map (\(a, b, c, d, e) -> b) xs,
738 map (\(a, b, c, d, e) -> c) xs,
739 map (\(a, b, c, d, e) -> d) xs,
740 map (\(a, b, c, d, e) -> e) xs)
741
742 unzip6 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
743 Vector v f, Vector v (a, b, c, d, e, f))
744 => v (a, b, c, d, e, f) -> (v a, v b, v c, v d, v e, v f)
745 {-# INLINE unzip6 #-}
746 unzip6 xs = (map (\(a, b, c, d, e, f) -> a) xs,
747 map (\(a, b, c, d, e, f) -> b) xs,
748 map (\(a, b, c, d, e, f) -> c) xs,
749 map (\(a, b, c, d, e, f) -> d) xs,
750 map (\(a, b, c, d, e, f) -> e) xs,
751 map (\(a, b, c, d, e, f) -> f) xs)
752
753 -- Comparisons
754 -- -----------
755
756 eq :: (Vector v a, Eq a) => v a -> v a -> Bool
757 {-# INLINE eq #-}
758 xs `eq` ys = stream xs == stream ys
759
760 cmp :: (Vector v a, Ord a) => v a -> v a -> Ordering
761 {-# INLINE cmp #-}
762 cmp xs ys = compare (stream xs) (stream ys)
763
764 -- Filtering
765 -- ---------
766
767 -- | Drop elements that do not satisfy the predicate
768 filter :: Vector v a => (a -> Bool) -> v a -> v a
769 {-# INLINE filter #-}
770 filter f = unstream . inplace (MStream.filter f) . stream
771
772 -- | Drop elements that do not satisfy the predicate (applied to values and
773 -- their indices)
774 ifilter :: Vector v a => (Int -> a -> Bool) -> v a -> v a
775 {-# INLINE ifilter #-}
776 ifilter f = unstream
777 . inplace (MStream.map snd . MStream.filter (uncurry f)
778 . MStream.indexed)
779 . stream
780
781 -- | Yield the longest prefix of elements satisfying the predicate.
782 takeWhile :: Vector v a => (a -> Bool) -> v a -> v a
783 {-# INLINE takeWhile #-}
784 takeWhile f = unstream . Stream.takeWhile f . stream
785
786 -- | Drop the longest prefix of elements that satisfy the predicate.
787 dropWhile :: Vector v a => (a -> Bool) -> v a -> v a
788 {-# INLINE dropWhile #-}
789 dropWhile f = unstream . Stream.dropWhile f . stream
790
791 -- | Split the vector in two parts, the first one containing those elements
792 -- that satisfy the predicate and the second one those that don't. The order
793 -- of the elements is not preserved.
794 unstablePartition :: Vector v a => (a -> Bool) -> v a -> (v a, v a)
795 {-# INLINE unstablePartition #-}
796 unstablePartition f = unstablePartition_stream f . stream
797
798 unstablePartition_stream
799 :: Vector v a => (a -> Bool) -> Stream a -> (v a, v a)
800 {-# INLINE_STREAM unstablePartition_stream #-}
801 unstablePartition_stream f s = s `seq` runST (
802 do
803 (mv1,mv2) <- M.unstablePartitionStream f s
804 v1 <- unsafeFreeze mv1
805 v2 <- unsafeFreeze mv2
806 return (v1,v2))
807
808 unstablePartition_new :: Vector v a => (a -> Bool) -> New a -> (v a, v a)
809 {-# INLINE_STREAM unstablePartition_new #-}
810 unstablePartition_new f (New.New p) = runST (
811 do
812 mv <- p
813 i <- M.unstablePartition f mv
814 v <- unsafeFreeze mv
815 return (take i v, drop i v))
816
817 {-# RULES
818
819 "unstablePartition" forall f v p.
820 unstablePartition_stream f (stream (new' v p))
821 = unstablePartition_new f p
822
823 #-}
824
825
826 -- FIXME: make span and break fusible
827
828 -- | Split the vector into the longest prefix of elements that satisfy the
829 -- predicate and the rest.
830 span :: Vector v a => (a -> Bool) -> v a -> (v a, v a)
831 {-# INLINE span #-}
832 span f = break (not . f)
833
834 -- | Split the vector into the longest prefix of elements that do not satisfy
835 -- the predicate and the rest.
836 break :: Vector v a => (a -> Bool) -> v a -> (v a, v a)
837 {-# INLINE break #-}
838 break f xs = case findIndex f xs of
839 Just i -> (unsafeSlice 0 i xs, unsafeSlice i (length xs - i) xs)
840 Nothing -> (xs, empty)
841
842
843 -- Searching
844 -- ---------
845
846 infix 4 `elem`
847 -- | Check whether the vector contains an element
848 elem :: (Vector v a, Eq a) => a -> v a -> Bool
849 {-# INLINE elem #-}
850 elem x = Stream.elem x . stream
851
852 infix 4 `notElem`
853 -- | Inverse of `elem`
854 notElem :: (Vector v a, Eq a) => a -> v a -> Bool
855 {-# INLINE notElem #-}
856 notElem x = Stream.notElem x . stream
857
858 -- | Yield 'Just' the first element matching the predicate or 'Nothing' if no
859 -- such element exists.
860 find :: Vector v a => (a -> Bool) -> v a -> Maybe a
861 {-# INLINE find #-}
862 find f = Stream.find f . stream
863
864 -- | Yield 'Just' the index of the first element matching the predicate or
865 -- 'Nothing' if no such element exists.
866 findIndex :: Vector v a => (a -> Bool) -> v a -> Maybe Int
867 {-# INLINE findIndex #-}
868 findIndex f = Stream.findIndex f . stream
869
870 -- | Yield the indices of elements satisfying the predicate
871 findIndices :: (Vector v a, Vector v Int) => (a -> Bool) -> v a -> v Int
872 {-# INLINE findIndices #-}
873 findIndices f = unstream
874 . Stream.map fst
875 . Stream.filter (f . snd)
876 . Stream.indexed
877 . stream
878
879 -- | Yield 'Just' the index of the first occurence of the given element or
880 -- 'Nothing' if the vector does not contain the element
881 elemIndex :: (Vector v a, Eq a) => a -> v a -> Maybe Int
882 {-# INLINE elemIndex #-}
883 elemIndex x = findIndex (x==)
884
885 -- | Yield the indices of all occurences of the given element
886 elemIndices :: (Vector v a, Vector v Int, Eq a) => a -> v a -> v Int
887 {-# INLINE elemIndices #-}
888 elemIndices x = findIndices (x==)
889
890 -- Folding
891 -- -------
892
893 -- | Left fold
894 foldl :: Vector v b => (a -> b -> a) -> a -> v b -> a
895 {-# INLINE foldl #-}
896 foldl f z = Stream.foldl f z . stream
897
898 -- | Left fold on non-empty vectors
899 foldl1 :: Vector v a => (a -> a -> a) -> v a -> a
900 {-# INLINE foldl1 #-}
901 foldl1 f = Stream.foldl1 f . stream
902
903 -- | Left fold with strict accumulator
904 foldl' :: Vector v b => (a -> b -> a) -> a -> v b -> a
905 {-# INLINE foldl' #-}
906 foldl' f z = Stream.foldl' f z . stream
907
908 -- | Left fold on non-empty vectors with strict accumulator
909 foldl1' :: Vector v a => (a -> a -> a) -> v a -> a
910 {-# INLINE foldl1' #-}
911 foldl1' f = Stream.foldl1' f . stream
912
913 -- | Right fold
914 foldr :: Vector v a => (a -> b -> b) -> b -> v a -> b
915 {-# INLINE foldr #-}
916 foldr f z = Stream.foldr f z . stream
917
918 -- | Right fold on non-empty vectors
919 foldr1 :: Vector v a => (a -> a -> a) -> v a -> a
920 {-# INLINE foldr1 #-}
921 foldr1 f = Stream.foldr1 f . stream
922
923 -- | Left fold (function applied to each element and its index)
924 ifoldl :: Vector v b => (a -> Int -> b -> a) -> a -> v b -> a
925 {-# INLINE ifoldl #-}
926 ifoldl f z = Stream.foldl (uncurry . f) z . Stream.indexed . stream
927
928 -- | Left fold with strict accumulator (function applied to each element and
929 -- its index)
930 ifoldl' :: Vector v b => (a -> Int -> b -> a) -> a -> v b -> a
931 {-# INLINE ifoldl' #-}
932 ifoldl' f z = Stream.foldl' (uncurry . f) z . Stream.indexed . stream
933
934 -- | Right fold (function applied to each element and its index)
935 ifoldr :: Vector v a => (Int -> a -> b -> b) -> b -> v a -> b
936 {-# INLINE ifoldr #-}
937 ifoldr f z = Stream.foldr (uncurry f) z . Stream.indexed . stream
938
939 -- Specialised folds
940 -- -----------------
941
942 all :: Vector v a => (a -> Bool) -> v a -> Bool
943 {-# INLINE all #-}
944 all f = Stream.and . Stream.map f . stream
945
946 any :: Vector v a => (a -> Bool) -> v a -> Bool
947 {-# INLINE any #-}
948 any f = Stream.or . Stream.map f . stream
949
950 and :: Vector v Bool => v Bool -> Bool
951 {-# INLINE and #-}
952 and = Stream.and . stream
953
954 or :: Vector v Bool => v Bool -> Bool
955 {-# INLINE or #-}
956 or = Stream.or . stream
957
958 sum :: (Vector v a, Num a) => v a -> a
959 {-# INLINE sum #-}
960 sum = Stream.foldl' (+) 0 . stream
961
962 product :: (Vector v a, Num a) => v a -> a
963 {-# INLINE product #-}
964 product = Stream.foldl' (*) 1 . stream
965
966 maximum :: (Vector v a, Ord a) => v a -> a
967 {-# INLINE maximum #-}
968 maximum = Stream.foldl1' max . stream
969
970 maximumBy :: Vector v a => (a -> a -> Ordering) -> v a -> a
971 {-# INLINE maximumBy #-}
972 maximumBy cmp = Stream.foldl1' maxBy . stream
973 where
974 {-# INLINE maxBy #-}
975 maxBy x y = case cmp x y of
976 LT -> y
977 _ -> x
978
979 minimum :: (Vector v a, Ord a) => v a -> a
980 {-# INLINE minimum #-}
981 minimum = Stream.foldl1' min . stream
982
983 minimumBy :: Vector v a => (a -> a -> Ordering) -> v a -> a
984 {-# INLINE minimumBy #-}
985 minimumBy cmp = Stream.foldl1' minBy . stream
986 where
987 {-# INLINE minBy #-}
988 minBy x y = case cmp x y of
989 GT -> y
990 _ -> x
991
992 maxIndex :: (Vector v a, Ord a) => v a -> Int
993 {-# INLINE maxIndex #-}
994 maxIndex = maxIndexBy compare
995
996 maxIndexBy :: Vector v a => (a -> a -> Ordering) -> v a -> Int
997 {-# INLINE maxIndexBy #-}
998 maxIndexBy cmp = fst . Stream.foldl1' imax . Stream.indexed . stream
999 where
1000 imax (i,x) (j,y) = case cmp x y of
1001 LT -> (j,y)
1002 _ -> (i,x)
1003
1004 minIndex :: (Vector v a, Ord a) => v a -> Int
1005 {-# INLINE minIndex #-}
1006 minIndex = minIndexBy compare
1007
1008 minIndexBy :: Vector v a => (a -> a -> Ordering) -> v a -> Int
1009 {-# INLINE minIndexBy #-}
1010 minIndexBy cmp = fst . Stream.foldl1' imin . Stream.indexed . stream
1011 where
1012 imin (i,x) (j,y) = case cmp x y of
1013 GT -> (j,y)
1014 _ -> (i,x)
1015
1016
1017 -- Unfolding
1018 -- ---------
1019
1020 unfoldr :: Vector v a => (b -> Maybe (a, b)) -> b -> v a
1021 {-# INLINE unfoldr #-}
1022 unfoldr f = unstream . Stream.unfoldr f
1023
1024 -- Scans
1025 -- -----
1026
1027 -- | Prefix scan
1028 prescanl :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1029 {-# INLINE prescanl #-}
1030 prescanl f z = unstream . inplace (MStream.prescanl f z) . stream
1031
1032 -- | Prefix scan with strict accumulator
1033 prescanl' :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1034 {-# INLINE prescanl' #-}
1035 prescanl' f z = unstream . inplace (MStream.prescanl' f z) . stream
1036
1037 -- | Suffix scan
1038 postscanl :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1039 {-# INLINE postscanl #-}
1040 postscanl f z = unstream . inplace (MStream.postscanl f z) . stream
1041
1042 -- | Suffix scan with strict accumulator
1043 postscanl' :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1044 {-# INLINE postscanl' #-}
1045 postscanl' f z = unstream . inplace (MStream.postscanl' f z) . stream
1046
1047 -- | Haskell-style scan
1048 scanl :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1049 {-# INLINE scanl #-}
1050 scanl f z = unstream . Stream.scanl f z . stream
1051
1052 -- | Haskell-style scan with strict accumulator
1053 scanl' :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1054 {-# INLINE scanl' #-}
1055 scanl' f z = unstream . Stream.scanl' f z . stream
1056
1057 -- | Scan over a non-empty vector
1058 scanl1 :: Vector v a => (a -> a -> a) -> v a -> v a
1059 {-# INLINE scanl1 #-}
1060 scanl1 f = unstream . inplace (MStream.scanl1 f) . stream
1061
1062 -- | Scan over a non-empty vector with a strict accumulator
1063 scanl1' :: Vector v a => (a -> a -> a) -> v a -> v a
1064 {-# INLINE scanl1' #-}
1065 scanl1' f = unstream . inplace (MStream.scanl1' f) . stream
1066
1067 -- Enumeration
1068 -- -----------
1069
1070 enumFromTo :: (Vector v a, Enum a) => a -> a -> v a
1071 {-# INLINE enumFromTo #-}
1072 enumFromTo x y = unstream (Stream.enumFromTo x y)
1073
1074 enumFromThenTo :: (Vector v a, Enum a) => a -> a -> a -> v a
1075 {-# INLINE enumFromThenTo #-}
1076 enumFromThenTo x y z = unstream (Stream.enumFromThenTo x y z)
1077
1078 -- | Convert a vector to a list
1079 toList :: Vector v a => v a -> [a]
1080 {-# INLINE toList #-}
1081 toList = Stream.toList . stream
1082
1083 -- | Convert a list to a vector
1084 fromList :: Vector v a => [a] -> v a
1085 {-# INLINE fromList #-}
1086 fromList = unstream . Stream.fromList
1087