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