Add monadic vector combinators
[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-2010
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, (++), force,
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 partition, unstablePartition, span, break,
57
58 -- * Searching
59 elem, notElem, find, findIndex, findIndices, elemIndex, elemIndices,
60
61 -- * Folding
62 foldl, foldl1, foldl', foldl1', foldr, foldr1, foldr', foldr1',
63 ifoldl, ifoldl', ifoldr, 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, unfoldrN,
73
74 -- * Scans
75 prescanl, prescanl',
76 postscanl, postscanl',
77 scanl, scanl', scanl1, scanl1',
78 prescanr, prescanr',
79 postscanr, postscanr',
80 scanr, scanr', scanr1, scanr1',
81
82 -- * Enumeration
83 enumFromN, enumFromStepN, enumFromTo, enumFromThenTo,
84
85 -- * Conversion to/from lists
86 toList, fromList, fromListN,
87
88 -- * Monadic operations
89 replicateM, mapM, mapM_, forM, forM_, zipWithM, zipWithM_, filterM,
90 foldM, foldM', fold1M, fold1M',
91
92 -- * Destructive operations
93 create, modify, copy, unsafeCopy,
94
95 -- * Conversion to/from Streams
96 stream, unstream, streamR, unstreamR,
97
98 -- * Recycling support
99 new, clone,
100
101 -- * Utilities for defining Data instances
102 gfoldl, dataCast, mkType
103 ) where
104
105 import Data.Vector.Generic.Base
106
107 import Data.Vector.Generic.Mutable ( MVector )
108 import qualified Data.Vector.Generic.Mutable as M
109
110 import qualified Data.Vector.Generic.New as New
111 import Data.Vector.Generic.New ( New )
112
113 import qualified Data.Vector.Fusion.Stream as Stream
114 import Data.Vector.Fusion.Stream ( Stream, MStream, inplace )
115 import qualified Data.Vector.Fusion.Stream.Monadic as MStream
116 import Data.Vector.Fusion.Stream.Size
117 import Data.Vector.Fusion.Util
118
119 import Control.Monad.ST ( ST, runST )
120 import Control.Monad.Primitive
121 import qualified Control.Monad as Monad
122 import Prelude hiding ( length, null,
123 replicate, (++),
124 head, last,
125 init, tail, take, drop, reverse,
126 map, concatMap,
127 zipWith, zipWith3, zip, zip3, unzip, unzip3,
128 filter, takeWhile, dropWhile, span, break,
129 elem, notElem,
130 foldl, foldl1, foldr, foldr1,
131 all, any, and, or, sum, product, maximum, minimum,
132 scanl, scanl1, scanr, scanr1,
133 enumFromTo, enumFromThenTo,
134 mapM, mapM_ )
135
136 import Data.Typeable ( Typeable1, gcast1 )
137 import Data.Data ( Data, DataType, mkNorepType )
138
139 #include "vector.h"
140
141 -- Fusion
142 -- ------
143
144 -- | Construct a pure vector from a monadic initialiser
145 new :: Vector v a => New v a -> v a
146 {-# INLINE_STREAM new #-}
147 new m = m `seq` runST (unsafeFreeze =<< New.run m)
148
149 clone :: Vector v a => v a -> New v a
150 {-# INLINE_STREAM clone #-}
151 clone v = v `seq` New.create (
152 do
153 mv <- M.new (length v)
154 unsafeCopy mv v
155 return mv)
156
157 -- | Convert a vector to a 'Stream'
158 stream :: Vector v a => v a -> Stream a
159 {-# INLINE_STREAM stream #-}
160 stream v = v `seq` (Stream.unfoldr get 0 `Stream.sized` Exact n)
161 where
162 n = length v
163
164 -- NOTE: the False case comes first in Core so making it the recursive one
165 -- makes the code easier to read
166 {-# INLINE get #-}
167 get i | i >= n = Nothing
168 | otherwise = case basicUnsafeIndexM v i of Box x -> Just (x, i+1)
169
170 -- | Create a vector from a 'Stream'
171 unstream :: Vector v a => Stream a -> v a
172 {-# INLINE unstream #-}
173 unstream s = new (New.unstream s)
174
175 {-# RULES
176
177 "stream/unstream [Vector]" forall s.
178 stream (new (New.unstream s)) = s
179
180 "New.unstream/stream [Vector]" forall v.
181 New.unstream (stream v) = clone v
182
183 "clone/new [Vector]" forall p.
184 clone (new p) = p
185
186 "inplace [Vector]"
187 forall (f :: forall m. Monad m => MStream m a -> MStream m a) m.
188 New.unstream (inplace f (stream (new m))) = New.transform f m
189
190 "uninplace [Vector]"
191 forall (f :: forall m. Monad m => MStream m a -> MStream m a) m.
192 stream (new (New.transform f m)) = inplace f (stream (new m))
193
194 #-}
195
196 -- | Convert a vector to a 'Stream'
197 streamR :: Vector v a => v a -> Stream a
198 {-# INLINE_STREAM streamR #-}
199 streamR v = v `seq` (Stream.unfoldr get n `Stream.sized` Exact n)
200 where
201 n = length v
202
203 {-# INLINE get #-}
204 get 0 = Nothing
205 get i = let i' = i-1
206 in
207 case basicUnsafeIndexM v i' of Box x -> Just (x, i')
208
209 -- | Create a vector from a 'Stream'
210 unstreamR :: Vector v a => Stream a -> v a
211 {-# INLINE unstreamR #-}
212 unstreamR s = new (New.unstreamR s)
213
214 {-# RULES
215
216 "streamR/unstreamR [Vector]" forall s.
217 streamR (new (New.unstreamR s)) = s
218
219 "New.unstreamR/streamR/new [Vector]" forall p.
220 New.unstreamR (streamR (new p)) = p
221
222 "inplace right [Vector]"
223 forall (f :: forall m. Monad m => MStream m a -> MStream m a) m.
224 New.unstreamR (inplace f (streamR (new m))) = New.transformR f m
225
226 "uninplace right [Vector]"
227 forall (f :: forall m. Monad m => MStream m a -> MStream m a) m.
228 streamR (new (New.transformR f m)) = inplace f (streamR (new m))
229
230 #-}
231
232 -- Length
233 -- ------
234
235 length :: Vector v a => v a -> Int
236 {-# INLINE_STREAM length #-}
237 length v = basicLength v
238
239 {-# RULES
240
241 "length/unstream [Vector]" forall s.
242 length (new (New.unstream s)) = Stream.length s
243
244 #-}
245
246 null :: Vector v a => v a -> Bool
247 {-# INLINE_STREAM null #-}
248 null v = basicLength v == 0
249
250 {-# RULES
251
252 "null/unstream [Vector]" forall s.
253 null (new (New.unstream s)) = Stream.null s
254
255 #-}
256
257 -- Construction
258 -- ------------
259
260 -- | Empty vector
261 empty :: Vector v a => v a
262 {-# INLINE empty #-}
263 empty = unstream Stream.empty
264
265 -- | Vector with exaclty one element
266 singleton :: forall v a. Vector v a => a -> v a
267 {-# INLINE singleton #-}
268 singleton x = elemseq (undefined :: v a) x
269 $ unstream (Stream.singleton x)
270
271 -- | Vector of the given length with the given value in each position
272 replicate :: forall v a. Vector v a => Int -> a -> v a
273 {-# INLINE replicate #-}
274 replicate n x = elemseq (undefined :: v a) x
275 $ unstream
276 $ Stream.replicate n x
277
278 -- | Generate a vector of the given length by applying the function to each
279 -- index
280 generate :: Vector v a => Int -> (Int -> a) -> v a
281 {-# INLINE generate #-}
282 generate n f = unstream (Stream.generate n f)
283
284 -- | Prepend an element
285 cons :: forall v a. Vector v a => a -> v a -> v a
286 {-# INLINE cons #-}
287 cons x v = elemseq (undefined :: v a) x
288 $ unstream
289 $ Stream.cons x
290 $ stream v
291
292 -- | Append an element
293 snoc :: forall v a. Vector v a => v a -> a -> v a
294 {-# INLINE snoc #-}
295 snoc v x = elemseq (undefined :: v a) x
296 $ unstream
297 $ Stream.snoc (stream v) x
298
299 infixr 5 ++
300 -- | Concatenate two vectors
301 (++) :: Vector v a => v a -> v a -> v a
302 {-# INLINE (++) #-}
303 v ++ w = unstream (stream v Stream.++ stream w)
304
305 -- | Create a copy of a vector. Useful when dealing with slices.
306 force :: Vector v a => v a -> v a
307 {-# INLINE_STREAM force #-}
308 force v = new (clone v)
309
310 -- Accessing individual elements
311 -- -----------------------------
312
313 -- | Indexing
314 (!) :: Vector v a => v a -> Int -> a
315 {-# INLINE_STREAM (!) #-}
316 v ! i = BOUNDS_CHECK(checkIndex) "(!)" i (length v)
317 $ unId (basicUnsafeIndexM v i)
318
319 -- | First element
320 head :: Vector v a => v a -> a
321 {-# INLINE_STREAM head #-}
322 head v = v ! 0
323
324 -- | Last element
325 last :: Vector v a => v a -> a
326 {-# INLINE_STREAM last #-}
327 last v = v ! (length v - 1)
328
329 -- | Unsafe indexing without bounds checking
330 unsafeIndex :: Vector v a => v a -> Int -> a
331 {-# INLINE_STREAM unsafeIndex #-}
332 unsafeIndex v i = UNSAFE_CHECK(checkIndex) "unsafeIndex" i (length v)
333 $ unId (basicUnsafeIndexM v i)
334
335 -- | Yield the first element of a vector without checking if the vector is
336 -- empty
337 unsafeHead :: Vector v a => v a -> a
338 {-# INLINE_STREAM unsafeHead #-}
339 unsafeHead v = unsafeIndex v 0
340
341 -- | Yield the last element of a vector without checking if the vector is
342 -- empty
343 unsafeLast :: Vector v a => v a -> a
344 {-# INLINE_STREAM unsafeLast #-}
345 unsafeLast v = unsafeIndex v (length v - 1)
346
347 {-# RULES
348
349 "(!)/unstream [Vector]" forall i s.
350 new (New.unstream s) ! i = s Stream.!! i
351
352 "head/unstream [Vector]" forall s.
353 head (new (New.unstream s)) = Stream.head s
354
355 "last/unstream [Vector]" forall s.
356 last (new (New.unstream s)) = Stream.last s
357
358 "unsafeIndex/unstream [Vector]" forall i s.
359 unsafeIndex (new (New.unstream s)) i = s Stream.!! i
360
361 "unsafeHead/unstream [Vector]" forall s.
362 unsafeHead (new (New.unstream s)) = Stream.head s
363
364 "unsafeLast/unstream [Vector]" forall s.
365 unsafeLast (new (New.unstream s)) = Stream.last s
366
367 #-}
368
369 -- | Monadic indexing which can be strict in the vector while remaining lazy in
370 -- the element.
371 indexM :: (Vector v a, Monad m) => v a -> Int -> m a
372 {-# INLINE_STREAM indexM #-}
373 indexM v i = BOUNDS_CHECK(checkIndex) "indexM" i (length v)
374 $ basicUnsafeIndexM v i
375
376 headM :: (Vector v a, Monad m) => v a -> m a
377 {-# INLINE_STREAM headM #-}
378 headM v = indexM v 0
379
380 lastM :: (Vector v a, Monad m) => v a -> m a
381 {-# INLINE_STREAM lastM #-}
382 lastM v = indexM v (length v - 1)
383
384 -- | Unsafe monadic indexing without bounds checks
385 unsafeIndexM :: (Vector v a, Monad m) => v a -> Int -> m a
386 {-# INLINE_STREAM unsafeIndexM #-}
387 unsafeIndexM v i = UNSAFE_CHECK(checkIndex) "unsafeIndexM" i (length v)
388 $ basicUnsafeIndexM v i
389
390 unsafeHeadM :: (Vector v a, Monad m) => v a -> m a
391 {-# INLINE_STREAM unsafeHeadM #-}
392 unsafeHeadM v = unsafeIndexM v 0
393
394 unsafeLastM :: (Vector v a, Monad m) => v a -> m a
395 {-# INLINE_STREAM unsafeLastM #-}
396 unsafeLastM v = unsafeIndexM v (length v - 1)
397
398 -- FIXME: the rhs of these rules are lazy in the stream which is WRONG
399 {- RULES
400
401 "indexM/unstream [Vector]" forall v i s.
402 indexM (new' v (New.unstream s)) i = return (s Stream.!! i)
403
404 "headM/unstream [Vector]" forall v s.
405 headM (new' v (New.unstream s)) = return (Stream.head s)
406
407 "lastM/unstream [Vector]" forall v s.
408 lastM (new' v (New.unstream s)) = return (Stream.last s)
409
410 -}
411
412 -- Subarrays
413 -- ---------
414
415 -- | Yield a part of the vector without copying it.
416 slice :: Vector v a => Int -- ^ starting index
417 -> Int -- ^ length
418 -> v a
419 -> v a
420 {-# INLINE_STREAM slice #-}
421 slice i n v = BOUNDS_CHECK(checkSlice) "slice" i n (length v)
422 $ basicUnsafeSlice i n v
423
424 -- | Yield all but the last element without copying.
425 init :: Vector v a => v a -> v a
426 {-# INLINE_STREAM init #-}
427 init v = slice 0 (length v - 1) v
428
429 -- | All but the first element (without copying).
430 tail :: Vector v a => v a -> v a
431 {-# INLINE_STREAM tail #-}
432 tail v = slice 1 (length v - 1) v
433
434 -- | Yield the first @n@ elements without copying.
435 take :: Vector v a => Int -> v a -> v a
436 {-# INLINE_STREAM take #-}
437 take n v = unsafeSlice 0 (delay_inline min n' (length v)) v
438 where n' = max n 0
439
440 -- | Yield all but the first @n@ elements without copying.
441 drop :: Vector v a => Int -> v a -> v a
442 {-# INLINE_STREAM drop #-}
443 drop n v = unsafeSlice (delay_inline min n' len)
444 (delay_inline max 0 (len - n')) v
445 where n' = max n 0
446 len = length v
447
448 -- | Unsafely yield a part of the vector without copying it and without
449 -- performing bounds checks.
450 unsafeSlice :: Vector v a => Int -- ^ starting index
451 -> Int -- ^ length
452 -> v a
453 -> v a
454 {-# INLINE_STREAM unsafeSlice #-}
455 unsafeSlice i n v = UNSAFE_CHECK(checkSlice) "unsafeSlice" i n (length v)
456 $ basicUnsafeSlice i n v
457
458 unsafeInit :: Vector v a => v a -> v a
459 {-# INLINE_STREAM unsafeInit #-}
460 unsafeInit v = unsafeSlice 0 (length v - 1) v
461
462 unsafeTail :: Vector v a => v a -> v a
463 {-# INLINE_STREAM unsafeTail #-}
464 unsafeTail v = unsafeSlice 1 (length v - 1) v
465
466 unsafeTake :: Vector v a => Int -> v a -> v a
467 {-# INLINE unsafeTake #-}
468 unsafeTake n v = unsafeSlice 0 n v
469
470 unsafeDrop :: Vector v a => Int -> v a -> v a
471 {-# INLINE unsafeDrop #-}
472 unsafeDrop n v = unsafeSlice n (length v - n) v
473
474 {-# RULES
475
476 "slice/new [Vector]" forall i n p.
477 slice i n (new p) = new (New.slice i n p)
478
479 "init/new [Vector]" forall p.
480 init (new p) = new (New.init p)
481
482 "tail/new [Vector]" forall p.
483 tail (new p) = new (New.tail p)
484
485 "take/new [Vector]" forall n p.
486 take n (new p) = new (New.take n p)
487
488 "drop/new [Vector]" forall n p.
489 drop n (new p) = new (New.drop n p)
490
491 "unsafeSlice/new [Vector]" forall i n p.
492 unsafeSlice i n (new p) = new (New.unsafeSlice i n p)
493
494 "unsafeInit/new [Vector]" forall p.
495 unsafeInit (new p) = new (New.unsafeInit p)
496
497 "unsafeTail/new [Vector]" forall p.
498 unsafeTail (new p) = new (New.unsafeTail p)
499
500 #-}
501
502 -- Permutations
503 -- ------------
504
505 unsafeAccum_stream
506 :: Vector v a => (a -> b -> a) -> v a -> Stream (Int,b) -> v a
507 {-# INLINE unsafeAccum_stream #-}
508 unsafeAccum_stream f v s = s `seq` modify (\mv -> M.unsafeAccum f mv s) v
509
510 unsafeAccum :: Vector v a => (a -> b -> a) -> v a -> [(Int,b)] -> v a
511 {-# INLINE unsafeAccum #-}
512 unsafeAccum f v us = unsafeAccum_stream f v (Stream.fromList us)
513
514 unsafeAccumulate :: (Vector v a, Vector v (Int, b))
515 => (a -> b -> a) -> v a -> v (Int,b) -> v a
516 {-# INLINE unsafeAccumulate #-}
517 unsafeAccumulate f v us = unsafeAccum_stream f v (stream us)
518
519 unsafeAccumulate_ :: (Vector v a, Vector v Int, Vector v b)
520 => (a -> b -> a) -> v a -> v Int -> v b -> v a
521 {-# INLINE unsafeAccumulate_ #-}
522 unsafeAccumulate_ f v is xs
523 = unsafeAccum_stream f v (Stream.zipWith (,) (stream is) (stream xs))
524
525 accum_stream :: Vector v a => (a -> b -> a) -> v a -> Stream (Int,b) -> v a
526 {-# INLINE accum_stream #-}
527 accum_stream f v s = s `seq` modify (\mv -> M.accum f mv s) v
528
529 accum :: Vector v a => (a -> b -> a) -> v a -> [(Int,b)] -> v a
530 {-# INLINE accum #-}
531 accum f v us = accum_stream f v (Stream.fromList us)
532
533 accumulate :: (Vector v a, Vector v (Int, b))
534 => (a -> b -> a) -> v a -> v (Int,b) -> v a
535 {-# INLINE accumulate #-}
536 accumulate f v us = accum_stream f v (stream us)
537
538 accumulate_ :: (Vector v a, Vector v Int, Vector v b)
539 => (a -> b -> a) -> v a -> v Int -> v b -> v a
540 {-# INLINE accumulate_ #-}
541 accumulate_ f v is xs = accum_stream f v (Stream.zipWith (,) (stream is)
542 (stream xs))
543
544
545 unsafeUpdate_stream :: Vector v a => v a -> Stream (Int,a) -> v a
546 {-# INLINE unsafeUpdate_stream #-}
547 unsafeUpdate_stream v s = s `seq` modify (\mv -> M.unsafeUpdate mv s) v
548
549 unsafeUpd :: Vector v a => v a -> [(Int, a)] -> v a
550 {-# INLINE unsafeUpd #-}
551 unsafeUpd v us = unsafeUpdate_stream v (Stream.fromList us)
552
553 unsafeUpdate :: (Vector v a, Vector v (Int, a)) => v a -> v (Int, a) -> v a
554 {-# INLINE unsafeUpdate #-}
555 unsafeUpdate v w = unsafeUpdate_stream v (stream w)
556
557 unsafeUpdate_ :: (Vector v a, Vector v Int) => v a -> v Int -> v a -> v a
558 {-# INLINE unsafeUpdate_ #-}
559 unsafeUpdate_ v is w
560 = unsafeUpdate_stream v (Stream.zipWith (,) (stream is) (stream w))
561
562 update_stream :: Vector v a => v a -> Stream (Int,a) -> v a
563 {-# INLINE update_stream #-}
564 update_stream v s = s `seq` modify (\mv -> M.update mv s) v
565
566 (//) :: Vector v a => v a -> [(Int, a)] -> v a
567 {-# INLINE (//) #-}
568 v // us = update_stream v (Stream.fromList us)
569
570 update :: (Vector v a, Vector v (Int, a)) => v a -> v (Int, a) -> v a
571 {-# INLINE update #-}
572 update v w = update_stream v (stream w)
573
574 update_ :: (Vector v a, Vector v Int) => v a -> v Int -> v a -> v a
575 {-# INLINE update_ #-}
576 update_ v is w = update_stream v (Stream.zipWith (,) (stream is) (stream w))
577
578 -- This somewhat non-intuitive definition ensures that the resulting vector
579 -- does not retain references to the original one even if it is lazy in its
580 -- elements. This would not be the case if we simply used
581 --
582 -- backpermute v is = map (v!) is
583 backpermute :: (Vector v a, Vector v Int) => v a -> v Int -> v a
584 {-# INLINE backpermute #-}
585 backpermute v is = seq v
586 $ unstream
587 $ Stream.unbox
588 $ Stream.map (indexM v)
589 $ stream is
590
591 unsafeBackpermute :: (Vector v a, Vector v Int) => v a -> v Int -> v a
592 {-# INLINE unsafeBackpermute #-}
593 unsafeBackpermute v is = seq v
594 $ unstream
595 $ Stream.unbox
596 $ Stream.map (unsafeIndexM v)
597 $ stream is
598
599 -- FIXME: make this fuse better, add support for recycling
600 reverse :: (Vector v a) => v a -> v a
601 {-# INLINE reverse #-}
602 reverse = unstream . streamR
603
604 -- Mapping
605 -- -------
606
607 -- | Map a function over a vector
608 map :: (Vector v a, Vector v b) => (a -> b) -> v a -> v b
609 {-# INLINE map #-}
610 map f = unstream . inplace (MStream.map f) . stream
611
612 -- | Apply a function to every index/value pair
613 imap :: (Vector v a, Vector v b) => (Int -> a -> b) -> v a -> v b
614 {-# INLINE imap #-}
615 imap f = unstream . inplace (MStream.map (uncurry f) . MStream.indexed)
616 . stream
617
618 concatMap :: (Vector v a, Vector v b) => (a -> v b) -> v a -> v b
619 {-# INLINE concatMap #-}
620 concatMap f = unstream . Stream.concatMap (stream . f) . stream
621
622 -- Zipping/unzipping
623 -- -----------------
624
625 -- | Zip two vectors with the given function.
626 zipWith :: (Vector v a, Vector v b, Vector v c)
627 => (a -> b -> c) -> v a -> v b -> v c
628 {-# INLINE zipWith #-}
629 zipWith f xs ys = unstream (Stream.zipWith f (stream xs) (stream ys))
630
631 -- | Zip three vectors with the given function.
632 zipWith3 :: (Vector v a, Vector v b, Vector v c, Vector v d)
633 => (a -> b -> c -> d) -> v a -> v b -> v c -> v d
634 {-# INLINE zipWith3 #-}
635 zipWith3 f as bs cs = unstream (Stream.zipWith3 f (stream as)
636 (stream bs)
637 (stream cs))
638
639 zipWith4 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e)
640 => (a -> b -> c -> d -> e) -> v a -> v b -> v c -> v d -> v e
641 {-# INLINE zipWith4 #-}
642 zipWith4 f as bs cs ds
643 = unstream (Stream.zipWith4 f (stream as)
644 (stream bs)
645 (stream cs)
646 (stream ds))
647
648 zipWith5 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
649 Vector v f)
650 => (a -> b -> c -> d -> e -> f) -> v a -> v b -> v c -> v d -> v e
651 -> v f
652 {-# INLINE zipWith5 #-}
653 zipWith5 f as bs cs ds es
654 = unstream (Stream.zipWith5 f (stream as)
655 (stream bs)
656 (stream cs)
657 (stream ds)
658 (stream es))
659
660 zipWith6 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
661 Vector v f, Vector v g)
662 => (a -> b -> c -> d -> e -> f -> g)
663 -> v a -> v b -> v c -> v d -> v e -> v f -> v g
664 {-# INLINE zipWith6 #-}
665 zipWith6 f as bs cs ds es fs
666 = unstream (Stream.zipWith6 f (stream as)
667 (stream bs)
668 (stream cs)
669 (stream ds)
670 (stream es)
671 (stream fs))
672
673 -- | Zip two vectors and their indices with the given function.
674 izipWith :: (Vector v a, Vector v b, Vector v c)
675 => (Int -> a -> b -> c) -> v a -> v b -> v c
676 {-# INLINE izipWith #-}
677 izipWith f xs ys = unstream
678 (Stream.zipWith (uncurry f) (Stream.indexed (stream xs))
679 (stream ys))
680
681 -- | Zip three vectors and their indices with the given function.
682 izipWith3 :: (Vector v a, Vector v b, Vector v c, Vector v d)
683 => (Int -> a -> b -> c -> d) -> v a -> v b -> v c -> v d
684 {-# INLINE izipWith3 #-}
685 izipWith3 f as bs cs
686 = unstream (Stream.zipWith3 (uncurry f) (Stream.indexed (stream as))
687 (stream bs)
688 (stream cs))
689
690 izipWith4 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e)
691 => (Int -> a -> b -> c -> d -> e) -> v a -> v b -> v c -> v d -> v e
692 {-# INLINE izipWith4 #-}
693 izipWith4 f as bs cs ds
694 = unstream (Stream.zipWith4 (uncurry f) (Stream.indexed (stream as))
695 (stream bs)
696 (stream cs)
697 (stream ds))
698
699 izipWith5 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
700 Vector v f)
701 => (Int -> a -> b -> c -> d -> e -> f) -> v a -> v b -> v c -> v d
702 -> v e -> v f
703 {-# INLINE izipWith5 #-}
704 izipWith5 f as bs cs ds es
705 = unstream (Stream.zipWith5 (uncurry f) (Stream.indexed (stream as))
706 (stream bs)
707 (stream cs)
708 (stream ds)
709 (stream es))
710
711 izipWith6 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
712 Vector v f, Vector v g)
713 => (Int -> a -> b -> c -> d -> e -> f -> g)
714 -> v a -> v b -> v c -> v d -> v e -> v f -> v g
715 {-# INLINE izipWith6 #-}
716 izipWith6 f as bs cs ds es fs
717 = unstream (Stream.zipWith6 (uncurry f) (Stream.indexed (stream as))
718 (stream bs)
719 (stream cs)
720 (stream ds)
721 (stream es)
722 (stream fs))
723
724 zip :: (Vector v a, Vector v b, Vector v (a,b)) => v a -> v b -> v (a, b)
725 {-# INLINE zip #-}
726 zip = zipWith (,)
727
728 zip3 :: (Vector v a, Vector v b, Vector v c, Vector v (a, b, c))
729 => v a -> v b -> v c -> v (a, b, c)
730 {-# INLINE zip3 #-}
731 zip3 = zipWith3 (,,)
732
733 zip4 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v (a, b, c, d))
734 => v a -> v b -> v c -> v d -> v (a, b, c, d)
735 {-# INLINE zip4 #-}
736 zip4 = zipWith4 (,,,)
737
738 zip5 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
739 Vector v (a, b, c, d, e))
740 => v a -> v b -> v c -> v d -> v e -> v (a, b, c, d, e)
741 {-# INLINE zip5 #-}
742 zip5 = zipWith5 (,,,,)
743
744 zip6 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
745 Vector v f, Vector v (a, b, c, d, e, f))
746 => v a -> v b -> v c -> v d -> v e -> v f -> v (a, b, c, d, e, f)
747 {-# INLINE zip6 #-}
748 zip6 = zipWith6 (,,,,,)
749
750 unzip :: (Vector v a, Vector v b, Vector v (a,b)) => v (a, b) -> (v a, v b)
751 {-# INLINE unzip #-}
752 unzip xs = (map fst xs, map snd xs)
753
754 unzip3 :: (Vector v a, Vector v b, Vector v c, Vector v (a, b, c))
755 => v (a, b, c) -> (v a, v b, v c)
756 {-# INLINE unzip3 #-}
757 unzip3 xs = (map (\(a, b, c) -> a) xs,
758 map (\(a, b, c) -> b) xs,
759 map (\(a, b, c) -> c) xs)
760
761 unzip4 :: (Vector v a, Vector v b, Vector v c, Vector v d,
762 Vector v (a, b, c, d))
763 => v (a, b, c, d) -> (v a, v b, v c, v d)
764 {-# INLINE unzip4 #-}
765 unzip4 xs = (map (\(a, b, c, d) -> a) xs,
766 map (\(a, b, c, d) -> b) xs,
767 map (\(a, b, c, d) -> c) xs,
768 map (\(a, b, c, d) -> d) xs)
769
770 unzip5 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
771 Vector v (a, b, c, d, e))
772 => v (a, b, c, d, e) -> (v a, v b, v c, v d, v e)
773 {-# INLINE unzip5 #-}
774 unzip5 xs = (map (\(a, b, c, d, e) -> a) xs,
775 map (\(a, b, c, d, e) -> b) xs,
776 map (\(a, b, c, d, e) -> c) xs,
777 map (\(a, b, c, d, e) -> d) xs,
778 map (\(a, b, c, d, e) -> e) xs)
779
780 unzip6 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
781 Vector v f, Vector v (a, b, c, d, e, f))
782 => v (a, b, c, d, e, f) -> (v a, v b, v c, v d, v e, v f)
783 {-# INLINE unzip6 #-}
784 unzip6 xs = (map (\(a, b, c, d, e, f) -> a) xs,
785 map (\(a, b, c, d, e, f) -> b) xs,
786 map (\(a, b, c, d, e, f) -> c) xs,
787 map (\(a, b, c, d, e, f) -> d) xs,
788 map (\(a, b, c, d, e, f) -> e) xs,
789 map (\(a, b, c, d, e, f) -> f) xs)
790
791 -- Comparisons
792 -- -----------
793
794 eq :: (Vector v a, Eq a) => v a -> v a -> Bool
795 {-# INLINE eq #-}
796 xs `eq` ys = stream xs == stream ys
797
798 cmp :: (Vector v a, Ord a) => v a -> v a -> Ordering
799 {-# INLINE cmp #-}
800 cmp xs ys = compare (stream xs) (stream ys)
801
802 -- Filtering
803 -- ---------
804
805 -- | Drop elements that do not satisfy the predicate
806 filter :: Vector v a => (a -> Bool) -> v a -> v a
807 {-# INLINE filter #-}
808 filter f = unstream . inplace (MStream.filter f) . stream
809
810 -- | Drop elements that do not satisfy the predicate (applied to values and
811 -- their indices)
812 ifilter :: Vector v a => (Int -> a -> Bool) -> v a -> v a
813 {-# INLINE ifilter #-}
814 ifilter f = unstream
815 . inplace (MStream.map snd . MStream.filter (uncurry f)
816 . MStream.indexed)
817 . stream
818
819 -- | Yield the longest prefix of elements satisfying the predicate.
820 takeWhile :: Vector v a => (a -> Bool) -> v a -> v a
821 {-# INLINE takeWhile #-}
822 takeWhile f = unstream . Stream.takeWhile f . stream
823
824 -- | Drop the longest prefix of elements that satisfy the predicate.
825 dropWhile :: Vector v a => (a -> Bool) -> v a -> v a
826 {-# INLINE dropWhile #-}
827 dropWhile f = unstream . Stream.dropWhile f . stream
828
829 -- | Split the vector in two parts, the first one containing those elements
830 -- that satisfy the predicate and the second one those that don't. The
831 -- relative order of the elements is preserved at the cost of a (sometimes)
832 -- reduced performance compared to 'unstablePartition'.
833 partition :: Vector v a => (a -> Bool) -> v a -> (v a, v a)
834 {-# INLINE partition #-}
835 partition f = partition_stream f . stream
836
837 -- FIXME: Make this inplace-fusible (look at how stable_partition is
838 -- implemented in C++)
839
840 partition_stream :: Vector v a => (a -> Bool) -> Stream a -> (v a, v a)
841 {-# INLINE_STREAM partition_stream #-}
842 partition_stream f s = s `seq` runST (
843 do
844 (mv1,mv2) <- M.partitionStream f s
845 v1 <- unsafeFreeze mv1
846 v2 <- unsafeFreeze mv2
847 return (v1,v2))
848
849 -- | Split the vector in two parts, the first one containing those elements
850 -- that satisfy the predicate and the second one those that don't. The order
851 -- of the elements is not preserved but the operation is often faster than
852 -- 'partition'.
853 unstablePartition :: Vector v a => (a -> Bool) -> v a -> (v a, v a)
854 {-# INLINE unstablePartition #-}
855 unstablePartition f = unstablePartition_stream f . stream
856
857 unstablePartition_stream
858 :: Vector v a => (a -> Bool) -> Stream a -> (v a, v a)
859 {-# INLINE_STREAM unstablePartition_stream #-}
860 unstablePartition_stream f s = s `seq` runST (
861 do
862 (mv1,mv2) <- M.unstablePartitionStream f s
863 v1 <- unsafeFreeze mv1
864 v2 <- unsafeFreeze mv2
865 return (v1,v2))
866
867 unstablePartition_new :: Vector v a => (a -> Bool) -> New v a -> (v a, v a)
868 {-# INLINE_STREAM unstablePartition_new #-}
869 unstablePartition_new f (New.New p) = runST (
870 do
871 mv <- p
872 i <- M.unstablePartition f mv
873 v <- unsafeFreeze mv
874 return (unsafeTake i v, unsafeDrop i v))
875
876 {-# RULES
877
878 "unstablePartition" forall f p.
879 unstablePartition_stream f (stream (new p))
880 = unstablePartition_new f p
881
882 #-}
883
884
885 -- FIXME: make span and break fusible
886
887 -- | Split the vector into the longest prefix of elements that satisfy the
888 -- predicate and the rest.
889 span :: Vector v a => (a -> Bool) -> v a -> (v a, v a)
890 {-# INLINE span #-}
891 span f = break (not . f)
892
893 -- | Split the vector into the longest prefix of elements that do not satisfy
894 -- the predicate and the rest.
895 break :: Vector v a => (a -> Bool) -> v a -> (v a, v a)
896 {-# INLINE break #-}
897 break f xs = case findIndex f xs of
898 Just i -> (unsafeSlice 0 i xs, unsafeSlice i (length xs - i) xs)
899 Nothing -> (xs, empty)
900
901
902 -- Searching
903 -- ---------
904
905 infix 4 `elem`
906 -- | Check whether the vector contains an element
907 elem :: (Vector v a, Eq a) => a -> v a -> Bool
908 {-# INLINE elem #-}
909 elem x = Stream.elem x . stream
910
911 infix 4 `notElem`
912 -- | Inverse of `elem`
913 notElem :: (Vector v a, Eq a) => a -> v a -> Bool
914 {-# INLINE notElem #-}
915 notElem x = Stream.notElem x . stream
916
917 -- | Yield 'Just' the first element matching the predicate or 'Nothing' if no
918 -- such element exists.
919 find :: Vector v a => (a -> Bool) -> v a -> Maybe a
920 {-# INLINE find #-}
921 find f = Stream.find f . stream
922
923 -- | Yield 'Just' the index of the first element matching the predicate or
924 -- 'Nothing' if no such element exists.
925 findIndex :: Vector v a => (a -> Bool) -> v a -> Maybe Int
926 {-# INLINE findIndex #-}
927 findIndex f = Stream.findIndex f . stream
928
929 -- | Yield the indices of elements satisfying the predicate
930 findIndices :: (Vector v a, Vector v Int) => (a -> Bool) -> v a -> v Int
931 {-# INLINE findIndices #-}
932 findIndices f = unstream
933 . inplace (MStream.map fst . MStream.filter (f . snd)
934 . MStream.indexed)
935 . stream
936
937 -- | Yield 'Just' the index of the first occurence of the given element or
938 -- 'Nothing' if the vector does not contain the element
939 elemIndex :: (Vector v a, Eq a) => a -> v a -> Maybe Int
940 {-# INLINE elemIndex #-}
941 elemIndex x = findIndex (x==)
942
943 -- | Yield the indices of all occurences of the given element
944 elemIndices :: (Vector v a, Vector v Int, Eq a) => a -> v a -> v Int
945 {-# INLINE elemIndices #-}
946 elemIndices x = findIndices (x==)
947
948 -- Folding
949 -- -------
950
951 -- | Left fold
952 foldl :: Vector v b => (a -> b -> a) -> a -> v b -> a
953 {-# INLINE foldl #-}
954 foldl f z = Stream.foldl f z . stream
955
956 -- | Left fold on non-empty vectors
957 foldl1 :: Vector v a => (a -> a -> a) -> v a -> a
958 {-# INLINE foldl1 #-}
959 foldl1 f = Stream.foldl1 f . stream
960
961 -- | Left fold with strict accumulator
962 foldl' :: Vector v b => (a -> b -> a) -> a -> v b -> a
963 {-# INLINE foldl' #-}
964 foldl' f z = Stream.foldl' f z . stream
965
966 -- | Left fold on non-empty vectors with strict accumulator
967 foldl1' :: Vector v a => (a -> a -> a) -> v a -> a
968 {-# INLINE foldl1' #-}
969 foldl1' f = Stream.foldl1' f . stream
970
971 -- | Right fold
972 foldr :: Vector v a => (a -> b -> b) -> b -> v a -> b
973 {-# INLINE foldr #-}
974 foldr f z = Stream.foldr f z . stream
975
976 -- | Right fold on non-empty vectors
977 foldr1 :: Vector v a => (a -> a -> a) -> v a -> a
978 {-# INLINE foldr1 #-}
979 foldr1 f = Stream.foldr1 f . stream
980
981 -- | Right fold with a strict accumulator
982 foldr' :: Vector v a => (a -> b -> b) -> b -> v a -> b
983 {-# INLINE foldr' #-}
984 foldr' f z = Stream.foldl' (flip f) z . streamR
985
986 -- | Right fold on non-empty vectors with strict accumulator
987 foldr1' :: Vector v a => (a -> a -> a) -> v a -> a
988 {-# INLINE foldr1' #-}
989 foldr1' f = Stream.foldl1' (flip f) . streamR
990
991 -- | Left fold (function applied to each element and its index)
992 ifoldl :: Vector v b => (a -> Int -> b -> a) -> a -> v b -> a
993 {-# INLINE ifoldl #-}
994 ifoldl f z = Stream.foldl (uncurry . f) z . Stream.indexed . stream
995
996 -- | Left fold with strict accumulator (function applied to each element and
997 -- its index)
998 ifoldl' :: Vector v b => (a -> Int -> b -> a) -> a -> v b -> a
999 {-# INLINE ifoldl' #-}
1000 ifoldl' f z = Stream.foldl' (uncurry . f) z . Stream.indexed . stream
1001
1002 -- | Right fold (function applied to each element and its index)
1003 ifoldr :: Vector v a => (Int -> a -> b -> b) -> b -> v a -> b
1004 {-# INLINE ifoldr #-}
1005 ifoldr f z = Stream.foldr (uncurry f) z . Stream.indexed . stream
1006
1007 -- | Right fold with strict accumulator (function applied to each element and
1008 -- its index)
1009 ifoldr' :: Vector v a => (Int -> a -> b -> b) -> b -> v a -> b
1010 {-# INLINE ifoldr' #-}
1011 ifoldr' f z xs = Stream.foldl' (flip (uncurry f)) z
1012 $ Stream.indexedR (length xs) $ streamR xs
1013
1014 -- Specialised folds
1015 -- -----------------
1016
1017 all :: Vector v a => (a -> Bool) -> v a -> Bool
1018 {-# INLINE all #-}
1019 all f = Stream.and . Stream.map f . stream
1020
1021 any :: Vector v a => (a -> Bool) -> v a -> Bool
1022 {-# INLINE any #-}
1023 any f = Stream.or . Stream.map f . stream
1024
1025 and :: Vector v Bool => v Bool -> Bool
1026 {-# INLINE and #-}
1027 and = Stream.and . stream
1028
1029 or :: Vector v Bool => v Bool -> Bool
1030 {-# INLINE or #-}
1031 or = Stream.or . stream
1032
1033 sum :: (Vector v a, Num a) => v a -> a
1034 {-# INLINE sum #-}
1035 sum = Stream.foldl' (+) 0 . stream
1036
1037 product :: (Vector v a, Num a) => v a -> a
1038 {-# INLINE product #-}
1039 product = Stream.foldl' (*) 1 . stream
1040
1041 maximum :: (Vector v a, Ord a) => v a -> a
1042 {-# INLINE maximum #-}
1043 maximum = Stream.foldl1' max . stream
1044
1045 maximumBy :: Vector v a => (a -> a -> Ordering) -> v a -> a
1046 {-# INLINE maximumBy #-}
1047 maximumBy cmp = Stream.foldl1' maxBy . stream
1048 where
1049 {-# INLINE maxBy #-}
1050 maxBy x y = case cmp x y of
1051 LT -> y
1052 _ -> x
1053
1054 minimum :: (Vector v a, Ord a) => v a -> a
1055 {-# INLINE minimum #-}
1056 minimum = Stream.foldl1' min . stream
1057
1058 minimumBy :: Vector v a => (a -> a -> Ordering) -> v a -> a
1059 {-# INLINE minimumBy #-}
1060 minimumBy cmp = Stream.foldl1' minBy . stream
1061 where
1062 {-# INLINE minBy #-}
1063 minBy x y = case cmp x y of
1064 GT -> y
1065 _ -> x
1066
1067 maxIndex :: (Vector v a, Ord a) => v a -> Int
1068 {-# INLINE maxIndex #-}
1069 maxIndex = maxIndexBy compare
1070
1071 maxIndexBy :: Vector v a => (a -> a -> Ordering) -> v a -> Int
1072 {-# INLINE maxIndexBy #-}
1073 maxIndexBy cmp = fst . Stream.foldl1' imax . Stream.indexed . stream
1074 where
1075 imax (i,x) (j,y) = case cmp x y of
1076 LT -> (j,y)
1077 _ -> (i,x)
1078
1079 minIndex :: (Vector v a, Ord a) => v a -> Int
1080 {-# INLINE minIndex #-}
1081 minIndex = minIndexBy compare
1082
1083 minIndexBy :: Vector v a => (a -> a -> Ordering) -> v a -> Int
1084 {-# INLINE minIndexBy #-}
1085 minIndexBy cmp = fst . Stream.foldl1' imin . Stream.indexed . stream
1086 where
1087 imin (i,x) (j,y) = case cmp x y of
1088 GT -> (j,y)
1089 _ -> (i,x)
1090
1091
1092 -- Unfolding
1093 -- ---------
1094
1095 -- | Unfold
1096 unfoldr :: Vector v a => (b -> Maybe (a, b)) -> b -> v a
1097 {-# INLINE unfoldr #-}
1098 unfoldr f = unstream . Stream.unfoldr f
1099
1100 -- | Unfoldr at most @n@ elements.
1101 unfoldrN :: Vector v a => Int -> (b -> Maybe (a, b)) -> b -> v a
1102 {-# INLINE unfoldrN #-}
1103 unfoldrN n f = unstream . Stream.unfoldrN n f
1104
1105 -- Scans
1106 -- -----
1107
1108 -- | Prefix scan
1109 prescanl :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1110 {-# INLINE prescanl #-}
1111 prescanl f z = unstream . inplace (MStream.prescanl f z) . stream
1112
1113 -- | Prefix scan with strict accumulator
1114 prescanl' :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1115 {-# INLINE prescanl' #-}
1116 prescanl' f z = unstream . inplace (MStream.prescanl' f z) . stream
1117
1118 -- | Suffix scan
1119 postscanl :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1120 {-# INLINE postscanl #-}
1121 postscanl f z = unstream . inplace (MStream.postscanl f z) . stream
1122
1123 -- | Suffix scan with strict accumulator
1124 postscanl' :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1125 {-# INLINE postscanl' #-}
1126 postscanl' f z = unstream . inplace (MStream.postscanl' f z) . stream
1127
1128 -- | Haskell-style scan
1129 scanl :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1130 {-# INLINE scanl #-}
1131 scanl f z = unstream . Stream.scanl f z . stream
1132
1133 -- | Haskell-style scan with strict accumulator
1134 scanl' :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
1135 {-# INLINE scanl' #-}
1136 scanl' f z = unstream . Stream.scanl' f z . stream
1137
1138 -- | Scan over a non-empty vector
1139 scanl1 :: Vector v a => (a -> a -> a) -> v a -> v a
1140 {-# INLINE scanl1 #-}
1141 scanl1 f = unstream . inplace (MStream.scanl1 f) . stream
1142
1143 -- | Scan over a non-empty vector with a strict accumulator
1144 scanl1' :: Vector v a => (a -> a -> a) -> v a -> v a
1145 {-# INLINE scanl1' #-}
1146 scanl1' f = unstream . inplace (MStream.scanl1' f) . stream
1147
1148
1149 -- | Prefix right-to-left scan
1150 prescanr :: (Vector v a, Vector v b) => (a -> b -> b) -> b -> v a -> v b
1151 {-# INLINE prescanr #-}
1152 prescanr f z = unstreamR . inplace (MStream.prescanl (flip f) z) . streamR
1153
1154 -- | Prefix right-to-left scan with strict accumulator
1155 prescanr' :: (Vector v a, Vector v b) => (a -> b -> b) -> b -> v a -> v b
1156 {-# INLINE prescanr' #-}
1157 prescanr' f z = unstreamR . inplace (MStream.prescanl' (flip f) z) . streamR
1158
1159 -- | Suffix right-to-left scan
1160 postscanr :: (Vector v a, Vector v b) => (a -> b -> b) -> b -> v a -> v b
1161 {-# INLINE postscanr #-}
1162 postscanr f z = unstreamR . inplace (MStream.postscanl (flip f) z) . streamR
1163
1164 -- | Suffix right-to-left scan with strict accumulator
1165 postscanr' :: (Vector v a, Vector v b) => (a -> b -> b) -> b -> v a -> v b
1166 {-# INLINE postscanr' #-}
1167 postscanr' f z = unstreamR . inplace (MStream.postscanl' (flip f) z) . streamR
1168
1169 -- | Haskell-style right-to-left scan
1170 scanr :: (Vector v a, Vector v b) => (a -> b -> b) -> b -> v a -> v b
1171 {-# INLINE scanr #-}
1172 scanr f z = unstreamR . Stream.scanl (flip f) z . streamR
1173
1174 -- | Haskell-style right-to-left scan with strict accumulator
1175 scanr' :: (Vector v a, Vector v b) => (a -> b -> b) -> b -> v a -> v b
1176 {-# INLINE scanr' #-}
1177 scanr' f z = unstreamR . Stream.scanl' (flip f) z . streamR
1178
1179 -- | Right-to-left scan over a non-empty vector
1180 scanr1 :: Vector v a => (a -> a -> a) -> v a -> v a
1181 {-# INLINE scanr1 #-}
1182 scanr1 f = unstreamR . inplace (MStream.scanl1 (flip f)) . streamR
1183
1184 -- | Right-to-left scan over a non-empty vector with a strict accumulator
1185 scanr1' :: Vector v a => (a -> a -> a) -> v a -> v a
1186 {-# INLINE scanr1' #-}
1187 scanr1' f = unstreamR . inplace (MStream.scanl1' (flip f)) . streamR
1188
1189 -- Enumeration
1190 -- -----------
1191
1192 -- | Yield a vector of the given length containing the values @x@, @x+1@ etc.
1193 -- This operation is usually more efficient than 'enumFromTo'.
1194 enumFromN :: (Vector v a, Num a) => a -> Int -> v a
1195 {-# INLINE enumFromN #-}
1196 enumFromN x n = enumFromStepN x 1 n
1197
1198 -- | Yield a vector of the given length containing the values @x@, @x+y@,
1199 -- @x+y+y@ etc. This operations is usually more efficient than
1200 -- 'enumFromThenTo'.
1201 enumFromStepN :: forall v a. (Vector v a, Num a) => a -> a -> Int -> v a
1202 {-# INLINE enumFromStepN #-}
1203 enumFromStepN x y n = elemseq (undefined :: v a) x
1204 $ elemseq (undefined :: v a) y
1205 $ unstream
1206 $ Stream.enumFromStepN x y n
1207
1208 -- | Enumerate values from @x@ to @y@.
1209 --
1210 -- /WARNING:/ This operation can be very inefficient. If at all possible, use
1211 -- 'enumFromN' instead.
1212 enumFromTo :: (Vector v a, Enum a) => a -> a -> v a
1213 {-# INLINE enumFromTo #-}
1214 enumFromTo x y = unstream (Stream.enumFromTo x y)
1215
1216 -- | Enumerate values from @x@ to @y@ with a specific step @z@.
1217 --
1218 -- /WARNING:/ This operation can be very inefficient. If at all possible, use
1219 -- 'enumFromStepN' instead.
1220 enumFromThenTo :: (Vector v a, Enum a) => a -> a -> a -> v a
1221 {-# INLINE enumFromThenTo #-}
1222 enumFromThenTo x y z = unstream (Stream.enumFromThenTo x y z)
1223
1224 -- Conversion to/from lists
1225 -- ------------------------
1226
1227 -- | Convert a vector to a list
1228 toList :: Vector v a => v a -> [a]
1229 {-# INLINE toList #-}
1230 toList = Stream.toList . stream
1231
1232 -- | Convert a list to a vector
1233 fromList :: Vector v a => [a] -> v a
1234 {-# INLINE fromList #-}
1235 fromList = unstream . Stream.fromList
1236
1237 -- | Convert the first @n@ elements of a list to a vector
1238 --
1239 -- > fromListN n xs = fromList (take n xs)
1240 fromListN :: Vector v a => Int -> [a] -> v a
1241 {-# INLINE fromListN #-}
1242 fromListN n = unstream . Stream.fromListN n
1243
1244 unstreamM :: (Vector v a, Monad m) => MStream m a -> m (v a)
1245 {-# INLINE_STREAM unstreamM #-}
1246 unstreamM s = do
1247 xs <- MStream.toList s
1248 return $ unstream $ Stream.unsafeFromList (MStream.size s) xs
1249
1250 -- Monadic operations
1251 -- ------------------
1252
1253 -- FIXME: specialise various combinators for ST and IO?
1254
1255 -- | Perform the monadic action the given number of times and store the
1256 -- results in a vector.
1257 replicateM :: (Monad m, Vector v a) => Int -> m a -> m (v a)
1258 {-# INLINE replicateM #-}
1259 replicateM n m = fromListN n `Monad.liftM` Monad.replicateM n m
1260
1261 -- | Apply the monadic action to all elements of the vector, yielding a vector
1262 -- of results
1263 mapM :: (Monad m, Vector v a, Vector v b) => (a -> m b) -> v a -> m (v b)
1264 {-# INLINE mapM #-}
1265 mapM f = unstreamM . Stream.mapM f . stream
1266
1267 -- | Apply the monadic action to all elements of a vector and ignore the
1268 -- results
1269 mapM_ :: (Monad m, Vector v a) => (a -> m b) -> v a -> m ()
1270 {-# INLINE mapM_ #-}
1271 mapM_ f = Stream.mapM_ f . stream
1272
1273 -- | Apply the monadic action to all elements of the vector, yielding a vector
1274 -- of results
1275 forM :: (Monad m, Vector v a, Vector v b) => v a -> (a -> m b) -> m (v b)
1276 {-# INLINE forM #-}
1277 forM as f = mapM f as
1278
1279 -- | Apply the monadic action to all elements of a vector and ignore the
1280 -- results
1281 forM_ :: (Monad m, Vector v a) => v a -> (a -> m b) -> m ()
1282 {-# INLINE forM_ #-}
1283 forM_ as f = mapM_ f as
1284
1285 -- | Zip the two vectors with the monadic action and yield a vector of results
1286 zipWithM :: (Monad m, Vector v a, Vector v b, Vector v c)
1287 => (a -> b -> m c) -> v a -> v b -> m (v c)
1288 {-# INLINE zipWithM #-}
1289 zipWithM f as bs = unstreamM $ Stream.zipWithM f (stream as) (stream bs)
1290
1291 -- | Zip the two vectors with the monadic action and ignore the results
1292 zipWithM_ :: (Monad m, Vector v a, Vector v b)
1293 => (a -> b -> m c) -> v a -> v b -> m ()
1294 {-# INLINE zipWithM_ #-}
1295 zipWithM_ f as bs = Stream.zipWithM_ f (stream as) (stream bs)
1296
1297 -- | Drop elements that do not satisfy the monadic predicate
1298 filterM :: (Monad m, Vector v a) => (a -> m Bool) -> v a -> m (v a)
1299 {-# INLINE filterM #-}
1300 filterM f = unstreamM . Stream.filterM f . stream
1301
1302 -- | Monadic fold
1303 foldM :: (Monad m, Vector v b) => (a -> b -> m a) -> a -> v b -> m a
1304 {-# INLINE foldM #-}
1305 foldM m z = Stream.foldM m z . stream
1306
1307 -- | Monadic fold over non-empty vectors
1308 fold1M :: (Monad m, Vector v a) => (a -> a -> m a) -> v a -> m a
1309 {-# INLINE fold1M #-}
1310 fold1M m = Stream.fold1M m . stream
1311
1312 -- | Monadic fold with strict accumulator
1313 foldM' :: (Monad m, Vector v b) => (a -> b -> m a) -> a -> v b -> m a
1314 {-# INLINE foldM' #-}
1315 foldM' m z = Stream.foldM' m z . stream
1316
1317 -- | Monad fold over non-empty vectors with strict accumulator
1318 fold1M' :: (Monad m, Vector v a) => (a -> a -> m a) -> v a -> m a
1319 {-# INLINE fold1M' #-}
1320 fold1M' m = Stream.fold1M' m . stream
1321
1322 -- Destructive operations
1323 -- ----------------------
1324
1325 -- | Destructively initialise a vector.
1326 create :: Vector v a => (forall s. ST s (Mutable v s a)) -> v a
1327 {-# INLINE create #-}
1328 create p = new (New.create p)
1329
1330 -- | Apply a destructive operation to a vector. The operation modifies a
1331 -- copy of the vector unless it can be safely performed in place.
1332 modify :: Vector v a => (forall s. Mutable v s a -> ST s ()) -> v a -> v a
1333 {-# INLINE modify #-}
1334 modify p = new . New.modify p . clone
1335
1336 -- | Copy an immutable vector into a mutable one. The two vectors must have
1337 -- the same length. This is not checked.
1338 unsafeCopy
1339 :: (PrimMonad m, Vector v a) => Mutable v (PrimState m) a -> v a -> m ()
1340 {-# INLINE unsafeCopy #-}
1341 unsafeCopy dst src = UNSAFE_CHECK(check) "unsafeCopy" "length mismatch"
1342 (M.length dst == length src)
1343 $ (dst `seq` src `seq` basicUnsafeCopy dst src)
1344
1345 -- | Copy an immutable vector into a mutable one. The two vectors must have the
1346 -- same length.
1347 copy
1348 :: (PrimMonad m, Vector v a) => Mutable v (PrimState m) a -> v a -> m ()
1349 {-# INLINE copy #-}
1350 copy dst src = BOUNDS_CHECK(check) "copy" "length mismatch"
1351 (M.length dst == length src)
1352 $ unsafeCopy dst src
1353
1354 -- Utilities for defining Data instances
1355 -- -------------------------------------
1356
1357 -- | Generic definion of 'Data.Data.gfoldl' that views a 'Vector' as a
1358 -- list.
1359 gfoldl :: (Vector v a, Data a)
1360 => (forall d b. Data d => c (d -> b) -> d -> c b)
1361 -> (forall g. g -> c g)
1362 -> v a
1363 -> c (v a)
1364 {-# INLINE gfoldl #-}
1365 gfoldl f z v = z fromList `f` toList v
1366
1367 mkType :: String -> DataType
1368 {-# INLINE mkType #-}
1369 mkType = mkNorepType
1370
1371 dataCast :: (Vector v a, Data a, Typeable1 v, Typeable1 t)
1372 => (forall d. Data d => c (t d)) -> Maybe (c (v a))
1373 {-# INLINE dataCast #-}
1374 dataCast f = gcast1 f
1375