MVector.clear
[darcs-mirrors/vector.git] / Data / Vector / MVector.hs
1 {-# LANGUAGE MultiParamTypeClasses #-}
2 -- |
3 -- Module : Data.Vector.MVector
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 mutable vectors
12 --
13
14 #include "phases.h"
15
16 module Data.Vector.MVector (
17 MVectorPure(..), MVector(..),
18
19 slice, new, newWith, read, write, copy, grow, unstream, update, reverse, map
20 ) where
21
22 import qualified Data.Vector.Stream as Stream
23 import Data.Vector.Stream ( Stream )
24 import Data.Vector.Stream.Size
25
26 import Control.Monad.ST ( ST )
27 import Control.Exception ( assert )
28
29 import GHC.Float (
30 double2Int, int2Double
31 )
32
33 import Prelude hiding ( length, reverse, map, read )
34
35 gROWTH_FACTOR :: Double
36 gROWTH_FACTOR = 1.5
37
38 -- | Basic pure functions on mutable vectors
39 class MVectorPure v a where
40 -- | Length of the mutable vector
41 length :: v a -> Int
42
43 -- | Yield a part of the mutable vector without copying it. No range checks!
44 unsafeSlice :: v a -> Int -- ^ starting index
45 -> Int -- ^ length of the slice
46 -> v a
47
48 -- Check whether two vectors overlap.
49 overlaps :: v a -> v a -> Bool
50
51 -- | Class of mutable vectors. The type @m@ is the monad in which the mutable
52 -- vector can be transformed and @a@ is the type of elements.
53 --
54 class (Monad m, MVectorPure v a) => MVector v m a where
55 -- | Create a mutable vector of the given length. Length is not checked!
56 unsafeNew :: Int -> m (v a)
57
58 -- | Create a mutable vector of the given length and fill it with an
59 -- initial value. Length is not checked!
60 unsafeNewWith :: Int -> a -> m (v a)
61
62 -- | Yield the element at the given position. Index is not checked!
63 unsafeRead :: v a -> Int -> m a
64
65 -- | Replace the element at the given position. Index is not checked!
66 unsafeWrite :: v a -> Int -> a -> m ()
67
68 -- | Clear all references to external objects
69 clear :: v a -> m ()
70
71 -- | Write the value at each position.
72 set :: v a -> a -> m ()
73
74 -- | Copy a vector. The two vectors may not overlap. This is not checked!
75 unsafeCopy :: v a -- ^ target
76 -> v a -- ^ source
77 -> m ()
78
79 -- | Grow a vector by the given number of elements. The length is not
80 -- checked!
81 unsafeGrow :: v a -> Int -> m (v a)
82
83 {-# INLINE unsafeNewWith #-}
84 unsafeNewWith n x = do
85 v <- unsafeNew n
86 set v x
87 return v
88
89 {-# INLINE set #-}
90 set v x = do_set 0
91 where
92 n = length v
93
94 do_set i | i < n = do
95 unsafeWrite v i x
96 do_set (i+1)
97 | otherwise = return ()
98
99 {-# INLINE unsafeCopy #-}
100 unsafeCopy dst src = do_copy 0
101 where
102 n = length src
103
104 do_copy i | i < n = do
105 x <- unsafeRead src i
106 unsafeWrite dst i x
107 do_copy (i+1)
108 | otherwise = return ()
109
110 {-# INLINE unsafeGrow #-}
111 unsafeGrow v by = do
112 v' <- unsafeNew (n+by)
113 unsafeCopy (unsafeSlice v' 0 n) v
114 return v'
115 where
116 n = length v
117
118 -- | Test whether the index is valid for the vector
119 inBounds :: MVectorPure v a => v a -> Int -> Bool
120 {-# INLINE inBounds #-}
121 inBounds v i = i >= 0 && i < length v
122
123 -- | Yield a part of the mutable vector without copying it. Safer version of
124 -- 'unsafeSlice'.
125 slice :: MVectorPure v a => v a -> Int -> Int -> v a
126 {-# INLINE slice #-}
127 slice v i n = assert (i >=0 && n >= 0 && i+n <= length v)
128 $ unsafeSlice v i n
129
130 -- | Create a mutable vector of the given length. Safer version of
131 -- 'unsafeNew'.
132 new :: MVector v m a => Int -> m (v a)
133 {-# INLINE new #-}
134 new n = assert (n >= 0) $ unsafeNew n
135
136 -- | Create a mutable vector of the given length and fill it with an
137 -- initial value. Safer version of 'unsafeNewWith'.
138 newWith :: MVector v m a => Int -> a -> m (v a)
139 {-# INLINE newWith #-}
140 newWith n x = assert (n >= 0) $ unsafeNewWith n x
141
142 -- | Yield the element at the given position. Safer version of 'unsafeRead'.
143 read :: MVector v m a => v a -> Int -> m a
144 {-# INLINE read #-}
145 read v i = assert (inBounds v i) $ unsafeRead v i
146
147 -- | Replace the element at the given position. Safer version of
148 -- 'unsafeWrite'.
149 write :: MVector v m a => v a -> Int -> a -> m ()
150 {-# INLINE write #-}
151 write v i x = assert (inBounds v i) $ unsafeWrite v i x
152
153 -- | Copy a vector. The two vectors may not overlap. Safer version of
154 -- 'unsafeCopy'.
155 copy :: MVector v m a => v a -> v a -> m ()
156 {-# INLINE copy #-}
157 copy dst src = assert (not (dst `overlaps` src) && length dst == length src)
158 $ unsafeCopy dst src
159
160 -- | Grow a vector by the given number of elements. Safer version of
161 -- 'unsafeGrow'.
162 grow :: MVector v m a => v a -> Int -> m (v a)
163 {-# INLINE grow #-}
164 grow v by = assert (by >= 0)
165 $ unsafeGrow v by
166
167
168 -- | Create a new mutable vector and fill it with elements from the 'Stream'.
169 -- The vector will grow logarithmically if the 'Size' hint of the 'Stream' is
170 -- inexact.
171 unstream :: MVector v m a => Stream a -> m (v a)
172 {-# INLINE_STREAM unstream #-}
173 unstream s = case upperBound (Stream.size s) of
174 Just n -> unstreamMax s n
175 Nothing -> unstreamUnknown s
176
177 unstreamMax :: MVector v m a => Stream a -> Int -> m (v a)
178 {-# INLINE unstreamMax #-}
179 unstreamMax s n
180 = do
181 v <- new n
182 let put i x = do { write v i x; return (i+1) }
183 n' <- Stream.foldM put 0 s
184 return $ slice v 0 n'
185
186 unstreamUnknown :: MVector v m a => Stream a -> m (v a)
187 {-# INLINE unstreamUnknown #-}
188 unstreamUnknown s
189 = do
190 v <- new 0
191 (v', n) <- Stream.foldM put (v, 0) s
192 return $ slice v' 0 n
193 where
194 {-# INLINE put #-}
195 put (v, i) x = do
196 v' <- enlarge v i
197 unsafeWrite v' i x
198 return (v', i+1)
199
200 {-# INLINE enlarge #-}
201 enlarge v i | i < length v = return v
202 | otherwise = unsafeGrow v
203 . max 1
204 . double2Int
205 $ int2Double (length v) * gROWTH_FACTOR
206
207 update :: MVector v m a => v a -> Stream (Int, a) -> m ()
208 {-# INLINE update #-}
209 update v s = Stream.mapM_ put s
210 where
211 {-# INLINE put #-}
212 put (i, x) = write v i x
213
214 reverse :: MVector v m a => v a -> m ()
215 {-# INLINE reverse #-}
216 reverse v = reverse_loop 0 (length v - 1)
217 where
218 reverse_loop i j | i < j = do
219 x <- unsafeRead v i
220 y <- unsafeRead v j
221 unsafeWrite v i y
222 unsafeWrite v j x
223 reverse_loop _ _ = return ()
224
225
226 map :: MVector v m a => (a -> a) -> v a -> m ()
227 {-# INLINE map #-}
228 map f v = map_loop 0
229 where
230 n = length v
231
232 map_loop i | i <= n = do
233 x <- read v i
234 write v i (f x)
235 | otherwise = return ()
236