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