Add D.V.IVector.{indexM,headM,lastM}
[darcs-mirrors/vector.git] / Data / Vector / IVector.hs
1 {-# LANGUAGE Rank2Types, MultiParamTypeClasses, FlexibleContexts,
2 ScopedTypeVariables #-}
3 -- |
4 -- Module : Data.Vector.IVector
5 -- Copyright : (c) Roman Leshchinskiy 2008
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 #include "phases.h"
16
17 module Data.Vector.IVector (
18 -- * Immutable vectors
19 IVector,
20
21 -- * Length information
22 length, null,
23
24 -- * Construction
25 empty, singleton, cons, snoc, replicate, (++), copy,
26
27 -- * Accessing individual elements
28 (!), head, last, indexM, headM, lastM,
29
30 -- * Subvectors
31 slice, init, tail, take, drop,
32
33 -- * Permutations
34 accum, (//), update, backpermute, reverse,
35
36 -- * Mapping
37 map, concatMap,
38
39 -- * Zipping and unzipping
40 zipWith, zipWith3, zip, zip3, unzip, unzip3,
41
42 -- * Comparisons
43 eq, cmp,
44
45 -- * Filtering
46 filter, takeWhile, dropWhile,
47
48 -- * Searching
49 elem, notElem, find, findIndex,
50
51 -- * Folding
52 foldl, foldl1, foldl', foldl1', foldr, foldr1,
53
54 -- * Specialised folds
55 and, or, sum, product, maximum, minimum,
56
57 -- * Enumeration
58 enumFromTo, enumFromThenTo,
59
60 -- * Unfolding
61 unfoldr,
62
63 -- * Scans
64 prescanl, prescanl',
65
66 -- * Conversion to/from lists
67 toList, fromList,
68
69 -- * Conversion to/from Streams
70 stream, unstream,
71
72 -- * MVector-based initialisation
73 new,
74
75 -- * Unsafe functions
76 unsafeSlice, unsafeIndex,
77
78 -- * Utility functions
79 vlength, vnew
80 ) where
81
82 import qualified Data.Vector.MVector as MVector
83 import Data.Vector.MVector ( MVector )
84
85 import qualified Data.Vector.MVector.New as New
86 import Data.Vector.MVector.New ( New )
87
88 import qualified Data.Vector.Fusion.Stream as Stream
89 import Data.Vector.Fusion.Stream ( Stream, MStream )
90 import qualified Data.Vector.Fusion.Stream.Monadic as MStream
91 import Data.Vector.Fusion.Stream.Size
92
93 import Control.Exception ( assert )
94
95 import Prelude hiding ( length, null,
96 replicate, (++),
97 head, last,
98 init, tail, take, drop, reverse,
99 map, concatMap,
100 zipWith, zipWith3, zip, zip3, unzip, unzip3,
101 filter, takeWhile, dropWhile,
102 elem, notElem,
103 foldl, foldl1, foldr, foldr1,
104 and, or, sum, product, maximum, minimum,
105 enumFromTo, enumFromThenTo )
106
107 -- | Class of immutable vectors.
108 --
109 class IVector v a where
110 -- | Construct a pure vector from a monadic initialiser (not fusible!)
111 vnew :: (forall mv m. MVector mv m a => m (mv a)) -> v a
112
113 -- | Length of the vector (not fusible!)
114 vlength :: v a -> Int
115
116 -- | Yield a part of the vector without copying it. No range checks!
117 unsafeSlice :: v a -> Int -> Int -> v a
118
119 -- | Apply the given function to the element at the given position. This
120 -- interface prevents us from being too lazy. Suppose we had
121 --
122 -- > unsafeIndex' :: v a -> Int -> a
123 --
124 -- instead. Now, if we wanted to copy a vector, we'd do something like
125 --
126 -- > copy mv v ... = ... unsafeWrite mv i (unsafeIndex' v i) ...
127 --
128 -- For lazy vectors, the indexing would not be evaluated which means that we
129 -- would retain a reference to the original vector in each element we write.
130 -- This would be bad!
131 --
132 -- With 'unsafeIndex', we can do
133 --
134 -- > copy mv v ... = ... unsafeIndex v i (unsafeWrite mv i) ...
135 --
136 -- which does not have this problem.
137 --
138 unsafeIndex :: v a -> Int -> (a -> b) -> b
139
140 -- Fusion
141 -- ------
142
143 -- | Construct a pure vector from a monadic initialiser
144 new :: IVector v a => New a -> v a
145 {-# INLINE new #-}
146 new m = new' undefined m
147
148 -- | Same as 'new' but with a dummy argument necessary for correctly typing
149 -- the rule @uninplace@.
150 --
151 -- See http://hackage.haskell.org/trac/ghc/ticket/2600
152 new' :: IVector v a => v a -> New a -> v a
153 {-# INLINE_STREAM new' #-}
154 new' _ m = vnew (New.run m)
155
156 -- | Convert a vector to a 'Stream'
157 stream :: IVector v a => v a -> Stream a
158 {-# INLINE_STREAM stream #-}
159 stream v = v `seq` (Stream.unfoldr get 0 `Stream.sized` Exact n)
160 where
161 n = length v
162
163 {-# INLINE get #-}
164 get i | i < n = unsafeIndex v i $ \x -> Just (x, i+1)
165 | otherwise = Nothing
166
167 -- | Create a vector from a 'Stream'
168 unstream :: IVector v a => Stream a -> v a
169 {-# INLINE unstream #-}
170 unstream s = new (New.unstream s)
171
172 {-# RULES
173
174 "stream/unstream [IVector]" forall v s.
175 stream (new' v (New.unstream s)) = s
176
177 "New.unstream/stream/new [IVector]" forall v p.
178 New.unstream (stream (new' v p)) = p
179
180 #-}
181
182 inplace :: (forall m. Monad m => MStream m a -> MStream m a)
183 -> Stream a -> Stream a
184 {-# INLINE_STREAM inplace #-}
185 inplace f s = f s
186
187 {-# RULES
188
189 "inplace [IVector]"
190 forall (f :: forall m. Monad m => MStream m a -> MStream m a) v m.
191 New.unstream (inplace f (stream (new' v m))) = New.transform f m
192
193 "uninplace [IVector]"
194 forall (f :: forall m. Monad m => MStream m a -> MStream m a) v m.
195 stream (new' v (New.transform f m)) = inplace f (stream (new' v m))
196
197 "inplace/inplace [IVector]"
198 forall (f :: forall m. Monad m => MStream m a -> MStream m a)
199 (g :: forall m. Monad m => MStream m a -> MStream m a)
200 s.
201 inplace f (inplace g s) = inplace (f . g) s
202
203 #-}
204
205 -- Length
206 -- ------
207
208 length :: IVector v a => v a -> Int
209 {-# INLINE_STREAM length #-}
210 length v = vlength v
211
212 {-# RULES
213
214 "length/unstream [IVector]" forall v s.
215 length (new' v (New.unstream s)) = Stream.length s
216
217 #-}
218
219 null :: IVector v a => v a -> Bool
220 {-# INLINE_STREAM null #-}
221 null v = vlength v == 0
222
223 {-# RULES
224
225 "null/unstream [IVector]" forall v s.
226 null (new' v (New.unstream s)) = Stream.null s
227
228 #-}
229
230 -- Construction
231 -- ------------
232
233 -- | Empty vector
234 empty :: IVector v a => v a
235 {-# INLINE empty #-}
236 empty = unstream Stream.empty
237
238 -- | Vector with exaclty one element
239 singleton :: IVector v a => a -> v a
240 {-# INLINE singleton #-}
241 singleton x = unstream (Stream.singleton x)
242
243 -- | Vector of the given length with the given value in each position
244 replicate :: IVector v a => Int -> a -> v a
245 {-# INLINE replicate #-}
246 replicate n = unstream . Stream.replicate n
247
248 -- | Prepend an element
249 cons :: IVector v a => a -> v a -> v a
250 {-# INLINE cons #-}
251 cons x = unstream . Stream.cons x . stream
252
253 -- | Append an element
254 snoc :: IVector v a => v a -> a -> v a
255 {-# INLINE snoc #-}
256 snoc v = unstream . Stream.snoc (stream v)
257
258 infixr 5 ++
259 -- | Concatenate two vectors
260 (++) :: IVector v a => v a -> v a -> v a
261 {-# INLINE (++) #-}
262 v ++ w = unstream (stream v Stream.++ stream w)
263
264 -- | Create a copy of a vector. Useful when dealing with slices.
265 copy :: IVector v a => v a -> v a
266 {-# INLINE_STREAM copy #-}
267 copy = unstream . stream
268
269 {-# RULES
270
271 "copy/unstream [IVector]" forall v s.
272 copy (new' v (New.unstream s)) = new' v (New.unstream s)
273
274 #-}
275
276 -- Accessing individual elements
277 -- -----------------------------
278
279 -- | Indexing
280 (!) :: IVector v a => v a -> Int -> a
281 {-# INLINE_STREAM (!) #-}
282 v ! i = assert (i >= 0 && i < length v)
283 $ unsafeIndex v i id
284
285 -- | First element
286 head :: IVector v a => v a -> a
287 {-# INLINE_STREAM head #-}
288 head v = v ! 0
289
290 -- | Last element
291 last :: IVector v a => v a -> a
292 {-# INLINE_STREAM last #-}
293 last v = v ! (length v - 1)
294
295 {-# RULES
296
297 "(!)/unstream [IVector]" forall v i s.
298 new' v (New.unstream s) ! i = s Stream.!! i
299
300 "head/unstream [IVector]" forall v s.
301 head (new' v (New.unstream s)) = Stream.head s
302
303 "last/unstream [IVector]" forall v s.
304 last (new' v (New.unstream s)) = Stream.last s
305
306 #-}
307
308 -- | Monadic indexing which can be strict in the array while remaining lazy in
309 -- the element.
310 indexM :: (IVector v a, Monad m) => v a -> Int -> m a
311 {-# INLINE_STREAM indexM #-}
312 indexM v i = assert (i >= 0 && i < length v)
313 $ unsafeIndex v i return
314
315 headM :: (IVector v a, Monad m) => v a -> m a
316 {-# INLINE_STREAM headM #-}
317 headM v = indexM v 0
318
319 lastM :: (IVector v a, Monad m) => v a -> m a
320 {-# INLINE_STREAM lastM #-}
321 lastM v = indexM v (length v - 1)
322
323 {-# RULES
324
325 "indexM/unstream [IVector]" forall v i s.
326 indexM (new' v (New.unstream s)) i = return (s Stream.!! i)
327
328 "headM/unstream [IVector]" forall v s.
329 headM (new' v (New.unstream s)) = return (Stream.head s)
330
331 "lastM/unstream [IVector]" forall v s.
332 lastM (new' v (New.unstream s)) = return (Stream.last s)
333
334 #-}
335
336 -- Subarrays
337 -- ---------
338
339 -- FIXME: slicing doesn't work with the inplace stuff at the moment
340
341 -- | Yield a part of the vector without copying it. Safer version of
342 -- 'unsafeSlice'.
343 slice :: IVector v a => v a -> Int -- ^ starting index
344 -> Int -- ^ length
345 -> v a
346 {-# INLINE_STREAM slice #-}
347 slice v i n = assert (i >= 0 && n >= 0 && i+n <= length v)
348 $ unsafeSlice v i n
349
350 -- | Yield all but the last element without copying.
351 init :: IVector v a => v a -> v a
352 {-# INLINE_STREAM init #-}
353 init v = slice v 0 (length v - 1)
354
355 -- | All but the first element (without copying).
356 tail :: IVector v a => v a -> v a
357 {-# INLINE_STREAM tail #-}
358 tail v = slice v 1 (length v - 1)
359
360 -- | Yield the first @n@ elements without copying.
361 take :: IVector v a => Int -> v a -> v a
362 {-# INLINE_STREAM take #-}
363 take n v = slice v 0 (min n' (length v))
364 where n' = max n 0
365
366 -- | Yield all but the first @n@ elements without copying.
367 drop :: IVector v a => Int -> v a -> v a
368 {-# INLINE_STREAM drop #-}
369 drop n v = slice v (min n' len) (max 0 (len - n'))
370 where n' = max n 0
371 len = length v
372
373 {-# RULES
374
375 "slice/new [IVector]" forall v p i n.
376 slice (new' v p) i n = new' v (New.slice p i n)
377
378 "init/new [IVector]" forall v p.
379 init (new' v p) = new' v (New.init p)
380
381 "tail/new [IVector]" forall v p.
382 tail (new' v p) = new' v (New.tail p)
383
384 "take/new [IVector]" forall n v p.
385 take n (new' v p) = new' v (New.take n p)
386
387 "drop/new [IVector]" forall n v p.
388 drop n (new' v p) = new' v (New.drop n p)
389
390 #-}
391
392 -- Permutations
393 -- ------------
394
395 accum :: IVector v a => (a -> b -> a) -> v a -> [(Int,b)] -> v a
396 {-# INLINE accum #-}
397 accum f v us = new (New.accum f (New.unstream (stream v))
398 (Stream.fromList us))
399
400 (//) :: IVector v a => v a -> [(Int, a)] -> v a
401 {-# INLINE (//) #-}
402 v // us = new (New.update (New.unstream (stream v))
403 (Stream.fromList us))
404
405 update :: (IVector v a, IVector v (Int, a)) => v a -> v (Int, a) -> v a
406 {-# INLINE update #-}
407 update v w = new (New.update (New.unstream (stream v)) (stream w))
408
409 backpermute :: (IVector v a, IVector v Int) => v a -> v Int -> v a
410 {-# INLINE backpermute #-}
411 backpermute v is = v `seq` map (v!) is
412
413 reverse :: (IVector v a) => v a -> v a
414 {-# INLINE reverse #-}
415 reverse = new . New.reverse . New.unstream . stream
416
417 -- Mapping
418 -- -------
419
420 -- | Map a function over a vector
421 map :: (IVector v a, IVector v b) => (a -> b) -> v a -> v b
422 {-# INLINE map #-}
423 map f = unstream . Stream.map f . stream
424
425 inplace_map :: IVector v a => (a -> a) -> v a -> v a
426 {-# INLINE inplace_map #-}
427 inplace_map f = unstream . inplace (MStream.map f) . stream
428
429 {-# RULES
430
431 "map->inplace_map [IVector]" map = inplace_map
432
433 #-}
434
435 concatMap :: (IVector v a, IVector v b) => (a -> v b) -> v a -> v b
436 {-# INLINE concatMap #-}
437 concatMap f = unstream . Stream.concatMap (stream . f) . stream
438
439 -- Zipping/unzipping
440 -- -----------------
441
442 -- | Zip two vectors with the given function.
443 zipWith :: (IVector v a, IVector v b, IVector v c) => (a -> b -> c) -> v a -> v b -> v c
444 {-# INLINE zipWith #-}
445 zipWith f xs ys = unstream (Stream.zipWith f (stream xs) (stream ys))
446
447 -- | Zip three vectors with the given function.
448 zipWith3 :: (IVector v a, IVector v b, IVector v c, IVector v d) => (a -> b -> c -> d) -> v a -> v b -> v c -> v d
449 {-# INLINE zipWith3 #-}
450 zipWith3 f xs ys zs = unstream (Stream.zipWith3 f (stream xs) (stream ys) (stream zs))
451
452 zip :: (IVector v a, IVector v b, IVector v (a,b)) => v a -> v b -> v (a, b)
453 {-# INLINE zip #-}
454 zip = zipWith (,)
455
456 zip3 :: (IVector v a, IVector v b, IVector v c, IVector v (a, b, c)) => v a -> v b -> v c -> v (a, b, c)
457 {-# INLINE zip3 #-}
458 zip3 = zipWith3 (,,)
459
460 unzip :: (IVector v a, IVector v b, IVector v (a,b)) => v (a, b) -> (v a, v b)
461 {-# INLINE unzip #-}
462 unzip xs = (map fst xs, map snd xs)
463
464 unzip3 :: (IVector v a, IVector v b, IVector v c, IVector v (a, b, c)) => v (a, b, c) -> (v a, v b, v c)
465 {-# INLINE unzip3 #-}
466 unzip3 xs = (map (\(a, b, c) -> a) xs, map (\(a, b, c) -> b) xs, map (\(a, b, c) -> c) xs)
467
468 -- Comparisons
469 -- -----------
470
471 eq :: (IVector v a, Eq a) => v a -> v a -> Bool
472 {-# INLINE eq #-}
473 xs `eq` ys = stream xs == stream ys
474
475 cmp :: (IVector v a, Ord a) => v a -> v a -> Ordering
476 {-# INLINE cmp #-}
477 cmp xs ys = compare (stream xs) (stream ys)
478
479 -- Filtering
480 -- ---------
481
482 -- | Drop elements which do not satisfy the predicate
483 filter :: IVector v a => (a -> Bool) -> v a -> v a
484 {-# INLINE filter #-}
485 filter f = unstream . inplace (MStream.filter f) . stream
486
487 -- | Yield the longest prefix of elements satisfying the predicate.
488 takeWhile :: IVector v a => (a -> Bool) -> v a -> v a
489 {-# INLINE takeWhile #-}
490 takeWhile f = unstream . Stream.takeWhile f . stream
491
492 -- | Drop the longest prefix of elements that satisfy the predicate.
493 dropWhile :: IVector v a => (a -> Bool) -> v a -> v a
494 {-# INLINE dropWhile #-}
495 dropWhile f = unstream . Stream.dropWhile f . stream
496
497 -- Searching
498 -- ---------
499
500 infix 4 `elem`
501 -- | Check whether the vector contains an element
502 elem :: (IVector v a, Eq a) => a -> v a -> Bool
503 {-# INLINE elem #-}
504 elem x = Stream.elem x . stream
505
506 infix 4 `notElem`
507 -- | Inverse of `elem`
508 notElem :: (IVector v a, Eq a) => a -> v a -> Bool
509 {-# INLINE notElem #-}
510 notElem x = Stream.notElem x . stream
511
512 -- | Yield 'Just' the first element matching the predicate or 'Nothing' if no
513 -- such element exists.
514 find :: IVector v a => (a -> Bool) -> v a -> Maybe a
515 {-# INLINE find #-}
516 find f = Stream.find f . stream
517
518 -- | Yield 'Just' the index of the first element matching the predicate or
519 -- 'Nothing' if no such element exists.
520 findIndex :: IVector v a => (a -> Bool) -> v a -> Maybe Int
521 {-# INLINE findIndex #-}
522 findIndex f = Stream.findIndex f . stream
523
524 -- Folding
525 -- -------
526
527 -- | Left fold
528 foldl :: IVector v b => (a -> b -> a) -> a -> v b -> a
529 {-# INLINE foldl #-}
530 foldl f z = Stream.foldl f z . stream
531
532 -- | Lefgt fold on non-empty vectors
533 foldl1 :: IVector v a => (a -> a -> a) -> v a -> a
534 {-# INLINE foldl1 #-}
535 foldl1 f = Stream.foldl1 f . stream
536
537 -- | Left fold with strict accumulator
538 foldl' :: IVector v b => (a -> b -> a) -> a -> v b -> a
539 {-# INLINE foldl' #-}
540 foldl' f z = Stream.foldl' f z . stream
541
542 -- | Left fold on non-empty vectors with strict accumulator
543 foldl1' :: IVector v a => (a -> a -> a) -> v a -> a
544 {-# INLINE foldl1' #-}
545 foldl1' f = Stream.foldl1' f . stream
546
547 -- | Right fold
548 foldr :: IVector v a => (a -> b -> b) -> b -> v a -> b
549 {-# INLINE foldr #-}
550 foldr f z = Stream.foldr f z . stream
551
552 -- | Right fold on non-empty vectors
553 foldr1 :: IVector v a => (a -> a -> a) -> v a -> a
554 {-# INLINE foldr1 #-}
555 foldr1 f = Stream.foldr1 f . stream
556
557 -- Specialised folds
558 -- -----------------
559
560 and :: IVector v Bool => v Bool -> Bool
561 {-# INLINE and #-}
562 and = Stream.and . stream
563
564 or :: IVector v Bool => v Bool -> Bool
565 {-# INLINE or #-}
566 or = Stream.or . stream
567
568 sum :: (IVector v a, Num a) => v a -> a
569 {-# INLINE sum #-}
570 sum = Stream.foldl' (+) 0 . stream
571
572 product :: (IVector v a, Num a) => v a -> a
573 {-# INLINE product #-}
574 product = Stream.foldl' (*) 1 . stream
575
576 maximum :: (IVector v a, Ord a) => v a -> a
577 {-# INLINE maximum #-}
578 maximum = Stream.foldl1' max . stream
579
580 minimum :: (IVector v a, Ord a) => v a -> a
581 {-# INLINE minimum #-}
582 minimum = Stream.foldl1' min . stream
583
584 -- Enumeration
585 -- -----------
586
587 enumFromTo :: (IVector v a, Enum a) => a -> a -> v a
588 {-# INLINE enumFromTo #-}
589 enumFromTo from to = from `seq` to `seq` unfoldr enumFromTo_go (fromEnum from)
590 where
591 to_i = fromEnum to
592 enumFromTo_go i | i <= to_i = Just (toEnum i, i + 1)
593 | otherwise = Nothing
594
595 enumFromThenTo :: (IVector v a, Enum a) => a -> a -> a -> v a
596 {-# INLINE enumFromThenTo #-}
597 enumFromThenTo from next to = from `seq` next `seq` to `seq` unfoldr enumFromThenTo_go from_i
598 where
599 from_i = fromEnum from
600 to_i = fromEnum to
601 step_i = fromEnum next - from_i
602 enumFromThenTo_go i | i <= to_i = Just (toEnum i, i + step_i)
603 | otherwise = Nothing
604
605 -- Unfolding
606 -- ---------
607
608 unfoldr :: IVector v a => (b -> Maybe (a, b)) -> b -> v a
609 {-# INLINE unfoldr #-}
610 unfoldr f = unstream . Stream.unfoldr f
611
612 -- Scans
613 -- -----
614
615 -- | Prefix scan
616 prescanl :: (IVector v a, IVector v b) => (a -> b -> a) -> a -> v b -> v a
617 {-# INLINE prescanl #-}
618 prescanl f z = unstream . Stream.prescanl f z . stream
619
620 inplace_prescanl :: IVector v a => (a -> a -> a) -> a -> v a -> v a
621 {-# INLINE inplace_prescanl #-}
622 inplace_prescanl f z = unstream . inplace (MStream.prescanl f z) . stream
623
624 {-# RULES
625
626 "prescanl -> inplace_prescanl [IVector]" prescanl = inplace_prescanl
627
628 #-}
629
630 -- | Prefix scan with strict accumulator
631 prescanl' :: (IVector v a, IVector v b) => (a -> b -> a) -> a -> v b -> v a
632 {-# INLINE prescanl' #-}
633 prescanl' f z = unstream . Stream.prescanl' f z . stream
634
635 inplace_prescanl' :: IVector v a => (a -> a -> a) -> a -> v a -> v a
636 {-# INLINE inplace_prescanl' #-}
637 inplace_prescanl' f z = unstream . inplace (MStream.prescanl' f z) . stream
638
639 {-# RULES
640
641 "prescanl' -> inplace_prescanl' [IVector]" prescanl' = inplace_prescanl'
642
643 #-}
644
645
646 -- | Convert a vector to a list
647 toList :: IVector v a => v a -> [a]
648 {-# INLINE toList #-}
649 toList = Stream.toList . stream
650
651 -- | Convert a list to a vector
652 fromList :: IVector v a => [a] -> v a
653 {-# INLINE fromList #-}
654 fromList = unstream . Stream.fromList
655