module Network.Socket.BufferPool.Buffer (
newBufferPool,
withBufferPool,
mallocBS,
copy,
) where
import qualified Data.ByteString as BS
import Data.ByteString.Internal (ByteString (..))
import Data.ByteString.Unsafe (unsafeDrop, unsafeTake)
import Data.IORef (newIORef, readIORef, writeIORef)
import Foreign.ForeignPtr
import Foreign.Marshal.Alloc (finalizerFree, mallocBytes)
import Foreign.Marshal.Utils (copyBytes)
import Foreign.Ptr (castPtr, plusPtr)
import Network.Socket.BufferPool.Types
newBufferPool :: Int -> Int -> IO BufferPool
newBufferPool :: Int -> Int -> IO BufferPool
newBufferPool Int
l Int
h = Int -> Int -> IORef ByteString -> BufferPool
BufferPool Int
l Int
h (IORef ByteString -> BufferPool)
-> IO (IORef ByteString) -> IO BufferPool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ByteString -> IO (IORef ByteString)
forall a. a -> IO (IORef a)
newIORef ByteString
BS.empty
withBufferPool :: BufferPool -> (Buffer -> BufSize -> IO Int) -> IO ByteString
withBufferPool :: BufferPool -> (Buffer -> Int -> IO Int) -> IO ByteString
withBufferPool (BufferPool Int
l Int
h IORef ByteString
ref) Buffer -> Int -> IO Int
f = do
buf0 <- IORef ByteString -> IO ByteString
forall a. IORef a -> IO a
readIORef IORef ByteString
ref
buf <-
if BS.length buf0 >= l
then return buf0
else mallocBS h
consumed <- withForeignBuffer buf f
writeIORef ref $ unsafeDrop consumed buf
return $ unsafeTake consumed buf
withForeignBuffer :: ByteString -> (Buffer -> BufSize -> IO Int) -> IO Int
withForeignBuffer :: ByteString -> (Buffer -> Int -> IO Int) -> IO Int
withForeignBuffer (PS ForeignPtr Word8
ps Int
s Int
l) Buffer -> Int -> IO Int
f = ForeignPtr Word8 -> (Buffer -> IO Int) -> IO Int
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr Word8
ps ((Buffer -> IO Int) -> IO Int) -> (Buffer -> IO Int) -> IO Int
forall a b. (a -> b) -> a -> b
$ \Buffer
p -> Buffer -> Int -> IO Int
f (Buffer -> Ptr (ZonkAny 0)
forall a b. Ptr a -> Ptr b
castPtr Buffer
p Ptr (ZonkAny 0) -> Int -> Buffer
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
s) Int
l
{-# INLINE withForeignBuffer #-}
mallocBS :: Int -> IO ByteString
mallocBS :: Int -> IO ByteString
mallocBS Int
size = do
ptr <- Int -> IO Buffer
forall a. Int -> IO (Ptr a)
mallocBytes Int
size
fptr <- newForeignPtr finalizerFree ptr
return $ PS fptr 0 size
{-# INLINE mallocBS #-}
copy :: Buffer -> ByteString -> IO Buffer
copy :: Buffer -> ByteString -> IO Buffer
copy Buffer
ptr (PS ForeignPtr Word8
fp Int
o Int
l) = ForeignPtr Word8 -> (Buffer -> IO Buffer) -> IO Buffer
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr Word8
fp ((Buffer -> IO Buffer) -> IO Buffer)
-> (Buffer -> IO Buffer) -> IO Buffer
forall a b. (a -> b) -> a -> b
$ \Buffer
p -> do
Buffer -> Buffer -> Int -> IO ()
forall a. Ptr a -> Ptr a -> Int -> IO ()
copyBytes Buffer
ptr (Buffer
p Buffer -> Int -> Buffer
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
o) (Int -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
l)
Buffer -> IO Buffer
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Buffer -> IO Buffer) -> Buffer -> IO Buffer
forall a b. (a -> b) -> a -> b
$ Buffer
ptr Buffer -> Int -> Buffer
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
l
{-# INLINE copy #-}