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