{-# LANGUAGE BangPatterns  #-}
{-# LANGUAGE MagicHash     #-}
{-# LANGUAGE TypeFamilies  #-}
{-# LANGUAGE UnboxedTuples #-}
{- |
Memory access primitives.

Includes code from the [store-core](https://hackage.haskell.org/package/store-core) package.
-}
module Flat.Memory
  ( chunksToByteString
  , chunksToByteArray
  , ByteArray
  , pokeByteArray
  , pokeByteString
  , unsafeCreateUptoN'
  , minusPtr
  , peekByteString
  )
where

import           Control.Monad            (foldM_, when)
import           Control.Monad.Primitive  (PrimMonad (..))
import qualified Data.ByteString          as B
import qualified Data.ByteString.Internal as BS
import           Data.Primitive.ByteArray (ByteArray, ByteArray#,
                                           MutableByteArray (..), newByteArray,
                                           unsafeFreezeByteArray)
import           Foreign                  (Ptr, Word8, minusPtr, plusPtr,
                                           withForeignPtr)
import           GHC.Prim                 (copyAddrToByteArray#,
                                           copyByteArrayToAddr#)
import           GHC.Ptr                  (Ptr (..))
import           GHC.Types                (IO (..), Int (..))
import           System.IO.Unsafe         (unsafeDupablePerformIO,
                                           unsafePerformIO)

unsafeCreateUptoN' :: Int -> (Ptr Word8 -> IO (Int, a)) -> (BS.ByteString, a)
unsafeCreateUptoN' :: forall a. Int -> (Ptr Word8 -> IO (Int, a)) -> (ByteString, a)
unsafeCreateUptoN' Int
l Ptr Word8 -> IO (Int, a)
f = forall a. IO a -> a
unsafeDupablePerformIO (forall a. Int -> (Ptr Word8 -> IO (Int, a)) -> IO (ByteString, a)
createUptoN' Int
l Ptr Word8 -> IO (Int, a)
f)
{-# INLINE unsafeCreateUptoN' #-}

createUptoN' :: Int -> (Ptr Word8 -> IO (Int, a)) -> IO (BS.ByteString, a)
createUptoN' :: forall a. Int -> (Ptr Word8 -> IO (Int, a)) -> IO (ByteString, a)
createUptoN' Int
l Ptr Word8 -> IO (Int, a)
f = do
  ForeignPtr Word8
fp        <- forall a. Int -> IO (ForeignPtr a)
BS.mallocByteString Int
l
  (Int
l', a
res) <- forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr Word8
fp forall a b. (a -> b) -> a -> b
$ \Ptr Word8
p -> Ptr Word8 -> IO (Int, a)
f Ptr Word8
p
  --print (unwords ["Buffer allocated:",show l,"bytes, used:",show l',"bytes"])
  forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
l' forall a. Ord a => a -> a -> Bool
> Int
l) forall a b. (a -> b) -> a -> b
$ forall a. HasCallStack => [Char] -> a
error
    ([[Char]] -> [Char]
unwords
      [[Char]
"Buffer overflow, allocated:", forall a. Show a => a -> [Char]
show Int
l, [Char]
"bytes, used:", forall a. Show a => a -> [Char]
show Int
l', [Char]
"bytes"]
    )
  forall (m :: * -> *) a. Monad m => a -> m a
return (ForeignPtr Word8 -> Int -> Int -> ByteString
BS.PS ForeignPtr Word8
fp Int
0 Int
l', a
res) -- , minusPtr l')
{-# INLINE createUptoN' #-}

-- |Copy bytestring to given pointer, returns new pointer
pokeByteString :: B.ByteString -> Ptr Word8 -> IO (Ptr Word8)
pokeByteString :: ByteString -> Ptr Word8 -> IO (Ptr Word8)
pokeByteString (BS.PS ForeignPtr Word8
foreignPointer Int
sourceOffset Int
sourceLength) Ptr Word8
destPointer =
  do
    forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr Word8
foreignPointer forall a b. (a -> b) -> a -> b
$ \Ptr Word8
sourcePointer -> Ptr Word8 -> Ptr Word8 -> Int -> IO ()
BS.memcpy
      Ptr Word8
destPointer
      (Ptr Word8
sourcePointer forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
sourceOffset)
      Int
sourceLength
    forall (m :: * -> *) a. Monad m => a -> m a
return (Ptr Word8
destPointer forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
sourceLength)

{-| Create a new bytestring, copying sourceLen bytes from sourcePtr

@since 0.6
-}
peekByteString ::
  Ptr Word8 -- ^ sourcePtr
  -> Int -- ^ sourceLen
  -> BS.ByteString
peekByteString :: Ptr Word8 -> Int -> ByteString
peekByteString Ptr Word8
sourcePtr Int
sourceLength = Int -> (Ptr Word8 -> IO ()) -> ByteString
BS.unsafeCreate Int
sourceLength forall a b. (a -> b) -> a -> b
$ \Ptr Word8
destPointer -> Ptr Word8 -> Ptr Word8 -> Int -> IO ()
BS.memcpy Ptr Word8
destPointer Ptr Word8
sourcePtr Int
sourceLength

-- |Copy ByteArray to given pointer, returns new pointer
pokeByteArray :: ByteArray# -> Int -> Int -> Ptr Word8 -> IO (Ptr Word8)
pokeByteArray :: ByteArray# -> Int -> Int -> Ptr Word8 -> IO (Ptr Word8)
pokeByteArray ByteArray#
sourceArr Int
sourceOffset Int
len Ptr Word8
dest = do
  forall a. ByteArray# -> Int -> Ptr a -> Int -> IO ()
copyByteArrayToAddr ByteArray#
sourceArr Int
sourceOffset Ptr Word8
dest Int
len
  let !dest' :: Ptr Word8
dest' = Ptr Word8
dest forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
len
  forall (m :: * -> *) a. Monad m => a -> m a
return Ptr Word8
dest'
{-# INLINE pokeByteArray #-}


-- | Wrapper around @copyByteArrayToAddr#@ primop.
--
-- Copied from the store-core package
copyByteArrayToAddr :: ByteArray# -> Int -> Ptr a -> Int -> IO ()
copyByteArrayToAddr :: forall a. ByteArray# -> Int -> Ptr a -> Int -> IO ()
copyByteArrayToAddr ByteArray#
arr (I# Int#
offset) (Ptr Addr#
addr) (I# Int#
len) =
  forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO (\State# RealWorld
s -> (# forall d.
ByteArray# -> Int# -> Addr# -> Int# -> State# d -> State# d
copyByteArrayToAddr# ByteArray#
arr Int#
offset Addr#
addr Int#
len State# RealWorld
s, () #))
{-# INLINE copyByteArrayToAddr #-}

chunksToByteString :: (Ptr Word8, [Int]) -> BS.ByteString
chunksToByteString :: (Ptr Word8, [Int]) -> ByteString
chunksToByteString (Ptr Word8
sourcePtr0, [Int]
lens) =
  Int -> (Ptr Word8 -> IO ()) -> ByteString
BS.unsafeCreate (forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Int]
lens) forall a b. (a -> b) -> a -> b
$ \Ptr Word8
destPtr0 -> forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m ()
foldM_
    (\(Ptr Word8
destPtr, Ptr Word8
sourcePtr) Int
sourceLength ->
      Ptr Word8 -> Ptr Word8 -> Int -> IO ()
BS.memcpy Ptr Word8
destPtr Ptr Word8
sourcePtr Int
sourceLength
        forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall (m :: * -> *) a. Monad m => a -> m a
return
             ( Ptr Word8
destPtr forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
sourceLength
             , Ptr Word8
sourcePtr forall a b. Ptr a -> Int -> Ptr b
`plusPtr` (Int
sourceLength forall a. Num a => a -> a -> a
+ Int
1)
             )
    )
    (Ptr Word8
destPtr0, Ptr Word8
sourcePtr0)
    [Int]
lens

chunksToByteArray :: (Ptr Word8, [Int]) -> (ByteArray, Int)
chunksToByteArray :: (Ptr Word8, [Int]) -> (ByteArray, Int)
chunksToByteArray (Ptr Word8
sourcePtr0, [Int]
lens) = forall a. IO a -> a
unsafePerformIO forall a b. (a -> b) -> a -> b
$ do
  let len :: Int
len = forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Int]
lens
  MutableByteArray RealWorld
arr <- forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableByteArray (PrimState m))
newByteArray Int
len
  forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m ()
foldM_
    (\(Int
destOff, Ptr Word8
sourcePtr) Int
sourceLength ->
      forall a.
Ptr a -> MutableByteArray (PrimState IO) -> Int -> Int -> IO ()
copyAddrToByteArray Ptr Word8
sourcePtr MutableByteArray RealWorld
arr Int
destOff Int
sourceLength forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall (m :: * -> *) a. Monad m => a -> m a
return
        (Int
destOff forall a. Num a => a -> a -> a
+ Int
sourceLength, Ptr Word8
sourcePtr forall a b. Ptr a -> Int -> Ptr b
`plusPtr` (Int
sourceLength forall a. Num a => a -> a -> a
+ Int
1))
    )
    (Int
0, Ptr Word8
sourcePtr0)
    [Int]
lens
  ByteArray
farr <- forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> m ByteArray
unsafeFreezeByteArray MutableByteArray RealWorld
arr
  forall (m :: * -> *) a. Monad m => a -> m a
return (ByteArray
farr, Int
len)


-- | Wrapper around @copyAddrToByteArray#@ primop.
--
-- Copied from the store-core package
copyAddrToByteArray
  :: Ptr a -> MutableByteArray (PrimState IO) -> Int -> Int -> IO ()
copyAddrToByteArray :: forall a.
Ptr a -> MutableByteArray (PrimState IO) -> Int -> Int -> IO ()
copyAddrToByteArray (Ptr Addr#
addr) (MutableByteArray MutableByteArray# (PrimState IO)
arr) (I# Int#
offset) (I# Int#
len) =
  forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO (\State# RealWorld
s -> (# forall d.
Addr#
-> MutableByteArray# d -> Int# -> Int# -> State# d -> State# d
copyAddrToByteArray# Addr#
addr MutableByteArray# (PrimState IO)
arr Int#
offset Int#
len State# RealWorld
s, () #))
{-# INLINE copyAddrToByteArray #-}