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