Add replicatePrimM and specialise replicateM
[darcs-mirrors/vector.git] / Data / Vector / Mutable.hs
1 {-# LANGUAGE MultiParamTypeClasses, FlexibleInstances, BangPatterns #-}
2
3 -- |
4 -- Module : Data.Vector.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 boxed vectors.
13 --
14
15 module Data.Vector.Mutable (
16 -- * Mutable boxed vectors
17 MVector(..), IOVector, STVector,
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 Control.Monad (when)
56 import qualified Data.Vector.Generic.Mutable as G
57 import Data.Primitive.Array
58 import Control.Monad.Primitive
59
60 import Prelude hiding ( length, null, replicate, reverse, map, read,
61 take, drop, splitAt, init, tail )
62
63 import Data.Typeable ( Typeable )
64
65 #include "vector.h"
66
67 -- | Mutable boxed vectors keyed on the monad they live in ('IO' or @'ST' s@).
68 data MVector s a = MVector {-# UNPACK #-} !Int
69 {-# UNPACK #-} !Int
70 {-# UNPACK #-} !(MutableArray s a)
71 deriving ( Typeable )
72
73 type IOVector = MVector RealWorld
74 type STVector s = MVector s
75
76 instance G.MVector MVector a where
77 {-# INLINE basicLength #-}
78 basicLength (MVector _ n _) = n
79
80 {-# INLINE basicUnsafeSlice #-}
81 basicUnsafeSlice j m (MVector i n arr) = MVector (i+j) m arr
82
83 {-# INLINE basicOverlaps #-}
84 basicOverlaps (MVector i m arr1) (MVector j n arr2)
85 = sameMutableArray arr1 arr2
86 && (between i j (j+n) || between j i (i+m))
87 where
88 between x y z = x >= y && x < z
89
90 {-# INLINE basicUnsafeNew #-}
91 basicUnsafeNew n
92 = do
93 arr <- newArray n uninitialised
94 return (MVector 0 n arr)
95
96 {-# INLINE basicUnsafeReplicate #-}
97 basicUnsafeReplicate n x
98 = do
99 arr <- newArray n x
100 return (MVector 0 n arr)
101
102 {-# INLINE basicUnsafeRead #-}
103 basicUnsafeRead (MVector i n arr) j = readArray arr (i+j)
104
105 {-# INLINE basicUnsafeWrite #-}
106 basicUnsafeWrite (MVector i n arr) j x = writeArray arr (i+j) x
107
108 basicUnsafeMove dst@(MVector iDst n arrDst) src@(MVector iSrc _ arrSrc)
109 = case n of
110 0 -> return ()
111 1 -> readArray arrSrc iSrc >>= writeArray arrDst iDst
112 2 -> do
113 x <- readArray arrSrc iSrc
114 y <- readArray arrSrc (iSrc + 1)
115 writeArray arrDst iDst x
116 writeArray arrDst (iDst + 1) y
117 _
118 | overlaps dst src
119 -> case compare iDst iSrc of
120 LT -> moveBackwards arrDst iDst iSrc n
121 EQ -> return ()
122 GT | (iDst - iSrc) * 2 < n
123 -> moveForwardsLargeOverlap arrDst iDst iSrc n
124 | otherwise
125 -> moveForwardsSmallOverlap arrDst iDst iSrc n
126 | otherwise -> G.basicUnsafeCopy dst src
127
128 {-# INLINE basicClear #-}
129 basicClear v = G.set v uninitialised
130
131 {-# INLINE moveBackwards #-}
132 moveBackwards :: PrimMonad m => MutableArray (PrimState m) a -> Int -> Int -> Int -> m ()
133 moveBackwards !arr !dstOff !srcOff !len =
134 INTERNAL_CHECK(check) "moveBackwards" "not a backwards move" (dstOff < srcOff)
135 $ loopM len $ \ i -> readArray arr (srcOff + i) >>= writeArray arr (dstOff + i)
136
137 {-# INLINE moveForwardsSmallOverlap #-}
138 -- Performs a move when dstOff > srcOff, optimized for when the overlap of the intervals is small.
139 moveForwardsSmallOverlap :: PrimMonad m => MutableArray (PrimState m) a -> Int -> Int -> Int -> m ()
140 moveForwardsSmallOverlap !arr !dstOff !srcOff !len =
141 INTERNAL_CHECK(check) "moveForwardsSmallOverlap" "not a forward move" (dstOff > srcOff)
142 $ do
143 tmp <- newArray overlap uninitialised
144 loopM overlap $ \ i -> readArray arr (dstOff + i) >>= writeArray tmp i
145 loopM nonOverlap $ \ i -> readArray arr (srcOff + i) >>= writeArray arr (dstOff + i)
146 loopM overlap $ \ i -> readArray tmp i >>= writeArray arr (dstOff + nonOverlap + i)
147 where nonOverlap = dstOff - srcOff; overlap = len - nonOverlap
148
149 -- Performs a move when dstOff > srcOff, optimized for when the overlap of the intervals is large.
150 moveForwardsLargeOverlap :: PrimMonad m => MutableArray (PrimState m) a -> Int -> Int -> Int -> m ()
151 moveForwardsLargeOverlap !arr !dstOff !srcOff !len =
152 INTERNAL_CHECK(check) "moveForwardsLargeOverlap" "not a forward move" (dstOff > srcOff)
153 $ do
154 queue <- newArray nonOverlap uninitialised
155 loopM nonOverlap $ \ i -> readArray arr (srcOff + i) >>= writeArray queue i
156 let mov !i !qTop = when (i < dstOff + len) $ do
157 x <- readArray arr i
158 y <- readArray queue qTop
159 writeArray arr i y
160 writeArray queue qTop x
161 mov (i+1) (if qTop + 1 >= nonOverlap then 0 else qTop + 1)
162 mov dstOff 0
163 where nonOverlap = dstOff - srcOff
164
165 {-# INLINE loopM #-}
166 loopM :: Monad m => Int -> (Int -> m a) -> m ()
167 loopM !n k = let
168 go i = when (i < n) (k i >> go (i+1))
169 in go 0
170
171 uninitialised :: a
172 uninitialised = error "Data.Vector.Mutable: uninitialised element"
173
174 -- Length information
175 -- ------------------
176
177 -- | Length of the mutable vector.
178 length :: MVector s a -> Int
179 {-# INLINE length #-}
180 length = G.length
181
182 -- | Check whether the vector is empty
183 null :: MVector s a -> Bool
184 {-# INLINE null #-}
185 null = G.null
186
187 -- Extracting subvectors
188 -- ---------------------
189
190 -- | Yield a part of the mutable vector without copying it.
191 slice :: Int -> Int -> MVector s a -> MVector s a
192 {-# INLINE slice #-}
193 slice = G.slice
194
195 take :: Int -> MVector s a -> MVector s a
196 {-# INLINE take #-}
197 take = G.take
198
199 drop :: Int -> MVector s a -> MVector s a
200 {-# INLINE drop #-}
201 drop = G.drop
202
203 {-# INLINE splitAt #-}
204 splitAt :: Int -> MVector s a -> (MVector s a, MVector s a)
205 splitAt = G.splitAt
206
207 init :: MVector s a -> MVector s a
208 {-# INLINE init #-}
209 init = G.init
210
211 tail :: MVector s a -> MVector s a
212 {-# INLINE tail #-}
213 tail = G.tail
214
215 -- | Yield a part of the mutable vector without copying it. No bounds checks
216 -- are performed.
217 unsafeSlice :: Int -- ^ starting index
218 -> Int -- ^ length of the slice
219 -> MVector s a
220 -> MVector s a
221 {-# INLINE unsafeSlice #-}
222 unsafeSlice = G.unsafeSlice
223
224 unsafeTake :: Int -> MVector s a -> MVector s a
225 {-# INLINE unsafeTake #-}
226 unsafeTake = G.unsafeTake
227
228 unsafeDrop :: Int -> MVector s a -> MVector s a
229 {-# INLINE unsafeDrop #-}
230 unsafeDrop = G.unsafeDrop
231
232 unsafeInit :: MVector s a -> MVector s a
233 {-# INLINE unsafeInit #-}
234 unsafeInit = G.unsafeInit
235
236 unsafeTail :: MVector s a -> MVector s a
237 {-# INLINE unsafeTail #-}
238 unsafeTail = G.unsafeTail
239
240 -- Overlapping
241 -- -----------
242
243 -- Check whether two vectors overlap.
244 overlaps :: MVector s a -> MVector s a -> Bool
245 {-# INLINE overlaps #-}
246 overlaps = G.overlaps
247
248 -- Initialisation
249 -- --------------
250
251 -- | Create a mutable vector of the given length.
252 new :: PrimMonad m => Int -> m (MVector (PrimState m) a)
253 {-# INLINE new #-}
254 new = G.new
255
256 -- | Create a mutable vector of the given length. The length is not checked.
257 unsafeNew :: PrimMonad m => Int -> m (MVector (PrimState m) a)
258 {-# INLINE unsafeNew #-}
259 unsafeNew = G.unsafeNew
260
261 -- | Create a mutable vector of the given length (0 if the length is negative)
262 -- and fill it with an initial value.
263 replicate :: PrimMonad m => Int -> a -> m (MVector (PrimState m) a)
264 {-# INLINE replicate #-}
265 replicate = G.replicate
266
267 -- | Create a mutable vector of the given length (0 if the length is negative)
268 -- and fill it with values produced by repeatedly executing the monadic action.
269 replicateM :: PrimMonad m => Int -> m a -> m (MVector (PrimState m) a)
270 {-# INLINE replicateM #-}
271 replicateM = G.replicateM
272
273 -- | Create a copy of a mutable vector.
274 clone :: PrimMonad m => MVector (PrimState m) a -> m (MVector (PrimState m) a)
275 {-# INLINE clone #-}
276 clone = G.clone
277
278 -- Growing
279 -- -------
280
281 -- | Grow a vector by the given number of elements. The number must be
282 -- positive.
283 grow :: PrimMonad m
284 => MVector (PrimState m) a -> Int -> m (MVector (PrimState m) a)
285 {-# INLINE grow #-}
286 grow = G.grow
287
288 -- | Grow a vector by the given number of elements. The number must be
289 -- positive but this is not checked.
290 unsafeGrow :: PrimMonad m
291 => MVector (PrimState m) a -> Int -> m (MVector (PrimState m) a)
292 {-# INLINE unsafeGrow #-}
293 unsafeGrow = G.unsafeGrow
294
295 -- Restricting memory usage
296 -- ------------------------
297
298 -- | Reset all elements of the vector to some undefined value, clearing all
299 -- references to external objects. This is usually a noop for unboxed vectors.
300 clear :: PrimMonad m => MVector (PrimState m) a -> m ()
301 {-# INLINE clear #-}
302 clear = G.clear
303
304 -- Accessing individual elements
305 -- -----------------------------
306
307 -- | Yield the element at the given position.
308 read :: PrimMonad m => MVector (PrimState m) a -> Int -> m a
309 {-# INLINE read #-}
310 read = G.read
311
312 -- | Replace the element at the given position.
313 write :: PrimMonad m => MVector (PrimState m) a -> Int -> a -> m ()
314 {-# INLINE write #-}
315 write = G.write
316
317 -- | Swap the elements at the given positions.
318 swap :: PrimMonad m => MVector (PrimState m) a -> Int -> Int -> m ()
319 {-# INLINE swap #-}
320 swap = G.swap
321
322
323 -- | Yield the element at the given position. No bounds checks are performed.
324 unsafeRead :: PrimMonad m => MVector (PrimState m) a -> Int -> m a
325 {-# INLINE unsafeRead #-}
326 unsafeRead = G.unsafeRead
327
328 -- | Replace the element at the given position. No bounds checks are performed.
329 unsafeWrite :: PrimMonad m => MVector (PrimState m) a -> Int -> a -> m ()
330 {-# INLINE unsafeWrite #-}
331 unsafeWrite = G.unsafeWrite
332
333 -- | Swap the elements at the given positions. No bounds checks are performed.
334 unsafeSwap :: PrimMonad m => MVector (PrimState m) a -> Int -> Int -> m ()
335 {-# INLINE unsafeSwap #-}
336 unsafeSwap = G.unsafeSwap
337
338 -- Filling and copying
339 -- -------------------
340
341 -- | Set all elements of the vector to the given value.
342 set :: PrimMonad m => MVector (PrimState m) a -> a -> m ()
343 {-# INLINE set #-}
344 set = G.set
345
346 -- | Copy a vector. The two vectors must have the same length and may not
347 -- overlap.
348 copy :: PrimMonad m
349 => MVector (PrimState m) a -> MVector (PrimState m) a -> m ()
350 {-# INLINE copy #-}
351 copy = G.copy
352
353 -- | Copy a vector. The two vectors must have the same length and may not
354 -- overlap. This is not checked.
355 unsafeCopy :: PrimMonad m => MVector (PrimState m) a -- ^ target
356 -> MVector (PrimState m) a -- ^ source
357 -> m ()
358 {-# INLINE unsafeCopy #-}
359 unsafeCopy = G.unsafeCopy
360
361 -- | Move the contents of a vector. The two vectors must have the same
362 -- length.
363 --
364 -- If the vectors do not overlap, then this is equivalent to 'copy'.
365 -- Otherwise, the copying is performed as if the source vector were
366 -- copied to a temporary vector and then the temporary vector was copied
367 -- to the target vector.
368 move :: PrimMonad m
369 => MVector (PrimState m) a -> MVector (PrimState m) a -> m ()
370 {-# INLINE move #-}
371 move = G.move
372
373 -- | Move the contents of a vector. The two vectors must have the same
374 -- length, but this is not checked.
375 --
376 -- If the vectors do not overlap, then this is equivalent to 'unsafeCopy'.
377 -- Otherwise, the copying is performed as if the source vector were
378 -- copied to a temporary vector and then the temporary vector was copied
379 -- to the target vector.
380 unsafeMove :: PrimMonad m => MVector (PrimState m) a -- ^ target
381 -> MVector (PrimState m) a -- ^ source
382 -> m ()
383 {-# INLINE unsafeMove #-}
384 unsafeMove = G.unsafeMove
385
386 -- Deprecated functions
387 -- --------------------
388
389 -- | /DEPRECATED/ Use 'replicate' instead
390 newWith :: PrimMonad m => Int -> a -> m (MVector (PrimState m) a)
391 {-# INLINE newWith #-}
392 newWith = G.replicate
393
394 -- | /DEPRECATED/ Use 'replicate' instead
395 unsafeNewWith :: PrimMonad m => Int -> a -> m (MVector (PrimState m) a)
396 {-# INLINE unsafeNewWith #-}
397 unsafeNewWith = G.replicate
398
399 {-# DEPRECATED newWith, unsafeNewWith "Use replicate instead" #-}
400