{-# LANGUAGE MagicHash #-}
{-# LANGUAGE UnboxedTuples #-}
module Foundation.Array.Unboxed.Mutable
( MUArray(..)
, sizeInMutableBytesOfContent
, mutableLength
, mutableSame
, new
, newPinned
, newNative
, mutableForeignMem
, copyAt
, unsafeWrite
, unsafeRead
, write
, read
) where
import GHC.Prim
import GHC.Types
import GHC.Ptr
import Foundation.Internal.Base
import qualified Foundation.Internal.Environment as Environment
import Foundation.Internal.Types
import Foundation.Internal.Primitive
import Foundation.Internal.Proxy
import Foundation.Primitive.Monad
import Foundation.Primitive.Types
import Foundation.Primitive.FinalPtr
import Foundation.Array.Common
import Foundation.Numerical
data MUArray ty st =
MUVecMA {-# UNPACK #-} !(Offset ty)
{-# UNPACK #-} !(Size ty)
{-# UNPACK #-} !PinnedStatus
(MutableByteArray# st)
| MUVecAddr {-# UNPACK #-} !(Offset ty)
{-# UNPACK #-} !(Size ty)
!(FinalPtr ty)
mutableArrayProxyTy :: MUArray ty st -> Proxy ty
mutableArrayProxyTy _ = Proxy
sizeInMutableBytesOfContent :: PrimType ty => MUArray ty s -> Size8
sizeInMutableBytesOfContent = primSizeInBytes . mutableArrayProxyTy
{-# INLINE sizeInMutableBytesOfContent #-}
read :: (PrimMonad prim, PrimType ty) => MUArray ty (PrimState prim) -> Int -> prim ty
read array n
| n < 0 || n >= len = primThrow (OutOfBound OOB_Read n len)
| otherwise = unsafeRead array n
where len = mutableLength array
{-# INLINE read #-}
unsafeRead :: (PrimMonad prim, PrimType ty) => MUArray ty (PrimState prim) -> Int -> prim ty
unsafeRead (MUVecMA start _ _ mba) i = primMbaRead mba (start+.i)
unsafeRead (MUVecAddr start _ fptr) i = withFinalPtr fptr $ \(Ptr addr) -> primAddrRead addr (start+.i)
{-# INLINE unsafeRead #-}
write :: (PrimMonad prim, PrimType ty) => MUArray ty (PrimState prim) -> Int -> ty -> prim ()
write array n val
| n < 0 || n >= len = primThrow (OutOfBound OOB_Write n len)
| otherwise = unsafeWrite array n val
where
len = mutableLength array
{-# INLINE write #-}
unsafeWrite :: (PrimMonad prim, PrimType ty) => MUArray ty (PrimState prim) -> Int -> ty -> prim ()
unsafeWrite (MUVecMA start _ _ mba) i v = primMbaWrite mba (start+.i) v
unsafeWrite (MUVecAddr start _ fptr) i v = withFinalPtr fptr $ \(Ptr addr) -> primAddrWrite addr (start+.i) v
{-# INLINE unsafeWrite #-}
newPinned :: (PrimMonad prim, PrimType ty) => Size ty -> prim (MUArray ty (PrimState prim))
newPinned n = newFake n Proxy
where newFake :: (PrimMonad prim, PrimType ty) => Size ty -> Proxy ty -> prim (MUArray ty (PrimState prim))
newFake sz ty = primitive $ \s1 ->
case newAlignedPinnedByteArray# bytes 8# s1 of
(# s2, mba #) -> (# s2, MUVecMA (Offset 0) sz pinned mba #)
where
!(Size (I# bytes)) = sizeOfE (primSizeInBytes ty) sz
{-# INLINE newFake #-}
newUnpinned :: (PrimMonad prim, PrimType ty) => Size ty -> prim (MUArray ty (PrimState prim))
newUnpinned n = newFake n Proxy
where newFake :: (PrimMonad prim, PrimType ty) => Size ty -> Proxy ty -> prim (MUArray ty (PrimState prim))
newFake sz ty = primitive $ \s1 ->
case newByteArray# bytes s1 of
(# s2, mba #) -> (# s2, MUVecMA (Offset 0) sz unpinned mba #)
where
!(Size (I# bytes)) = sizeOfE (primSizeInBytes ty) sz
{-# INLINE newFake #-}
new :: (PrimMonad prim, PrimType ty) => Size ty -> prim (MUArray ty (PrimState prim))
new sz
| sizeRecast sz <= maxSizeUnpinned = newUnpinned sz
| otherwise = newPinned sz
where
maxSizeUnpinned = Environment.unsafeUArrayUnpinnedMaxSize
{-# INLINE new #-}
mutableSame :: MUArray ty st -> MUArray ty st -> Bool
mutableSame (MUVecMA sa ea _ ma) (MUVecMA sb eb _ mb) = and [ sa == sb, ea == eb, bool# (sameMutableByteArray# ma mb)]
mutableSame (MUVecAddr s1 e1 f1) (MUVecAddr s2 e2 f2) = and [ s1 == s2, e1 == e2, finalPtrSameMemory f1 f2 ]
mutableSame (MUVecMA {}) (MUVecAddr {}) = False
mutableSame (MUVecAddr {}) (MUVecMA {}) = False
newNative :: (PrimMonad prim, PrimType ty) => Size ty -> (MutableByteArray# (PrimState prim) -> prim ()) -> prim (MUArray ty (PrimState prim))
newNative n f = do
muvec <- new n
case muvec of
(MUVecMA _ _ _ mba) -> f mba >> return muvec
(MUVecAddr {}) -> error "internal error: unboxed new only supposed to allocate natively"
mutableForeignMem :: (PrimMonad prim, PrimType ty)
=> FinalPtr ty
-> Int
-> prim (MUArray ty (PrimState prim))
mutableForeignMem fptr nb = return $ MUVecAddr (Offset 0) (Size nb) fptr
copyAt :: (PrimMonad prim, PrimType ty)
=> MUArray ty (PrimState prim)
-> Offset ty
-> MUArray ty (PrimState prim)
-> Offset ty
-> Size ty
-> prim ()
copyAt (MUVecMA dstStart _ _ dstMba) ed uvec@(MUVecMA srcStart _ _ srcBa) es n =
primitive $ \st -> (# copyMutableByteArray# srcBa os dstMba od nBytes st, () #)
where
!sz = primSizeInBytes (mutableArrayProxyTy uvec)
!(Offset (I# os)) = offsetOfE sz (srcStart + es)
!(Offset (I# od)) = offsetOfE sz (dstStart + ed)
!(Size (I# nBytes)) = sizeOfE sz n
copyAt (MUVecMA dstStart _ _ dstMba) ed muvec@(MUVecAddr srcStart _ srcFptr) es n =
withFinalPtr srcFptr $ \srcPtr ->
let !(Ptr srcAddr) = srcPtr `plusPtr` os
in primitive $ \s -> (# compatCopyAddrToByteArray# srcAddr dstMba od nBytes s, () #)
where
!sz = primSizeInBytes (mutableArrayProxyTy muvec)
!(Offset os) = offsetOfE sz (srcStart + es)
!(Offset (I# od)) = offsetOfE sz (dstStart + ed)
!(Size (I# nBytes)) = sizeOfE sz n
copyAt dst od src os n = loop od os
where
!(Offset endIndex) = os `offsetPlusE` n
loop !(Offset d) !(Offset i)
| i == endIndex = return ()
| otherwise = unsafeRead src i >>= unsafeWrite dst d >> loop (Offset $ d+1) (Offset $ i+1)
mutableLength :: PrimType ty => MUArray ty st -> Int
mutableLength (MUVecMA _ (Size end) _ _) = end
mutableLength (MUVecAddr _ (Size end) _) = end