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