Adapt fusion rules
[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 -- | Zip two vectors with the given function.
300 zipWith :: (IVector v a, IVector v b, IVector v c) => (a -> b -> c) -> v a -> v b -> v c
301 {-# INLINE zipWith #-}
302 zipWith f xs ys = unstream (Stream.zipWith f (stream xs) (stream ys))
303
304 -- Filtering
305 -- ---------
306
307 -- | Drop elements which do not satisfy the predicate
308 filter :: IVector v a => (a -> Bool) -> v a -> v a
309 {-# INLINE filter #-}
310 filter f = unstream . Stream.filter f . stream
311
312 -- | Yield the longest prefix of elements satisfying the predicate without
313 -- copying.
314 takeWhileSlice :: IVector v a => (a -> Bool) -> v a -> v a
315 {-# INLINE takeWhileSlice #-}
316 takeWhileSlice f v = case findIndex (not . f) v of
317 Just n -> takeSlice n v
318 Nothing -> v
319
320 -- | Copy the longest prefix of elements satisfying the predicate to a new
321 -- vector
322 takeWhile :: IVector v a => (a -> Bool) -> v a -> v a
323 {-# INLINE takeWhile #-}
324 takeWhile f = unstream . Stream.takeWhile f . stream
325
326 -- | Drop the longest prefix of elements that satisfy the predicate without
327 -- copying
328 dropWhileSlice :: IVector v a => (a -> Bool) -> v a -> v a
329 {-# INLINE dropWhileSlice #-}
330 dropWhileSlice f v = case findIndex (not . f) v of
331 Just n -> dropSlice n v
332 Nothing -> v
333
334 -- | Drop the longest prefix of elements that satisfy the predicate and copy
335 -- the rest to a new vector.
336 dropWhile :: IVector v a => (a -> Bool) -> v a -> v a
337 {-# INLINE dropWhile #-}
338 dropWhile f = unstream . Stream.dropWhile f . stream
339
340 {-# RULES
341
342 "takeWhileSlice/unstream" forall f s.
343 takeWhileSlice f (new (Mut.unstream s)) = takeWhile f (new (Mut.unstream s))
344
345 "dropWhileSlice/unstream" forall f s.
346 dropWhileSlice f (new (Mut.unstream s)) = dropWhile f (new (Mut.unstream s))
347
348 #-}
349
350 -- Searching
351 -- ---------
352
353 infix 4 `elem`
354 -- | Check whether the vector contains an element
355 elem :: (IVector v a, Eq a) => a -> v a -> Bool
356 {-# INLINE elem #-}
357 elem x = Stream.elem x . stream
358
359 infix 4 `notElem`
360 -- | Inverse of `elem`
361 notElem :: (IVector v a, Eq a) => a -> v a -> Bool
362 {-# INLINE notElem #-}
363 notElem x = Stream.notElem x . stream
364
365 -- | Yield 'Just' the first element matching the predicate or 'Nothing' if no
366 -- such element exists.
367 find :: IVector v a => (a -> Bool) -> v a -> Maybe a
368 {-# INLINE find #-}
369 find f = Stream.find f . stream
370
371 -- | Yield 'Just' the index of the first element matching the predicate or
372 -- 'Nothing' if no such element exists.
373 findIndex :: IVector v a => (a -> Bool) -> v a -> Maybe Int
374 {-# INLINE findIndex #-}
375 findIndex f = Stream.findIndex f . stream
376
377 -- Folding
378 -- -------
379
380 -- | Left fold
381 foldl :: IVector v b => (a -> b -> a) -> a -> v b -> a
382 {-# INLINE foldl #-}
383 foldl f z = Stream.foldl f z . stream
384
385 -- | Lefgt fold on non-empty vectors
386 foldl1 :: IVector v a => (a -> a -> a) -> v a -> a
387 {-# INLINE foldl1 #-}
388 foldl1 f = Stream.foldl1 f . stream
389
390 -- | Left fold with strict accumulator
391 foldl' :: IVector v b => (a -> b -> a) -> a -> v b -> a
392 {-# INLINE foldl' #-}
393 foldl' f z = Stream.foldl' f z . stream
394
395 -- | Left fold on non-empty vectors with strict accumulator
396 foldl1' :: IVector v a => (a -> a -> a) -> v a -> a
397 {-# INLINE foldl1' #-}
398 foldl1' f = Stream.foldl1' f . stream
399
400 -- | Right fold
401 foldr :: IVector v a => (a -> b -> b) -> b -> v a -> b
402 {-# INLINE foldr #-}
403 foldr f z = Stream.foldr f z . stream
404
405 -- | Right fold on non-empty vectors
406 foldr1 :: IVector v a => (a -> a -> a) -> v a -> a
407 {-# INLINE foldr1 #-}
408 foldr1 f = Stream.foldr1 f . stream
409
410 -- | Convert a vector to a list
411 toList :: IVector v a => v a -> [a]
412 {-# INLINE toList #-}
413 toList = Stream.toList . stream
414
415 -- | Convert a list to a vector
416 fromList :: IVector v a => [a] -> v a
417 {-# INLINE fromList #-}
418 fromList = unstream . Stream.fromList
419