content: {-# LANGUAGE MultiParamTypeClasses, FlexibleInstances, ScopedTypeVariables #-}
  
  -- |
  -- Module      : Data.Vector.Storable.Mutable
  -- Copyright   : (c) Roman Leshchinskiy 2009-2010
  -- License     : BSD-style
  --
  -- Maintainer  : Roman Leshchinskiy <rl@cse.unsw.edu.au>
  -- Stability   : experimental
  -- Portability : non-portable
  -- 
  -- Mutable vectors based on Storable.
  --
  
  module Data.Vector.Storable.Mutable(
    -- * Mutable vectors of 'Storable' types
    MVector(..), IOVector, STVector, Storable,
  
    -- * Accessors
  
    -- ** Length information
    length, null,
  
    -- ** Extracting subvectors
    slice, init, tail, take, drop, splitAt,
    unsafeSlice, unsafeInit, unsafeTail, unsafeTake, unsafeDrop,
  
    -- ** Overlapping
    overlaps,
  
    -- * Construction
  
    -- ** Initialisation
    new, unsafeNew, replicate, replicateM, clone,
  
    -- ** Growing
    grow, unsafeGrow,
  
    -- ** Restricting memory usage
    clear,
  
    -- * Accessing individual elements
    read, write, swap,
    unsafeRead, unsafeWrite, unsafeSwap,
  
    -- * Modifying vectors
  
    -- ** Filling and copying
    set, copy, move, unsafeCopy, unsafeMove,
  
    -- * Unsafe conversions
    unsafeCast,
  
    -- * Raw pointers
    unsafeFromForeignPtr, unsafeFromForeignPtr0,
    unsafeToForeignPtr,   unsafeToForeignPtr0,
    unsafeWith
  ) where
  
  import Control.DeepSeq ( NFData(rnf) )
  
  import qualified Data.Vector.Generic.Mutable as G
  import Data.Vector.Storable.Internal
  
  import Foreign.Storable
  import Foreign.ForeignPtr
  
  #if __GLASGOW_HASKELL__ >= 605
  import GHC.ForeignPtr (mallocPlainForeignPtrBytes)
  #endif
  
  import Foreign.Ptr
  import Foreign.Marshal.Array ( advancePtr, copyArray, moveArray )
  import Foreign.C.Types ( CInt )
  
  import Control.Monad.Primitive
  import Data.Primitive.Addr
  import Data.Primitive.Types (Prim)
  
  import GHC.Word (Word8, Word16, Word32, Word64)
  import GHC.Ptr (Ptr(..))
  
  import Prelude hiding ( length, null, replicate, reverse, map, read,
                          take, drop, splitAt, init, tail )
  
  import Data.Typeable ( Typeable )
  
  #include "vector.h"
  
  -- | Mutable 'Storable'-based vectors
  data MVector s a = MVector {-# UNPACK #-} !Int
                             {-# UNPACK #-} !(ForeignPtr a)
          deriving ( Typeable )
  
  type IOVector = MVector RealWorld
  type STVector s = MVector s
  
  instance NFData (MVector s a) where
    rnf (MVector _ _) = ()
  
  instance Storable a => G.MVector MVector a where
    {-# INLINE basicLength #-}
    basicLength (MVector n _) = n
  
    {-# INLINE basicUnsafeSlice #-}
    basicUnsafeSlice j m (MVector n fp) = MVector m (updPtr (`advancePtr` j) fp)
  
    -- FIXME: this relies on non-portable pointer comparisons
    {-# INLINE basicOverlaps #-}
    basicOverlaps (MVector m fp) (MVector n fq)
      = between p q (q `advancePtr` n) || between q p (p `advancePtr` m)
      where
        between x y z = x >= y && x < z
        p = getPtr fp
        q = getPtr fq
  
    {-# INLINE basicUnsafeNew #-}
    basicUnsafeNew n
      = unsafePrimToPrim
      $ do
          fp <- mallocVector n
          return $ MVector n fp
  
    {-# INLINE basicUnsafeRead #-}
    basicUnsafeRead (MVector _ fp) i
      = unsafePrimToPrim
      $ withForeignPtr fp (`peekElemOff` i)
  
    {-# INLINE basicUnsafeWrite #-}
    basicUnsafeWrite (MVector _ fp) i x
      = unsafePrimToPrim
      $ withForeignPtr fp $ \p -> pokeElemOff p i x
  
    {-# INLINE basicSet #-}
    basicSet = storableSet
  
    {-# INLINE basicUnsafeCopy #-}
    basicUnsafeCopy (MVector n fp) (MVector _ fq)
      = unsafePrimToPrim
      $ withForeignPtr fp $ \p ->
        withForeignPtr fq $ \q ->
        copyArray p q n
    
    {-# INLINE basicUnsafeMove #-}
    basicUnsafeMove (MVector n fp) (MVector _ fq)
      = unsafePrimToPrim
      $ withForeignPtr fp $ \p ->
        withForeignPtr fq $ \q ->
        moveArray p q n
  
  storableSet :: (Storable a, PrimMonad m) => MVector (PrimState m) a -> a -> m ()
  {-# INLINE storableSet #-}
  storableSet v@(MVector n fp) x
    | n == 0 = return ()
    | otherwise = unsafePrimToPrim $
                  case sizeOf x of
                    1 -> storableSetAsPrim n fp x (undefined :: Word8)
                    2 -> storableSetAsPrim n fp x (undefined :: Word16)
                    4 -> storableSetAsPrim n fp x (undefined :: Word32)
                    8 -> storableSetAsPrim n fp x (undefined :: Word64)
                    _ -> withForeignPtr fp $ \p -> do
                         poke p x
  
                         let do_set i
                               | 2*i < n = do
                                   copyArray (p `advancePtr` i) p i
                                   do_set (2*i)
                               | otherwise = copyArray (p `advancePtr` i) p (n-i)
  
                         do_set 1
  
  storableSetAsPrim
    :: (Storable a, Prim b) => Int -> ForeignPtr a -> a -> b -> IO ()
  {-# INLINE [0] storableSetAsPrim #-}
  storableSetAsPrim n fp x y = withForeignPtr fp $ \(Ptr p) -> do
    poke (Ptr p) x
    let q = Addr p
    w <- readOffAddr q 0
    setAddr (q `plusAddr` sizeOf x) (n-1) (w `asTypeOf` y)
  
  {-# INLINE mallocVector #-}
  mallocVector :: Storable a => Int -> IO (ForeignPtr a)
  mallocVector =
  #if __GLASGOW_HASKELL__ >= 605
      doMalloc undefined
          where
            doMalloc :: Storable b => b -> Int -> IO (ForeignPtr b)
            doMalloc dummy size = mallocPlainForeignPtrBytes (size * sizeOf dummy)
  #else
      mallocForeignPtrArray
  #endif
  
  -- Length information
  -- ------------------
  
  -- | Length of the mutable vector.
  length :: Storable a => MVector s a -> Int
  {-# INLINE length #-}
  length = G.length
  
  -- | Check whether the vector is empty
  null :: Storable a => MVector s a -> Bool
  {-# INLINE null #-}
  null = G.null
  
  -- Extracting subvectors
  -- ---------------------
  
  -- | Yield a part of the mutable vector without copying it.
  slice :: Storable a => Int -> Int -> MVector s a -> MVector s a
  {-# INLINE slice #-}
  slice = G.slice
  
  take :: Storable a => Int -> MVector s a -> MVector s a
  {-# INLINE take #-}
  take = G.take
  
  drop :: Storable a => Int -> MVector s a -> MVector s a
  {-# INLINE drop #-}
  drop = G.drop
  
  splitAt :: Storable a => Int -> MVector s a -> (MVector s a, MVector s a)
  {-# INLINE splitAt #-}
  splitAt = G.splitAt
  
  init :: Storable a => MVector s a -> MVector s a
  {-# INLINE init #-}
  init = G.init
  
  tail :: Storable a => MVector s a -> MVector s a
  {-# INLINE tail #-}
  tail = G.tail
  
  -- | Yield a part of the mutable vector without copying it. No bounds checks
  -- are performed.
  unsafeSlice :: Storable a
              => Int  -- ^ starting index
              -> Int  -- ^ length of the slice
              -> MVector s a
              -> MVector s a
  {-# INLINE unsafeSlice #-}
  unsafeSlice = G.unsafeSlice
  
  unsafeTake :: Storable a => Int -> MVector s a -> MVector s a
  {-# INLINE unsafeTake #-}
  unsafeTake = G.unsafeTake
  
  unsafeDrop :: Storable a => Int -> MVector s a -> MVector s a
  {-# INLINE unsafeDrop #-}
  unsafeDrop = G.unsafeDrop
  
  unsafeInit :: Storable a => MVector s a -> MVector s a
  {-# INLINE unsafeInit #-}
  unsafeInit = G.unsafeInit
  
  unsafeTail :: Storable a => MVector s a -> MVector s a
  {-# INLINE unsafeTail #-}
  unsafeTail = G.unsafeTail
  
  -- Overlapping
  -- -----------
  
  -- Check whether two vectors overlap.
  overlaps :: Storable a => MVector s a -> MVector s a -> Bool
  {-# INLINE overlaps #-}
  overlaps = G.overlaps
  
  -- Initialisation
  -- --------------
  
  -- | Create a mutable vector of the given length.
  new :: (PrimMonad m, Storable a) => Int -> m (MVector (PrimState m) a)
  {-# INLINE new #-}
  new = G.new
  
  -- | Create a mutable vector of the given length. The length is not checked.
  unsafeNew :: (PrimMonad m, Storable a) => Int -> m (MVector (PrimState m) a)
  {-# INLINE unsafeNew #-}
  unsafeNew = G.unsafeNew
  
  -- | Create a mutable vector of the given length (0 if the length is negative)
  -- and fill it with an initial value.
  replicate :: (PrimMonad m, Storable a) => Int -> a -> m (MVector (PrimState m) a)
  {-# INLINE replicate #-}
  replicate = G.replicate
  
  -- | Create a mutable vector of the given length (0 if the length is negative)
  -- and fill it with values produced by repeatedly executing the monadic action.
  replicateM :: (PrimMonad m, Storable a) => Int -> m a -> m (MVector (PrimState m) a)
  {-# INLINE replicateM #-}
  replicateM = G.replicateM
  
  -- | Create a copy of a mutable vector.
  clone :: (PrimMonad m, Storable a)
        => MVector (PrimState m) a -> m (MVector (PrimState m) a)
  {-# INLINE clone #-}
  clone = G.clone
  
  -- Growing
  -- -------
  
  -- | Grow a vector by the given number of elements. The number must be
  -- positive.
  grow :: (PrimMonad m, Storable a)  
                => MVector (PrimState m) a -> Int -> m (MVector (PrimState m) a)
  {-# INLINE grow #-}
  grow = G.grow
  
  -- | Grow a vector by the given number of elements. The number must be
  -- positive but this is not checked.
  unsafeGrow :: (PrimMonad m, Storable a)
                 => MVector (PrimState m) a -> Int -> m (MVector (PrimState m) a)
  {-# INLINE unsafeGrow #-}
  unsafeGrow = G.unsafeGrow
  
  -- Restricting memory usage
  -- ------------------------
  
  -- | Reset all elements of the vector to some undefined value, clearing all
  -- references to external objects. This is usually a noop for unboxed vectors. 
  clear :: (PrimMonad m, Storable a) => MVector (PrimState m) a -> m ()
  {-# INLINE clear #-}
  clear = G.clear
  
  -- Accessing individual elements
  -- -----------------------------
  
  -- | Yield the element at the given position.
  read :: (PrimMonad m, Storable a) => MVector (PrimState m) a -> Int -> m a
  {-# INLINE read #-}
  read = G.read
  
  -- | Replace the element at the given position.
  write
      :: (PrimMonad m, Storable a) => MVector (PrimState m) a -> Int -> a -> m ()
  {-# INLINE write #-}
  write = G.write
  
  -- | Swap the elements at the given positions.
  swap
      :: (PrimMonad m, Storable a) => MVector (PrimState m) a -> Int -> Int -> m ()
  {-# INLINE swap #-}
  swap = G.swap
  
  
  -- | Yield the element at the given position. No bounds checks are performed.
  unsafeRead :: (PrimMonad m, Storable a) => MVector (PrimState m) a -> Int -> m a
  {-# INLINE unsafeRead #-}
  unsafeRead = G.unsafeRead
  
  -- | Replace the element at the given position. No bounds checks are performed.
  unsafeWrite
      :: (PrimMonad m, Storable a) =>  MVector (PrimState m) a -> Int -> a -> m ()
  {-# INLINE unsafeWrite #-}
  unsafeWrite = G.unsafeWrite
  
  -- | Swap the elements at the given positions. No bounds checks are performed.
  unsafeSwap
      :: (PrimMonad m, Storable a) => MVector (PrimState m) a -> Int -> Int -> m ()
  {-# INLINE unsafeSwap #-}
  unsafeSwap = G.unsafeSwap
  
  -- Filling and copying
  -- -------------------
  
  -- | Set all elements of the vector to the given value.
  set :: (PrimMonad m, Storable a) => MVector (PrimState m) a -> a -> m ()
  {-# INLINE set #-}
  set = G.set
  
  -- | Copy a vector. The two vectors must have the same length and may not
  -- overlap.
  copy :: (PrimMonad m, Storable a) 
                   => MVector (PrimState m) a -> MVector (PrimState m) a -> m ()
  {-# INLINE copy #-}
  copy = G.copy
  
  -- | Copy a vector. The two vectors must have the same length and may not
  -- overlap. This is not checked.
  unsafeCopy :: (PrimMonad m, Storable a)
             => MVector (PrimState m) a   -- ^ target
             -> MVector (PrimState m) a   -- ^ source
             -> m ()
  {-# INLINE unsafeCopy #-}
  unsafeCopy = G.unsafeCopy
  
  -- | Move the contents of a vector. The two vectors must have the same
  -- length.
  -- 
  -- If the vectors do not overlap, then this is equivalent to 'copy'.
  -- Otherwise, the copying is performed as if the source vector were
  -- copied to a temporary vector and then the temporary vector was copied
  -- to the target vector.
  move :: (PrimMonad m, Storable a)
                   => MVector (PrimState m) a -> MVector (PrimState m) a -> m ()
  {-# INLINE move #-}
  move = G.move
  
  -- | Move the contents of a vector. The two vectors must have the same
  -- length, but this is not checked.
  -- 
  -- If the vectors do not overlap, then this is equivalent to 'unsafeCopy'.
  -- Otherwise, the copying is performed as if the source vector were
  -- copied to a temporary vector and then the temporary vector was copied
  -- to the target vector.
  unsafeMove :: (PrimMonad m, Storable a)
                            => MVector (PrimState m) a   -- ^ target
                            -> MVector (PrimState m) a   -- ^ source
                            -> m ()
  {-# INLINE unsafeMove #-}
  unsafeMove = G.unsafeMove
  
  -- Unsafe conversions
  -- ------------------
  
  -- | /O(1)/ Unsafely cast a mutable vector from one element type to another.
  -- The operation just changes the type of the underlying pointer and does not
  -- modify the elements.
  --
  -- The resulting vector contains as many elements as can fit into the
  -- underlying memory block.
  --
  unsafeCast :: forall a b s.
                (Storable a, Storable b) => MVector s a -> MVector s b
  {-# INLINE unsafeCast #-}
  unsafeCast (MVector n fp)
    = MVector ((n * sizeOf (undefined :: a)) `div` sizeOf (undefined :: b))
              (castForeignPtr fp)
  
  -- Raw pointers
  -- ------------
  
  -- | Create a mutable vector from a 'ForeignPtr' with an offset and a length.
  --
  -- Modifying data through the 'ForeignPtr' afterwards is unsafe if the vector
  -- could have been frozen before the modification.
  --
  --  If your offset is 0 it is more efficient to use 'unsafeFromForeignPtr0'.
  unsafeFromForeignPtr :: Storable a
                       => ForeignPtr a    -- ^ pointer
                       -> Int             -- ^ offset
                       -> Int             -- ^ length
                       -> MVector s a
  {-# INLINE unsafeFromForeignPtr #-}
  unsafeFromForeignPtr fp i n = unsafeFromForeignPtr0 fp' n
      where
        fp' = updPtr (`advancePtr` i) fp
  
  {-# RULES
  "unsafeFromForeignPtr fp 0 n -> unsafeFromForeignPtr0 fp n " forall fp n.
    unsafeFromForeignPtr fp 0 n = unsafeFromForeignPtr0 fp n
    #-}
  
  -- | /O(1)/ Create a mutable vector from a 'ForeignPtr' and a length.
  --
  -- It is assumed the pointer points directly to the data (no offset).
  -- Use `unsafeFromForeignPtr` if you need to specify an offset.
  --
  -- Modifying data through the 'ForeignPtr' afterwards is unsafe if the vector
  -- could have been frozen before the modification.
  unsafeFromForeignPtr0 :: Storable a
                        => ForeignPtr a    -- ^ pointer
                        -> Int             -- ^ length
                        -> MVector s a
  {-# INLINE unsafeFromForeignPtr0 #-}
  unsafeFromForeignPtr0 fp n = MVector n fp
  
  -- | Yield the underlying 'ForeignPtr' together with the offset to the data
  -- and its length. Modifying the data through the 'ForeignPtr' is
  -- unsafe if the vector could have frozen before the modification.
  unsafeToForeignPtr :: Storable a => MVector s a -> (ForeignPtr a, Int, Int)
  {-# INLINE unsafeToForeignPtr #-}
  unsafeToForeignPtr (MVector n fp) = (fp, 0, n)
  
  -- | /O(1)/ Yield the underlying 'ForeignPtr' together with its length.
  --
  -- You can assume the pointer points directly to the data (no offset).
  --
  -- Modifying the data through the 'ForeignPtr' is unsafe if the vector could
  -- have frozen before the modification.
  unsafeToForeignPtr0 :: Storable a => MVector s a -> (ForeignPtr a, Int)
  {-# INLINE unsafeToForeignPtr0 #-}
  unsafeToForeignPtr0 (MVector n fp) = (fp, n)
  
  -- | Pass a pointer to the vector's data to the IO action. Modifying data
  -- through the pointer is unsafe if the vector could have been frozen before
  -- the modification.
  unsafeWith :: Storable a => IOVector a -> (Ptr a -> IO b) -> IO b
  {-# INLINE unsafeWith #-}
  unsafeWith (MVector n fp) = withForeignPtr fp
  
  