{-# language BangPatterns #-}
{-# language BlockArguments #-}
{-# language MagicHash #-}
{-# language NamedFieldPuns #-}
{-# language TypeApplications #-}
{-# language UnboxedTuples #-}

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

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

-- | 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 -> 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
  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 :: MutableByteArray RealWorld
arr@(PM.MutableByteArray MutableByteArray# RealWorld
arr#) <- forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableByteArray (PrimState m))
PM.newPinnedByteArray Int
maxSz
  Int
sz <- Ptr Word8 -> IO Int
f (forall s. MutableByteArray s -> Ptr Word8
PM.mutableByteArrayContents MutableByteArray RealWorld
arr)
  forall s. MutableByteArray s -> IO ()
touchMutableByteArrayIO MutableByteArray RealWorld
arr
  forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a -> Int -> m ()
PM.shrinkMutablePrimArray (forall s a. MutableByteArray# s -> MutablePrimArray s a
PM.MutablePrimArray @Exts.RealWorld @Word8 MutableByteArray# RealWorld
arr#) Int
sz
  ByteArray
r <- forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> m ByteArray
PM.unsafeFreezeByteArray MutableByteArray RealWorld
arr
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteArray -> Int -> Int -> Bytes
Bytes ByteArray
r Int
0 Int
sz)

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

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