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