    {-# 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 = 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 = 16384

    minBufferSize :: Int
    minBufferSize :: BufSize
    minBufferSize = 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 size :: 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 0 BufSize
    size
    {-# INLINE mallocBS #-}

    usefulBuffer :: ByteString -> Bool
    usefulBuffer :: ByteString -> Bool
    usefulBuffer buffer :: 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 pool :: 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 pool :: BufferPool
    pool buffer :: 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 ps :: ForeignPtr Word8
    ps s :: BufSize
    s l :: BufSize
    l) f :: (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
    $ \p :: 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 pool :: BufferPool
    pool f :: (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 ptr :: Buffer
    ptr size :: 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 fp :: ForeignPtr Word8
    fp o :: BufSize
    o l :: 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
    $ \p :: 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 ptr :: Buffer
    ptr siz :: BufSize
    siz io :: 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 0 BufSize
    siz