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