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