module Network.Socket.BufferPool.Buffer (
    newBufferPool,
    withBufferPool,
    tryWithBufferPool,
    eitherWithBufferPool,
    mallocBS,
    copy,
) where

import qualified Data.ByteString as BS
import Data.ByteString.Internal (ByteString (..))
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

----------------------------------------------------------------

-- | Creating a buffer pool.
--   The first argument is the lower limit.
--   When the size of the buffer in the poll is lower than this limit,
--   the buffer is thrown awany (and is eventually freed).
--   Then a new buffer is allocated.
--   The second argument is the size for the new allocation.
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

----------------------------------------------------------------

-- | Using a buffer pool.
--   The second argument is a function which returns
--   how many bytes are filled in the buffer.
--   This function should return non negative 'Int'.
--   The buffer in the buffer pool is automatically managed.
--
-- /Caveats of 'BufferPool' apply./
withBufferPool :: BufferPool -> (Buffer -> BufSize -> IO Int) -> IO ByteString
withBufferPool :: BufferPool -> (Buffer -> Int -> IO Int) -> IO ByteString
withBufferPool pool :: BufferPool
pool@(BufferPool Int
_ Int
_ IORef ByteString
ref) Buffer -> Int -> IO Int
f = do
    (buf, consumed) <- BufferPool -> (Buffer -> Int -> IO Int) -> IO (ByteString, Int)
forall r.
BufferPool -> (Buffer -> Int -> IO r) -> IO (ByteString, r)
applyBufferPool BufferPool
pool Buffer -> Int -> IO Int
f
    writeIORef ref $ BS.drop consumed buf
    return $ BS.take consumed buf

-- | L ike 'withBufferPool' for fillers that can decline to fill:
--   a negative return value from the filler leaves the pool untouched
--   and produces 'Nothing'.
--
-- /Caveats of 'BufferPool' apply./
tryWithBufferPool
    :: BufferPool -> (Buffer -> BufSize -> IO Int) -> IO (Maybe ByteString)
tryWithBufferPool :: BufferPool -> (Buffer -> Int -> IO Int) -> IO (Maybe ByteString)
tryWithBufferPool pool :: BufferPool
pool@(BufferPool Int
_ Int
_ IORef ByteString
ref) Buffer -> Int -> IO Int
f = do
    (buf, consumed) <- BufferPool -> (Buffer -> Int -> IO Int) -> IO (ByteString, Int)
forall r.
BufferPool -> (Buffer -> Int -> IO r) -> IO (ByteString, r)
applyBufferPool BufferPool
pool Buffer -> Int -> IO Int
f
    if consumed < 0
        then do
            writeIORef ref buf
            return Nothing
        else do
            writeIORef ref $ BS.drop consumed buf
            return $ Just $ BS.take consumed buf

-- | Like 'withBufferPool' for fillers which can give up: 'Left' leaves
--   the pool untouched.
--
-- /Caveats of 'BufferPool' apply./
eitherWithBufferPool
    :: BufferPool
    -> (Buffer -> BufSize -> IO (Either a Int))
    -> IO (Either a ByteString)
eitherWithBufferPool :: forall a.
BufferPool
-> (Buffer -> Int -> IO (Either a Int)) -> IO (Either a ByteString)
eitherWithBufferPool pool :: BufferPool
pool@(BufferPool Int
_ Int
_ IORef ByteString
ref) Buffer -> Int -> IO (Either a Int)
f = do
    (buf, r) <- BufferPool
-> (Buffer -> Int -> IO (Either a Int))
-> IO (ByteString, Either a Int)
forall r.
BufferPool -> (Buffer -> Int -> IO r) -> IO (ByteString, r)
applyBufferPool BufferPool
pool Buffer -> Int -> IO (Either a Int)
f
    case r of
        Left a
a -> do
            IORef ByteString -> ByteString -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef ByteString
ref ByteString
buf
            Either a ByteString -> IO (Either a ByteString)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either a ByteString -> IO (Either a ByteString))
-> Either a ByteString -> IO (Either a ByteString)
forall a b. (a -> b) -> a -> b
$ a -> Either a ByteString
forall a b. a -> Either a b
Left a
a
        Right Int
consumed -> do
            IORef ByteString -> ByteString -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef ByteString
ref (ByteString -> IO ()) -> ByteString -> IO ()
forall a b. (a -> b) -> a -> b
$ Int -> ByteString -> ByteString
BS.drop Int
consumed ByteString
buf
            Either a ByteString -> IO (Either a ByteString)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either a ByteString -> IO (Either a ByteString))
-> Either a ByteString -> IO (Either a ByteString)
forall a b. (a -> b) -> a -> b
$ ByteString -> Either a ByteString
forall a b. b -> Either a b
Right (ByteString -> Either a ByteString)
-> ByteString -> Either a ByteString
forall a b. (a -> b) -> a -> b
$ Int -> ByteString -> ByteString
BS.take Int
consumed ByteString
buf

applyBufferPool
    :: BufferPool -> (Buffer -> BufSize -> IO r) -> IO (ByteString, r)
applyBufferPool :: forall r.
BufferPool -> (Buffer -> Int -> IO r) -> IO (ByteString, r)
applyBufferPool (BufferPool Int
l Int
h IORef ByteString
ref) Buffer -> Int -> IO r
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
    return (buf, consumed)

withForeignBuffer :: ByteString -> (Buffer -> BufSize -> IO r) -> IO r
withForeignBuffer :: forall r. ByteString -> (Buffer -> Int -> IO r) -> IO r
withForeignBuffer (PS ForeignPtr Word8
ps Int
s Int
l) Buffer -> Int -> IO r
f = ForeignPtr Word8 -> (Buffer -> IO r) -> IO r
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr Word8
ps ((Buffer -> IO r) -> IO r) -> (Buffer -> IO r) -> IO r
forall a b. (a -> b) -> a -> b
$ \Buffer
p -> Buffer -> Int -> IO r
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 #-}

----------------------------------------------------------------

-- | Allocating a byte string.
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 #-}

-- | Copying the bytestring to the buffer.
--   This function returns the point where the next copy should start.
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 #-}