Add unsafeInit and unsafeTail
[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,
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 {-# RULES
460
461 "slice/new [Vector]" forall i n v p.
462 slice i n (new' v p) = new' v (New.slice i n p)
463
464 "init/new [Vector]" forall v p.
465 init (new' v p) = new' v (New.init p)
466
467 "tail/new [Vector]" forall v p.
468 tail (new' v p) = new' v (New.tail p)
469
470 "take/new [Vector]" forall n v p.
471 take n (new' v p) = new' v (New.take n p)
472
473 "drop/new [Vector]" forall n v p.
474 drop n (new' v p) = new' v (New.drop n p)
475
476 "unsafeSlice/new [Vector]" forall i n v p.
477 unsafeSlice i n (new' v p) = new' v (New.unsafeSlice i n p)
478
479 "unsafeInit/new [Vector]" forall v p.
480 unsafeInit (new' v p) = new' v (New.unsafeInit p)
481
482 "unsafeTail/new [Vector]" forall v p.
483 unsafeTail (new' v p) = new' v (New.unsafeTail p)
484
485 #-}
486
487 -- Permutations
488 -- ------------
489
490 unsafeAccum_stream
491 :: Vector v a => (a -> b -> a) -> v a -> Stream (Int,b) -> v a
492 {-# INLINE unsafeAccum_stream #-}
493 unsafeAccum_stream f v s = new (New.accum f (New.unstream (stream v)) s)
494
495 unsafeAccum :: Vector v a => (a -> b -> a) -> v a -> [(Int,b)] -> v a
496 {-# INLINE unsafeAccum #-}
497 unsafeAccum f v us = unsafeAccum_stream f v (Stream.fromList us)
498
499 unsafeAccumulate :: (Vector v a, Vector v (Int, b))
500 => (a -> b -> a) -> v a -> v (Int,b) -> v a
501 {-# INLINE unsafeAccumulate #-}
502 unsafeAccumulate f v us = unsafeAccum_stream f v (stream us)
503
504 unsafeAccumulate_ :: (Vector v a, Vector v Int, Vector v b)
505 => (a -> b -> a) -> v a -> v Int -> v b -> v a
506 {-# INLINE unsafeAccumulate_ #-}
507 unsafeAccumulate_ f v is xs
508 = unsafeAccum_stream f v (Stream.zipWith (,) (stream is) (stream xs))
509
510 accum_stream :: Vector v a => (a -> b -> a) -> v a -> Stream (Int,b) -> v a
511 {-# INLINE accum_stream #-}
512 accum_stream f v s = new (New.accum f (New.unstream (stream v)) s)
513
514 accum :: Vector v a => (a -> b -> a) -> v a -> [(Int,b)] -> v a
515 {-# INLINE accum #-}
516 accum f v us = accum_stream f v (Stream.fromList us)
517
518 accumulate :: (Vector v a, Vector v (Int, b))
519 => (a -> b -> a) -> v a -> v (Int,b) -> v a
520 {-# INLINE accumulate #-}
521 accumulate f v us = accum_stream f v (stream us)
522
523 accumulate_ :: (Vector v a, Vector v Int, Vector v b)
524 => (a -> b -> a) -> v a -> v Int -> v b -> v a
525 {-# INLINE accumulate_ #-}
526 accumulate_ f v is xs = accum_stream f v (Stream.zipWith (,) (stream is)
527 (stream xs))
528
529
530 unsafeUpdate_stream :: Vector v a => v a -> Stream (Int,a) -> v a
531 {-# INLINE unsafeUpdate_stream #-}
532 unsafeUpdate_stream v s = new (New.unsafeUpdate (New.unstream (stream v)) s)
533
534 unsafeUpd :: Vector v a => v a -> [(Int, a)] -> v a
535 {-# INLINE unsafeUpd #-}
536 unsafeUpd v us = unsafeUpdate_stream v (Stream.fromList us)
537
538 unsafeUpdate :: (Vector v a, Vector v (Int, a)) => v a -> v (Int, a) -> v a
539 {-# INLINE unsafeUpdate #-}
540 unsafeUpdate v w = unsafeUpdate_stream v (stream w)
541
542 unsafeUpdate_ :: (Vector v a, Vector v Int) => v a -> v Int -> v a -> v a
543 {-# INLINE unsafeUpdate_ #-}
544 unsafeUpdate_ v is w
545 = unsafeUpdate_stream v (Stream.zipWith (,) (stream is) (stream w))
546
547 update_stream :: Vector v a => v a -> Stream (Int,a) -> v a
548 {-# INLINE update_stream #-}
549 update_stream v s = new (New.update (New.unstream (stream v)) s)
550
551 (//) :: Vector v a => v a -> [(Int, a)] -> v a
552 {-# INLINE (//) #-}
553 v // us = update_stream v (Stream.fromList us)
554
555 update :: (Vector v a, Vector v (Int, a)) => v a -> v (Int, a) -> v a
556 {-# INLINE update #-}
557 update v w = update_stream v (stream w)
558
559 update_ :: (Vector v a, Vector v Int) => v a -> v Int -> v a -> v a
560 {-# INLINE update_ #-}
561 update_ v is w = update_stream v (Stream.zipWith (,) (stream is) (stream w))
562
563 -- This somewhat non-intuitive definition ensures that the resulting vector
564 -- does not retain references to the original one even if it is lazy in its
565 -- elements. This would not be the case if we simply used
566 --
567 -- backpermute v is = map (v!) is
568 backpermute :: (Vector v a, Vector v Int) => v a -> v Int -> v a
569 {-# INLINE backpermute #-}
570 backpermute v is = seq v
571 $ unstream
572 $ Stream.unbox
573 $ Stream.map (indexM v)
574 $ stream is
575
576 reverse :: (Vector v a) => v a -> v a
577 {-# INLINE reverse #-}
578 reverse = new . New.reverse . New.unstream . stream
579
580 -- Mapping
581 -- -------
582
583 -- | Map a function over a vector
584 map :: (Vector v a, Vector v b) => (a -> b) -> v a -> v b
585 {-# INLINE map #-}
586 map f = unstream . inplace (MStream.map f) . stream
587
588 -- | Apply a function to every index/value pair
589 imap :: (Vector v a, Vector v b) => (Int -> a -> b) -> v a -> v b
590 {-# INLINE imap #-}
591 imap f = unstream . inplace (MStream.map (uncurry f) . MStream.indexed)
592 . stream
593
594 concatMap :: (Vector v a, Vector v b) => (a -> v b) -> v a -> v b
595 {-# INLINE concatMap #-}
596 concatMap f = unstream . Stream.concatMap (stream . f) . stream
597
598 -- Zipping/unzipping
599 -- -----------------
600
601 -- | Zip two vectors with the given function.
602 zipWith :: (Vector v a, Vector v b, Vector v c)
603 => (a -> b -> c) -> v a -> v b -> v c
604 {-# INLINE zipWith #-}
605 zipWith f xs ys = unstream (Stream.zipWith f (stream xs) (stream ys))
606
607 -- | Zip three vectors with the given function.
608 zipWith3 :: (Vector v a, Vector v b, Vector v c, Vector v d)
609 => (a -> b -> c -> d) -> v a -> v b -> v c -> v d
610 {-# INLINE zipWith3 #-}
611 zipWith3 f as bs cs = unstream (Stream.zipWith3 f (stream as)
612 (stream bs)
613 (stream cs))
614
615 zipWith4 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e)
616 => (a -> b -> c -> d -> e) -> v a -> v b -> v c -> v d -> v e
617 {-# INLINE zipWith4 #-}
618 zipWith4 f as bs cs ds
619 = unstream (Stream.zipWith4 f (stream as)
620 (stream bs)
621 (stream cs)
622 (stream ds))
623
624 zipWith5 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
625 Vector v f)
626 => (a -> b -> c -> d -> e -> f) -> v a -> v b -> v c -> v d -> v e
627 -> v f
628 {-# INLINE zipWith5 #-}
629 zipWith5 f as bs cs ds es
630 = unstream (Stream.zipWith5 f (stream as)
631 (stream bs)
632 (stream cs)
633 (stream ds)
634 (stream es))
635
636 zipWith6 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
637 Vector v f, Vector v g)
638 => (a -> b -> c -> d -> e -> f -> g)
639 -> v a -> v b -> v c -> v d -> v e -> v f -> v g
640 {-# INLINE zipWith6 #-}
641 zipWith6 f as bs cs ds es fs
642 = unstream (Stream.zipWith6 f (stream as)
643 (stream bs)
644 (stream cs)
645 (stream ds)
646 (stream es)
647 (stream fs))
648
649 -- | Zip two vectors and their indices with the given function.
650 izipWith :: (Vector v a, Vector v b, Vector v c)
651 => (Int -> a -> b -> c) -> v a -> v b -> v c
652 {-# INLINE izipWith #-}
653 izipWith f xs ys = unstream
654 (Stream.zipWith (uncurry f) (Stream.indexed (stream xs))
655 (stream ys))
656
657 -- | Zip three vectors and their indices with the given function.
658 izipWith3 :: (Vector v a, Vector v b, Vector v c, Vector v d)
659 => (Int -> a -> b -> c -> d) -> v a -> v b -> v c -> v d
660 {-# INLINE izipWith3 #-}
661 izipWith3 f as bs cs
662 = unstream (Stream.zipWith3 (uncurry f) (Stream.indexed (stream as))
663 (stream bs)
664 (stream cs))
665
666 izipWith4 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e)
667 => (Int -> a -> b -> c -> d -> e) -> v a -> v b -> v c -> v d -> v e
668 {-# INLINE izipWith4 #-}
669 izipWith4 f as bs cs ds
670 = unstream (Stream.zipWith4 (uncurry f) (Stream.indexed (stream as))
671 (stream bs)
672 (stream cs)
673 (stream ds))
674
675 izipWith5 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
676 Vector v f)
677 => (Int -> a -> b -> c -> d -> e -> f) -> v a -> v b -> v c -> v d
678 -> v e -> v f
679 {-# INLINE izipWith5 #-}
680 izipWith5 f as bs cs ds es
681 = unstream (Stream.zipWith5 (uncurry f) (Stream.indexed (stream as))
682 (stream bs)
683 (stream cs)
684 (stream ds)
685 (stream es))
686
687 izipWith6 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
688 Vector v f, Vector v g)
689 => (Int -> a -> b -> c -> d -> e -> f -> g)
690 -> v a -> v b -> v c -> v d -> v e -> v f -> v g
691 {-# INLINE izipWith6 #-}
692 izipWith6 f as bs cs ds es fs
693 = unstream (Stream.zipWith6 (uncurry f) (Stream.indexed (stream as))
694 (stream bs)
695 (stream cs)
696 (stream ds)
697 (stream es)
698 (stream fs))
699
700 zip :: (Vector v a, Vector v b, Vector v (a,b)) => v a -> v b -> v (a, b)
701 {-# INLINE zip #-}
702 zip = zipWith (,)
703
704 zip3 :: (Vector v a, Vector v b, Vector v c, Vector v (a, b, c))
705 => v a -> v b -> v c -> v (a, b, c)
706 {-# INLINE zip3 #-}
707 zip3 = zipWith3 (,,)
708
709 zip4 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v (a, b, c, d))
710 => v a -> v b -> v c -> v d -> v (a, b, c, d)
711 {-# INLINE zip4 #-}
712 zip4 = zipWith4 (,,,)
713
714 zip5 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
715 Vector v (a, b, c, d, e))
716 => v a -> v b -> v c -> v d -> v e -> v (a, b, c, d, e)
717 {-# INLINE zip5 #-}
718 zip5 = zipWith5 (,,,,)
719
720 zip6 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
721 Vector v f, Vector v (a, b, c, d, e, f))
722 => v a -> v b -> v c -> v d -> v e -> v f -> v (a, b, c, d, e, f)
723 {-# INLINE zip6 #-}
724 zip6 = zipWith6 (,,,,,)
725
726 unzip :: (Vector v a, Vector v b, Vector v (a,b)) => v (a, b) -> (v a, v b)
727 {-# INLINE unzip #-}
728 unzip xs = (map fst xs, map snd xs)
729
730 unzip3 :: (Vector v a, Vector v b, Vector v c, Vector v (a, b, c))
731 => v (a, b, c) -> (v a, v b, v c)
732 {-# INLINE unzip3 #-}
733 unzip3 xs = (map (\(a, b, c) -> a) xs,
734 map (\(a, b, c) -> b) xs,
735 map (\(a, b, c) -> c) xs)
736
737 unzip4 :: (Vector v a, Vector v b, Vector v c, Vector v d,
738 Vector v (a, b, c, d))
739 => v (a, b, c, d) -> (v a, v b, v c, v d)
740 {-# INLINE unzip4 #-}
741 unzip4 xs = (map (\(a, b, c, d) -> a) xs,
742 map (\(a, b, c, d) -> b) xs,
743 map (\(a, b, c, d) -> c) xs,
744 map (\(a, b, c, d) -> d) xs)
745
746 unzip5 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
747 Vector v (a, b, c, d, e))
748 => v (a, b, c, d, e) -> (v a, v b, v c, v d, v e)
749 {-# INLINE unzip5 #-}
750 unzip5 xs = (map (\(a, b, c, d, e) -> a) xs,
751 map (\(a, b, c, d, e) -> b) xs,
752 map (\(a, b, c, d, e) -> c) xs,
753 map (\(a, b, c, d, e) -> d) xs,
754 map (\(a, b, c, d, e) -> e) xs)
755
756 unzip6 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
757 Vector v f, Vector v (a, b, c, d, e, f))
758 => v (a, b, c, d, e, f) -> (v a, v b, v c, v d, v e, v f)
759 {-# INLINE unzip6 #-}
760 unzip6 xs = (map (\(a, b, c, d, e, f) -> a) xs,
761 map (\(a, b, c, d, e, f) -> b) xs,
762 map (\(a, b, c, d, e, f) -> c) xs,
763 map (\(a, b, c, d, e, f) -> d) xs,
764 map (\(a, b, c, d, e, f) -> e) xs,
765 map (\(a, b, c, d, e, f) -> f) xs)
766
767 -- Comparisons
768 -- -----------
769
770 eq :: (Vector v a, Eq a) => v a -> v a -> Bool
771 {-# INLINE eq #-}
772 xs `eq` ys = stream xs == stream ys
773
774 cmp :: (Vector v a, Ord a) => v a -> v a -> Ordering
775 {-# INLINE cmp #-}
776 cmp xs ys = compare (stream xs) (stream ys)
777
778 -- Filtering
779 -- ---------
780
781 -- | Drop elements that do not satisfy the predicate
782 filter :: Vector v a => (a -> Bool) -> v a -> v a
783 {-# INLINE filter #-}
784 filter f = unstream . inplace (MStream.filter f) . stream
785
786 -- | Drop elements that do not satisfy the predicate (applied to values and
787 -- their indices)
788 ifilter :: Vector v a => (Int -> a -> Bool) -> v a -> v a
789 {-# INLINE ifilter #-}
790 ifilter f = unstream
791 . inplace (MStream.map snd . MStream.filter (uncurry f)
792 . MStream.indexed)
793 . stream
794
795 -- | Yield the longest prefix of elements satisfying the predicate.
796 takeWhile :: Vector v a => (a -> Bool) -> v a -> v a
797 {-# INLINE takeWhile #-}
798 takeWhile f = unstream . Stream.takeWhile f . stream
799
800 -- | Drop the longest prefix of elements that satisfy the predicate.
801 dropWhile :: Vector v a => (a -> Bool) -> v a -> v a
802 {-# INLINE dropWhile #-}
803 dropWhile f = unstream . Stream.dropWhile f . stream
804
805 -- | Split the vector in two parts, the first one containing those elements
806 -- that satisfy the predicate and the second one those that don't. The order
807 -- of the elements is not preserved.
808 unstablePartition :: Vector v a => (a -> Bool) -> v a -> (v a, v a)
809 {-# INLINE unstablePartition #-}
810 unstablePartition f = unstablePartition_stream f . stream
811
812 unstablePartition_stream
813 :: Vector v a => (a -> Bool) -> Stream a -> (v a, v a)
814 {-# INLINE_STREAM unstablePartition_stream #-}
815 unstablePartition_stream f s = s `seq` runST (
816 do
817 (mv1,mv2) <- M.unstablePartitionStream f s
818 v1 <- unsafeFreeze mv1
819 v2 <- unsafeFreeze mv2
820 return (v1,v2))
821
822 unstablePartition_new :: Vector v a => (a -> Bool) -> New a -> (v a, v a)
823 {-# INLINE_STREAM unstablePartition_new #-}
824 unstablePartition_new f (New.New p) = runST (
825 do
826 mv <- p
827 i <- M.unstablePartition f mv
828 v <- unsafeFreeze mv
829 return (take i v, drop i v))
830
831 {-# RULES
832
833 "unstablePartition" forall f v p.
834 unstablePartition_stream f (stream (new' v p))
835 = unstablePartition_new f p
836
837 #-}
838
839
840 -- FIXME: make span and break fusible
841
842 -- | Split the vector into the longest prefix of elements that satisfy the
843 -- predicate and the rest.
844 span :: Vector v a => (a -> Bool) -> v a -> (v a, v a)
845 {-# INLINE span #-}
846 span f = break (not . f)
847
848 -- | Split the vector into the longest prefix of elements that do not satisfy
849 -- the predicate and the rest.
850 break :: Vector v a => (a -> Bool) -> v a -> (v a, v a)
851 {-# INLINE break #-}
852 break f xs = case findIndex f xs of
853 Just i -> (unsafeSlice 0 i xs, unsafeSlice i (length xs - i) xs)
854 Nothing -> (xs, empty)
855
856
857 -- Searching
858 -- ---------
859
860 infix 4 `elem`
861 -- | Check whether the vector contains an element
862 elem :: (Vector v a, Eq a) => a -> v a -> Bool
863 {-# INLINE elem #-}
864 elem x = Stream.elem x . stream
865
866 infix 4 `notElem`
867 -- | Inverse of `elem`
868 notElem :: (Vector v a, Eq a) => a -> v a -> Bool
869 {-# INLINE notElem #-}
870 notElem x = Stream.notElem x . stream
871
872 -- | Yield 'Just' the first element matching the predicate or 'Nothing' if no
873 -- such element exists.
874 find :: Vector v a => (a -> Bool) -> v a -> Maybe a
875 {-# INLINE find #-}
876 find f = Stream.find f . stream
877
878 -- | Yield 'Just' the index of the first element matching the predicate or
879 -- 'Nothing' if no such element exists.
880 findIndex :: Vector v a => (a -> Bool) -> v a -> Maybe Int
881 {-# INLINE findIndex #-}
882 findIndex f = Stream.findIndex f . stream
883
884 -- | Yield the indices of elements satisfying the predicate
885 findIndices :: (Vector v a, Vector v Int) => (a -> Bool) -> v a -> v Int
886 {-# INLINE findIndices #-}
887 findIndices f = unstream
888 . Stream.map fst
889 . Stream.filter (f . snd)
890 . Stream.indexed
891 . stream
892
893 -- | Yield 'Just' the index of the first occurence of the given element or
894 -- 'Nothing' if the vector does not contain the element
895 elemIndex :: (Vector v a, Eq a) => a -> v a -> Maybe Int
896 {-# INLINE elemIndex #-}
897 elemIndex x = findIndex (x==)
898
899 -- | Yield the indices of all occurences of the given element
900 elemIndices :: (Vector v a, Vector v Int, Eq a) => a -> v a -> v Int
901 {-# INLINE elemIndices #-}
902 elemIndices x = findIndices (x==)
903
904 -- Folding
905 -- -------
906
907 -- | Left fold
908 foldl :: Vector v b => (a -> b -> a) -> a -> v b -> a
909 {-# INLINE foldl #-}
910 foldl f z = Stream.foldl f z . stream
911
912 -- | Left fold on non-empty vectors
913 foldl1 :: Vector v a => (a -> a -> a) -> v a -> a
914 {-# INLINE foldl1 #-}
915 foldl1 f = Stream.foldl1 f . stream
916
917 -- | Left fold with strict accumulator
918 foldl' :: Vector v b => (a -> b -> a) -> a -> v b -> a
919 {-# INLINE foldl' #-}
920 foldl' f z = Stream.foldl' f z . stream
921
922 -- | Left fold on non-empty vectors with strict accumulator
923 foldl1' :: Vector v a => (a -> a -> a) -> v a -> a
924 {-# INLINE foldl1' #-}
925 foldl1' f = Stream.foldl1' f . stream
926
927 -- | Right fold
928 foldr :: Vector v a => (a -> b -> b) -> b -> v a -> b
929 {-# INLINE foldr #-}
930 foldr f z = Stream.foldr f z . stream
931
932 -- | Right fold on non-empty vectors
933 foldr1 :: Vector v a => (a -> a -> a) -> v a -> a
934 {-# INLINE foldr1 #-}
935 foldr1 f = Stream.foldr1 f . stream
936
937 -- | Left fold (function applied to each element and its index)
938 ifoldl :: Vector v b => (a -> Int -> b -> a) -> a -> v b -> a
939 {-# INLINE ifoldl #-}
940 ifoldl f z = Stream.foldl (uncurry . f) z . Stream.indexed . stream
941
942 -- | Left fold with strict accumulator (function applied to each element and
943 -- its index)
944 ifoldl' :: Vector v b => (a -> Int -> b -> a) -> a -> v b -> a
945 {-# INLINE ifoldl' #-}
946 ifoldl' f z = Stream.foldl' (uncurry . f) z . Stream.indexed . stream
947
948 -- | Right fold (function applied to each element and its index)
949 ifoldr :: Vector v a => (Int -> a -> b -> b) -> b -> v a -> b
950 {-# INLINE ifoldr #-}
951 ifoldr f z = Stream.foldr (uncurry f) z . Stream.indexed . stream
952
953 -- Specialised folds
954 -- -----------------
955
956 all :: Vector v a => (a -> Bool) -> v a -> Bool
957 {-# INLINE all #-}
958 all f = Stream.and . Stream.map f . stream
959
960 any :: Vector v a => (a -> Bool) -> v a -> Bool
961 {-# INLINE any #-}
962 any f = Stream.or . Stream.map f . stream
963
964 and :: Vector v Bool => v Bool -> Bool
965 {-# INLINE and #-}
966 and = Stream.and . stream
967
968 or :: Vector v Bool => v Bool -> Bool
969 {-# INLINE or #-}
970 or = Stream.or . stream
971
972 sum :: (Vector v a, Num a) => v a -> a
973 {-# INLINE sum #-}
974 sum = Stream.foldl' (+) 0 . stream
975
976 product :: (Vector v a, Num a) => v a -> a
977 {-# INLINE product #-}
978 product = Stream.foldl' (*) 1 . stream
979
980 maximum :: (Vector v a, Ord a) => v a -> a
981 {-# INLINE maximum #-}
982 maximum = Stream.foldl1' max . stream
983
984 maximumBy :: Vector v a => (a -> a -> Ordering) -> v a -> a
985 {-# INLINE maximumBy #-}
986 maximumBy cmp = Stream.foldl1' maxBy . stream
987 where
988 {-# INLINE maxBy #-}
989 maxBy x y = case cmp x y of
990 LT -> y
991 _ -> x
992
993 minimum :: (Vector v a, Ord a) => v a -> a
994 {-# INLINE minimum #-}
995 minimum = Stream.foldl1' min . stream
996
997 minimumBy :: Vector v a => (a -> a -> Ordering) -> v a -> a
998 {-# INLINE minimumBy #-}
999 minimumBy cmp = Stream.foldl1' minBy . stream
1000 where
1001 {-# INLINE minBy #-}
1002 minBy x y = case cmp x y of
1003 GT -> y
1004 _ -> x
1005
1006 maxIndex :: (Vector v a, Ord a) => v a -> Int
1007 {-# INLINE maxIndex #-}
1008 maxIndex = maxIndexBy compare
1009
1010 maxIndexBy :: Vector v a => (a -> a -> Ordering) -> v a -> Int
1011 {-# INLINE maxIndexBy #-}
1012 maxIndexBy cmp = fst . Stream.foldl1' imax . Stream.indexed . stream
1013 where
1014 imax (i,x) (j,y) = case cmp x y of
1015 LT -> (j,y)
1016 _ -> (i,x)
1017
1018 minIndex :: (Vector v a, Ord a) => v a -> Int
1019 {-# INLINE minIndex #-}
1020 minIndex = minIndexBy compare
1021
1022 minIndexBy :: Vector v a => (a -> a -> Ordering) -> v a -> Int
1023 {-# INLINE minIndexBy #-}
1024 minIndexBy cmp = fst . Stream.foldl1' imin . Stream.indexed . stream
1025 where
1026 imin (i,x) (j,y) = case cmp x y of
1027 GT -> (j,y)
1028 _ -> (i,x)
1029
1030
1031 -- Unfolding
1032 -- ---------
1033
1034 unfoldr :: Vector v a => (b -> Maybe (a, b)) -> b -> v a
1035 {-# INLINE unfoldr #-}
1036 unfoldr f = unstream . Stream.unfoldr f
1037
1038 -- Scans
1039 -- -----
1040
1041 -- | Prefix scan
1042 prescanl :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1043 {-# INLINE prescanl #-}
1044 prescanl f z = unstream . inplace (MStream.prescanl f z) . stream
1045
1046 -- | Prefix scan with strict accumulator
1047 prescanl' :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1048 {-# INLINE prescanl' #-}
1049 prescanl' f z = unstream . inplace (MStream.prescanl' f z) . stream
1050
1051 -- | Suffix scan
1052 postscanl :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1053 {-# INLINE postscanl #-}
1054 postscanl f z = unstream . inplace (MStream.postscanl f z) . stream
1055
1056 -- | Suffix scan with strict accumulator
1057 postscanl' :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1058 {-# INLINE postscanl' #-}
1059 postscanl' f z = unstream . inplace (MStream.postscanl' f z) . stream
1060
1061 -- | Haskell-style scan
1062 scanl :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1063 {-# INLINE scanl #-}
1064 scanl f z = unstream . Stream.scanl f z . stream
1065
1066 -- | Haskell-style scan with strict accumulator
1067 scanl' :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1068 {-# INLINE scanl' #-}
1069 scanl' f z = unstream . Stream.scanl' f z . stream
1070
1071 -- | Scan over a non-empty vector
1072 scanl1 :: Vector v a => (a -> a -> a) -> v a -> v a
1073 {-# INLINE scanl1 #-}
1074 scanl1 f = unstream . inplace (MStream.scanl1 f) . stream
1075
1076 -- | Scan over a non-empty vector with a strict accumulator
1077 scanl1' :: Vector v a => (a -> a -> a) -> v a -> v a
1078 {-# INLINE scanl1' #-}
1079 scanl1' f = unstream . inplace (MStream.scanl1' f) . stream
1080
1081 -- Enumeration
1082 -- -----------
1083
1084 enumFromTo :: (Vector v a, Enum a) => a -> a -> v a
1085 {-# INLINE enumFromTo #-}
1086 enumFromTo x y = unstream (Stream.enumFromTo x y)
1087
1088 enumFromThenTo :: (Vector v a, Enum a) => a -> a -> a -> v a
1089 {-# INLINE enumFromThenTo #-}
1090 enumFromThenTo x y z = unstream (Stream.enumFromThenTo x y z)
1091
1092 -- | Convert a vector to a list
1093 toList :: Vector v a => v a -> [a]
1094 {-# INLINE toList #-}
1095 toList = Stream.toList . stream
1096
1097 -- | Convert a list to a vector
1098 fromList :: Vector v a => [a] -> v a
1099 {-# INLINE fromList #-}
1100 fromList = unstream . Stream.fromList
1101