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