Add findIndices, elemIndex, elemIndices
[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, 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 :: forall v a. Vector v a => a -> v a
246 {-# INLINE singleton #-}
247 singleton x = elemseq (undefined :: v a) x
248 $ unstream (Stream.singleton x)
249
250 -- | Vector of the given length with the given value in each position
251 replicate :: forall v a. Vector v a => Int -> a -> v a
252 {-# INLINE replicate #-}
253 replicate n x = elemseq (undefined :: v a) x
254 $ unstream
255 $ Stream.replicate n x
256
257 -- | Prepend an element
258 cons :: forall v a. Vector v a => a -> v a -> v a
259 {-# INLINE cons #-}
260 cons x v = elemseq (undefined :: v a) x
261 $ unstream
262 $ Stream.cons x
263 $ stream v
264
265 -- | Append an element
266 snoc :: forall v a. Vector v a => v a -> a -> v a
267 {-# INLINE snoc #-}
268 snoc v x = elemseq (undefined :: v a) x
269 $ unstream
270 $ Stream.snoc (stream v) x
271
272 infixr 5 ++
273 -- | Concatenate two vectors
274 (++) :: Vector v a => v a -> v a -> v a
275 {-# INLINE (++) #-}
276 v ++ w = unstream (stream v Stream.++ stream w)
277
278 -- | Create a copy of a vector. Useful when dealing with slices.
279 copy :: Vector v a => v a -> v a
280 {-# INLINE_STREAM copy #-}
281 copy = unstream . stream
282
283 {-# RULES
284
285 "copy/unstream [Vector]" forall v s.
286 copy (new' v (New.unstream s)) = new' v (New.unstream s)
287
288 #-}
289
290 -- Accessing individual elements
291 -- -----------------------------
292
293 -- | Indexing
294 (!) :: Vector v a => v a -> Int -> a
295 {-# INLINE_STREAM (!) #-}
296 v ! i = BOUNDS_CHECK(checkIndex) "(!)" i (length v)
297 $ unId (basicUnsafeIndexM v i)
298
299 -- | Unsafe indexing without bounds checking
300 unsafeIndex :: Vector v a => v a -> Int -> a
301 {-# INLINE_STREAM unsafeIndex #-}
302 unsafeIndex v i = UNSAFE_CHECK(checkIndex) "unsafeIndex" i (length v)
303 $ unId (basicUnsafeIndexM v i)
304
305 -- | First element
306 head :: Vector v a => v a -> a
307 {-# INLINE_STREAM head #-}
308 head v = v ! 0
309
310 -- | Last element
311 last :: Vector v a => v a -> a
312 {-# INLINE_STREAM last #-}
313 last v = v ! (length v - 1)
314
315 {-# RULES
316
317 "(!)/unstream [Vector]" forall v i s.
318 new' v (New.unstream s) ! i = s Stream.!! i
319
320 "unsafeIndex/unstream [Vector]" forall v i s.
321 unsafeIndex (new' v (New.unstream s)) i = s Stream.!! i
322
323 "head/unstream [Vector]" forall v s.
324 head (new' v (New.unstream s)) = Stream.head s
325
326 "last/unstream [Vector]" forall v s.
327 last (new' v (New.unstream s)) = Stream.last s
328
329 #-}
330
331 -- | Monadic indexing which can be strict in the vector while remaining lazy in
332 -- the element.
333 indexM :: (Vector v a, Monad m) => v a -> Int -> m a
334 {-# INLINE_STREAM indexM #-}
335 indexM v i = BOUNDS_CHECK(checkIndex) "indexM" i (length v)
336 $ basicUnsafeIndexM v i
337
338 -- | Unsafe monadic indexing without bounds checks
339 unsafeIndexM :: (Vector v a, Monad m) => v a -> Int -> m a
340 {-# INLINE_STREAM unsafeIndexM #-}
341 unsafeIndexM v i = UNSAFE_CHECK(checkIndex) "unsafeIndexM" i (length v)
342 $ basicUnsafeIndexM v i
343
344 headM :: (Vector v a, Monad m) => v a -> m a
345 {-# INLINE_STREAM headM #-}
346 headM v = indexM v 0
347
348 lastM :: (Vector v a, Monad m) => v a -> m a
349 {-# INLINE_STREAM lastM #-}
350 lastM v = indexM v (length v - 1)
351
352 -- FIXME: the rhs of these rules are lazy in the stream which is WRONG
353 {- RULES
354
355 "indexM/unstream [Vector]" forall v i s.
356 indexM (new' v (New.unstream s)) i = return (s Stream.!! i)
357
358 "headM/unstream [Vector]" forall v s.
359 headM (new' v (New.unstream s)) = return (Stream.head s)
360
361 "lastM/unstream [Vector]" forall v s.
362 lastM (new' v (New.unstream s)) = return (Stream.last s)
363
364 -}
365
366 -- Subarrays
367 -- ---------
368
369 -- FIXME: slicing doesn't work with the inplace stuff at the moment
370
371 -- | Yield a part of the vector without copying it.
372 slice :: Vector v a => v a -> Int -- ^ starting index
373 -> Int -- ^ length
374 -> v a
375 {-# INLINE_STREAM slice #-}
376 slice v i n = BOUNDS_CHECK(checkSlice) "slice" i n (length v)
377 $ basicUnsafeSlice v i n
378
379 -- | Unsafely yield a part of the vector without copying it and without
380 -- performing bounds checks.
381 unsafeSlice :: Vector v a => v a -> Int -- ^ starting index
382 -> Int -- ^ length
383 -> v a
384 {-# INLINE_STREAM unsafeSlice #-}
385 unsafeSlice v i n = UNSAFE_CHECK(checkSlice) "unsafeSlice" i n (length v)
386 $ basicUnsafeSlice v i n
387
388 -- | Yield all but the last element without copying.
389 init :: Vector v a => v a -> v a
390 {-# INLINE_STREAM init #-}
391 init v = slice v 0 (length v - 1)
392
393 -- | All but the first element (without copying).
394 tail :: Vector v a => v a -> v a
395 {-# INLINE_STREAM tail #-}
396 tail v = slice v 1 (length v - 1)
397
398 -- | Yield the first @n@ elements without copying.
399 take :: Vector v a => Int -> v a -> v a
400 {-# INLINE_STREAM take #-}
401 take n v = slice v 0 (min n' (length v))
402 where n' = max n 0
403
404 -- | Yield all but the first @n@ elements without copying.
405 drop :: Vector v a => Int -> v a -> v a
406 {-# INLINE_STREAM drop #-}
407 drop n v = slice v (min n' len) (max 0 (len - n'))
408 where n' = max n 0
409 len = length v
410
411 {-# RULES
412
413 "slice/new [Vector]" forall v p i n.
414 slice (new' v p) i n = new' v (New.slice p i n)
415
416 "unsafeSlice/new [Vector]" forall v p i n.
417 unsafeSlice (new' v p) i n = new' v (New.unsafeSlice p i n)
418
419 "init/new [Vector]" forall v p.
420 init (new' v p) = new' v (New.init p)
421
422 "tail/new [Vector]" forall v p.
423 tail (new' v p) = new' v (New.tail p)
424
425 "take/new [Vector]" forall n v p.
426 take n (new' v p) = new' v (New.take n p)
427
428 "drop/new [Vector]" forall n v p.
429 drop n (new' v p) = new' v (New.drop n p)
430
431 #-}
432
433 -- Permutations
434 -- ------------
435
436 unsafeAccum_stream
437 :: Vector v a => (a -> b -> a) -> v a -> Stream (Int,b) -> v a
438 {-# INLINE unsafeAccum_stream #-}
439 unsafeAccum_stream f v s = new (New.accum f (New.unstream (stream v)) s)
440
441 unsafeAccum :: Vector v a => (a -> b -> a) -> v a -> [(Int,b)] -> v a
442 {-# INLINE unsafeAccum #-}
443 unsafeAccum f v us = unsafeAccum_stream f v (Stream.fromList us)
444
445 unsafeAccumulate :: (Vector v a, Vector v (Int, b))
446 => (a -> b -> a) -> v a -> v (Int,b) -> v a
447 {-# INLINE unsafeAccumulate #-}
448 unsafeAccumulate f v us = unsafeAccum_stream f v (stream us)
449
450 unsafeAccumulate_ :: (Vector v a, Vector v Int, Vector v b)
451 => (a -> b -> a) -> v a -> v Int -> v b -> v a
452 {-# INLINE unsafeAccumulate_ #-}
453 unsafeAccumulate_ f v is xs
454 = unsafeAccum_stream f v (Stream.zipWith (,) (stream is) (stream xs))
455
456 accum_stream :: Vector v a => (a -> b -> a) -> v a -> Stream (Int,b) -> v a
457 {-# INLINE accum_stream #-}
458 accum_stream f v s = new (New.accum f (New.unstream (stream v)) s)
459
460 accum :: Vector v a => (a -> b -> a) -> v a -> [(Int,b)] -> v a
461 {-# INLINE accum #-}
462 accum f v us = accum_stream f v (Stream.fromList us)
463
464 accumulate :: (Vector v a, Vector v (Int, b))
465 => (a -> b -> a) -> v a -> v (Int,b) -> v a
466 {-# INLINE accumulate #-}
467 accumulate f v us = accum_stream f v (stream us)
468
469 accumulate_ :: (Vector v a, Vector v Int, Vector v b)
470 => (a -> b -> a) -> v a -> v Int -> v b -> v a
471 {-# INLINE accumulate_ #-}
472 accumulate_ f v is xs = accum_stream f v (Stream.zipWith (,) (stream is)
473 (stream xs))
474
475
476 unsafeUpdate_stream :: Vector v a => v a -> Stream (Int,a) -> v a
477 {-# INLINE unsafeUpdate_stream #-}
478 unsafeUpdate_stream v s = new (New.unsafeUpdate (New.unstream (stream v)) s)
479
480 unsafeUpd :: Vector v a => v a -> [(Int, a)] -> v a
481 {-# INLINE unsafeUpd #-}
482 unsafeUpd v us = unsafeUpdate_stream v (Stream.fromList us)
483
484 unsafeUpdate :: (Vector v a, Vector v (Int, a)) => v a -> v (Int, a) -> v a
485 {-# INLINE unsafeUpdate #-}
486 unsafeUpdate v w = unsafeUpdate_stream v (stream w)
487
488 unsafeUpdate_ :: (Vector v a, Vector v Int) => v a -> v Int -> v a -> v a
489 {-# INLINE unsafeUpdate_ #-}
490 unsafeUpdate_ v is w
491 = unsafeUpdate_stream v (Stream.zipWith (,) (stream is) (stream w))
492
493 update_stream :: Vector v a => v a -> Stream (Int,a) -> v a
494 {-# INLINE update_stream #-}
495 update_stream v s = new (New.update (New.unstream (stream v)) s)
496
497 (//) :: Vector v a => v a -> [(Int, a)] -> v a
498 {-# INLINE (//) #-}
499 v // us = update_stream v (Stream.fromList us)
500
501 update :: (Vector v a, Vector v (Int, a)) => v a -> v (Int, a) -> v a
502 {-# INLINE update #-}
503 update v w = update_stream v (stream w)
504
505 update_ :: (Vector v a, Vector v Int) => v a -> v Int -> v a -> v a
506 {-# INLINE update_ #-}
507 update_ v is w = update_stream v (Stream.zipWith (,) (stream is) (stream w))
508
509 -- This somewhat non-intuitive definition ensures that the resulting vector
510 -- does not retain references to the original one even if it is lazy in its
511 -- elements. This would not be the case if we simply used
512 --
513 -- backpermute v is = map (v!) is
514 backpermute :: (Vector v a, Vector v Int) => v a -> v Int -> v a
515 {-# INLINE backpermute #-}
516 backpermute v is = seq v
517 $ unstream
518 $ Stream.unbox
519 $ Stream.map (indexM v)
520 $ stream is
521
522 reverse :: (Vector v a) => v a -> v a
523 {-# INLINE reverse #-}
524 reverse = new . New.reverse . New.unstream . stream
525
526 -- Mapping
527 -- -------
528
529 -- | Map a function over a vector
530 map :: (Vector v a, Vector v b) => (a -> b) -> v a -> v b
531 {-# INLINE map #-}
532 map f = unstream . inplace (MStream.map f) . stream
533
534 -- | Apply a function to every index/value pair
535 imap :: (Vector v a, Vector v b) => (Int -> a -> b) -> v a -> v b
536 {-# INLINE imap #-}
537 imap f = unstream . inplace (MStream.map (uncurry f) . MStream.indexed)
538 . stream
539
540 concatMap :: (Vector v a, Vector v b) => (a -> v b) -> v a -> v b
541 {-# INLINE concatMap #-}
542 concatMap f = unstream . Stream.concatMap (stream . f) . stream
543
544 -- Zipping/unzipping
545 -- -----------------
546
547 -- | Zip two vectors with the given function.
548 zipWith :: (Vector v a, Vector v b, Vector v c)
549 => (a -> b -> c) -> v a -> v b -> v c
550 {-# INLINE zipWith #-}
551 zipWith f xs ys = unstream (Stream.zipWith f (stream xs) (stream ys))
552
553 -- | Zip three vectors with the given function.
554 zipWith3 :: (Vector v a, Vector v b, Vector v c, Vector v d)
555 => (a -> b -> c -> d) -> v a -> v b -> v c -> v d
556 {-# INLINE zipWith3 #-}
557 zipWith3 f as bs cs = unstream (Stream.zipWith3 f (stream as)
558 (stream bs)
559 (stream cs))
560
561 zipWith4 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e)
562 => (a -> b -> c -> d -> e) -> v a -> v b -> v c -> v d -> v e
563 {-# INLINE zipWith4 #-}
564 zipWith4 f as bs cs ds
565 = unstream (Stream.zipWith4 f (stream as)
566 (stream bs)
567 (stream cs)
568 (stream ds))
569
570 zipWith5 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
571 Vector v f)
572 => (a -> b -> c -> d -> e -> f) -> v a -> v b -> v c -> v d -> v e
573 -> v f
574 {-# INLINE zipWith5 #-}
575 zipWith5 f as bs cs ds es
576 = unstream (Stream.zipWith5 f (stream as)
577 (stream bs)
578 (stream cs)
579 (stream ds)
580 (stream es))
581
582 zipWith6 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
583 Vector v f, Vector v g)
584 => (a -> b -> c -> d -> e -> f -> g)
585 -> v a -> v b -> v c -> v d -> v e -> v f -> v g
586 {-# INLINE zipWith6 #-}
587 zipWith6 f as bs cs ds es fs
588 = unstream (Stream.zipWith6 f (stream as)
589 (stream bs)
590 (stream cs)
591 (stream ds)
592 (stream es)
593 (stream fs))
594
595 -- | Zip two vectors and their indices with the given function.
596 izipWith :: (Vector v a, Vector v b, Vector v c)
597 => (Int -> a -> b -> c) -> v a -> v b -> v c
598 {-# INLINE izipWith #-}
599 izipWith f xs ys = unstream
600 (Stream.zipWith (uncurry f) (Stream.indexed (stream xs))
601 (stream ys))
602
603 -- | Zip three vectors and their indices with the given function.
604 izipWith3 :: (Vector v a, Vector v b, Vector v c, Vector v d)
605 => (Int -> a -> b -> c -> d) -> v a -> v b -> v c -> v d
606 {-# INLINE izipWith3 #-}
607 izipWith3 f as bs cs
608 = unstream (Stream.zipWith3 (uncurry f) (Stream.indexed (stream as))
609 (stream bs)
610 (stream cs))
611
612 izipWith4 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e)
613 => (Int -> a -> b -> c -> d -> e) -> v a -> v b -> v c -> v d -> v e
614 {-# INLINE izipWith4 #-}
615 izipWith4 f as bs cs ds
616 = unstream (Stream.zipWith4 (uncurry f) (Stream.indexed (stream as))
617 (stream bs)
618 (stream cs)
619 (stream ds))
620
621 izipWith5 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
622 Vector v f)
623 => (Int -> a -> b -> c -> d -> e -> f) -> v a -> v b -> v c -> v d
624 -> v e -> v f
625 {-# INLINE izipWith5 #-}
626 izipWith5 f as bs cs ds es
627 = unstream (Stream.zipWith5 (uncurry f) (Stream.indexed (stream as))
628 (stream bs)
629 (stream cs)
630 (stream ds)
631 (stream es))
632
633 izipWith6 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
634 Vector v f, Vector v g)
635 => (Int -> a -> b -> c -> d -> e -> f -> g)
636 -> v a -> v b -> v c -> v d -> v e -> v f -> v g
637 {-# INLINE izipWith6 #-}
638 izipWith6 f as bs cs ds es fs
639 = unstream (Stream.zipWith6 (uncurry f) (Stream.indexed (stream as))
640 (stream bs)
641 (stream cs)
642 (stream ds)
643 (stream es)
644 (stream fs))
645
646 zip :: (Vector v a, Vector v b, Vector v (a,b)) => v a -> v b -> v (a, b)
647 {-# INLINE zip #-}
648 zip = zipWith (,)
649
650 zip3 :: (Vector v a, Vector v b, Vector v c, Vector v (a, b, c))
651 => v a -> v b -> v c -> v (a, b, c)
652 {-# INLINE zip3 #-}
653 zip3 = zipWith3 (,,)
654
655 zip4 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v (a, b, c, d))
656 => v a -> v b -> v c -> v d -> v (a, b, c, d)
657 {-# INLINE zip4 #-}
658 zip4 = zipWith4 (,,,)
659
660 zip5 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
661 Vector v (a, b, c, d, e))
662 => v a -> v b -> v c -> v d -> v e -> v (a, b, c, d, e)
663 {-# INLINE zip5 #-}
664 zip5 = zipWith5 (,,,,)
665
666 zip6 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
667 Vector v f, Vector v (a, b, c, d, e, f))
668 => v a -> v b -> v c -> v d -> v e -> v f -> v (a, b, c, d, e, f)
669 {-# INLINE zip6 #-}
670 zip6 = zipWith6 (,,,,,)
671
672 unzip :: (Vector v a, Vector v b, Vector v (a,b)) => v (a, b) -> (v a, v b)
673 {-# INLINE unzip #-}
674 unzip xs = (map fst xs, map snd xs)
675
676 unzip3 :: (Vector v a, Vector v b, Vector v c, Vector v (a, b, c))
677 => v (a, b, c) -> (v a, v b, v c)
678 {-# INLINE unzip3 #-}
679 unzip3 xs = (map (\(a, b, c) -> a) xs,
680 map (\(a, b, c) -> b) xs,
681 map (\(a, b, c) -> c) xs)
682
683 unzip4 :: (Vector v a, Vector v b, Vector v c, Vector v d,
684 Vector v (a, b, c, d))
685 => v (a, b, c, d) -> (v a, v b, v c, v d)
686 {-# INLINE unzip4 #-}
687 unzip4 xs = (map (\(a, b, c, d) -> a) xs,
688 map (\(a, b, c, d) -> b) xs,
689 map (\(a, b, c, d) -> c) xs,
690 map (\(a, b, c, d) -> d) xs)
691
692 unzip5 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
693 Vector v (a, b, c, d, e))
694 => v (a, b, c, d, e) -> (v a, v b, v c, v d, v e)
695 {-# INLINE unzip5 #-}
696 unzip5 xs = (map (\(a, b, c, d, e) -> a) xs,
697 map (\(a, b, c, d, e) -> b) xs,
698 map (\(a, b, c, d, e) -> c) xs,
699 map (\(a, b, c, d, e) -> d) xs,
700 map (\(a, b, c, d, e) -> e) xs)
701
702 unzip6 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
703 Vector v f, Vector v (a, b, c, d, e, f))
704 => v (a, b, c, d, e, f) -> (v a, v b, v c, v d, v e, v f)
705 {-# INLINE unzip6 #-}
706 unzip6 xs = (map (\(a, b, c, d, e, f) -> a) xs,
707 map (\(a, b, c, d, e, f) -> b) xs,
708 map (\(a, b, c, d, e, f) -> c) xs,
709 map (\(a, b, c, d, e, f) -> d) xs,
710 map (\(a, b, c, d, e, f) -> e) xs,
711 map (\(a, b, c, d, e, f) -> f) xs)
712
713 -- Comparisons
714 -- -----------
715
716 eq :: (Vector v a, Eq a) => v a -> v a -> Bool
717 {-# INLINE eq #-}
718 xs `eq` ys = stream xs == stream ys
719
720 cmp :: (Vector v a, Ord a) => v a -> v a -> Ordering
721 {-# INLINE cmp #-}
722 cmp xs ys = compare (stream xs) (stream ys)
723
724 -- Filtering
725 -- ---------
726
727 -- | Drop elements that do not satisfy the predicate
728 filter :: Vector v a => (a -> Bool) -> v a -> v a
729 {-# INLINE filter #-}
730 filter f = unstream . inplace (MStream.filter f) . stream
731
732 -- | Drop elements that do not satisfy the predicate (applied to values and
733 -- their indices)
734 ifilter :: Vector v a => (Int -> a -> Bool) -> v a -> v a
735 {-# INLINE ifilter #-}
736 ifilter f = unstream
737 . inplace (MStream.map snd . MStream.filter (uncurry f)
738 . MStream.indexed)
739 . stream
740
741 -- | Yield the longest prefix of elements satisfying the predicate.
742 takeWhile :: Vector v a => (a -> Bool) -> v a -> v a
743 {-# INLINE takeWhile #-}
744 takeWhile f = unstream . Stream.takeWhile f . stream
745
746 -- | Drop the longest prefix of elements that satisfy the predicate.
747 dropWhile :: Vector v a => (a -> Bool) -> v a -> v a
748 {-# INLINE dropWhile #-}
749 dropWhile f = unstream . Stream.dropWhile f . stream
750
751 -- Searching
752 -- ---------
753
754 infix 4 `elem`
755 -- | Check whether the vector contains an element
756 elem :: (Vector v a, Eq a) => a -> v a -> Bool
757 {-# INLINE elem #-}
758 elem x = Stream.elem x . stream
759
760 infix 4 `notElem`
761 -- | Inverse of `elem`
762 notElem :: (Vector v a, Eq a) => a -> v a -> Bool
763 {-# INLINE notElem #-}
764 notElem x = Stream.notElem x . stream
765
766 -- | Yield 'Just' the first element matching the predicate or 'Nothing' if no
767 -- such element exists.
768 find :: Vector v a => (a -> Bool) -> v a -> Maybe a
769 {-# INLINE find #-}
770 find f = Stream.find f . stream
771
772 -- | Yield 'Just' the index of the first element matching the predicate or
773 -- 'Nothing' if no such element exists.
774 findIndex :: Vector v a => (a -> Bool) -> v a -> Maybe Int
775 {-# INLINE findIndex #-}
776 findIndex f = Stream.findIndex f . stream
777
778 -- | Yield the indices of elements satisfying the predicate
779 findIndices :: (Vector v a, Vector v Int) => (a -> Bool) -> v a -> v Int
780 {-# INLINE findIndices #-}
781 findIndices f = unstream
782 . Stream.map fst
783 . Stream.filter (f . snd)
784 . Stream.indexed
785 . stream
786
787 -- | Yield 'Just' the index of the first occurence of the given element or
788 -- 'Nothing' if the vector does not contain the element
789 elemIndex :: (Vector v a, Eq a) => a -> v a -> Maybe Int
790 {-# INLINE elemIndex #-}
791 elemIndex x = findIndex (x==)
792
793 -- | Yield the indices of all occurences of the given element
794 elemIndices :: (Vector v a, Vector v Int, Eq a) => a -> v a -> v Int
795 {-# INLINE elemIndices #-}
796 elemIndices x = findIndices (x==)
797
798 -- Folding
799 -- -------
800
801 -- | Left fold
802 foldl :: Vector v b => (a -> b -> a) -> a -> v b -> a
803 {-# INLINE foldl #-}
804 foldl f z = Stream.foldl f z . stream
805
806 -- | Left fold on non-empty vectors
807 foldl1 :: Vector v a => (a -> a -> a) -> v a -> a
808 {-# INLINE foldl1 #-}
809 foldl1 f = Stream.foldl1 f . stream
810
811 -- | Left fold with strict accumulator
812 foldl' :: Vector v b => (a -> b -> a) -> a -> v b -> a
813 {-# INLINE foldl' #-}
814 foldl' f z = Stream.foldl' f z . stream
815
816 -- | Left fold on non-empty vectors with strict accumulator
817 foldl1' :: Vector v a => (a -> a -> a) -> v a -> a
818 {-# INLINE foldl1' #-}
819 foldl1' f = Stream.foldl1' f . stream
820
821 -- | Right fold
822 foldr :: Vector v a => (a -> b -> b) -> b -> v a -> b
823 {-# INLINE foldr #-}
824 foldr f z = Stream.foldr f z . stream
825
826 -- | Right fold on non-empty vectors
827 foldr1 :: Vector v a => (a -> a -> a) -> v a -> a
828 {-# INLINE foldr1 #-}
829 foldr1 f = Stream.foldr1 f . stream
830
831 -- | Left fold (function applied to each element and its index)
832 ifoldl :: Vector v b => (a -> Int -> b -> a) -> a -> v b -> a
833 {-# INLINE ifoldl #-}
834 ifoldl f z = Stream.foldl (uncurry . f) z . Stream.indexed . stream
835
836 -- | Left fold with strict accumulator (function applied to each element and
837 -- its index)
838 ifoldl' :: Vector v b => (a -> Int -> b -> a) -> a -> v b -> a
839 {-# INLINE ifoldl' #-}
840 ifoldl' f z = Stream.foldl' (uncurry . f) z . Stream.indexed . stream
841
842 -- | Right fold (function applied to each element and its index)
843 ifoldr :: Vector v a => (Int -> a -> b -> b) -> b -> v a -> b
844 {-# INLINE ifoldr #-}
845 ifoldr f z = Stream.foldr (uncurry f) z . Stream.indexed . stream
846
847 -- Specialised folds
848 -- -----------------
849
850 and :: Vector v Bool => v Bool -> Bool
851 {-# INLINE and #-}
852 and = Stream.and . stream
853
854 or :: Vector v Bool => v Bool -> Bool
855 {-# INLINE or #-}
856 or = Stream.or . stream
857
858 sum :: (Vector v a, Num a) => v a -> a
859 {-# INLINE sum #-}
860 sum = Stream.foldl' (+) 0 . stream
861
862 product :: (Vector v a, Num a) => v a -> a
863 {-# INLINE product #-}
864 product = Stream.foldl' (*) 1 . stream
865
866 maximum :: (Vector v a, Ord a) => v a -> a
867 {-# INLINE maximum #-}
868 maximum = Stream.foldl1' max . stream
869
870 minimum :: (Vector v a, Ord a) => v a -> a
871 {-# INLINE minimum #-}
872 minimum = Stream.foldl1' min . stream
873
874 minIndex :: (Vector v a, Ord a) => v a -> Int
875 {-# INLINE minIndex #-}
876 minIndex = fst . Stream.foldl1' imin . Stream.indexed . stream
877 where
878 imin (i,x) (j,y) | x <= y = (i,x)
879 | otherwise = (j,y)
880
881 maxIndex :: (Vector v a, Ord a) => v a -> Int
882 {-# INLINE maxIndex #-}
883 maxIndex = fst . Stream.foldl1' imax . Stream.indexed . stream
884 where
885 imax (i,x) (j,y) | x >= y = (i,x)
886 | otherwise = (j,y)
887
888
889 -- Unfolding
890 -- ---------
891
892 unfoldr :: Vector v a => (b -> Maybe (a, b)) -> b -> v a
893 {-# INLINE unfoldr #-}
894 unfoldr f = unstream . Stream.unfoldr f
895
896 -- Scans
897 -- -----
898
899 -- | Prefix scan
900 prescanl :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
901 {-# INLINE prescanl #-}
902 prescanl f z = unstream . inplace (MStream.prescanl f z) . stream
903
904 -- | Prefix scan with strict accumulator
905 prescanl' :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
906 {-# INLINE prescanl' #-}
907 prescanl' f z = unstream . inplace (MStream.prescanl' f z) . stream
908
909 -- | Suffix scan
910 postscanl :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
911 {-# INLINE postscanl #-}
912 postscanl f z = unstream . inplace (MStream.postscanl f z) . stream
913
914 -- | Suffix scan with strict accumulator
915 postscanl' :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
916 {-# INLINE postscanl' #-}
917 postscanl' f z = unstream . inplace (MStream.postscanl' f z) . stream
918
919 -- | Haskell-style scan
920 scanl :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
921 {-# INLINE scanl #-}
922 scanl f z = unstream . Stream.scanl f z . stream
923
924 -- | Haskell-style scan with strict accumulator
925 scanl' :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
926 {-# INLINE scanl' #-}
927 scanl' f z = unstream . Stream.scanl' f z . stream
928
929 -- | Scan over a non-empty vector
930 scanl1 :: Vector v a => (a -> a -> a) -> v a -> v a
931 {-# INLINE scanl1 #-}
932 scanl1 f = unstream . inplace (MStream.scanl1 f) . stream
933
934 -- | Scan over a non-empty vector with a strict accumulator
935 scanl1' :: Vector v a => (a -> a -> a) -> v a -> v a
936 {-# INLINE scanl1' #-}
937 scanl1' f = unstream . inplace (MStream.scanl1' f) . stream
938
939 -- Enumeration
940 -- -----------
941
942 enumFromTo :: (Vector v a, Enum a) => a -> a -> v a
943 {-# INLINE enumFromTo #-}
944 enumFromTo x y = unstream (Stream.enumFromTo x y)
945
946 enumFromThenTo :: (Vector v a, Enum a) => a -> a -> a -> v a
947 {-# INLINE enumFromThenTo #-}
948 enumFromThenTo x y z = unstream (Stream.enumFromThenTo x y z)
949
950 -- | Convert a vector to a list
951 toList :: Vector v a => v a -> [a]
952 {-# INLINE toList #-}
953 toList = Stream.toList . stream
954
955 -- | Convert a list to a vector
956 fromList :: Vector v a => [a] -> v a
957 {-# INLINE fromList #-}
958 fromList = unstream . Stream.fromList
959