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