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