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