Add replicatePrimM and specialise replicateM
[darcs-mirrors/vector.git] / Data / Vector / Primitive / Mutable.hs
1 {-# LANGUAGE MultiParamTypeClasses, FlexibleInstances, ScopedTypeVariables #-}
2
3 -- |
4 -- Module : Data.Vector.Primitive.Mutable
5 -- Copyright : (c) Roman Leshchinskiy 2008-2010
6 -- License : BSD-style
7 --
8 -- Maintainer : Roman Leshchinskiy <rl@cse.unsw.edu.au>
9 -- Stability : experimental
10 -- Portability : non-portable
11 --
12 -- Mutable primitive vectors.
13 --
14
15 module Data.Vector.Primitive.Mutable (
16 -- * Mutable vectors of primitive types
17 MVector(..), IOVector, STVector, Prim,
18
19 -- * Accessors
20
21 -- ** Length information
22 length, null,
23
24 -- ** Extracting subvectors
25 slice, init, tail, take, drop, splitAt,
26 unsafeSlice, unsafeInit, unsafeTail, unsafeTake, unsafeDrop,
27
28 -- ** Overlapping
29 overlaps,
30
31 -- * Construction
32
33 -- ** Initialisation
34 new, unsafeNew, replicate, replicateM, clone,
35
36 -- ** Growing
37 grow, unsafeGrow,
38
39 -- ** Restricting memory usage
40 clear,
41
42 -- * Accessing individual elements
43 read, write, swap,
44 unsafeRead, unsafeWrite, unsafeSwap,
45
46 -- * Modifying vectors
47
48 -- ** Filling and copying
49 set, copy, move, unsafeCopy, unsafeMove,
50
51 -- * Deprecated operations
52 newWith, unsafeNewWith
53 ) where
54
55 import qualified Data.Vector.Generic.Mutable as G
56 import Data.Primitive.ByteArray
57 import Data.Primitive ( Prim, sizeOf )
58 import Control.Monad.Primitive
59 import Control.Monad ( liftM )
60
61 import Prelude hiding ( length, null, replicate, reverse, map, read,
62 take, drop, splitAt, init, tail )
63
64 import Data.Typeable ( Typeable )
65
66 #include "vector.h"
67
68 -- | Mutable vectors of primitive types.
69 data MVector s a = MVector {-# UNPACK #-} !Int
70 {-# UNPACK #-} !Int
71 {-# UNPACK #-} !(MutableByteArray s)
72 deriving ( Typeable )
73
74 type IOVector = MVector RealWorld
75 type STVector s = MVector s
76
77 instance Prim a => G.MVector MVector a where
78 basicLength (MVector _ n _) = n
79 basicUnsafeSlice j m (MVector i n arr)
80 = MVector (i+j) m arr
81
82 {-# INLINE basicOverlaps #-}
83 basicOverlaps (MVector i m arr1) (MVector j n arr2)
84 = sameMutableByteArray arr1 arr2
85 && (between i j (j+n) || between j i (i+m))
86 where
87 between x y z = x >= y && x < z
88
89 {-# INLINE basicUnsafeNew #-}
90 basicUnsafeNew n = MVector 0 n
91 `liftM` newByteArray (n * sizeOf (undefined :: a))
92
93 {-# INLINE basicUnsafeRead #-}
94 basicUnsafeRead (MVector i n arr) j = readByteArray arr (i+j)
95
96 {-# INLINE basicUnsafeWrite #-}
97 basicUnsafeWrite (MVector i n arr) j x = writeByteArray arr (i+j) x
98
99 {-# INLINE basicUnsafeCopy #-}
100 basicUnsafeCopy (MVector i n dst) (MVector j _ src)
101 = memcpyByteArray dst (i * sz) src (j * sz) (n * sz)
102 where
103 sz = sizeOf (undefined :: a)
104
105 {-# INLINE basicUnsafeMove #-}
106 basicUnsafeMove (MVector i n dst) (MVector j _ src)
107 = memmoveByteArray dst (i * sz) src (j * sz) (n * sz)
108 where
109 sz = sizeOf (undefined :: a)
110
111 -- Length information
112 -- ------------------
113
114 -- | Length of the mutable vector.
115 length :: Prim a => MVector s a -> Int
116 {-# INLINE length #-}
117 length = G.length
118
119 -- | Check whether the vector is empty
120 null :: Prim a => MVector s a -> Bool
121 {-# INLINE null #-}
122 null = G.null
123
124 -- Extracting subvectors
125 -- ---------------------
126
127 -- | Yield a part of the mutable vector without copying it.
128 slice :: Prim a => Int -> Int -> MVector s a -> MVector s a
129 {-# INLINE slice #-}
130 slice = G.slice
131
132 take :: Prim a => Int -> MVector s a -> MVector s a
133 {-# INLINE take #-}
134 take = G.take
135
136 drop :: Prim a => Int -> MVector s a -> MVector s a
137 {-# INLINE drop #-}
138 drop = G.drop
139
140 splitAt :: Prim a => Int -> MVector s a -> (MVector s a, MVector s a)
141 {-# INLINE splitAt #-}
142 splitAt = G.splitAt
143
144 init :: Prim a => MVector s a -> MVector s a
145 {-# INLINE init #-}
146 init = G.init
147
148 tail :: Prim a => MVector s a -> MVector s a
149 {-# INLINE tail #-}
150 tail = G.tail
151
152 -- | Yield a part of the mutable vector without copying it. No bounds checks
153 -- are performed.
154 unsafeSlice :: Prim a
155 => Int -- ^ starting index
156 -> Int -- ^ length of the slice
157 -> MVector s a
158 -> MVector s a
159 {-# INLINE unsafeSlice #-}
160 unsafeSlice = G.unsafeSlice
161
162 unsafeTake :: Prim a => Int -> MVector s a -> MVector s a
163 {-# INLINE unsafeTake #-}
164 unsafeTake = G.unsafeTake
165
166 unsafeDrop :: Prim a => Int -> MVector s a -> MVector s a
167 {-# INLINE unsafeDrop #-}
168 unsafeDrop = G.unsafeDrop
169
170 unsafeInit :: Prim a => MVector s a -> MVector s a
171 {-# INLINE unsafeInit #-}
172 unsafeInit = G.unsafeInit
173
174 unsafeTail :: Prim a => MVector s a -> MVector s a
175 {-# INLINE unsafeTail #-}
176 unsafeTail = G.unsafeTail
177
178 -- Overlapping
179 -- -----------
180
181 -- Check whether two vectors overlap.
182 overlaps :: Prim a => MVector s a -> MVector s a -> Bool
183 {-# INLINE overlaps #-}
184 overlaps = G.overlaps
185
186 -- Initialisation
187 -- --------------
188
189 -- | Create a mutable vector of the given length.
190 new :: (PrimMonad m, Prim a) => Int -> m (MVector (PrimState m) a)
191 {-# INLINE new #-}
192 new = G.new
193
194 -- | Create a mutable vector of the given length. The length is not checked.
195 unsafeNew :: (PrimMonad m, Prim a) => Int -> m (MVector (PrimState m) a)
196 {-# INLINE unsafeNew #-}
197 unsafeNew = G.unsafeNew
198
199 -- | Create a mutable vector of the given length (0 if the length is negative)
200 -- and fill it with an initial value.
201 replicate :: (PrimMonad m, Prim a) => Int -> a -> m (MVector (PrimState m) a)
202 {-# INLINE replicate #-}
203 replicate = G.replicate
204
205 -- | Create a mutable vector of the given length (0 if the length is negative)
206 -- and fill it with values produced by repeatedly executing the monadic action.
207 replicateM :: (PrimMonad m, Prim a) => Int -> m a -> m (MVector (PrimState m) a)
208 {-# INLINE replicateM #-}
209 replicateM = G.replicateM
210
211 -- | Create a copy of a mutable vector.
212 clone :: (PrimMonad m, Prim a)
213 => MVector (PrimState m) a -> m (MVector (PrimState m) a)
214 {-# INLINE clone #-}
215 clone = G.clone
216
217 -- Growing
218 -- -------
219
220 -- | Grow a vector by the given number of elements. The number must be
221 -- positive.
222 grow :: (PrimMonad m, Prim a)
223 => MVector (PrimState m) a -> Int -> m (MVector (PrimState m) a)
224 {-# INLINE grow #-}
225 grow = G.grow
226
227 -- | Grow a vector by the given number of elements. The number must be
228 -- positive but this is not checked.
229 unsafeGrow :: (PrimMonad m, Prim a)
230 => MVector (PrimState m) a -> Int -> m (MVector (PrimState m) a)
231 {-# INLINE unsafeGrow #-}
232 unsafeGrow = G.unsafeGrow
233
234 -- Restricting memory usage
235 -- ------------------------
236
237 -- | Reset all elements of the vector to some undefined value, clearing all
238 -- references to external objects. This is usually a noop for unboxed vectors.
239 clear :: (PrimMonad m, Prim a) => MVector (PrimState m) a -> m ()
240 {-# INLINE clear #-}
241 clear = G.clear
242
243 -- Accessing individual elements
244 -- -----------------------------
245
246 -- | Yield the element at the given position.
247 read :: (PrimMonad m, Prim a) => MVector (PrimState m) a -> Int -> m a
248 {-# INLINE read #-}
249 read = G.read
250
251 -- | Replace the element at the given position.
252 write :: (PrimMonad m, Prim a) => MVector (PrimState m) a -> Int -> a -> m ()
253 {-# INLINE write #-}
254 write = G.write
255
256 -- | Swap the elements at the given positions.
257 swap :: (PrimMonad m, Prim a) => MVector (PrimState m) a -> Int -> Int -> m ()
258 {-# INLINE swap #-}
259 swap = G.swap
260
261
262 -- | Yield the element at the given position. No bounds checks are performed.
263 unsafeRead :: (PrimMonad m, Prim a) => MVector (PrimState m) a -> Int -> m a
264 {-# INLINE unsafeRead #-}
265 unsafeRead = G.unsafeRead
266
267 -- | Replace the element at the given position. No bounds checks are performed.
268 unsafeWrite
269 :: (PrimMonad m, Prim a) => MVector (PrimState m) a -> Int -> a -> m ()
270 {-# INLINE unsafeWrite #-}
271 unsafeWrite = G.unsafeWrite
272
273 -- | Swap the elements at the given positions. No bounds checks are performed.
274 unsafeSwap
275 :: (PrimMonad m, Prim a) => MVector (PrimState m) a -> Int -> Int -> m ()
276 {-# INLINE unsafeSwap #-}
277 unsafeSwap = G.unsafeSwap
278
279 -- Filling and copying
280 -- -------------------
281
282 -- | Set all elements of the vector to the given value.
283 set :: (PrimMonad m, Prim a) => MVector (PrimState m) a -> a -> m ()
284 {-# INLINE set #-}
285 set = G.set
286
287 -- | Copy a vector. The two vectors must have the same length and may not
288 -- overlap.
289 copy :: (PrimMonad m, Prim a)
290 => MVector (PrimState m) a -> MVector (PrimState m) a -> m ()
291 {-# INLINE copy #-}
292 copy = G.copy
293
294 -- | Copy a vector. The two vectors must have the same length and may not
295 -- overlap. This is not checked.
296 unsafeCopy :: (PrimMonad m, Prim a)
297 => MVector (PrimState m) a -- ^ target
298 -> MVector (PrimState m) a -- ^ source
299 -> m ()
300 {-# INLINE unsafeCopy #-}
301 unsafeCopy = G.unsafeCopy
302
303 -- | Move the contents of a vector. The two vectors must have the same
304 -- length.
305 --
306 -- If the vectors do not overlap, then this is equivalent to 'copy'.
307 -- Otherwise, the copying is performed as if the source vector were
308 -- copied to a temporary vector and then the temporary vector was copied
309 -- to the target vector.
310 move :: (PrimMonad m, Prim a)
311 => MVector (PrimState m) a -> MVector (PrimState m) a -> m ()
312 {-# INLINE move #-}
313 move = G.move
314
315 -- | Move the contents of a vector. The two vectors must have the same
316 -- length, but this is not checked.
317 --
318 -- If the vectors do not overlap, then this is equivalent to 'unsafeCopy'.
319 -- Otherwise, the copying is performed as if the source vector were
320 -- copied to a temporary vector and then the temporary vector was copied
321 -- to the target vector.
322 unsafeMove :: (PrimMonad m, Prim a)
323 => MVector (PrimState m) a -- ^ target
324 -> MVector (PrimState m) a -- ^ source
325 -> m ()
326 {-# INLINE unsafeMove #-}
327 unsafeMove = G.unsafeMove
328
329 -- Deprecated functions
330 -- --------------------
331
332 -- | /DEPRECATED/ Use 'replicate' instead
333 newWith :: (PrimMonad m, Prim a) => Int -> a -> m (MVector (PrimState m) a)
334 {-# INLINE newWith #-}
335 newWith = G.replicate
336
337 -- | /DEPRECATED/ Use 'replicate' instead
338 unsafeNewWith :: (PrimMonad m, Prim a) => Int -> a -> m (MVector (PrimState m) a)
339 {-# INLINE unsafeNewWith #-}
340 unsafeNewWith = G.replicate
341
342 {-# DEPRECATED newWith, unsafeNewWith "Use replicate instead" #-}
343