Add unstablePartition, span, break
[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
28 -- * Subvectors
29 slice, init, tail, take, drop,
30 unsafeSlice,
31
32 -- * Permutations
33 accum, accumulate, accumulate_,
34 (//), update, update_,
35 backpermute, reverse,
36
37 -- * Mapping
38 map, imap, concatMap,
39
40 -- * Zipping and unzipping
41 zipWith, zipWith3, zipWith4, zipWith5, zipWith6,
42 izipWith, izipWith3, izipWith4, izipWith5, izipWith6,
43 zip, zip3, zip4, zip5, zip6,
44 unzip, unzip3, unzip4, unzip5, unzip6,
45
46 -- * Comparisons
47 eq, cmp,
48
49 -- * Filtering
50 filter, ifilter, takeWhile, dropWhile,
51 unstablePartition, span, break,
52
53 -- * Searching
54 elem, notElem, find, findIndex, findIndices, elemIndex, elemIndices,
55
56 -- * Folding
57 foldl, foldl1, foldl', foldl1', foldr, foldr1,
58 ifoldl, ifoldl', ifoldr,
59
60 -- * Specialised folds
61 all, any, and, or,
62 sum, product,
63 maximum, maximumBy, minimum, minimumBy,
64 minIndex, minIndexBy, maxIndex, maxIndexBy,
65
66 -- * Unfolding
67 unfoldr,
68
69 -- * Scans
70 prescanl, prescanl',
71 postscanl, postscanl',
72 scanl, scanl', scanl1, scanl1',
73
74 -- * Enumeration
75 enumFromTo, enumFromThenTo,
76
77 -- * Conversion to/from lists
78 toList, fromList,
79
80 -- * Conversion to/from Streams
81 stream, unstream,
82
83 -- * MVector-based initialisation
84 new,
85
86 -- * Unsafe operations
87 unsafeIndex, unsafeIndexM,
88 unsafeAccum, unsafeAccumulate, unsafeAccumulate_,
89 unsafeUpd, unsafeUpdate, unsafeUpdate_
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 :: v a -> Int -> Int -> 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 -- | Unsafe indexing without bounds checking
311 unsafeIndex :: Vector v a => v a -> Int -> a
312 {-# INLINE_STREAM unsafeIndex #-}
313 unsafeIndex v i = UNSAFE_CHECK(checkIndex) "unsafeIndex" i (length v)
314 $ unId (basicUnsafeIndexM v i)
315
316 -- | First element
317 head :: Vector v a => v a -> a
318 {-# INLINE_STREAM head #-}
319 head v = v ! 0
320
321 -- | Last element
322 last :: Vector v a => v a -> a
323 {-# INLINE_STREAM last #-}
324 last v = v ! (length v - 1)
325
326 {-# RULES
327
328 "(!)/unstream [Vector]" forall v i s.
329 new' v (New.unstream s) ! i = s Stream.!! i
330
331 "unsafeIndex/unstream [Vector]" forall v i s.
332 unsafeIndex (new' v (New.unstream s)) i = s Stream.!! i
333
334 "head/unstream [Vector]" forall v s.
335 head (new' v (New.unstream s)) = Stream.head s
336
337 "last/unstream [Vector]" forall v s.
338 last (new' v (New.unstream s)) = Stream.last s
339
340 #-}
341
342 -- | Monadic indexing which can be strict in the vector while remaining lazy in
343 -- the element.
344 indexM :: (Vector v a, Monad m) => v a -> Int -> m a
345 {-# INLINE_STREAM indexM #-}
346 indexM v i = BOUNDS_CHECK(checkIndex) "indexM" i (length v)
347 $ basicUnsafeIndexM v i
348
349 -- | Unsafe monadic indexing without bounds checks
350 unsafeIndexM :: (Vector v a, Monad m) => v a -> Int -> m a
351 {-# INLINE_STREAM unsafeIndexM #-}
352 unsafeIndexM v i = UNSAFE_CHECK(checkIndex) "unsafeIndexM" i (length v)
353 $ basicUnsafeIndexM v i
354
355 headM :: (Vector v a, Monad m) => v a -> m a
356 {-# INLINE_STREAM headM #-}
357 headM v = indexM v 0
358
359 lastM :: (Vector v a, Monad m) => v a -> m a
360 {-# INLINE_STREAM lastM #-}
361 lastM v = indexM v (length v - 1)
362
363 -- FIXME: the rhs of these rules are lazy in the stream which is WRONG
364 {- RULES
365
366 "indexM/unstream [Vector]" forall v i s.
367 indexM (new' v (New.unstream s)) i = return (s Stream.!! i)
368
369 "headM/unstream [Vector]" forall v s.
370 headM (new' v (New.unstream s)) = return (Stream.head s)
371
372 "lastM/unstream [Vector]" forall v s.
373 lastM (new' v (New.unstream s)) = return (Stream.last s)
374
375 -}
376
377 -- Subarrays
378 -- ---------
379
380 -- FIXME: slicing doesn't work with the inplace stuff at the moment
381
382 -- | Yield a part of the vector without copying it.
383 slice :: Vector v a => v a -> Int -- ^ starting index
384 -> Int -- ^ length
385 -> v a
386 {-# INLINE_STREAM slice #-}
387 slice v i n = BOUNDS_CHECK(checkSlice) "slice" i n (length v)
388 $ basicUnsafeSlice v i n
389
390 -- | Unsafely yield a part of the vector without copying it and without
391 -- performing bounds checks.
392 unsafeSlice :: Vector v a => v a -> Int -- ^ starting index
393 -> Int -- ^ length
394 -> v a
395 {-# INLINE_STREAM unsafeSlice #-}
396 unsafeSlice v i n = UNSAFE_CHECK(checkSlice) "unsafeSlice" i n (length v)
397 $ basicUnsafeSlice v i n
398
399 -- | Yield all but the last element without copying.
400 init :: Vector v a => v a -> v a
401 {-# INLINE_STREAM init #-}
402 init v = slice v 0 (length v - 1)
403
404 -- | All but the first element (without copying).
405 tail :: Vector v a => v a -> v a
406 {-# INLINE_STREAM tail #-}
407 tail v = slice v 1 (length v - 1)
408
409 -- | Yield the first @n@ elements without copying.
410 take :: Vector v a => Int -> v a -> v a
411 {-# INLINE_STREAM take #-}
412 take n v = slice v 0 (min n' (length v))
413 where n' = max n 0
414
415 -- | Yield all but the first @n@ elements without copying.
416 drop :: Vector v a => Int -> v a -> v a
417 {-# INLINE_STREAM drop #-}
418 drop n v = slice v (min n' len) (max 0 (len - n'))
419 where n' = max n 0
420 len = length v
421
422 {-# RULES
423
424 "slice/new [Vector]" forall v p i n.
425 slice (new' v p) i n = new' v (New.slice p i n)
426
427 "unsafeSlice/new [Vector]" forall v p i n.
428 unsafeSlice (new' v p) i n = new' v (New.unsafeSlice p i n)
429
430 "init/new [Vector]" forall v p.
431 init (new' v p) = new' v (New.init p)
432
433 "tail/new [Vector]" forall v p.
434 tail (new' v p) = new' v (New.tail p)
435
436 "take/new [Vector]" forall n v p.
437 take n (new' v p) = new' v (New.take n p)
438
439 "drop/new [Vector]" forall n v p.
440 drop n (new' v p) = new' v (New.drop n p)
441
442 #-}
443
444 -- Permutations
445 -- ------------
446
447 unsafeAccum_stream
448 :: Vector v a => (a -> b -> a) -> v a -> Stream (Int,b) -> v a
449 {-# INLINE unsafeAccum_stream #-}
450 unsafeAccum_stream f v s = new (New.accum f (New.unstream (stream v)) s)
451
452 unsafeAccum :: Vector v a => (a -> b -> a) -> v a -> [(Int,b)] -> v a
453 {-# INLINE unsafeAccum #-}
454 unsafeAccum f v us = unsafeAccum_stream f v (Stream.fromList us)
455
456 unsafeAccumulate :: (Vector v a, Vector v (Int, b))
457 => (a -> b -> a) -> v a -> v (Int,b) -> v a
458 {-# INLINE unsafeAccumulate #-}
459 unsafeAccumulate f v us = unsafeAccum_stream f v (stream us)
460
461 unsafeAccumulate_ :: (Vector v a, Vector v Int, Vector v b)
462 => (a -> b -> a) -> v a -> v Int -> v b -> v a
463 {-# INLINE unsafeAccumulate_ #-}
464 unsafeAccumulate_ f v is xs
465 = unsafeAccum_stream f v (Stream.zipWith (,) (stream is) (stream xs))
466
467 accum_stream :: Vector v a => (a -> b -> a) -> v a -> Stream (Int,b) -> v a
468 {-# INLINE accum_stream #-}
469 accum_stream f v s = new (New.accum f (New.unstream (stream v)) s)
470
471 accum :: Vector v a => (a -> b -> a) -> v a -> [(Int,b)] -> v a
472 {-# INLINE accum #-}
473 accum f v us = accum_stream f v (Stream.fromList us)
474
475 accumulate :: (Vector v a, Vector v (Int, b))
476 => (a -> b -> a) -> v a -> v (Int,b) -> v a
477 {-# INLINE accumulate #-}
478 accumulate f v us = accum_stream f v (stream us)
479
480 accumulate_ :: (Vector v a, Vector v Int, Vector v b)
481 => (a -> b -> a) -> v a -> v Int -> v b -> v a
482 {-# INLINE accumulate_ #-}
483 accumulate_ f v is xs = accum_stream f v (Stream.zipWith (,) (stream is)
484 (stream xs))
485
486
487 unsafeUpdate_stream :: Vector v a => v a -> Stream (Int,a) -> v a
488 {-# INLINE unsafeUpdate_stream #-}
489 unsafeUpdate_stream v s = new (New.unsafeUpdate (New.unstream (stream v)) s)
490
491 unsafeUpd :: Vector v a => v a -> [(Int, a)] -> v a
492 {-# INLINE unsafeUpd #-}
493 unsafeUpd v us = unsafeUpdate_stream v (Stream.fromList us)
494
495 unsafeUpdate :: (Vector v a, Vector v (Int, a)) => v a -> v (Int, a) -> v a
496 {-# INLINE unsafeUpdate #-}
497 unsafeUpdate v w = unsafeUpdate_stream v (stream w)
498
499 unsafeUpdate_ :: (Vector v a, Vector v Int) => v a -> v Int -> v a -> v a
500 {-# INLINE unsafeUpdate_ #-}
501 unsafeUpdate_ v is w
502 = unsafeUpdate_stream v (Stream.zipWith (,) (stream is) (stream w))
503
504 update_stream :: Vector v a => v a -> Stream (Int,a) -> v a
505 {-# INLINE update_stream #-}
506 update_stream v s = new (New.update (New.unstream (stream v)) s)
507
508 (//) :: Vector v a => v a -> [(Int, a)] -> v a
509 {-# INLINE (//) #-}
510 v // us = update_stream v (Stream.fromList us)
511
512 update :: (Vector v a, Vector v (Int, a)) => v a -> v (Int, a) -> v a
513 {-# INLINE update #-}
514 update v w = update_stream v (stream w)
515
516 update_ :: (Vector v a, Vector v Int) => v a -> v Int -> v a -> v a
517 {-# INLINE update_ #-}
518 update_ v is w = update_stream v (Stream.zipWith (,) (stream is) (stream w))
519
520 -- This somewhat non-intuitive definition ensures that the resulting vector
521 -- does not retain references to the original one even if it is lazy in its
522 -- elements. This would not be the case if we simply used
523 --
524 -- backpermute v is = map (v!) is
525 backpermute :: (Vector v a, Vector v Int) => v a -> v Int -> v a
526 {-# INLINE backpermute #-}
527 backpermute v is = seq v
528 $ unstream
529 $ Stream.unbox
530 $ Stream.map (indexM v)
531 $ stream is
532
533 reverse :: (Vector v a) => v a -> v a
534 {-# INLINE reverse #-}
535 reverse = new . New.reverse . New.unstream . stream
536
537 -- Mapping
538 -- -------
539
540 -- | Map a function over a vector
541 map :: (Vector v a, Vector v b) => (a -> b) -> v a -> v b
542 {-# INLINE map #-}
543 map f = unstream . inplace (MStream.map f) . stream
544
545 -- | Apply a function to every index/value pair
546 imap :: (Vector v a, Vector v b) => (Int -> a -> b) -> v a -> v b
547 {-# INLINE imap #-}
548 imap f = unstream . inplace (MStream.map (uncurry f) . MStream.indexed)
549 . stream
550
551 concatMap :: (Vector v a, Vector v b) => (a -> v b) -> v a -> v b
552 {-# INLINE concatMap #-}
553 concatMap f = unstream . Stream.concatMap (stream . f) . stream
554
555 -- Zipping/unzipping
556 -- -----------------
557
558 -- | Zip two vectors with the given function.
559 zipWith :: (Vector v a, Vector v b, Vector v c)
560 => (a -> b -> c) -> v a -> v b -> v c
561 {-# INLINE zipWith #-}
562 zipWith f xs ys = unstream (Stream.zipWith f (stream xs) (stream ys))
563
564 -- | Zip three vectors with the given function.
565 zipWith3 :: (Vector v a, Vector v b, Vector v c, Vector v d)
566 => (a -> b -> c -> d) -> v a -> v b -> v c -> v d
567 {-# INLINE zipWith3 #-}
568 zipWith3 f as bs cs = unstream (Stream.zipWith3 f (stream as)
569 (stream bs)
570 (stream cs))
571
572 zipWith4 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e)
573 => (a -> b -> c -> d -> e) -> v a -> v b -> v c -> v d -> v e
574 {-# INLINE zipWith4 #-}
575 zipWith4 f as bs cs ds
576 = unstream (Stream.zipWith4 f (stream as)
577 (stream bs)
578 (stream cs)
579 (stream ds))
580
581 zipWith5 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
582 Vector v f)
583 => (a -> b -> c -> d -> e -> f) -> v a -> v b -> v c -> v d -> v e
584 -> v f
585 {-# INLINE zipWith5 #-}
586 zipWith5 f as bs cs ds es
587 = unstream (Stream.zipWith5 f (stream as)
588 (stream bs)
589 (stream cs)
590 (stream ds)
591 (stream es))
592
593 zipWith6 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
594 Vector v f, Vector v g)
595 => (a -> b -> c -> d -> e -> f -> g)
596 -> v a -> v b -> v c -> v d -> v e -> v f -> v g
597 {-# INLINE zipWith6 #-}
598 zipWith6 f as bs cs ds es fs
599 = unstream (Stream.zipWith6 f (stream as)
600 (stream bs)
601 (stream cs)
602 (stream ds)
603 (stream es)
604 (stream fs))
605
606 -- | Zip two vectors and their indices with the given function.
607 izipWith :: (Vector v a, Vector v b, Vector v c)
608 => (Int -> a -> b -> c) -> v a -> v b -> v c
609 {-# INLINE izipWith #-}
610 izipWith f xs ys = unstream
611 (Stream.zipWith (uncurry f) (Stream.indexed (stream xs))
612 (stream ys))
613
614 -- | Zip three vectors and their indices with the given function.
615 izipWith3 :: (Vector v a, Vector v b, Vector v c, Vector v d)
616 => (Int -> a -> b -> c -> d) -> v a -> v b -> v c -> v d
617 {-# INLINE izipWith3 #-}
618 izipWith3 f as bs cs
619 = unstream (Stream.zipWith3 (uncurry f) (Stream.indexed (stream as))
620 (stream bs)
621 (stream cs))
622
623 izipWith4 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e)
624 => (Int -> a -> b -> c -> d -> e) -> v a -> v b -> v c -> v d -> v e
625 {-# INLINE izipWith4 #-}
626 izipWith4 f as bs cs ds
627 = unstream (Stream.zipWith4 (uncurry f) (Stream.indexed (stream as))
628 (stream bs)
629 (stream cs)
630 (stream ds))
631
632 izipWith5 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
633 Vector v f)
634 => (Int -> a -> b -> c -> d -> e -> f) -> v a -> v b -> v c -> v d
635 -> v e -> v f
636 {-# INLINE izipWith5 #-}
637 izipWith5 f as bs cs ds es
638 = unstream (Stream.zipWith5 (uncurry f) (Stream.indexed (stream as))
639 (stream bs)
640 (stream cs)
641 (stream ds)
642 (stream es))
643
644 izipWith6 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
645 Vector v f, Vector v g)
646 => (Int -> a -> b -> c -> d -> e -> f -> g)
647 -> v a -> v b -> v c -> v d -> v e -> v f -> v g
648 {-# INLINE izipWith6 #-}
649 izipWith6 f as bs cs ds es fs
650 = unstream (Stream.zipWith6 (uncurry f) (Stream.indexed (stream as))
651 (stream bs)
652 (stream cs)
653 (stream ds)
654 (stream es)
655 (stream fs))
656
657 zip :: (Vector v a, Vector v b, Vector v (a,b)) => v a -> v b -> v (a, b)
658 {-# INLINE zip #-}
659 zip = zipWith (,)
660
661 zip3 :: (Vector v a, Vector v b, Vector v c, Vector v (a, b, c))
662 => v a -> v b -> v c -> v (a, b, c)
663 {-# INLINE zip3 #-}
664 zip3 = zipWith3 (,,)
665
666 zip4 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v (a, b, c, d))
667 => v a -> v b -> v c -> v d -> v (a, b, c, d)
668 {-# INLINE zip4 #-}
669 zip4 = zipWith4 (,,,)
670
671 zip5 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
672 Vector v (a, b, c, d, e))
673 => v a -> v b -> v c -> v d -> v e -> v (a, b, c, d, e)
674 {-# INLINE zip5 #-}
675 zip5 = zipWith5 (,,,,)
676
677 zip6 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
678 Vector v f, Vector v (a, b, c, d, e, f))
679 => v a -> v b -> v c -> v d -> v e -> v f -> v (a, b, c, d, e, f)
680 {-# INLINE zip6 #-}
681 zip6 = zipWith6 (,,,,,)
682
683 unzip :: (Vector v a, Vector v b, Vector v (a,b)) => v (a, b) -> (v a, v b)
684 {-# INLINE unzip #-}
685 unzip xs = (map fst xs, map snd xs)
686
687 unzip3 :: (Vector v a, Vector v b, Vector v c, Vector v (a, b, c))
688 => v (a, b, c) -> (v a, v b, v c)
689 {-# INLINE unzip3 #-}
690 unzip3 xs = (map (\(a, b, c) -> a) xs,
691 map (\(a, b, c) -> b) xs,
692 map (\(a, b, c) -> c) xs)
693
694 unzip4 :: (Vector v a, Vector v b, Vector v c, Vector v d,
695 Vector v (a, b, c, d))
696 => v (a, b, c, d) -> (v a, v b, v c, v d)
697 {-# INLINE unzip4 #-}
698 unzip4 xs = (map (\(a, b, c, d) -> a) xs,
699 map (\(a, b, c, d) -> b) xs,
700 map (\(a, b, c, d) -> c) xs,
701 map (\(a, b, c, d) -> d) xs)
702
703 unzip5 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
704 Vector v (a, b, c, d, e))
705 => v (a, b, c, d, e) -> (v a, v b, v c, v d, v e)
706 {-# INLINE unzip5 #-}
707 unzip5 xs = (map (\(a, b, c, d, e) -> a) xs,
708 map (\(a, b, c, d, e) -> b) xs,
709 map (\(a, b, c, d, e) -> c) xs,
710 map (\(a, b, c, d, e) -> d) xs,
711 map (\(a, b, c, d, e) -> e) xs)
712
713 unzip6 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
714 Vector v f, Vector v (a, b, c, d, e, f))
715 => v (a, b, c, d, e, f) -> (v a, v b, v c, v d, v e, v f)
716 {-# INLINE unzip6 #-}
717 unzip6 xs = (map (\(a, b, c, d, e, f) -> a) xs,
718 map (\(a, b, c, d, e, f) -> b) xs,
719 map (\(a, b, c, d, e, f) -> c) xs,
720 map (\(a, b, c, d, e, f) -> d) xs,
721 map (\(a, b, c, d, e, f) -> e) xs,
722 map (\(a, b, c, d, e, f) -> f) xs)
723
724 -- Comparisons
725 -- -----------
726
727 eq :: (Vector v a, Eq a) => v a -> v a -> Bool
728 {-# INLINE eq #-}
729 xs `eq` ys = stream xs == stream ys
730
731 cmp :: (Vector v a, Ord a) => v a -> v a -> Ordering
732 {-# INLINE cmp #-}
733 cmp xs ys = compare (stream xs) (stream ys)
734
735 -- Filtering
736 -- ---------
737
738 -- | Drop elements that do not satisfy the predicate
739 filter :: Vector v a => (a -> Bool) -> v a -> v a
740 {-# INLINE filter #-}
741 filter f = unstream . inplace (MStream.filter f) . stream
742
743 -- | Drop elements that do not satisfy the predicate (applied to values and
744 -- their indices)
745 ifilter :: Vector v a => (Int -> a -> Bool) -> v a -> v a
746 {-# INLINE ifilter #-}
747 ifilter f = unstream
748 . inplace (MStream.map snd . MStream.filter (uncurry f)
749 . MStream.indexed)
750 . stream
751
752 -- | Yield the longest prefix of elements satisfying the predicate.
753 takeWhile :: Vector v a => (a -> Bool) -> v a -> v a
754 {-# INLINE takeWhile #-}
755 takeWhile f = unstream . Stream.takeWhile f . stream
756
757 -- | Drop the longest prefix of elements that satisfy the predicate.
758 dropWhile :: Vector v a => (a -> Bool) -> v a -> v a
759 {-# INLINE dropWhile #-}
760 dropWhile f = unstream . Stream.dropWhile f . stream
761
762 -- | Split the vector in two parts, the first one containing those elements
763 -- that satisfy the predicate and the second one those that don't. The order
764 -- of the elements is not preserved.
765 unstablePartition :: Vector v a => (a -> Bool) -> v a -> (v a, v a)
766 {-# INLINE unstablePartition #-}
767 unstablePartition f = unstablePartition_stream f . stream
768
769 unstablePartition_stream
770 :: Vector v a => (a -> Bool) -> Stream a -> (v a, v a)
771 {-# INLINE_STREAM unstablePartition_stream #-}
772 unstablePartition_stream f s = s `seq` runST (
773 do
774 (mv1,mv2) <- M.unstablePartitionStream f s
775 v1 <- unsafeFreeze mv1
776 v2 <- unsafeFreeze mv2
777 return (v1,v2))
778
779 unstablePartition_new :: Vector v a => (a -> Bool) -> New a -> (v a, v a)
780 {-# INLINE_STREAM unstablePartition_new #-}
781 unstablePartition_new f (New.New p) = runST (
782 do
783 mv <- p
784 i <- M.unstablePartition f mv
785 v <- unsafeFreeze mv
786 return (take i v, drop i v))
787
788 {-# RULES
789
790 "unstablePartition" forall f v p.
791 unstablePartition_stream f (stream (new' v p))
792 = unstablePartition_new f p
793
794 #-}
795
796
797 -- FIXME: make span and break fusible
798
799 -- | Split the vector into the longest prefix of elements that satisfy the
800 -- predicate and the rest.
801 span :: Vector v a => (a -> Bool) -> v a -> (v a, v a)
802 {-# INLINE span #-}
803 span f = break (not . f)
804
805 -- | Split the vector into the longest prefix of elements that do not satisfy
806 -- the predicate and the rest.
807 break :: Vector v a => (a -> Bool) -> v a -> (v a, v a)
808 {-# INLINE break #-}
809 break f xs = case findIndex f xs of
810 Just i -> (unsafeSlice xs 0 i, unsafeSlice xs i (length xs - i))
811 Nothing -> (xs, empty)
812
813
814 -- Searching
815 -- ---------
816
817 infix 4 `elem`
818 -- | Check whether the vector contains an element
819 elem :: (Vector v a, Eq a) => a -> v a -> Bool
820 {-# INLINE elem #-}
821 elem x = Stream.elem x . stream
822
823 infix 4 `notElem`
824 -- | Inverse of `elem`
825 notElem :: (Vector v a, Eq a) => a -> v a -> Bool
826 {-# INLINE notElem #-}
827 notElem x = Stream.notElem x . stream
828
829 -- | Yield 'Just' the first element matching the predicate or 'Nothing' if no
830 -- such element exists.
831 find :: Vector v a => (a -> Bool) -> v a -> Maybe a
832 {-# INLINE find #-}
833 find f = Stream.find f . stream
834
835 -- | Yield 'Just' the index of the first element matching the predicate or
836 -- 'Nothing' if no such element exists.
837 findIndex :: Vector v a => (a -> Bool) -> v a -> Maybe Int
838 {-# INLINE findIndex #-}
839 findIndex f = Stream.findIndex f . stream
840
841 -- | Yield the indices of elements satisfying the predicate
842 findIndices :: (Vector v a, Vector v Int) => (a -> Bool) -> v a -> v Int
843 {-# INLINE findIndices #-}
844 findIndices f = unstream
845 . Stream.map fst
846 . Stream.filter (f . snd)
847 . Stream.indexed
848 . stream
849
850 -- | Yield 'Just' the index of the first occurence of the given element or
851 -- 'Nothing' if the vector does not contain the element
852 elemIndex :: (Vector v a, Eq a) => a -> v a -> Maybe Int
853 {-# INLINE elemIndex #-}
854 elemIndex x = findIndex (x==)
855
856 -- | Yield the indices of all occurences of the given element
857 elemIndices :: (Vector v a, Vector v Int, Eq a) => a -> v a -> v Int
858 {-# INLINE elemIndices #-}
859 elemIndices x = findIndices (x==)
860
861 -- Folding
862 -- -------
863
864 -- | Left fold
865 foldl :: Vector v b => (a -> b -> a) -> a -> v b -> a
866 {-# INLINE foldl #-}
867 foldl f z = Stream.foldl f z . stream
868
869 -- | Left fold on non-empty vectors
870 foldl1 :: Vector v a => (a -> a -> a) -> v a -> a
871 {-# INLINE foldl1 #-}
872 foldl1 f = Stream.foldl1 f . stream
873
874 -- | Left fold with strict accumulator
875 foldl' :: Vector v b => (a -> b -> a) -> a -> v b -> a
876 {-# INLINE foldl' #-}
877 foldl' f z = Stream.foldl' f z . stream
878
879 -- | Left fold on non-empty vectors with strict accumulator
880 foldl1' :: Vector v a => (a -> a -> a) -> v a -> a
881 {-# INLINE foldl1' #-}
882 foldl1' f = Stream.foldl1' f . stream
883
884 -- | Right fold
885 foldr :: Vector v a => (a -> b -> b) -> b -> v a -> b
886 {-# INLINE foldr #-}
887 foldr f z = Stream.foldr f z . stream
888
889 -- | Right fold on non-empty vectors
890 foldr1 :: Vector v a => (a -> a -> a) -> v a -> a
891 {-# INLINE foldr1 #-}
892 foldr1 f = Stream.foldr1 f . stream
893
894 -- | Left fold (function applied to each element and its index)
895 ifoldl :: Vector v b => (a -> Int -> b -> a) -> a -> v b -> a
896 {-# INLINE ifoldl #-}
897 ifoldl f z = Stream.foldl (uncurry . f) z . Stream.indexed . stream
898
899 -- | Left fold with strict accumulator (function applied to each element and
900 -- its index)
901 ifoldl' :: Vector v b => (a -> Int -> b -> a) -> a -> v b -> a
902 {-# INLINE ifoldl' #-}
903 ifoldl' f z = Stream.foldl' (uncurry . f) z . Stream.indexed . stream
904
905 -- | Right fold (function applied to each element and its index)
906 ifoldr :: Vector v a => (Int -> a -> b -> b) -> b -> v a -> b
907 {-# INLINE ifoldr #-}
908 ifoldr f z = Stream.foldr (uncurry f) z . Stream.indexed . stream
909
910 -- Specialised folds
911 -- -----------------
912
913 all :: Vector v a => (a -> Bool) -> v a -> Bool
914 {-# INLINE all #-}
915 all f = Stream.and . Stream.map f . stream
916
917 any :: Vector v a => (a -> Bool) -> v a -> Bool
918 {-# INLINE any #-}
919 any f = Stream.or . Stream.map f . stream
920
921 and :: Vector v Bool => v Bool -> Bool
922 {-# INLINE and #-}
923 and = Stream.and . stream
924
925 or :: Vector v Bool => v Bool -> Bool
926 {-# INLINE or #-}
927 or = Stream.or . stream
928
929 sum :: (Vector v a, Num a) => v a -> a
930 {-# INLINE sum #-}
931 sum = Stream.foldl' (+) 0 . stream
932
933 product :: (Vector v a, Num a) => v a -> a
934 {-# INLINE product #-}
935 product = Stream.foldl' (*) 1 . stream
936
937 maximum :: (Vector v a, Ord a) => v a -> a
938 {-# INLINE maximum #-}
939 maximum = Stream.foldl1' max . stream
940
941 maximumBy :: Vector v a => (a -> a -> Ordering) -> v a -> a
942 {-# INLINE maximumBy #-}
943 maximumBy cmp = Stream.foldl1' maxBy . stream
944 where
945 {-# INLINE maxBy #-}
946 maxBy x y = case cmp x y of
947 LT -> y
948 _ -> x
949
950 minimum :: (Vector v a, Ord a) => v a -> a
951 {-# INLINE minimum #-}
952 minimum = Stream.foldl1' min . stream
953
954 minimumBy :: Vector v a => (a -> a -> Ordering) -> v a -> a
955 {-# INLINE minimumBy #-}
956 minimumBy cmp = Stream.foldl1' minBy . stream
957 where
958 {-# INLINE minBy #-}
959 minBy x y = case cmp x y of
960 GT -> y
961 _ -> x
962
963 maxIndex :: (Vector v a, Ord a) => v a -> Int
964 {-# INLINE maxIndex #-}
965 maxIndex = maxIndexBy compare
966
967 maxIndexBy :: Vector v a => (a -> a -> Ordering) -> v a -> Int
968 {-# INLINE maxIndexBy #-}
969 maxIndexBy cmp = fst . Stream.foldl1' imax . Stream.indexed . stream
970 where
971 imax (i,x) (j,y) = case cmp x y of
972 LT -> (j,y)
973 _ -> (i,x)
974
975 minIndex :: (Vector v a, Ord a) => v a -> Int
976 {-# INLINE minIndex #-}
977 minIndex = minIndexBy compare
978
979 minIndexBy :: Vector v a => (a -> a -> Ordering) -> v a -> Int
980 {-# INLINE minIndexBy #-}
981 minIndexBy cmp = fst . Stream.foldl1' imin . Stream.indexed . stream
982 where
983 imin (i,x) (j,y) = case cmp x y of
984 GT -> (j,y)
985 _ -> (i,x)
986
987
988 -- Unfolding
989 -- ---------
990
991 unfoldr :: Vector v a => (b -> Maybe (a, b)) -> b -> v a
992 {-# INLINE unfoldr #-}
993 unfoldr f = unstream . Stream.unfoldr f
994
995 -- Scans
996 -- -----
997
998 -- | Prefix scan
999 prescanl :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1000 {-# INLINE prescanl #-}
1001 prescanl f z = unstream . inplace (MStream.prescanl f z) . stream
1002
1003 -- | Prefix scan with strict accumulator
1004 prescanl' :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1005 {-# INLINE prescanl' #-}
1006 prescanl' f z = unstream . inplace (MStream.prescanl' f z) . stream
1007
1008 -- | Suffix scan
1009 postscanl :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1010 {-# INLINE postscanl #-}
1011 postscanl f z = unstream . inplace (MStream.postscanl f z) . stream
1012
1013 -- | Suffix scan with strict accumulator
1014 postscanl' :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1015 {-# INLINE postscanl' #-}
1016 postscanl' f z = unstream . inplace (MStream.postscanl' f z) . stream
1017
1018 -- | Haskell-style scan
1019 scanl :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1020 {-# INLINE scanl #-}
1021 scanl f z = unstream . Stream.scanl f z . stream
1022
1023 -- | Haskell-style scan with strict accumulator
1024 scanl' :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1025 {-# INLINE scanl' #-}
1026 scanl' f z = unstream . Stream.scanl' f z . stream
1027
1028 -- | Scan over a non-empty vector
1029 scanl1 :: Vector v a => (a -> a -> a) -> v a -> v a
1030 {-# INLINE scanl1 #-}
1031 scanl1 f = unstream . inplace (MStream.scanl1 f) . stream
1032
1033 -- | Scan over a non-empty vector with a strict accumulator
1034 scanl1' :: Vector v a => (a -> a -> a) -> v a -> v a
1035 {-# INLINE scanl1' #-}
1036 scanl1' f = unstream . inplace (MStream.scanl1' f) . stream
1037
1038 -- Enumeration
1039 -- -----------
1040
1041 enumFromTo :: (Vector v a, Enum a) => a -> a -> v a
1042 {-# INLINE enumFromTo #-}
1043 enumFromTo x y = unstream (Stream.enumFromTo x y)
1044
1045 enumFromThenTo :: (Vector v a, Enum a) => a -> a -> a -> v a
1046 {-# INLINE enumFromThenTo #-}
1047 enumFromThenTo x y z = unstream (Stream.enumFromThenTo x y z)
1048
1049 -- | Convert a vector to a list
1050 toList :: Vector v a => v a -> [a]
1051 {-# INLINE toList #-}
1052 toList = Stream.toList . stream
1053
1054 -- | Convert a list to a vector
1055 fromList :: Vector v a => [a] -> v a
1056 {-# INLINE fromList #-}
1057 fromList = unstream . Stream.fromList
1058