-- |
-- Module      : Basement.UArray.Mutable -- License     : BSD-style
-- Maintainer  : Vincent Hanquez <vincent@snarc.org>
-- Stability   : experimental
-- Portability : portable
--
-- A simple array abstraction that allow to use typed
-- array of bytes where the array is pinned in memory
-- to allow easy use with Foreign interfaces, ByteString
-- and always aligned to 64 bytes.
--
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Basement.UArray.Mutable
    ( MUArray(..)
    -- * Property queries
    , sizeInMutableBytesOfContent
    , mutableLength
    , mutableOffset
    , mutableSame
    , onMutableBackend
    -- * Allocation & Copy
    , new
    , newPinned
    , newNative
    , newNative_
    , mutableForeignMem
    , copyAt
    , copyFromPtr
    , copyToPtr
    , sub
    -- , copyAddr
    -- * Reading and Writing cells
    , unsafeWrite
    , unsafeRead
    , write
    , read
    , withMutablePtr
    , withMutablePtrHint
    ) where

import           GHC.Prim
import           GHC.Exts
import           GHC.Types
import           GHC.Ptr
import           Basement.Compat.Base
import           Basement.Compat.Primitive
import           Data.Proxy
import           Basement.Types.OffsetSize
import           Basement.Monad
import           Basement.PrimType
import           Basement.FinalPtr
import           Basement.Exception
import qualified Basement.Block         as BLK
import qualified Basement.Block.Mutable as MBLK
import           Basement.Block         (MutableBlock(..))
import           Basement.UArray.Base hiding (empty)
import           Basement.Numerical.Subtractive
import           Foreign.Marshal.Utils (copyBytes)

sizeInMutableBytesOfContent :: forall ty s . PrimType ty => MUArray ty s -> CountOf Word8
sizeInMutableBytesOfContent :: forall ty s. PrimType ty => MUArray ty s -> CountOf Word8
sizeInMutableBytesOfContent MUArray ty s
_ = forall ty. PrimType ty => Proxy ty -> CountOf Word8
primSizeInBytes (forall {k} (t :: k). Proxy t
Proxy :: Proxy ty)
{-# INLINE sizeInMutableBytesOfContent #-}

-- | read a cell in a mutable array.
--
-- If the index is out of bounds, an error is raised.
read :: (PrimMonad prim, PrimType ty) => MUArray ty (PrimState prim) -> Offset ty -> prim ty
read :: forall (prim :: * -> *) ty.
(PrimMonad prim, PrimType ty) =>
MUArray ty (PrimState prim) -> Offset ty -> prim ty
read MUArray ty (PrimState prim)
array Offset ty
n
    | forall ty. Offset ty -> CountOf ty -> Bool
isOutOfBound Offset ty
n CountOf ty
len = forall (prim :: * -> *) ty a.
PrimMonad prim =>
OutOfBoundOperation -> Offset ty -> CountOf ty -> prim a
primOutOfBound OutOfBoundOperation
OOB_Read Offset ty
n CountOf ty
len
    | Bool
otherwise          = forall (prim :: * -> *) ty.
(PrimMonad prim, PrimType ty) =>
MUArray ty (PrimState prim) -> Offset ty -> prim ty
unsafeRead MUArray ty (PrimState prim)
array Offset ty
n
  where len :: CountOf ty
len = forall ty st. PrimType ty => MUArray ty st -> CountOf ty
mutableLength MUArray ty (PrimState prim)
array
{-# INLINE read #-}

-- | Write to a cell in a mutable array.
--
-- If the index is out of bounds, an error is raised.
write :: (PrimMonad prim, PrimType ty) => MUArray ty (PrimState prim) -> Offset ty -> ty -> prim ()
write :: forall (prim :: * -> *) ty.
(PrimMonad prim, PrimType ty) =>
MUArray ty (PrimState prim) -> Offset ty -> ty -> prim ()
write MUArray ty (PrimState prim)
array Offset ty
n ty
val
    | forall ty. Offset ty -> CountOf ty -> Bool
isOutOfBound Offset ty
n CountOf ty
len = forall (prim :: * -> *) ty a.
PrimMonad prim =>
OutOfBoundOperation -> Offset ty -> CountOf ty -> prim a
primOutOfBound OutOfBoundOperation
OOB_Write Offset ty
n CountOf ty
len
    | Bool
otherwise          = forall (prim :: * -> *) ty.
(PrimMonad prim, PrimType ty) =>
MUArray ty (PrimState prim) -> Offset ty -> ty -> prim ()
unsafeWrite MUArray ty (PrimState prim)
array Offset ty
n ty
val
  where
    len :: CountOf ty
len = forall ty st. PrimType ty => MUArray ty st -> CountOf ty
mutableLength MUArray ty (PrimState prim)
array
{-# INLINE write #-}

empty :: (PrimType ty, PrimMonad prim) => prim (MUArray ty (PrimState prim))
empty :: forall ty (prim :: * -> *).
(PrimType ty, PrimMonad prim) =>
prim (MUArray ty (PrimState prim))
empty = forall ty st.
Offset ty -> CountOf ty -> MUArrayBackend ty st -> MUArray ty st
MUArray Offset ty
0 CountOf ty
0 forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall ty st. MutableBlock ty st -> MUArrayBackend ty st
MUArrayMBA forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (prim :: * -> *) ty.
PrimMonad prim =>
prim (MutableBlock ty (PrimState prim))
MBLK.mutableEmpty

mutableSame :: MUArray ty st -> MUArray ty st -> Bool
mutableSame :: forall ty st. MUArray ty st -> MUArray ty st -> Bool
mutableSame (MUArray Offset ty
sa CountOf ty
ea (MUArrayMBA (MutableBlock MutableByteArray# st
ma))) (MUArray Offset ty
sb CountOf ty
eb (MUArrayMBA (MutableBlock MutableByteArray# st
mb))) = (Offset ty
sa forall a. Eq a => a -> a -> Bool
== Offset ty
sb) Bool -> Bool -> Bool
&& (CountOf ty
ea forall a. Eq a => a -> a -> Bool
== CountOf ty
eb) Bool -> Bool -> Bool
&& Int# -> Bool
bool# (forall d. MutableByteArray# d -> MutableByteArray# d -> Int#
sameMutableByteArray# MutableByteArray# st
ma MutableByteArray# st
mb)
mutableSame (MUArray Offset ty
s1 CountOf ty
e1 (MUArrayAddr FinalPtr ty
f1)) (MUArray Offset ty
s2 CountOf ty
e2 (MUArrayAddr FinalPtr ty
f2)) = (Offset ty
s1 forall a. Eq a => a -> a -> Bool
== Offset ty
s2) Bool -> Bool -> Bool
&& (CountOf ty
e1 forall a. Eq a => a -> a -> Bool
== CountOf ty
e2) Bool -> Bool -> Bool
&& forall a b. FinalPtr a -> FinalPtr b -> Bool
finalPtrSameMemory FinalPtr ty
f1 FinalPtr ty
f2
mutableSame MUArray ty st
_ MUArray ty st
_ = Bool
False

mutableForeignMem :: (PrimMonad prim, PrimType ty)
                  => FinalPtr ty -- ^ the start pointer with a finalizer
                  -> Int         -- ^ the number of elements (in elements, not bytes)
                  -> prim (MUArray ty (PrimState prim))
mutableForeignMem :: forall (prim :: * -> *) ty.
(PrimMonad prim, PrimType ty) =>
FinalPtr ty -> Int -> prim (MUArray ty (PrimState prim))
mutableForeignMem FinalPtr ty
fptr Int
nb = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall ty st.
Offset ty -> CountOf ty -> MUArrayBackend ty st -> MUArray ty st
MUArray (forall ty. Int -> Offset ty
Offset Int
0) (forall ty. Int -> CountOf ty
CountOf Int
nb) (forall ty st. FinalPtr ty -> MUArrayBackend ty st
MUArrayAddr FinalPtr ty
fptr)

sub :: (PrimMonad prim, PrimType ty)
    => MUArray ty (PrimState prim)
    -> Int -- The number of elements to drop ahead
    -> Int -- Then the number of element to retain
    -> prim (MUArray ty (PrimState prim))
sub :: forall (prim :: * -> *) ty.
(PrimMonad prim, PrimType ty) =>
MUArray ty (PrimState prim)
-> Int -> Int -> prim (MUArray ty (PrimState prim))
sub (MUArray Offset ty
start CountOf ty
sz MUArrayBackend ty (PrimState prim)
back) Int
dropElems' Int
takeElems
    | Int
takeElems forall a. Ord a => a -> a -> Bool
<= Int
0 = forall ty (prim :: * -> *).
(PrimType ty, PrimMonad prim) =>
prim (MUArray ty (PrimState prim))
empty
    | Just CountOf ty
keepElems <- CountOf ty
sz forall a. Subtractive a => a -> a -> Difference a
- CountOf ty
dropElems, CountOf ty
keepElems forall a. Ord a => a -> a -> Bool
> CountOf ty
0
                     = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall ty st.
Offset ty -> CountOf ty -> MUArrayBackend ty st -> MUArray ty st
MUArray (Offset ty
start forall ty. Offset ty -> CountOf ty -> Offset ty
`offsetPlusE` CountOf ty
dropElems) (forall a. Ord a => a -> a -> a
min (forall ty. Int -> CountOf ty
CountOf Int
takeElems) CountOf ty
keepElems) MUArrayBackend ty (PrimState prim)
back
    | Bool
otherwise      = forall ty (prim :: * -> *).
(PrimType ty, PrimMonad prim) =>
prim (MUArray ty (PrimState prim))
empty
  where
    dropElems :: CountOf ty
dropElems = forall a. Ord a => a -> a -> a
max CountOf ty
0 (forall ty. Int -> CountOf ty
CountOf Int
dropElems')


-- | return the numbers of elements in a mutable array
mutableLength :: PrimType ty => MUArray ty st -> CountOf ty
mutableLength :: forall ty st. PrimType ty => MUArray ty st -> CountOf ty
mutableLength (MUArray Offset ty
_ CountOf ty
end MUArrayBackend ty st
_)   = CountOf ty
end

withMutablePtrHint :: forall ty prim a . (PrimMonad prim, PrimType ty)
                   => Bool
                   -> Bool
                   -> MUArray ty (PrimState prim)
                   -> (Ptr ty -> prim a)
                   -> prim a
withMutablePtrHint :: forall ty (prim :: * -> *) a.
(PrimMonad prim, PrimType ty) =>
Bool
-> Bool
-> MUArray ty (PrimState prim)
-> (Ptr ty -> prim a)
-> prim a
withMutablePtrHint Bool
skipCopy Bool
skipCopyBack (MUArray Offset ty
start CountOf ty
_ MUArrayBackend ty (PrimState prim)
back) Ptr ty -> prim a
f =
    case MUArrayBackend ty (PrimState prim)
back of
        MUArrayAddr FinalPtr ty
fptr -> forall (prim :: * -> *) p a.
PrimMonad prim =>
FinalPtr p -> (Ptr p -> prim a) -> prim a
withFinalPtr FinalPtr ty
fptr (\Ptr ty
ptr -> Ptr ty -> prim a
f (Ptr ty
ptr forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
os))
        MUArrayMBA MutableBlock ty (PrimState prim)
mb    -> forall ty (prim :: * -> *) a.
PrimMonad prim =>
Bool
-> Bool
-> MutableBlock ty (PrimState prim)
-> (Ptr ty -> prim a)
-> prim a
MBLK.withMutablePtrHint Bool
skipCopy Bool
skipCopyBack MutableBlock ty (PrimState prim)
mb forall a b. (a -> b) -> a -> b
$ \Ptr ty
ptr -> Ptr ty -> prim a
f (Ptr ty
ptr forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
os)
  where
    sz :: CountOf Word8
sz           = forall ty. PrimType ty => Proxy ty -> CountOf Word8
primSizeInBytes (forall {k} (t :: k). Proxy t
Proxy :: Proxy ty)
    !(Offset Int
os) = forall ty. CountOf Word8 -> Offset ty -> Offset Word8
offsetOfE CountOf Word8
sz Offset ty
start

-- | Create a pointer on the beginning of the mutable array
-- and call a function 'f'.
--
-- The mutable buffer can be mutated by the 'f' function
-- and the change will be reflected in the mutable array
--
-- If the mutable array is unpinned, a trampoline buffer
-- is created and the data is only copied when 'f' return.
withMutablePtr :: (PrimMonad prim, PrimType ty)
               => MUArray ty (PrimState prim)
               -> (Ptr ty -> prim a)
               -> prim a
withMutablePtr :: forall (prim :: * -> *) ty a.
(PrimMonad prim, PrimType ty) =>
MUArray ty (PrimState prim) -> (Ptr ty -> prim a) -> prim a
withMutablePtr = forall ty (prim :: * -> *) a.
(PrimMonad prim, PrimType ty) =>
Bool
-> Bool
-> MUArray ty (PrimState prim)
-> (Ptr ty -> prim a)
-> prim a
withMutablePtrHint Bool
False Bool
False

-- | Copy from a pointer, @count@ elements, into the mutable array
copyFromPtr :: forall prim ty . (PrimMonad prim, PrimType ty)
            => Ptr ty -> CountOf ty -> MUArray ty (PrimState prim) -> prim ()
copyFromPtr :: forall (prim :: * -> *) ty.
(PrimMonad prim, PrimType ty) =>
Ptr ty -> CountOf ty -> MUArray ty (PrimState prim) -> prim ()
copyFromPtr src :: Ptr ty
src@(Ptr Addr#
src#) CountOf ty
count MUArray ty (PrimState prim)
marr
    | CountOf ty
count forall a. Ord a => a -> a -> Bool
> CountOf ty
arrSz = forall (prim :: * -> *) ty a.
PrimMonad prim =>
OutOfBoundOperation -> Offset ty -> CountOf ty -> prim a
primOutOfBound OutOfBoundOperation
OOB_MemCopy (forall a. CountOf a -> Offset a
sizeAsOffset CountOf ty
count) CountOf ty
arrSz
    | Bool
otherwise     = forall (prim :: * -> *) ty a.
PrimMonad prim =>
(MutableBlock ty (PrimState prim) -> prim a)
-> (FinalPtr ty -> prim a) -> MUArray ty (PrimState prim) -> prim a
onMutableBackend MutableBlock ty (PrimState prim) -> prim ()
copyNative FinalPtr ty -> prim ()
copyPtr MUArray ty (PrimState prim)
marr
  where
    arrSz :: CountOf ty
arrSz = forall ty st. PrimType ty => MUArray ty st -> CountOf ty
mutableLength MUArray ty (PrimState prim)
marr
    ofs :: Offset ty
ofs = forall ty st. MUArray ty st -> Offset ty
mutableOffset MUArray ty (PrimState prim)
marr

    sz :: CountOf Word8
sz = forall ty. PrimType ty => Proxy ty -> CountOf Word8
primSizeInBytes (forall {k} (t :: k). Proxy t
Proxy :: Proxy ty)
    !count' :: CountOf Word8
count'@(CountOf bytes :: Int
bytes@(I# Int#
bytes#)) = forall ty. CountOf Word8 -> CountOf ty -> CountOf Word8
sizeOfE CountOf Word8
sz CountOf ty
count
    !off' :: Offset Word8
off'@(Offset od :: Int
od@(I# Int#
od#)) = forall ty. CountOf Word8 -> Offset ty -> Offset Word8
offsetOfE CountOf Word8
sz Offset ty
ofs

    copyNative :: MutableBlock ty (PrimState prim) -> prim ()
copyNative MutableBlock ty (PrimState prim)
mba = forall (prim :: * -> *) ty.
PrimMonad prim =>
MutableBlock ty (PrimState prim)
-> Offset Word8 -> Ptr ty -> CountOf Word8 -> prim ()
MBLK.unsafeCopyBytesPtr MutableBlock ty (PrimState prim)
mba Offset Word8
off' Ptr ty
src CountOf Word8
count'
    copyPtr :: FinalPtr ty -> prim ()
copyPtr FinalPtr ty
fptr = forall (prim :: * -> *) p a.
PrimMonad prim =>
FinalPtr p -> (Ptr p -> prim a) -> prim a
withFinalPtr FinalPtr ty
fptr forall a b. (a -> b) -> a -> b
$ \Ptr ty
dst ->
        forall (prim :: * -> *) a. PrimMonad prim => IO a -> prim a
unsafePrimFromIO forall a b. (a -> b) -> a -> b
$ forall a. Ptr a -> Ptr a -> Int -> IO ()
copyBytes (Ptr ty
dst forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
od) Ptr ty
src Int
bytes

-- | Copy all the block content to the memory starting at the destination address
copyToPtr :: forall ty prim . (PrimType ty, PrimMonad prim)
          => MUArray ty (PrimState prim) -- ^ the source mutable array to copy
          -> Ptr ty                      -- ^ The destination address where the copy is going to start
          -> prim ()
copyToPtr :: forall ty (prim :: * -> *).
(PrimType ty, PrimMonad prim) =>
MUArray ty (PrimState prim) -> Ptr ty -> prim ()
copyToPtr MUArray ty (PrimState prim)
marr dst :: Ptr ty
dst@(Ptr Addr#
dst#) = forall (prim :: * -> *) ty a.
PrimMonad prim =>
(MutableBlock ty (PrimState prim) -> prim a)
-> (FinalPtr ty -> prim a) -> MUArray ty (PrimState prim) -> prim a
onMutableBackend MutableBlock ty (PrimState prim) -> prim ()
copyNative FinalPtr ty -> prim ()
copyPtr MUArray ty (PrimState prim)
marr
  where
    copyNative :: MutableBlock ty (PrimState prim) -> prim ()
copyNative (MutableBlock MutableByteArray# (PrimState prim)
mba) = forall (m :: * -> *) a.
PrimMonad m =>
(State# (PrimState m) -> (# State# (PrimState m), a #)) -> m a
primitive forall a b. (a -> b) -> a -> b
$ \State# (PrimState prim)
s1 ->
        case forall d.
MutableByteArray# d -> State# d -> (# State# d, ByteArray# #)
unsafeFreezeByteArray# MutableByteArray# (PrimState prim)
mba State# (PrimState prim)
s1 of
            (# State# (PrimState prim)
s2, ByteArray#
ba #) -> (# forall d.
ByteArray# -> Int# -> Addr# -> Int# -> State# d -> State# d
copyByteArrayToAddr# ByteArray#
ba Int#
os# Addr#
dst# Int#
szBytes# State# (PrimState prim)
s2, () #)
    copyPtr :: FinalPtr ty -> prim ()
copyPtr FinalPtr ty
fptr = forall (prim :: * -> *) a. PrimMonad prim => IO a -> prim a
unsafePrimFromIO forall a b. (a -> b) -> a -> b
$ forall (prim :: * -> *) p a.
PrimMonad prim =>
FinalPtr p -> (Ptr p -> prim a) -> prim a
withFinalPtr FinalPtr ty
fptr forall a b. (a -> b) -> a -> b
$ \Ptr ty
ptr ->
        forall a. Ptr a -> Ptr a -> Int -> IO ()
copyBytes Ptr ty
dst (Ptr ty
ptr forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
os) Int
szBytes

    !(Offset os :: Int
os@(I# Int#
os#)) = forall a. PrimType a => Offset a -> Offset Word8
offsetInBytes forall a b. (a -> b) -> a -> b
$ forall ty st. MUArray ty st -> Offset ty
mutableOffset MUArray ty (PrimState prim)
marr
    !(CountOf szBytes :: Int
szBytes@(I# Int#
szBytes#)) = forall a. PrimType a => CountOf a -> CountOf Word8
sizeInBytes forall a b. (a -> b) -> a -> b
$ forall ty st. PrimType ty => MUArray ty st -> CountOf ty
mutableLength MUArray ty (PrimState prim)
marr

mutableOffset :: MUArray ty st -> Offset ty
mutableOffset :: forall ty st. MUArray ty st -> Offset ty
mutableOffset (MUArray Offset ty
ofs CountOf ty
_ MUArrayBackend ty st
_) = Offset ty
ofs