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