    {-# LANGUAGE BangPatterns #-}
    {-# LANGUAGE CPP #-}

    module Network.Wai.Handler.Warp.IO where

    import Control.Exception (mask_)
    import Data.ByteString.Builder (Builder)
    import Data.ByteString.Builder.Extra (Next (Chunk, Done, More), runBuilder)
    import Data.IORef (IORef, readIORef, writeIORef)
    import Network.Wai.Handler.Warp.Buffer
    import Network.Wai.Handler.Warp.Imports
    import Network.Wai.Handler.Warp.Types

    toBufIOWith :: Int -> IORef WriteBuffer -> (ByteString -> IO ()) -> Builder -> IO ()
    toBufIOWith :: Int
    -> IORef WriteBuffer -> (ByteString -> IO ()) -> Builder -> IO ()
    toBufIOWith Int
    maxRspBufSize IORef WriteBuffer
    writeBufferRef ByteString -> IO ()
    io Builder
    builder = do
      WriteBuffer
    writeBuffer <- IORef WriteBuffer -> IO WriteBuffer
    forall a. IORef a -> IO a
    readIORef IORef WriteBuffer
    writeBufferRef
      WriteBuffer -> BufferWriter -> IO ()
    loop WriteBuffer
    writeBuffer BufferWriter
    firstWriter
      where
        firstWriter :: BufferWriter
    firstWriter = Builder -> BufferWriter
    runBuilder Builder
    builder
        loop :: WriteBuffer -> BufferWriter -> IO ()
    loop WriteBuffer
    writeBuffer BufferWriter
    writer = do
          let buf :: Buffer
    buf = WriteBuffer -> Buffer
    bufBuffer WriteBuffer
    writeBuffer
              size :: Int
    size = WriteBuffer -> Int
    bufSize WriteBuffer
    writeBuffer
          (Int
    len, Next
    signal) <- BufferWriter
    writer Buffer
    buf Int
    size
          Buffer -> Int -> (ByteString -> IO ()) -> IO ()
    bufferIO Buffer
    buf Int
    len ByteString -> IO ()
    io
          case Next
    signal of
            Next
    Done -> () -> IO ()
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return ()
            More Int
    minSize BufferWriter
    next
              | Int
    size Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    < Int
    minSize -> do
                  Bool -> IO () -> IO ()
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
    when (Int
    minSize Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    > Int
    maxRspBufSize) (IO () -> IO ()) -> IO () -> IO ()
    forall a b. (a -> b) -> a -> b
    $
                    [Char] -> IO ()
    forall a. HasCallStack => [Char] -> a
    error ([Char] -> IO ()) -> [Char] -> IO ()
    forall a b. (a -> b) -> a -> b
    $ [Char]
    "Sending a Builder response required a buffer of size "
                              [Char] -> [Char] -> [Char]
    forall a. [a] -> [a] -> [a]
    ++ Int -> [Char]
    forall a. Show a => a -> [Char]
    show Int
    minSize [Char] -> [Char] -> [Char]
    forall a. [a] -> [a] -> [a]
    ++ [Char]
    " which is bigger than the specified maximum of "
                              [Char] -> [Char] -> [Char]
    forall a. [a] -> [a] -> [a]
    ++ Int -> [Char]
    forall a. Show a => a -> [Char]
    show Int
    maxRspBufSize [Char] -> [Char] -> [Char]
    forall a. [a] -> [a] -> [a]
    ++ [Char]
    "!"
                  -- The current WriteBuffer is too small to fit the next
                  -- batch of bytes from the Builder so we free it and
                  -- create a new bigger one. Freeing the current buffer,
                  -- creating a new one and writing it to the IORef need
                  -- to be performed atomically to prevent both double
                  -- frees and missed frees. So we mask async exceptions:
                  WriteBuffer
    biggerWriteBuffer <- IO WriteBuffer -> IO WriteBuffer
    forall a. IO a -> IO a
    mask_ (IO WriteBuffer -> IO WriteBuffer)
    -> IO WriteBuffer -> IO WriteBuffer
    forall a b. (a -> b) -> a -> b
    $ do
                    WriteBuffer -> IO ()
    bufFree WriteBuffer
    writeBuffer
                    WriteBuffer
    biggerWriteBuffer <- Int -> IO WriteBuffer
    createWriteBuffer Int
    minSize
                    IORef WriteBuffer -> WriteBuffer -> IO ()
    forall a. IORef a -> a -> IO ()
    writeIORef IORef WriteBuffer
    writeBufferRef WriteBuffer
    biggerWriteBuffer
                    WriteBuffer -> IO WriteBuffer
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return WriteBuffer
    biggerWriteBuffer
                  WriteBuffer -> BufferWriter -> IO ()
    loop WriteBuffer
    biggerWriteBuffer BufferWriter
    next
              | Bool
    otherwise -> WriteBuffer -> BufferWriter -> IO ()
    loop WriteBuffer
    writeBuffer BufferWriter
    next
            Chunk ByteString
    bs BufferWriter
    next -> do
              ByteString -> IO ()
    io ByteString
    bs
              WriteBuffer -> BufferWriter -> IO ()
    loop WriteBuffer
    writeBuffer BufferWriter
    next