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