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