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