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