{-# LANGUAGE BangPatterns #-}

module Network.Wai.Handler.Warp.Buffer (
    bufferSize
  , allocateBuffer
  , freeBuffer
  , mallocBS
  , newBufferPool
  , withBufferPool
  , toBuilderBuffer
  , copy
  , bufferIO
  ) where

import qualified Data.ByteString as BS
import Data.ByteString.Internal (memcpy)
import Data.ByteString.Unsafe (unsafeTake, unsafeDrop)
import Data.IORef (newIORef, readIORef, writeIORef)
import qualified Data.Streaming.ByteString.Builder.Buffer as B (Buffer (..))
import Foreign.ForeignPtr
import Foreign.Marshal.Alloc (mallocBytes, free, finalizerFree)
import Foreign.Ptr (castPtr, plusPtr)

import Network.Wai.Handler.Warp.Imports
import Network.Wai.Handler.Warp.Types

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

-- | The default size of the write buffer: 16384 (2^14 = 1024 * 16).
--   This is the maximum size of TLS record.
--   This is also the maximum size of HTTP/2 frame payload
--   (excluding frame header).
bufferSize :: BufSize
bufferSize :: BufSize
bufferSize = BufSize
16384

-- | Allocating a buffer with malloc().
allocateBuffer :: Int -> IO Buffer
allocateBuffer :: BufSize -> IO Buffer
allocateBuffer = BufSize -> IO Buffer
forall a. BufSize -> IO (Ptr a)
mallocBytes

-- | Releasing a buffer with free().
freeBuffer :: Buffer -> IO ()
freeBuffer :: Buffer -> IO ()
freeBuffer = Buffer -> IO ()
forall a. Ptr a -> IO ()
free

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

largeBufferSize :: Int
largeBufferSize :: BufSize
largeBufferSize = BufSize
16384

minBufferSize :: Int
minBufferSize :: BufSize
minBufferSize = BufSize
2048

newBufferPool :: IO BufferPool
newBufferPool :: IO BufferPool
newBufferPool = ByteString -> IO BufferPool
forall a. a -> IO (IORef a)
newIORef ByteString
BS.empty

mallocBS :: Int -> IO ByteString
mallocBS :: BufSize -> IO ByteString
mallocBS BufSize
size = do
    Buffer
ptr <- BufSize -> IO Buffer
allocateBuffer BufSize
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 (m :: * -> *) a. Monad m => a -> m a
return (ByteString -> IO ByteString) -> ByteString -> IO ByteString
forall a b. (a -> b) -> a -> b
$! ForeignPtr Word8 -> BufSize -> BufSize -> ByteString
PS ForeignPtr Word8
fptr BufSize
0 BufSize
size
{-# INLINE mallocBS #-}

usefulBuffer :: ByteString -> Bool
usefulBuffer :: ByteString -> Bool
usefulBuffer ByteString
buffer = ByteString -> BufSize
BS.length ByteString
buffer BufSize -> BufSize -> Bool
forall a. Ord a => a -> a -> Bool
>= BufSize
minBufferSize
{-# INLINE usefulBuffer #-}

getBuffer :: BufferPool -> IO ByteString
getBuffer :: BufferPool -> IO ByteString
getBuffer BufferPool
pool = do
    ByteString
buffer <- BufferPool -> IO ByteString
forall a. IORef a -> IO a
readIORef BufferPool
pool
    if ByteString -> Bool
usefulBuffer ByteString
buffer then ByteString -> IO ByteString
forall (m :: * -> *) a. Monad m => a -> m a
return ByteString
buffer else BufSize -> IO ByteString
mallocBS BufSize
largeBufferSize
{-# INLINE getBuffer #-}

putBuffer :: BufferPool -> ByteString -> IO ()
putBuffer :: BufferPool -> ByteString -> IO ()
putBuffer BufferPool
pool ByteString
buffer = BufferPool -> ByteString -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef BufferPool
pool ByteString
buffer
{-# INLINE putBuffer #-}

withForeignBuffer :: ByteString -> ((Buffer, BufSize) -> IO Int) -> IO Int
withForeignBuffer :: ByteString -> ((Buffer, BufSize) -> IO BufSize) -> IO BufSize
withForeignBuffer (PS ForeignPtr Word8
ps BufSize
s BufSize
l) (Buffer, BufSize) -> IO BufSize
f = ForeignPtr Word8 -> (Buffer -> IO BufSize) -> IO BufSize
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr Word8
ps ((Buffer -> IO BufSize) -> IO BufSize)
-> (Buffer -> IO BufSize) -> IO BufSize
forall a b. (a -> b) -> a -> b
$ \Buffer
p -> (Buffer, BufSize) -> IO BufSize
f (Buffer -> Ptr Any
forall a b. Ptr a -> Ptr b
castPtr Buffer
p Ptr Any -> BufSize -> Buffer
forall a b. Ptr a -> BufSize -> Ptr b
`plusPtr` BufSize
s, BufSize
l)
{-# INLINE withForeignBuffer #-}

withBufferPool :: BufferPool -> ((Buffer, BufSize) -> IO Int) -> IO ByteString
withBufferPool :: BufferPool -> ((Buffer, BufSize) -> IO BufSize) -> IO ByteString
withBufferPool BufferPool
pool (Buffer, BufSize) -> IO BufSize
f = do
    ByteString
buffer <- BufferPool -> IO ByteString
getBuffer BufferPool
pool
    BufSize
consumed <- ByteString -> ((Buffer, BufSize) -> IO BufSize) -> IO BufSize
withForeignBuffer ByteString
buffer (Buffer, BufSize) -> IO BufSize
f
    BufferPool -> ByteString -> IO ()
putBuffer BufferPool
pool (ByteString -> IO ()) -> ByteString -> IO ()
forall a b. (a -> b) -> a -> b
$! BufSize -> ByteString -> ByteString
unsafeDrop BufSize
consumed ByteString
buffer
    ByteString -> IO ByteString
forall (m :: * -> *) a. Monad m => a -> m a
return (ByteString -> IO ByteString) -> ByteString -> IO ByteString
forall a b. (a -> b) -> a -> b
$! BufSize -> ByteString -> ByteString
unsafeTake BufSize
consumed ByteString
buffer
{-# INLINE withBufferPool #-}

----------------------------------------------------------------
--
-- Utilities
--

toBuilderBuffer :: Buffer -> BufSize -> IO B.Buffer
toBuilderBuffer :: Buffer -> BufSize -> IO Buffer
toBuilderBuffer Buffer
ptr BufSize
size = do
    ForeignPtr Word8
fptr <- Buffer -> IO (ForeignPtr Word8)
forall a. Ptr a -> IO (ForeignPtr a)
newForeignPtr_ Buffer
ptr
    Buffer -> IO Buffer
forall (m :: * -> *) a. Monad m => a -> m a
return (Buffer -> IO Buffer) -> Buffer -> IO Buffer
forall a b. (a -> b) -> a -> b
$ ForeignPtr Word8 -> Buffer -> Buffer -> Buffer -> Buffer
B.Buffer ForeignPtr Word8
fptr Buffer
ptr Buffer
ptr (Buffer
ptr Buffer -> BufSize -> Buffer
forall a b. Ptr a -> BufSize -> Ptr b
`plusPtr` BufSize
size)

-- | 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 BufSize
o BufSize
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 -> BufSize -> IO ()
memcpy Buffer
ptr (Buffer
p Buffer -> BufSize -> Buffer
forall a b. Ptr a -> BufSize -> Ptr b
`plusPtr` BufSize
o) (BufSize -> BufSize
forall a b. (Integral a, Num b) => a -> b
fromIntegral BufSize
l)
    Buffer -> IO Buffer
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 -> BufSize -> Buffer
forall a b. Ptr a -> BufSize -> Ptr b
`plusPtr` BufSize
l
{-# INLINE copy #-}

bufferIO :: Buffer -> Int -> (ByteString -> IO ()) -> IO ()
bufferIO :: Buffer -> BufSize -> (ByteString -> IO ()) -> IO ()
bufferIO Buffer
ptr BufSize
siz ByteString -> IO ()
io = do
    ForeignPtr Word8
fptr <- Buffer -> IO (ForeignPtr Word8)
forall a. Ptr a -> IO (ForeignPtr a)
newForeignPtr_ Buffer
ptr
    ByteString -> IO ()
io (ByteString -> IO ()) -> ByteString -> IO ()
forall a b. (a -> b) -> a -> b
$ ForeignPtr Word8 -> BufSize -> BufSize -> ByteString
PS ForeignPtr Word8
fptr BufSize
0 BufSize
siz