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