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