Add unsafeIndex, unsafeLast etc.
[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 :: v a -> Int -> Int -> 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 => v a -> Int -- ^ starting index
411 -> Int -- ^ length
412 -> v a
413 {-# INLINE_STREAM slice #-}
414 slice v i n = BOUNDS_CHECK(checkSlice) "slice" i n (length v)
415 $ basicUnsafeSlice v i n
416
417 -- | Unsafely yield a part of the vector without copying it and without
418 -- performing bounds checks.
419 unsafeSlice :: Vector v a => v a -> Int -- ^ starting index
420 -> Int -- ^ length
421 -> v a
422 {-# INLINE_STREAM unsafeSlice #-}
423 unsafeSlice v i n = UNSAFE_CHECK(checkSlice) "unsafeSlice" i n (length v)
424 $ basicUnsafeSlice v i n
425
426 -- | Yield all but the last element without copying.
427 init :: Vector v a => v a -> v a
428 {-# INLINE_STREAM init #-}
429 init v = slice v 0 (length v - 1)
430
431 -- | All but the first element (without copying).
432 tail :: Vector v a => v a -> v a
433 {-# INLINE_STREAM tail #-}
434 tail v = slice v 1 (length v - 1)
435
436 -- | Yield the first @n@ elements without copying.
437 take :: Vector v a => Int -> v a -> v a
438 {-# INLINE_STREAM take #-}
439 take n v = slice v 0 (min n' (length v))
440 where n' = max n 0
441
442 -- | Yield all but the first @n@ elements without copying.
443 drop :: Vector v a => Int -> v a -> v a
444 {-# INLINE_STREAM drop #-}
445 drop n v = slice v (min n' len) (max 0 (len - n'))
446 where n' = max n 0
447 len = length v
448
449 {-# RULES
450
451 "slice/new [Vector]" forall v p i n.
452 slice (new' v p) i n = new' v (New.slice p i n)
453
454 "unsafeSlice/new [Vector]" forall v p i n.
455 unsafeSlice (new' v p) i n = new' v (New.unsafeSlice p i n)
456
457 "init/new [Vector]" forall v p.
458 init (new' v p) = new' v (New.init p)
459
460 "tail/new [Vector]" forall v p.
461 tail (new' v p) = new' v (New.tail p)
462
463 "take/new [Vector]" forall n v p.
464 take n (new' v p) = new' v (New.take n p)
465
466 "drop/new [Vector]" forall n v p.
467 drop n (new' v p) = new' v (New.drop n p)
468
469 #-}
470
471 -- Permutations
472 -- ------------
473
474 unsafeAccum_stream
475 :: Vector v a => (a -> b -> a) -> v a -> Stream (Int,b) -> v a
476 {-# INLINE unsafeAccum_stream #-}
477 unsafeAccum_stream f v s = new (New.accum f (New.unstream (stream v)) s)
478
479 unsafeAccum :: Vector v a => (a -> b -> a) -> v a -> [(Int,b)] -> v a
480 {-# INLINE unsafeAccum #-}
481 unsafeAccum f v us = unsafeAccum_stream f v (Stream.fromList us)
482
483 unsafeAccumulate :: (Vector v a, Vector v (Int, b))
484 => (a -> b -> a) -> v a -> v (Int,b) -> v a
485 {-# INLINE unsafeAccumulate #-}
486 unsafeAccumulate f v us = unsafeAccum_stream f v (stream us)
487
488 unsafeAccumulate_ :: (Vector v a, Vector v Int, Vector v b)
489 => (a -> b -> a) -> v a -> v Int -> v b -> v a
490 {-# INLINE unsafeAccumulate_ #-}
491 unsafeAccumulate_ f v is xs
492 = unsafeAccum_stream f v (Stream.zipWith (,) (stream is) (stream xs))
493
494 accum_stream :: Vector v a => (a -> b -> a) -> v a -> Stream (Int,b) -> v a
495 {-# INLINE accum_stream #-}
496 accum_stream f v s = new (New.accum f (New.unstream (stream v)) s)
497
498 accum :: Vector v a => (a -> b -> a) -> v a -> [(Int,b)] -> v a
499 {-# INLINE accum #-}
500 accum f v us = accum_stream f v (Stream.fromList us)
501
502 accumulate :: (Vector v a, Vector v (Int, b))
503 => (a -> b -> a) -> v a -> v (Int,b) -> v a
504 {-# INLINE accumulate #-}
505 accumulate f v us = accum_stream f v (stream us)
506
507 accumulate_ :: (Vector v a, Vector v Int, Vector v b)
508 => (a -> b -> a) -> v a -> v Int -> v b -> v a
509 {-# INLINE accumulate_ #-}
510 accumulate_ f v is xs = accum_stream f v (Stream.zipWith (,) (stream is)
511 (stream xs))
512
513
514 unsafeUpdate_stream :: Vector v a => v a -> Stream (Int,a) -> v a
515 {-# INLINE unsafeUpdate_stream #-}
516 unsafeUpdate_stream v s = new (New.unsafeUpdate (New.unstream (stream v)) s)
517
518 unsafeUpd :: Vector v a => v a -> [(Int, a)] -> v a
519 {-# INLINE unsafeUpd #-}
520 unsafeUpd v us = unsafeUpdate_stream v (Stream.fromList us)
521
522 unsafeUpdate :: (Vector v a, Vector v (Int, a)) => v a -> v (Int, a) -> v a
523 {-# INLINE unsafeUpdate #-}
524 unsafeUpdate v w = unsafeUpdate_stream v (stream w)
525
526 unsafeUpdate_ :: (Vector v a, Vector v Int) => v a -> v Int -> v a -> v a
527 {-# INLINE unsafeUpdate_ #-}
528 unsafeUpdate_ v is w
529 = unsafeUpdate_stream v (Stream.zipWith (,) (stream is) (stream w))
530
531 update_stream :: Vector v a => v a -> Stream (Int,a) -> v a
532 {-# INLINE update_stream #-}
533 update_stream v s = new (New.update (New.unstream (stream v)) s)
534
535 (//) :: Vector v a => v a -> [(Int, a)] -> v a
536 {-# INLINE (//) #-}
537 v // us = update_stream v (Stream.fromList us)
538
539 update :: (Vector v a, Vector v (Int, a)) => v a -> v (Int, a) -> v a
540 {-# INLINE update #-}
541 update v w = update_stream v (stream w)
542
543 update_ :: (Vector v a, Vector v Int) => v a -> v Int -> v a -> v a
544 {-# INLINE update_ #-}
545 update_ v is w = update_stream v (Stream.zipWith (,) (stream is) (stream w))
546
547 -- This somewhat non-intuitive definition ensures that the resulting vector
548 -- does not retain references to the original one even if it is lazy in its
549 -- elements. This would not be the case if we simply used
550 --
551 -- backpermute v is = map (v!) is
552 backpermute :: (Vector v a, Vector v Int) => v a -> v Int -> v a
553 {-# INLINE backpermute #-}
554 backpermute v is = seq v
555 $ unstream
556 $ Stream.unbox
557 $ Stream.map (indexM v)
558 $ stream is
559
560 reverse :: (Vector v a) => v a -> v a
561 {-# INLINE reverse #-}
562 reverse = new . New.reverse . New.unstream . stream
563
564 -- Mapping
565 -- -------
566
567 -- | Map a function over a vector
568 map :: (Vector v a, Vector v b) => (a -> b) -> v a -> v b
569 {-# INLINE map #-}
570 map f = unstream . inplace (MStream.map f) . stream
571
572 -- | Apply a function to every index/value pair
573 imap :: (Vector v a, Vector v b) => (Int -> a -> b) -> v a -> v b
574 {-# INLINE imap #-}
575 imap f = unstream . inplace (MStream.map (uncurry f) . MStream.indexed)
576 . stream
577
578 concatMap :: (Vector v a, Vector v b) => (a -> v b) -> v a -> v b
579 {-# INLINE concatMap #-}
580 concatMap f = unstream . Stream.concatMap (stream . f) . stream
581
582 -- Zipping/unzipping
583 -- -----------------
584
585 -- | Zip two vectors with the given function.
586 zipWith :: (Vector v a, Vector v b, Vector v c)
587 => (a -> b -> c) -> v a -> v b -> v c
588 {-# INLINE zipWith #-}
589 zipWith f xs ys = unstream (Stream.zipWith f (stream xs) (stream ys))
590
591 -- | Zip three vectors with the given function.
592 zipWith3 :: (Vector v a, Vector v b, Vector v c, Vector v d)
593 => (a -> b -> c -> d) -> v a -> v b -> v c -> v d
594 {-# INLINE zipWith3 #-}
595 zipWith3 f as bs cs = unstream (Stream.zipWith3 f (stream as)
596 (stream bs)
597 (stream cs))
598
599 zipWith4 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e)
600 => (a -> b -> c -> d -> e) -> v a -> v b -> v c -> v d -> v e
601 {-# INLINE zipWith4 #-}
602 zipWith4 f as bs cs ds
603 = unstream (Stream.zipWith4 f (stream as)
604 (stream bs)
605 (stream cs)
606 (stream ds))
607
608 zipWith5 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
609 Vector v f)
610 => (a -> b -> c -> d -> e -> f) -> v a -> v b -> v c -> v d -> v e
611 -> v f
612 {-# INLINE zipWith5 #-}
613 zipWith5 f as bs cs ds es
614 = unstream (Stream.zipWith5 f (stream as)
615 (stream bs)
616 (stream cs)
617 (stream ds)
618 (stream es))
619
620 zipWith6 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
621 Vector v f, Vector v g)
622 => (a -> b -> c -> d -> e -> f -> g)
623 -> v a -> v b -> v c -> v d -> v e -> v f -> v g
624 {-# INLINE zipWith6 #-}
625 zipWith6 f as bs cs ds es fs
626 = unstream (Stream.zipWith6 f (stream as)
627 (stream bs)
628 (stream cs)
629 (stream ds)
630 (stream es)
631 (stream fs))
632
633 -- | Zip two vectors and their indices with the given function.
634 izipWith :: (Vector v a, Vector v b, Vector v c)
635 => (Int -> a -> b -> c) -> v a -> v b -> v c
636 {-# INLINE izipWith #-}
637 izipWith f xs ys = unstream
638 (Stream.zipWith (uncurry f) (Stream.indexed (stream xs))
639 (stream ys))
640
641 -- | Zip three vectors and their indices with the given function.
642 izipWith3 :: (Vector v a, Vector v b, Vector v c, Vector v d)
643 => (Int -> a -> b -> c -> d) -> v a -> v b -> v c -> v d
644 {-# INLINE izipWith3 #-}
645 izipWith3 f as bs cs
646 = unstream (Stream.zipWith3 (uncurry f) (Stream.indexed (stream as))
647 (stream bs)
648 (stream cs))
649
650 izipWith4 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e)
651 => (Int -> a -> b -> c -> d -> e) -> v a -> v b -> v c -> v d -> v e
652 {-# INLINE izipWith4 #-}
653 izipWith4 f as bs cs ds
654 = unstream (Stream.zipWith4 (uncurry f) (Stream.indexed (stream as))
655 (stream bs)
656 (stream cs)
657 (stream ds))
658
659 izipWith5 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
660 Vector v f)
661 => (Int -> a -> b -> c -> d -> e -> f) -> v a -> v b -> v c -> v d
662 -> v e -> v f
663 {-# INLINE izipWith5 #-}
664 izipWith5 f as bs cs ds es
665 = unstream (Stream.zipWith5 (uncurry f) (Stream.indexed (stream as))
666 (stream bs)
667 (stream cs)
668 (stream ds)
669 (stream es))
670
671 izipWith6 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
672 Vector v f, Vector v g)
673 => (Int -> a -> b -> c -> d -> e -> f -> g)
674 -> v a -> v b -> v c -> v d -> v e -> v f -> v g
675 {-# INLINE izipWith6 #-}
676 izipWith6 f as bs cs ds es fs
677 = unstream (Stream.zipWith6 (uncurry f) (Stream.indexed (stream as))
678 (stream bs)
679 (stream cs)
680 (stream ds)
681 (stream es)
682 (stream fs))
683
684 zip :: (Vector v a, Vector v b, Vector v (a,b)) => v a -> v b -> v (a, b)
685 {-# INLINE zip #-}
686 zip = zipWith (,)
687
688 zip3 :: (Vector v a, Vector v b, Vector v c, Vector v (a, b, c))
689 => v a -> v b -> v c -> v (a, b, c)
690 {-# INLINE zip3 #-}
691 zip3 = zipWith3 (,,)
692
693 zip4 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v (a, b, c, d))
694 => v a -> v b -> v c -> v d -> v (a, b, c, d)
695 {-# INLINE zip4 #-}
696 zip4 = zipWith4 (,,,)
697
698 zip5 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
699 Vector v (a, b, c, d, e))
700 => v a -> v b -> v c -> v d -> v e -> v (a, b, c, d, e)
701 {-# INLINE zip5 #-}
702 zip5 = zipWith5 (,,,,)
703
704 zip6 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
705 Vector v f, Vector v (a, b, c, d, e, f))
706 => v a -> v b -> v c -> v d -> v e -> v f -> v (a, b, c, d, e, f)
707 {-# INLINE zip6 #-}
708 zip6 = zipWith6 (,,,,,)
709
710 unzip :: (Vector v a, Vector v b, Vector v (a,b)) => v (a, b) -> (v a, v b)
711 {-# INLINE unzip #-}
712 unzip xs = (map fst xs, map snd xs)
713
714 unzip3 :: (Vector v a, Vector v b, Vector v c, Vector v (a, b, c))
715 => v (a, b, c) -> (v a, v b, v c)
716 {-# INLINE unzip3 #-}
717 unzip3 xs = (map (\(a, b, c) -> a) xs,
718 map (\(a, b, c) -> b) xs,
719 map (\(a, b, c) -> c) xs)
720
721 unzip4 :: (Vector v a, Vector v b, Vector v c, Vector v d,
722 Vector v (a, b, c, d))
723 => v (a, b, c, d) -> (v a, v b, v c, v d)
724 {-# INLINE unzip4 #-}
725 unzip4 xs = (map (\(a, b, c, d) -> a) xs,
726 map (\(a, b, c, d) -> b) xs,
727 map (\(a, b, c, d) -> c) xs,
728 map (\(a, b, c, d) -> d) xs)
729
730 unzip5 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
731 Vector v (a, b, c, d, e))
732 => v (a, b, c, d, e) -> (v a, v b, v c, v d, v e)
733 {-# INLINE unzip5 #-}
734 unzip5 xs = (map (\(a, b, c, d, e) -> a) xs,
735 map (\(a, b, c, d, e) -> b) xs,
736 map (\(a, b, c, d, e) -> c) xs,
737 map (\(a, b, c, d, e) -> d) xs,
738 map (\(a, b, c, d, e) -> e) xs)
739
740 unzip6 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
741 Vector v f, Vector v (a, b, c, d, e, f))
742 => v (a, b, c, d, e, f) -> (v a, v b, v c, v d, v e, v f)
743 {-# INLINE unzip6 #-}
744 unzip6 xs = (map (\(a, b, c, d, e, f) -> a) xs,
745 map (\(a, b, c, d, e, f) -> b) xs,
746 map (\(a, b, c, d, e, f) -> c) xs,
747 map (\(a, b, c, d, e, f) -> d) xs,
748 map (\(a, b, c, d, e, f) -> e) xs,
749 map (\(a, b, c, d, e, f) -> f) xs)
750
751 -- Comparisons
752 -- -----------
753
754 eq :: (Vector v a, Eq a) => v a -> v a -> Bool
755 {-# INLINE eq #-}
756 xs `eq` ys = stream xs == stream ys
757
758 cmp :: (Vector v a, Ord a) => v a -> v a -> Ordering
759 {-# INLINE cmp #-}
760 cmp xs ys = compare (stream xs) (stream ys)
761
762 -- Filtering
763 -- ---------
764
765 -- | Drop elements that do not satisfy the predicate
766 filter :: Vector v a => (a -> Bool) -> v a -> v a
767 {-# INLINE filter #-}
768 filter f = unstream . inplace (MStream.filter f) . stream
769
770 -- | Drop elements that do not satisfy the predicate (applied to values and
771 -- their indices)
772 ifilter :: Vector v a => (Int -> a -> Bool) -> v a -> v a
773 {-# INLINE ifilter #-}
774 ifilter f = unstream
775 . inplace (MStream.map snd . MStream.filter (uncurry f)
776 . MStream.indexed)
777 . stream
778
779 -- | Yield the longest prefix of elements satisfying the predicate.
780 takeWhile :: Vector v a => (a -> Bool) -> v a -> v a
781 {-# INLINE takeWhile #-}
782 takeWhile f = unstream . Stream.takeWhile f . stream
783
784 -- | Drop the longest prefix of elements that satisfy the predicate.
785 dropWhile :: Vector v a => (a -> Bool) -> v a -> v a
786 {-# INLINE dropWhile #-}
787 dropWhile f = unstream . Stream.dropWhile f . stream
788
789 -- | Split the vector in two parts, the first one containing those elements
790 -- that satisfy the predicate and the second one those that don't. The order
791 -- of the elements is not preserved.
792 unstablePartition :: Vector v a => (a -> Bool) -> v a -> (v a, v a)
793 {-# INLINE unstablePartition #-}
794 unstablePartition f = unstablePartition_stream f . stream
795
796 unstablePartition_stream
797 :: Vector v a => (a -> Bool) -> Stream a -> (v a, v a)
798 {-# INLINE_STREAM unstablePartition_stream #-}
799 unstablePartition_stream f s = s `seq` runST (
800 do
801 (mv1,mv2) <- M.unstablePartitionStream f s
802 v1 <- unsafeFreeze mv1
803 v2 <- unsafeFreeze mv2
804 return (v1,v2))
805
806 unstablePartition_new :: Vector v a => (a -> Bool) -> New a -> (v a, v a)
807 {-# INLINE_STREAM unstablePartition_new #-}
808 unstablePartition_new f (New.New p) = runST (
809 do
810 mv <- p
811 i <- M.unstablePartition f mv
812 v <- unsafeFreeze mv
813 return (take i v, drop i v))
814
815 {-# RULES
816
817 "unstablePartition" forall f v p.
818 unstablePartition_stream f (stream (new' v p))
819 = unstablePartition_new f p
820
821 #-}
822
823
824 -- FIXME: make span and break fusible
825
826 -- | Split the vector into the longest prefix of elements that satisfy the
827 -- predicate and the rest.
828 span :: Vector v a => (a -> Bool) -> v a -> (v a, v a)
829 {-# INLINE span #-}
830 span f = break (not . f)
831
832 -- | Split the vector into the longest prefix of elements that do not satisfy
833 -- the predicate and the rest.
834 break :: Vector v a => (a -> Bool) -> v a -> (v a, v a)
835 {-# INLINE break #-}
836 break f xs = case findIndex f xs of
837 Just i -> (unsafeSlice xs 0 i, unsafeSlice xs i (length xs - i))
838 Nothing -> (xs, empty)
839
840
841 -- Searching
842 -- ---------
843
844 infix 4 `elem`
845 -- | Check whether the vector contains an element
846 elem :: (Vector v a, Eq a) => a -> v a -> Bool
847 {-# INLINE elem #-}
848 elem x = Stream.elem x . stream
849
850 infix 4 `notElem`
851 -- | Inverse of `elem`
852 notElem :: (Vector v a, Eq a) => a -> v a -> Bool
853 {-# INLINE notElem #-}
854 notElem x = Stream.notElem x . stream
855
856 -- | Yield 'Just' the first element matching the predicate or 'Nothing' if no
857 -- such element exists.
858 find :: Vector v a => (a -> Bool) -> v a -> Maybe a
859 {-# INLINE find #-}
860 find f = Stream.find f . stream
861
862 -- | Yield 'Just' the index of the first element matching the predicate or
863 -- 'Nothing' if no such element exists.
864 findIndex :: Vector v a => (a -> Bool) -> v a -> Maybe Int
865 {-# INLINE findIndex #-}
866 findIndex f = Stream.findIndex f . stream
867
868 -- | Yield the indices of elements satisfying the predicate
869 findIndices :: (Vector v a, Vector v Int) => (a -> Bool) -> v a -> v Int
870 {-# INLINE findIndices #-}
871 findIndices f = unstream
872 . Stream.map fst
873 . Stream.filter (f . snd)
874 . Stream.indexed
875 . stream
876
877 -- | Yield 'Just' the index of the first occurence of the given element or
878 -- 'Nothing' if the vector does not contain the element
879 elemIndex :: (Vector v a, Eq a) => a -> v a -> Maybe Int
880 {-# INLINE elemIndex #-}
881 elemIndex x = findIndex (x==)
882
883 -- | Yield the indices of all occurences of the given element
884 elemIndices :: (Vector v a, Vector v Int, Eq a) => a -> v a -> v Int
885 {-# INLINE elemIndices #-}
886 elemIndices x = findIndices (x==)
887
888 -- Folding
889 -- -------
890
891 -- | Left fold
892 foldl :: Vector v b => (a -> b -> a) -> a -> v b -> a
893 {-# INLINE foldl #-}
894 foldl f z = Stream.foldl f z . stream
895
896 -- | Left fold on non-empty vectors
897 foldl1 :: Vector v a => (a -> a -> a) -> v a -> a
898 {-# INLINE foldl1 #-}
899 foldl1 f = Stream.foldl1 f . stream
900
901 -- | Left fold with strict accumulator
902 foldl' :: Vector v b => (a -> b -> a) -> a -> v b -> a
903 {-# INLINE foldl' #-}
904 foldl' f z = Stream.foldl' f z . stream
905
906 -- | Left fold on non-empty vectors with strict accumulator
907 foldl1' :: Vector v a => (a -> a -> a) -> v a -> a
908 {-# INLINE foldl1' #-}
909 foldl1' f = Stream.foldl1' f . stream
910
911 -- | Right fold
912 foldr :: Vector v a => (a -> b -> b) -> b -> v a -> b
913 {-# INLINE foldr #-}
914 foldr f z = Stream.foldr f z . stream
915
916 -- | Right fold on non-empty vectors
917 foldr1 :: Vector v a => (a -> a -> a) -> v a -> a
918 {-# INLINE foldr1 #-}
919 foldr1 f = Stream.foldr1 f . stream
920
921 -- | Left fold (function applied to each element and its index)
922 ifoldl :: Vector v b => (a -> Int -> b -> a) -> a -> v b -> a
923 {-# INLINE ifoldl #-}
924 ifoldl f z = Stream.foldl (uncurry . f) z . Stream.indexed . stream
925
926 -- | Left fold with strict accumulator (function applied to each element and
927 -- its index)
928 ifoldl' :: Vector v b => (a -> Int -> b -> a) -> a -> v b -> a
929 {-# INLINE ifoldl' #-}
930 ifoldl' f z = Stream.foldl' (uncurry . f) z . Stream.indexed . stream
931
932 -- | Right fold (function applied to each element and its index)
933 ifoldr :: Vector v a => (Int -> a -> b -> b) -> b -> v a -> b
934 {-# INLINE ifoldr #-}
935 ifoldr f z = Stream.foldr (uncurry f) z . Stream.indexed . stream
936
937 -- Specialised folds
938 -- -----------------
939
940 all :: Vector v a => (a -> Bool) -> v a -> Bool
941 {-# INLINE all #-}
942 all f = Stream.and . Stream.map f . stream
943
944 any :: Vector v a => (a -> Bool) -> v a -> Bool
945 {-# INLINE any #-}
946 any f = Stream.or . Stream.map f . stream
947
948 and :: Vector v Bool => v Bool -> Bool
949 {-# INLINE and #-}
950 and = Stream.and . stream
951
952 or :: Vector v Bool => v Bool -> Bool
953 {-# INLINE or #-}
954 or = Stream.or . stream
955
956 sum :: (Vector v a, Num a) => v a -> a
957 {-# INLINE sum #-}
958 sum = Stream.foldl' (+) 0 . stream
959
960 product :: (Vector v a, Num a) => v a -> a
961 {-# INLINE product #-}
962 product = Stream.foldl' (*) 1 . stream
963
964 maximum :: (Vector v a, Ord a) => v a -> a
965 {-# INLINE maximum #-}
966 maximum = Stream.foldl1' max . stream
967
968 maximumBy :: Vector v a => (a -> a -> Ordering) -> v a -> a
969 {-# INLINE maximumBy #-}
970 maximumBy cmp = Stream.foldl1' maxBy . stream
971 where
972 {-# INLINE maxBy #-}
973 maxBy x y = case cmp x y of
974 LT -> y
975 _ -> x
976
977 minimum :: (Vector v a, Ord a) => v a -> a
978 {-# INLINE minimum #-}
979 minimum = Stream.foldl1' min . stream
980
981 minimumBy :: Vector v a => (a -> a -> Ordering) -> v a -> a
982 {-# INLINE minimumBy #-}
983 minimumBy cmp = Stream.foldl1' minBy . stream
984 where
985 {-# INLINE minBy #-}
986 minBy x y = case cmp x y of
987 GT -> y
988 _ -> x
989
990 maxIndex :: (Vector v a, Ord a) => v a -> Int
991 {-# INLINE maxIndex #-}
992 maxIndex = maxIndexBy compare
993
994 maxIndexBy :: Vector v a => (a -> a -> Ordering) -> v a -> Int
995 {-# INLINE maxIndexBy #-}
996 maxIndexBy cmp = fst . Stream.foldl1' imax . Stream.indexed . stream
997 where
998 imax (i,x) (j,y) = case cmp x y of
999 LT -> (j,y)
1000 _ -> (i,x)
1001
1002 minIndex :: (Vector v a, Ord a) => v a -> Int
1003 {-# INLINE minIndex #-}
1004 minIndex = minIndexBy compare
1005
1006 minIndexBy :: Vector v a => (a -> a -> Ordering) -> v a -> Int
1007 {-# INLINE minIndexBy #-}
1008 minIndexBy cmp = fst . Stream.foldl1' imin . Stream.indexed . stream
1009 where
1010 imin (i,x) (j,y) = case cmp x y of
1011 GT -> (j,y)
1012 _ -> (i,x)
1013
1014
1015 -- Unfolding
1016 -- ---------
1017
1018 unfoldr :: Vector v a => (b -> Maybe (a, b)) -> b -> v a
1019 {-# INLINE unfoldr #-}
1020 unfoldr f = unstream . Stream.unfoldr f
1021
1022 -- Scans
1023 -- -----
1024
1025 -- | Prefix scan
1026 prescanl :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1027 {-# INLINE prescanl #-}
1028 prescanl f z = unstream . inplace (MStream.prescanl f z) . stream
1029
1030 -- | Prefix scan with strict accumulator
1031 prescanl' :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1032 {-# INLINE prescanl' #-}
1033 prescanl' f z = unstream . inplace (MStream.prescanl' f z) . stream
1034
1035 -- | Suffix scan
1036 postscanl :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1037 {-# INLINE postscanl #-}
1038 postscanl f z = unstream . inplace (MStream.postscanl f z) . stream
1039
1040 -- | Suffix scan with strict accumulator
1041 postscanl' :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1042 {-# INLINE postscanl' #-}
1043 postscanl' f z = unstream . inplace (MStream.postscanl' f z) . stream
1044
1045 -- | Haskell-style scan
1046 scanl :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1047 {-# INLINE scanl #-}
1048 scanl f z = unstream . Stream.scanl f z . stream
1049
1050 -- | Haskell-style scan with strict accumulator
1051 scanl' :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1052 {-# INLINE scanl' #-}
1053 scanl' f z = unstream . Stream.scanl' f z . stream
1054
1055 -- | Scan over a non-empty vector
1056 scanl1 :: Vector v a => (a -> a -> a) -> v a -> v a
1057 {-# INLINE scanl1 #-}
1058 scanl1 f = unstream . inplace (MStream.scanl1 f) . stream
1059
1060 -- | Scan over a non-empty vector with a strict accumulator
1061 scanl1' :: Vector v a => (a -> a -> a) -> v a -> v a
1062 {-# INLINE scanl1' #-}
1063 scanl1' f = unstream . inplace (MStream.scanl1' f) . stream
1064
1065 -- Enumeration
1066 -- -----------
1067
1068 enumFromTo :: (Vector v a, Enum a) => a -> a -> v a
1069 {-# INLINE enumFromTo #-}
1070 enumFromTo x y = unstream (Stream.enumFromTo x y)
1071
1072 enumFromThenTo :: (Vector v a, Enum a) => a -> a -> a -> v a
1073 {-# INLINE enumFromThenTo #-}
1074 enumFromThenTo x y z = unstream (Stream.enumFromThenTo x y z)
1075
1076 -- | Convert a vector to a list
1077 toList :: Vector v a => v a -> [a]
1078 {-# INLINE toList #-}
1079 toList = Stream.toList . stream
1080
1081 -- | Convert a list to a vector
1082 fromList :: Vector v a => [a] -> v a
1083 {-# INLINE fromList #-}
1084 fromList = unstream . Stream.fromList
1085