{-# LANGUAGE BangPatterns #-}

module Network.Wai.Handler.Warp.Buffer (
    createWriteBuffer
  , allocateBuffer
  , freeBuffer
  , toBuilderBuffer
  , bufferIO
  ) where

import Data.IORef (IORef, readIORef)
import qualified Data.Streaming.ByteString.Builder.Buffer as B (Buffer (..))
import Foreign.ForeignPtr
import Foreign.Marshal.Alloc (mallocBytes, free)
import Foreign.Ptr (plusPtr)
import Network.Socket.BufferPool

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

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

-- | Allocate a buffer of the given size and wrap it in a 'WriteBuffer'
-- containing that size and a finalizer.
createWriteBuffer :: BufSize -> IO WriteBuffer
createWriteBuffer :: BufSize -> IO WriteBuffer
createWriteBuffer BufSize
size = do
  Buffer
bytes <- BufSize -> IO Buffer
allocateBuffer BufSize
size
  WriteBuffer -> IO WriteBuffer
forall (m :: * -> *) a. Monad m => a -> m a
return
    WriteBuffer :: Buffer -> BufSize -> IO () -> WriteBuffer
WriteBuffer
      { bufBuffer :: Buffer
bufBuffer = Buffer
bytes,
        bufSize :: BufSize
bufSize = BufSize
size,
        bufFree :: IO ()
bufFree = Buffer -> IO ()
freeBuffer Buffer
bytes
      }

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

-- | 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

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

toBuilderBuffer :: IORef WriteBuffer -> IO B.Buffer
toBuilderBuffer :: IORef WriteBuffer -> IO Buffer
toBuilderBuffer IORef WriteBuffer
writeBufferRef = do
    WriteBuffer
writeBuffer <- IORef WriteBuffer -> IO WriteBuffer
forall a. IORef a -> IO a
readIORef IORef WriteBuffer
writeBufferRef
    let ptr :: Buffer
ptr = WriteBuffer -> Buffer
bufBuffer WriteBuffer
writeBuffer
        size :: BufSize
size = WriteBuffer -> BufSize
bufSize WriteBuffer
writeBuffer
    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)

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