IVector.zip
[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 : 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,
23
24 -- * Construction
25 empty, singleton, cons, snoc, replicate, (++),
26
27 -- * Accessing individual elements
28 (!), head, last,
29
30 -- * Subvectors
31 slice, extract, takeSlice, take, dropSlice, drop,
32
33 -- * Permutations
34 (//), update, bpermute,
35
36 -- * Mapping and zipping
37 map, zipWith, zip,
38
39 -- * Comparisons
40 eq, cmp,
41
42 -- * Filtering
43 filter, takeWhileSlice, takeWhile, dropWhileSlice, dropWhile,
44
45 -- * Searching
46 elem, notElem, find, findIndex,
47
48 -- * Folding
49 foldl, foldl1, foldl', foldl1', foldr, foldr1,
50
51 -- * Conversion to/from lists
52 toList, fromList,
53
54 -- * Conversion to/from Streams
55 stream, unstream,
56
57 -- * MVector-based initialisation
58 new,
59
60 -- * Unsafe functions
61 unsafeSlice, unsafeIndex,
62
63 -- * Utility functions
64 vlength, vnew
65 ) where
66
67 import qualified Data.Vector.MVector as MVector
68 import Data.Vector.MVector ( MVector )
69
70 import qualified Data.Vector.MVector.Mut as Mut
71 import Data.Vector.MVector.Mut ( Mut )
72
73 import qualified Data.Vector.Stream as Stream
74 import Data.Vector.Stream ( Stream )
75 import qualified Data.Vector.MStream as MStream
76 import Data.Vector.MStream ( MStream )
77 import Data.Vector.Stream.Size
78
79 import Control.Exception ( assert )
80
81 import Prelude hiding ( length,
82 replicate, (++),
83 head, last,
84 init, tail, take, drop,
85 map, zipWith, zip,
86 filter, takeWhile, dropWhile,
87 elem, notElem,
88 foldl, foldl1, foldr, foldr1 )
89
90 -- | Class of immutable vectors.
91 --
92 class IVector v a where
93 -- | Construct a pure vector from a monadic initialiser (not fusible!)
94 vnew :: (forall mv m. MVector mv m a => m (mv a)) -> v a
95
96 -- | Length of the vector (not fusible!)
97 vlength :: v a -> Int
98
99 -- | Yield a part of the vector without copying it. No range checks!
100 unsafeSlice :: v a -> Int -> Int -> v a
101
102 -- | Apply the given function to the element at the given position. This
103 -- interface prevents us from being too lazy. Suppose we had
104 --
105 -- > unsafeIndex' :: v a -> Int -> a
106 --
107 -- instead. Now, if we wanted to copy a vector, we'd do something like
108 --
109 -- > copy mv v ... = ... unsafeWrite mv i (unsafeIndex' v i) ...
110 --
111 -- For lazy vectors, the indexing would not be evaluated which means that we
112 -- would retain a reference to the original vector in each element we write.
113 -- This would be bad!
114 --
115 -- With 'unsafeIndex', we can do
116 --
117 -- > copy mv v ... = ... unsafeIndex v i (unsafeWrite mv i) ...
118 --
119 -- which does not have this problem.
120 --
121 unsafeIndex :: v a -> Int -> (a -> b) -> b
122
123 -- Fusion
124 -- ------
125
126 -- | Construct a pure vector from a monadic initialiser
127 new :: IVector v a => Mut a -> v a
128 {-# INLINE_STREAM new #-}
129 new m = vnew (Mut.run m)
130
131 -- | Convert a vector to a 'Stream'
132 stream :: IVector v a => v a -> Stream a
133 {-# INLINE_STREAM stream #-}
134 stream v = v `seq` (Stream.unfold get 0 `Stream.sized` Exact n)
135 where
136 n = length v
137
138 {-# INLINE get #-}
139 get i | i < n = unsafeIndex v i $ \x -> Just (x, i+1)
140 | otherwise = Nothing
141
142 -- | Create a vector from a 'Stream'
143 unstream :: IVector v a => Stream a -> v a
144 {-# INLINE unstream #-}
145 unstream s = new (Mut.unstream s)
146
147 {-# RULES
148
149 "stream/unstream [IVector]" forall s.
150 stream (new (Mut.unstream s)) = s
151
152 "Mut.unstream/stream/new [IVector]" forall p.
153 Mut.unstream (stream (new p)) = p
154
155 #-}
156
157 inplace :: (Stream a -> Stream a)
158 -> (forall m. Monad m => MStream m a -> MStream m a)
159 -> Stream a -> Stream a
160 {-# INLINE_STREAM inplace #-}
161 inplace f _ s = f s
162
163 {-# RULES
164
165 "inplace [IVector]"
166 forall f (mf :: forall m. Monad m => MStream m a -> MStream m a) m.
167 Mut.unstream (inplace f mf (stream (new m))) = Mut.inplace mf m
168
169 "inplace/inplace [IVector]"
170 forall f (mf :: forall m. Monad m => MStream m a -> MStream m a)
171 g (mg :: forall m. Monad m => MStream m a -> MStream m a) s.
172 inplace f mf (inplace g mg s) = inplace (f . g) (mf . mg) s
173
174 #-}
175
176 -- Length
177 -- ------
178
179 length :: IVector v a => v a -> Int
180 {-# INLINE_STREAM length #-}
181 length v = vlength v
182
183 {-# RULES
184
185 "length/unstream [IVector]" forall s.
186 length (new (Mut.unstream s)) = Stream.length s
187
188 #-}
189
190 -- Construction
191 -- ------------
192
193 -- | Empty vector
194 empty :: IVector v a => v a
195 {-# INLINE empty #-}
196 empty = unstream Stream.empty
197
198 -- | Vector with exaclty one element
199 singleton :: IVector v a => a -> v a
200 {-# INLINE singleton #-}
201 singleton x = unstream (Stream.singleton x)
202
203 -- | Vector of the given length with the given value in each position
204 replicate :: IVector v a => Int -> a -> v a
205 {-# INLINE replicate #-}
206 replicate n = unstream . Stream.replicate n
207
208 -- | Prepend an element
209 cons :: IVector v a => a -> v a -> v a
210 {-# INLINE cons #-}
211 cons x = unstream . Stream.cons x . stream
212
213 -- | Append an element
214 snoc :: IVector v a => v a -> a -> v a
215 {-# INLINE snoc #-}
216 snoc v = unstream . Stream.snoc (stream v)
217
218 infixr 5 ++
219 -- | Concatenate two vectors
220 (++) :: IVector v a => v a -> v a -> v a
221 {-# INLINE (++) #-}
222 v ++ w = unstream (stream v Stream.++ stream w)
223
224 -- Accessing individual elements
225 -- -----------------------------
226
227 -- | Indexing
228 (!) :: IVector v a => v a -> Int -> a
229 {-# INLINE_STREAM (!) #-}
230 v ! i = assert (i >= 0 && i < length v)
231 $ unsafeIndex v i id
232
233 -- | First element
234 head :: IVector v a => v a -> a
235 {-# INLINE_STREAM head #-}
236 head v = v ! 0
237
238 -- | Last element
239 last :: IVector v a => v a -> a
240 {-# INLINE_STREAM last #-}
241 last v = v ! (length v - 1)
242
243 {-# RULES
244
245 "(!)/unstream [IVector]" forall i s.
246 new (Mut.unstream s) ! i = s Stream.!! i
247
248 "head/unstream [IVector]" forall s.
249 head (new (Mut.unstream s)) = Stream.head s
250
251 "last/unstream [IVector]" forall s.
252 last (new (Mut.unstream s)) = Stream.last s
253
254 #-}
255
256 -- Subarrays
257 -- ---------
258
259 -- | Yield a part of the vector without copying it. Safer version of
260 -- 'unsafeSlice'.
261 slice :: IVector v a => v a -> Int -- ^ starting index
262 -> Int -- ^ length
263 -> v a
264 {-# INLINE slice #-}
265 slice v i n = assert (i >= 0 && n >= 0 && i+n <= length v)
266 $ unsafeSlice v i n
267
268 -- | Copy @n@ elements starting at the given position to a new vector.
269 extract :: IVector v a => v a -> Int -- ^ starting index
270 -> Int -- ^ length
271 -> v a
272 {-# INLINE extract #-}
273 extract v i n = unstream (Stream.extract (stream v) i n)
274
275 -- | Yield the first @n@ elements without copying.
276 takeSlice :: IVector v a => Int -> v a -> v a
277 {-# INLINE takeSlice #-}
278 takeSlice n v = slice v 0 n
279
280 -- | Copy the first @n@ elements to a new vector.
281 take :: IVector v a => Int -> v a -> v a
282 {-# INLINE take #-}
283 take n = unstream . Stream.take n . stream
284
285 -- | Yield all but the first @n@ elements without copying.
286 dropSlice :: IVector v a => Int -> v a -> v a
287 {-# INLINE dropSlice #-}
288 dropSlice n v = slice v n (length v - n)
289
290 -- | Copy all but the first @n@ elements to a new vectors.
291 drop :: IVector v a => Int -> v a -> v a
292 {-# INLINE drop #-}
293 drop n = unstream . Stream.drop n . stream
294
295 {-# RULES
296
297 "slice/extract [IVector]" forall i n s.
298 slice (new (Mut.unstream s)) i n = extract (new (Mut.unstream s)) i n
299
300 "takeSlice/unstream [IVector]" forall n s.
301 takeSlice n (new (Mut.unstream s)) = take n (new (Mut.unstream s))
302
303 "dropSlice/unstream [IVector]" forall n s.
304 dropSlice n (new (Mut.unstream s)) = drop n (new (Mut.unstream s))
305
306 #-}
307
308 -- Permutations
309 -- ------------
310
311 (//) :: IVector v a => v a -> [(Int, a)] -> v a
312 {-# INLINE (//) #-}
313 v // us = new (Mut.update (Mut.unstream (stream v))
314 (Stream.fromList us))
315
316 update :: (IVector v a, IVector v (Int, a)) => v a -> v (Int, a) -> v a
317 {-# INLINE update #-}
318 update v w = new (Mut.update (Mut.unstream (stream v)) (stream w))
319
320 bpermute :: (IVector v a, IVector v Int) => v a -> v Int -> v a
321 {-# INLINE bpermute #-}
322 bpermute v is = is `seq` map (v!) is
323
324 -- Mapping/zipping
325 -- ---------------
326
327 -- | Map a function over a vector
328 map :: (IVector v a, IVector v b) => (a -> b) -> v a -> v b
329 {-# INLINE map #-}
330 map f = unstream . Stream.map f . stream
331
332 inplace_map :: IVector v a => (a -> a) -> v a -> v a
333 {-# INLINE inplace_map #-}
334 inplace_map f = unstream . inplace (Stream.map f) (MStream.map f) . stream
335
336 {-# RULES
337
338 "map->inplace_map [IVector]" map = inplace_map
339
340 #-}
341
342 -- | Zip two vectors with the given function.
343 zipWith :: (IVector v a, IVector v b, IVector v c) => (a -> b -> c) -> v a -> v b -> v c
344 {-# INLINE zipWith #-}
345 zipWith f xs ys = unstream (Stream.zipWith f (stream xs) (stream ys))
346
347 zip :: (IVector v a, IVector v b, IVector v (a,b)) => v a -> v b -> v (a, b)
348 {-# INLINE zip #-}
349 zip = zipWith (,)
350
351 -- Comparisons
352 -- -----------
353
354 eq :: (IVector v a, Eq a) => v a -> v a -> Bool
355 {-# INLINE eq #-}
356 xs `eq` ys = stream xs == stream ys
357
358 cmp :: (IVector v a, Ord a) => v a -> v a -> Ordering
359 {-# INLINE cmp #-}
360 cmp xs ys = compare (stream xs) (stream ys)
361
362 -- Filtering
363 -- ---------
364
365 -- | Drop elements which do not satisfy the predicate
366 filter :: IVector v a => (a -> Bool) -> v a -> v a
367 {-# INLINE filter #-}
368 filter f = unstream . inplace (Stream.filter f) (MStream.filter f) . stream
369
370 -- | Yield the longest prefix of elements satisfying the predicate without
371 -- copying.
372 takeWhileSlice :: IVector v a => (a -> Bool) -> v a -> v a
373 {-# INLINE takeWhileSlice #-}
374 takeWhileSlice f v = case findIndex (not . f) v of
375 Just n -> takeSlice n v
376 Nothing -> v
377
378 -- | Copy the longest prefix of elements satisfying the predicate to a new
379 -- vector
380 takeWhile :: IVector v a => (a -> Bool) -> v a -> v a
381 {-# INLINE takeWhile #-}
382 takeWhile f = unstream . Stream.takeWhile f . stream
383
384 -- | Drop the longest prefix of elements that satisfy the predicate without
385 -- copying
386 dropWhileSlice :: IVector v a => (a -> Bool) -> v a -> v a
387 {-# INLINE dropWhileSlice #-}
388 dropWhileSlice f v = case findIndex (not . f) v of
389 Just n -> dropSlice n v
390 Nothing -> v
391
392 -- | Drop the longest prefix of elements that satisfy the predicate and copy
393 -- the rest to a new vector.
394 dropWhile :: IVector v a => (a -> Bool) -> v a -> v a
395 {-# INLINE dropWhile #-}
396 dropWhile f = unstream . Stream.dropWhile f . stream
397
398 {-# RULES
399
400 "takeWhileSlice/unstream" forall f s.
401 takeWhileSlice f (new (Mut.unstream s)) = takeWhile f (new (Mut.unstream s))
402
403 "dropWhileSlice/unstream" forall f s.
404 dropWhileSlice f (new (Mut.unstream s)) = dropWhile f (new (Mut.unstream s))
405
406 #-}
407
408 -- Searching
409 -- ---------
410
411 infix 4 `elem`
412 -- | Check whether the vector contains an element
413 elem :: (IVector v a, Eq a) => a -> v a -> Bool
414 {-# INLINE elem #-}
415 elem x = Stream.elem x . stream
416
417 infix 4 `notElem`
418 -- | Inverse of `elem`
419 notElem :: (IVector v a, Eq a) => a -> v a -> Bool
420 {-# INLINE notElem #-}
421 notElem x = Stream.notElem x . stream
422
423 -- | Yield 'Just' the first element matching the predicate or 'Nothing' if no
424 -- such element exists.
425 find :: IVector v a => (a -> Bool) -> v a -> Maybe a
426 {-# INLINE find #-}
427 find f = Stream.find f . stream
428
429 -- | Yield 'Just' the index of the first element matching the predicate or
430 -- 'Nothing' if no such element exists.
431 findIndex :: IVector v a => (a -> Bool) -> v a -> Maybe Int
432 {-# INLINE findIndex #-}
433 findIndex f = Stream.findIndex f . stream
434
435 -- Folding
436 -- -------
437
438 -- | Left fold
439 foldl :: IVector v b => (a -> b -> a) -> a -> v b -> a
440 {-# INLINE foldl #-}
441 foldl f z = Stream.foldl f z . stream
442
443 -- | Lefgt fold on non-empty vectors
444 foldl1 :: IVector v a => (a -> a -> a) -> v a -> a
445 {-# INLINE foldl1 #-}
446 foldl1 f = Stream.foldl1 f . stream
447
448 -- | Left fold with strict accumulator
449 foldl' :: IVector v b => (a -> b -> a) -> a -> v b -> a
450 {-# INLINE foldl' #-}
451 foldl' f z = Stream.foldl' f z . stream
452
453 -- | Left fold on non-empty vectors with strict accumulator
454 foldl1' :: IVector v a => (a -> a -> a) -> v a -> a
455 {-# INLINE foldl1' #-}
456 foldl1' f = Stream.foldl1' f . stream
457
458 -- | Right fold
459 foldr :: IVector v a => (a -> b -> b) -> b -> v a -> b
460 {-# INLINE foldr #-}
461 foldr f z = Stream.foldr f z . stream
462
463 -- | Right fold on non-empty vectors
464 foldr1 :: IVector v a => (a -> a -> a) -> v a -> a
465 {-# INLINE foldr1 #-}
466 foldr1 f = Stream.foldr1 f . stream
467
468 -- | Convert a vector to a list
469 toList :: IVector v a => v a -> [a]
470 {-# INLINE toList #-}
471 toList = Stream.toList . stream
472
473 -- | Convert a list to a vector
474 fromList :: IVector v a => [a] -> v a
475 {-# INLINE fromList #-}
476 fromList = unstream . Stream.fromList
477