module Dahdit.Mem
  ( IxPtr (..)
  , ReadMem (..)
  , readSBSMem
  , viewSBSMem
  , viewBSMem
  , viewVecMem
  , WriteMem (..)
  , writeSBSMem
  , allocArrayMem
  , allocPtrMem
  , freezeSBSMem
  , freezeBSMem
  , freezeVecMem
  )
where

import Control.Monad.ST (ST, runST)
import Control.Monad.ST.Unsafe (unsafeIOToST)
import Dahdit.LiftedPrim (LiftedPrim (..), setByteArrayLifted)
import Dahdit.Proxy (proxyFor)
import Dahdit.Sizes (ByteCount (..), staticByteSize)
import Data.ByteString (ByteString)
import qualified Data.ByteString.Internal as BSI
import Data.ByteString.Short.Internal (ShortByteString (..))
import qualified Data.ByteString.Unsafe as BSU
import Data.Coerce (coerce)
import Data.Foldable (for_)
import Data.Primitive.ByteArray (ByteArray (..), MutableByteArray, cloneByteArray, copyByteArray, copyByteArrayToPtr, freezeByteArray, newByteArray, unsafeFreezeByteArray)
import Data.Primitive.Ptr (copyPtrToMutableByteArray)
import Data.Vector.Storable (Vector)
import qualified Data.Vector.Storable as VS
import Data.Word (Word8)
import Foreign.ForeignPtr (newForeignPtr)
import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
import Foreign.Marshal.Alloc (callocBytes, finalizerFree, free)
import Foreign.Ptr (Ptr, plusPtr)

-- | A wrapper over 'Ptr' with an additional phantom type index to align with 'ST' state.
newtype IxPtr s = IxPtr {forall s. IxPtr s -> Ptr Word8
unIxPtr :: Ptr Word8}
  deriving stock (Int -> IxPtr s -> ShowS
forall s. Int -> IxPtr s -> ShowS
forall s. [IxPtr s] -> ShowS
forall s. IxPtr s -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [IxPtr s] -> ShowS
$cshowList :: forall s. [IxPtr s] -> ShowS
show :: IxPtr s -> String
$cshow :: forall s. IxPtr s -> String
showsPrec :: Int -> IxPtr s -> ShowS
$cshowsPrec :: forall s. Int -> IxPtr s -> ShowS
Show)
  deriving newtype (IxPtr s -> IxPtr s -> Bool
forall s. IxPtr s -> IxPtr s -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: IxPtr s -> IxPtr s -> Bool
$c/= :: forall s. IxPtr s -> IxPtr s -> Bool
== :: IxPtr s -> IxPtr s -> Bool
$c== :: forall s. IxPtr s -> IxPtr s -> Bool
Eq, IxPtr s -> IxPtr s -> Bool
IxPtr s -> IxPtr s -> Ordering
IxPtr s -> IxPtr s -> IxPtr s
forall s. Eq (IxPtr s)
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall s. IxPtr s -> IxPtr s -> Bool
forall s. IxPtr s -> IxPtr s -> Ordering
forall s. IxPtr s -> IxPtr s -> IxPtr s
min :: IxPtr s -> IxPtr s -> IxPtr s
$cmin :: forall s. IxPtr s -> IxPtr s -> IxPtr s
max :: IxPtr s -> IxPtr s -> IxPtr s
$cmax :: forall s. IxPtr s -> IxPtr s -> IxPtr s
>= :: IxPtr s -> IxPtr s -> Bool
$c>= :: forall s. IxPtr s -> IxPtr s -> Bool
> :: IxPtr s -> IxPtr s -> Bool
$c> :: forall s. IxPtr s -> IxPtr s -> Bool
<= :: IxPtr s -> IxPtr s -> Bool
$c<= :: forall s. IxPtr s -> IxPtr s -> Bool
< :: IxPtr s -> IxPtr s -> Bool
$c< :: forall s. IxPtr s -> IxPtr s -> Bool
compare :: IxPtr s -> IxPtr s -> Ordering
$ccompare :: forall s. IxPtr s -> IxPtr s -> Ordering
Ord)

class ReadMem r where
  indexMemInBytes :: LiftedPrim a => r -> ByteCount -> a
  cloneArrayMemInBytes :: r -> ByteCount -> ByteCount -> ByteArray

instance ReadMem ByteArray where
  indexMemInBytes :: forall a. LiftedPrim a => ByteArray -> ByteCount -> a
indexMemInBytes = forall a. LiftedPrim a => ByteArray -> ByteCount -> a
indexArrayLiftedInBytes
  cloneArrayMemInBytes :: ByteArray -> ByteCount -> ByteCount -> ByteArray
cloneArrayMemInBytes ByteArray
arr ByteCount
off ByteCount
len = ByteArray -> Int -> Int -> ByteArray
cloneByteArray ByteArray
arr (coerce :: forall a b. Coercible a b => a -> b
coerce ByteCount
off) (coerce :: forall a b. Coercible a b => a -> b
coerce ByteCount
len)

clonePtr :: Ptr Word8 -> ByteCount -> ByteCount -> ByteArray
clonePtr :: Ptr Word8 -> ByteCount -> ByteCount -> ByteArray
clonePtr Ptr Word8
ptr ByteCount
off ByteCount
len = forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
  let wptr :: Ptr Word8
wptr = coerce :: forall a b. Coercible a b => a -> b
coerce (forall a b. Ptr a -> Int -> Ptr b
plusPtr Ptr Word8
ptr (coerce :: forall a b. Coercible a b => a -> b
coerce ByteCount
off)) :: Ptr Word8
  MutableByteArray s
marr <- forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableByteArray (PrimState m))
newByteArray (coerce :: forall a b. Coercible a b => a -> b
coerce ByteCount
len)
  forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MutableByteArray (PrimState m) -> Int -> Ptr a -> Int -> m ()
copyPtrToMutableByteArray MutableByteArray s
marr Int
0 Ptr Word8
wptr (coerce :: forall a b. Coercible a b => a -> b
coerce ByteCount
len)
  forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> m ByteArray
unsafeFreezeByteArray MutableByteArray s
marr

instance ReadMem (Ptr Word8) where
  indexMemInBytes :: forall a. LiftedPrim a => Ptr Word8 -> ByteCount -> a
indexMemInBytes = forall a. LiftedPrim a => Ptr Word8 -> ByteCount -> a
indexPtrLiftedInBytes
  cloneArrayMemInBytes :: Ptr Word8 -> ByteCount -> ByteCount -> ByteArray
cloneArrayMemInBytes = Ptr Word8 -> ByteCount -> ByteCount -> ByteArray
clonePtr

readSBSMem :: ReadMem r => r -> ByteCount -> ByteCount -> ShortByteString
readSBSMem :: forall r.
ReadMem r =>
r -> ByteCount -> ByteCount -> ShortByteString
readSBSMem r
mem ByteCount
off ByteCount
len = let !(ByteArray ByteArray#
frozArr) = forall r. ReadMem r => r -> ByteCount -> ByteCount -> ByteArray
cloneArrayMemInBytes r
mem ByteCount
off ByteCount
len in ByteArray# -> ShortByteString
SBS ByteArray#
frozArr

viewSBSMem :: ShortByteString -> ByteArray
viewSBSMem :: ShortByteString -> ByteArray
viewSBSMem (SBS ByteArray#
harr) = ByteArray# -> ByteArray
ByteArray ByteArray#
harr

viewBSMem :: ByteString -> Ptr Word8
viewBSMem :: ByteString -> Ptr Word8
viewBSMem ByteString
bs =
  let (ForeignPtr Word8
fp, Int
_) = ByteString -> (ForeignPtr Word8, Int)
BSI.toForeignPtr0 ByteString
bs
  in  forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr Word8
fp

viewVecMem :: Vector Word8 -> Ptr Word8
viewVecMem :: Vector Word8 -> Ptr Word8
viewVecMem Vector Word8
vec =
  let (ForeignPtr Word8
fp, Int
_) = forall a. Storable a => Vector a -> (ForeignPtr a, Int)
VS.unsafeToForeignPtr0 Vector Word8
vec
  in  forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr Word8
fp

class WriteMem q where
  writeMemInBytes :: LiftedPrim a => a -> q s -> ByteCount -> ST s ()
  copyArrayMemInBytes :: ByteArray -> ByteCount -> ByteCount -> q s -> ByteCount -> ST s ()
  setMemInBytes :: LiftedPrim a => ByteCount -> a -> q s -> ByteCount -> ST s ()
  releaseMem :: q s -> Maybe (IO ())

instance WriteMem MutableByteArray where
  writeMemInBytes :: forall a s.
LiftedPrim a =>
a -> MutableByteArray s -> ByteCount -> ST s ()
writeMemInBytes a
val MutableByteArray s
mem ByteCount
off = forall a (m :: * -> *).
(LiftedPrim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> ByteCount -> a -> m ()
writeArrayLiftedInBytes MutableByteArray s
mem ByteCount
off a
val
  copyArrayMemInBytes :: forall s.
ByteArray
-> ByteCount
-> ByteCount
-> MutableByteArray s
-> ByteCount
-> ST s ()
copyArrayMemInBytes ByteArray
arr ByteCount
arrOff ByteCount
arrLen MutableByteArray s
mem ByteCount
off = forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m)
-> Int -> ByteArray -> Int -> Int -> m ()
copyByteArray MutableByteArray s
mem (coerce :: forall a b. Coercible a b => a -> b
coerce ByteCount
off) ByteArray
arr (coerce :: forall a b. Coercible a b => a -> b
coerce ByteCount
arrOff) (coerce :: forall a b. Coercible a b => a -> b
coerce ByteCount
arrLen)
  setMemInBytes :: forall a s.
LiftedPrim a =>
ByteCount -> a -> MutableByteArray s -> ByteCount -> ST s ()
setMemInBytes ByteCount
len a
val MutableByteArray s
mem ByteCount
off = forall (m :: * -> *) a.
(PrimMonad m, LiftedPrim a) =>
MutableByteArray (PrimState m)
-> ByteCount -> ByteCount -> a -> m ()
setByteArrayLifted MutableByteArray s
mem ByteCount
off ByteCount
len a
val
  releaseMem :: forall s. MutableByteArray s -> Maybe (IO ())
releaseMem = forall a b. a -> b -> a
const forall a. Maybe a
Nothing

copyPtr :: ByteArray -> ByteCount -> ByteCount -> Ptr Word8 -> ByteCount -> ST s ()
copyPtr :: forall s.
ByteArray
-> ByteCount -> ByteCount -> Ptr Word8 -> ByteCount -> ST s ()
copyPtr ByteArray
arr ByteCount
arrOff ByteCount
arrLen Ptr Word8
ptr ByteCount
off =
  let wptr :: Ptr Word8
wptr = coerce :: forall a b. Coercible a b => a -> b
coerce (forall a b. Ptr a -> Int -> Ptr b
plusPtr Ptr Word8
ptr (coerce :: forall a b. Coercible a b => a -> b
coerce ByteCount
off)) :: Ptr Word8
  in  forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Ptr a -> ByteArray -> Int -> Int -> m ()
copyByteArrayToPtr Ptr Word8
wptr ByteArray
arr (coerce :: forall a b. Coercible a b => a -> b
coerce ByteCount
arrOff) (coerce :: forall a b. Coercible a b => a -> b
coerce ByteCount
arrLen)

setPtr :: LiftedPrim a => ByteCount -> a -> Ptr Word8 -> ByteCount -> ST s ()
setPtr :: forall a s.
LiftedPrim a =>
ByteCount -> a -> Ptr Word8 -> ByteCount -> ST s ()
setPtr ByteCount
len a
val Ptr Word8
ptr ByteCount
off = do
  let elemSize :: ByteCount
elemSize = forall a. StaticByteSized a => Proxy a -> ByteCount
staticByteSize (forall a. a -> Proxy a
proxyFor a
val)
      elemLen :: ByteCount
elemLen = forall a. Integral a => a -> a -> a
div (coerce :: forall a b. Coercible a b => a -> b
coerce ByteCount
len) ByteCount
elemSize
  forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [ByteCount
0 .. ByteCount
elemLen forall a. Num a => a -> a -> a
- ByteCount
1] forall a b. (a -> b) -> a -> b
$ \ByteCount
pos ->
    forall a (m :: * -> *).
(LiftedPrim a, PrimMonad m) =>
Ptr Word8 -> ByteCount -> a -> m ()
writePtrLiftedInBytes Ptr Word8
ptr (ByteCount
off forall a. Num a => a -> a -> a
+ ByteCount
pos forall a. Num a => a -> a -> a
* ByteCount
elemSize) a
val

instance WriteMem IxPtr where
  writeMemInBytes :: forall a s. LiftedPrim a => a -> IxPtr s -> ByteCount -> ST s ()
writeMemInBytes a
val IxPtr s
mem ByteCount
off = forall a (m :: * -> *).
(LiftedPrim a, PrimMonad m) =>
Ptr Word8 -> ByteCount -> a -> m ()
writePtrLiftedInBytes (forall s. IxPtr s -> Ptr Word8
unIxPtr IxPtr s
mem) ByteCount
off a
val
  copyArrayMemInBytes :: forall s.
ByteArray
-> ByteCount -> ByteCount -> IxPtr s -> ByteCount -> ST s ()
copyArrayMemInBytes ByteArray
arr ByteCount
arrOff ByteCount
arrLen = forall s.
ByteArray
-> ByteCount -> ByteCount -> Ptr Word8 -> ByteCount -> ST s ()
copyPtr ByteArray
arr ByteCount
arrOff ByteCount
arrLen forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall s. IxPtr s -> Ptr Word8
unIxPtr
  setMemInBytes :: forall a s.
LiftedPrim a =>
ByteCount -> a -> IxPtr s -> ByteCount -> ST s ()
setMemInBytes ByteCount
len a
val = forall a s.
LiftedPrim a =>
ByteCount -> a -> Ptr Word8 -> ByteCount -> ST s ()
setPtr ByteCount
len a
val forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall s. IxPtr s -> Ptr Word8
unIxPtr
  releaseMem :: forall s. IxPtr s -> Maybe (IO ())
releaseMem = forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Ptr a -> IO ()
free forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall s. IxPtr s -> Ptr Word8
unIxPtr

writeSBSMem :: WriteMem q => ShortByteString -> ByteCount -> q s -> ByteCount -> ST s ()
writeSBSMem :: forall (q :: * -> *) s.
WriteMem q =>
ShortByteString -> ByteCount -> q s -> ByteCount -> ST s ()
writeSBSMem (SBS ByteArray#
harr) = forall (q :: * -> *) s.
WriteMem q =>
ByteArray -> ByteCount -> ByteCount -> q s -> ByteCount -> ST s ()
copyArrayMemInBytes (ByteArray# -> ByteArray
ByteArray ByteArray#
harr) ByteCount
0

guardedFreeze :: (q s -> ByteCount -> ST s z) -> q s -> ByteCount -> ByteCount -> ST s z
guardedFreeze :: forall (q :: * -> *) s z.
(q s -> ByteCount -> ST s z)
-> q s -> ByteCount -> ByteCount -> ST s z
guardedFreeze q s -> ByteCount -> ST s z
freeze q s
arr ByteCount
len ByteCount
off =
  -- This is a sanity check - if it goes wrong then there's a bug in the library
  if ByteCount
off forall a. Eq a => a -> a -> Bool
/= ByteCount
len
    then forall a. HasCallStack => String -> a
error (String
"Invalid put length: (given " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show ByteCount
len forall a. [a] -> [a] -> [a]
++ String
", used " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show ByteCount
off forall a. [a] -> [a] -> [a]
++ String
")")
    else q s -> ByteCount -> ST s z
freeze q s
arr ByteCount
len

freezeSBSMem :: MutableByteArray s -> ByteCount -> ByteCount -> ST s ShortByteString
freezeSBSMem :: forall s.
MutableByteArray s
-> ByteCount -> ByteCount -> ST s ShortByteString
freezeSBSMem MutableByteArray s
marr ByteCount
cap ByteCount
len = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(ByteArray ByteArray#
harr) -> ByteArray# -> ShortByteString
SBS ByteArray#
harr) (if ByteCount
cap forall a. Eq a => a -> a -> Bool
== ByteCount
len then forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> m ByteArray
unsafeFreezeByteArray MutableByteArray s
marr else forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> Int -> Int -> m ByteArray
freezeByteArray MutableByteArray s
marr Int
0 (coerce :: forall a b. Coercible a b => a -> b
coerce ByteCount
len))

freezeBSMem :: IxPtr s -> ByteCount -> ByteCount -> ST s ByteString
freezeBSMem :: forall s. IxPtr s -> ByteCount -> ByteCount -> ST s ByteString
freezeBSMem (IxPtr Ptr Word8
ptr) ByteCount
_ ByteCount
len =
  forall a s. IO a -> ST s a
unsafeIOToST (Ptr Word8 -> Int -> IO () -> IO ByteString
BSU.unsafePackCStringFinalizer Ptr Word8
ptr (coerce :: forall a b. Coercible a b => a -> b
coerce ByteCount
len) (forall a. Ptr a -> IO ()
free Ptr Word8
ptr))

freezeVecMem :: IxPtr s -> ByteCount -> ByteCount -> ST s (Vector Word8)
freezeVecMem :: forall s. IxPtr s -> ByteCount -> ByteCount -> ST s (Vector Word8)
freezeVecMem (IxPtr Ptr Word8
ptr) ByteCount
_ ByteCount
len = forall a s. IO a -> ST s a
unsafeIOToST (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\ForeignPtr Word8
fp -> forall a. Storable a => ForeignPtr a -> Int -> Vector a
VS.unsafeFromForeignPtr0 ForeignPtr Word8
fp (coerce :: forall a b. Coercible a b => a -> b
coerce ByteCount
len)) (forall a. FinalizerPtr a -> Ptr a -> IO (ForeignPtr a)
newForeignPtr forall a. FinalizerPtr a
finalizerFree Ptr Word8
ptr))

allocPtrMem :: ByteCount -> ST s (IxPtr s)
allocPtrMem :: forall s. ByteCount -> ST s (IxPtr s)
allocPtrMem = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall s. Ptr Word8 -> IxPtr s
IxPtr forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a s. IO a -> ST s a
unsafeIOToST forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Int -> IO (Ptr a)
callocBytes forall b c a. (b -> c) -> (a -> b) -> a -> c
. coerce :: forall a b. Coercible a b => a -> b
coerce

allocArrayMem :: ByteCount -> ST s (MutableByteArray s)
allocArrayMem :: forall s. ByteCount -> ST s (MutableByteArray s)
allocArrayMem = forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableByteArray (PrimState m))
newByteArray forall b c a. (b -> c) -> (a -> b) -> a -> c
. coerce :: forall a b. Coercible a b => a -> b
coerce