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