{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE UnboxedTuples #-}

module Data.Bytes.IO
  ( hGet
  , hPut
  ) where

import Data.Bytes.Pure (contents, pin)
import Data.Bytes.Types (Bytes (Bytes))
import Data.Primitive (ByteArray (..), MutableByteArray)
import qualified Data.Primitive as PM
import Data.Word (Word8)
import Foreign.Ptr (Ptr)
import qualified GHC.Exts as Exts
import GHC.IO (IO (IO))
import System.IO (Handle)
import qualified System.IO as IO

{- | Read 'Bytes' directly from the specified 'Handle'. The resulting
'Bytes' are pinned. This is implemented with 'IO.hGetBuf'.
-}
hGet :: Handle -> Int -> IO Bytes
hGet :: Handle -> Int -> IO Bytes
hGet Handle
h Int
i = Int -> (Ptr Word8 -> IO Int) -> IO Bytes
createPinnedAndTrim Int
i (\Ptr Word8
p -> Handle -> Ptr Word8 -> Int -> IO Int
forall a. Handle -> Ptr a -> Int -> IO Int
IO.hGetBuf Handle
h Ptr Word8
p Int
i)

{- | Outputs 'Bytes' to the specified 'Handle'. This is implemented
with 'IO.hPutBuf'.
-}
hPut :: Handle -> Bytes -> IO ()
hPut :: Handle -> Bytes -> IO ()
hPut Handle
h Bytes
b0 = do
  let b1 :: Bytes
b1@(Bytes ByteArray
arr Int
_ Int
len) = Bytes -> Bytes
pin Bytes
b0
  Handle -> Ptr Word8 -> Int -> IO ()
forall a. Handle -> Ptr a -> Int -> IO ()
IO.hPutBuf Handle
h (Bytes -> Ptr Word8
contents Bytes
b1) Int
len
  ByteArray -> IO ()
touchByteArrayIO ByteArray
arr

-- Only used internally.
createPinnedAndTrim :: Int -> (Ptr Word8 -> IO Int) -> IO Bytes
{-# INLINE createPinnedAndTrim #-}
createPinnedAndTrim :: Int -> (Ptr Word8 -> IO Int) -> IO Bytes
createPinnedAndTrim Int
maxSz Ptr Word8 -> IO Int
f = do
  arr@(PM.MutableByteArray arr#) <- Int -> IO (MutableByteArray (PrimState IO))
forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableByteArray (PrimState m))
PM.newPinnedByteArray Int
maxSz
  sz <- f (PM.mutableByteArrayContents arr)
  touchMutableByteArrayIO arr
  PM.shrinkMutablePrimArray (PM.MutablePrimArray @Exts.RealWorld @Word8 arr#) sz
  r <- PM.unsafeFreezeByteArray arr
  pure (Bytes r 0 sz)

touchMutableByteArrayIO :: MutableByteArray s -> IO ()
touchMutableByteArrayIO :: forall s. MutableByteArray s -> IO ()
touchMutableByteArrayIO (PM.MutableByteArray MutableByteArray# s
x) =
  (State# RealWorld -> (# State# RealWorld, () #)) -> IO ()
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO (\State# RealWorld
s -> (# MutableByteArray# s -> State# RealWorld -> State# RealWorld
forall a d. a -> State# d -> State# d
Exts.touch# MutableByteArray# s
x State# RealWorld
s, () #))

touchByteArrayIO :: ByteArray -> IO ()
touchByteArrayIO :: ByteArray -> IO ()
touchByteArrayIO (ByteArray ByteArray#
x) =
  (State# RealWorld -> (# State# RealWorld, () #)) -> IO ()
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO (\State# RealWorld
s -> (# ByteArray# -> State# RealWorld -> State# RealWorld
forall a d. a -> State# d -> State# d
Exts.touch# ByteArray#
x State# RealWorld
s, () #))