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

import qualified Data.ByteString as BS
import Data.ByteString.Internal (ByteString(..), memcpy)
import Data.ByteString.Unsafe (unsafeTake, unsafeDrop)
import Data.IORef (newIORef, readIORef, writeIORef)
import Foreign.ForeignPtr
import Foreign.Marshal.Alloc (mallocBytes, finalizerFree)
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.
--   The buffer in the buffer pool is automatically managed.
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
    ByteString
buf0 <- IORef ByteString -> IO ByteString
forall a. IORef a -> IO a
readIORef IORef ByteString
ref
    ByteString
buf  <- if ByteString -> Int
BS.length ByteString
buf0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
l then ByteString -> IO ByteString
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ByteString
buf0
                                   else Int -> IO ByteString
mallocBS Int
h
    Int
consumed <- ByteString -> (Buffer -> Int -> IO Int) -> IO Int
withForeignBuffer ByteString
buf Buffer -> Int -> IO Int
f
    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
unsafeDrop Int
consumed ByteString
buf
    ByteString -> IO ByteString
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (ByteString -> IO ByteString) -> ByteString -> IO ByteString
forall a b. (a -> b) -> a -> b
$ Int -> ByteString -> ByteString
unsafeTake Int
consumed ByteString
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 Any
forall a b. Ptr a -> Ptr b
castPtr Buffer
p Ptr Any -> 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
    Buffer
ptr <- Int -> IO Buffer
forall a. Int -> IO (Ptr a)
mallocBytes Int
size
    ForeignPtr Word8
fptr <- FinalizerPtr Word8 -> Buffer -> IO (ForeignPtr Word8)
forall a. FinalizerPtr a -> Ptr a -> IO (ForeignPtr a)
newForeignPtr FinalizerPtr Word8
forall a. FinalizerPtr a
finalizerFree Buffer
ptr
    ByteString -> IO ByteString
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (ByteString -> IO ByteString) -> ByteString -> IO ByteString
forall a b. (a -> b) -> a -> b
$ ForeignPtr Word8 -> Int -> Int -> ByteString
PS ForeignPtr Word8
fptr Int
0 Int
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 ()
memcpy 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 #-}