    {-# LANGUAGE NamedFieldPuns #-}
    {-# LANGUAGE OverloadedStrings #-}
    {-# LANGUAGE RecordWildCards #-}
    {-# LANGUAGE ScopedTypeVariables #-}

    module Network.HTTP.Semantics.FillBuf (
        -- * Filling a buffer
        Next (..),
        DynaNext,
        BytesFilled,
        StreamingChunk (..),
        IsEndOfStream (..),
        CleanupStream,
        fillBuilderBodyGetNext,
        fillFileBodyGetNext,
        fillStreamBodyGetNext,
    ) where

    import Control.Exception (SomeException)
    import Control.Monad
    import qualified Data.ByteString as BS
    import Data.ByteString.Builder (Builder)
    import qualified Data.ByteString.Builder.Extra as B
    import Data.ByteString.Internal
    import Data.Int (Int64)
    import Data.Maybe
    import Foreign.Ptr (plusPtr)
    import Network.ByteOrder
    import Network.HTTP.Semantics.Client

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

    -- | Write part of a streaming response to the write buffer
    --
    -- In @http2@ this will be used to construct a single HTTP2 @DATA@ frame
    -- (see discussion of the maximum number of bytes, below).
    type DynaNext =
        Buffer
        -- ^ Write buffer
        -> Int
        -- ^ Maximum number of bytes we are allowed to write
        --
        -- In @http2@, this maximum will be set to the space left in the write
        -- buffer. Implicitly this also means that this maximum cannot exceed the
        -- maximum size of a HTTP2 frame, since in @http2@ the size of the write
        -- buffer is also used to set @SETTINGS_MAX_FRAME_SIZE@ (see
        -- @confBufferSize@).
        -> IO Next
        -- ^ Information on the data written, and on how to continue if not all data
        -- was written

    type BytesFilled = Int

    data Next
        = Next
            BytesFilled -- payload length
            Bool -- require flushing
            (Maybe DynaNext)
        | CancelNext (Maybe SomeException)

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

    data StreamingChunk
        = -- | Indicate that the stream is finished
          StreamingFinished (Maybe CleanupStream)
        | -- | Indicate that the stream is cancelled
          StreamingCancelled (Maybe SomeException)
        | -- | Flush the stream
          --
          -- This will cause the write buffer to be written to the network socket,
          -- without waiting for more data.
          StreamingFlush
        | -- | Construct a DATA frame, optionally terminating the stream
          StreamingBuilder Builder IsEndOfStream

    -- | Action to run prior to terminating the stream
    type CleanupStream = IO ()

    data IsEndOfStream
        = -- | The stream is not yet terminated
          NotEndOfStream
        | -- | The stream is terminated
          --
          -- In addition to indicating that the stream is terminated, we can also
          -- specify an optional `Cleanup` handler to be run.
          EndOfStream (Maybe CleanupStream)

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

    fillBuilderBodyGetNext :: Builder -> DynaNext
    fillBuilderBodyGetNext :: Builder -> DynaNext
    fillBuilderBodyGetNext Builder
    bb Buffer
    buf Int
    room = do
        (len, signal) <- Builder -> BufferWriter
    B.runBuilder Builder
    bb Buffer
    buf Int
    room
        return $ nextForBuilder len signal

    fillFileBodyGetNext
        :: PositionRead -> FileOffset -> ByteCount -> Sentinel -> DynaNext
    fillFileBodyGetNext :: PositionRead -> Int64 -> Int64 -> Sentinel -> DynaNext
    fillFileBodyGetNext PositionRead
    pread Int64
    start Int64
    bytecount Sentinel
    sentinel Buffer
    buf Int
    room = do
        len <- PositionRead
    pread Int64
    start (Int -> Int64 -> Int64
    mini Int
    room Int64
    bytecount) Buffer
    buf
        let len' = Int64 -> Int
    forall a b. (Integral a, Num b) => a -> b
    fromIntegral Int64
    len
        nextForFile len' pread (start + len) (bytecount - len) sentinel

    fillStreamBodyGetNext :: IO (Maybe StreamingChunk) -> DynaNext
    fillStreamBodyGetNext :: IO (Maybe StreamingChunk) -> DynaNext
    fillStreamBodyGetNext IO (Maybe StreamingChunk)
    takeQ = NextWithTotal
    loop Int
    0
      where
        loop :: NextWithTotal
        loop :: NextWithTotal
    loop Int
    total Buffer
    buf Int
    room = do
            mChunk <- IO (Maybe StreamingChunk)
    takeQ
            case mChunk of
                Just StreamingChunk
    chunk -> StreamingChunk -> NextWithTotal -> NextWithTotal
    runStreamingChunk StreamingChunk
    chunk NextWithTotal
    loop Int
    total Buffer
    buf Int
    room
                Maybe StreamingChunk
    Nothing -> Next -> IO Next
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return (Next -> IO Next) -> Next -> IO Next
    forall a b. (a -> b) -> a -> b
    $ Int -> Bool -> Maybe DynaNext -> Next
    Next Int
    total Bool
    False (DynaNext -> Maybe DynaNext
    forall a. a -> Maybe a
    Just (DynaNext -> Maybe DynaNext) -> DynaNext -> Maybe DynaNext
    forall a b. (a -> b) -> a -> b
    $ NextWithTotal
    loop Int
    0)

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

    fillBufBuilderOne :: Int -> B.BufferWriter -> DynaNext
    fillBufBuilderOne :: Int -> BufferWriter -> DynaNext
    fillBufBuilderOne Int
    minReq BufferWriter
    writer Buffer
    buf0 Int
    room = do
        if Int
    room Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    >= Int
    minReq
            then do
                (len, signal) <- BufferWriter
    writer Buffer
    buf0 Int
    room
                return $ nextForBuilder len signal
            else do
                Next -> IO Next
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return (Next -> IO Next) -> Next -> IO Next
    forall a b. (a -> b) -> a -> b
    $ Int -> Bool -> Maybe DynaNext -> Next
    Next Int
    0 Bool
    True (DynaNext -> Maybe DynaNext
    forall a. a -> Maybe a
    Just (DynaNext -> Maybe DynaNext) -> DynaNext -> Maybe DynaNext
    forall a b. (a -> b) -> a -> b
    $ Int -> BufferWriter -> DynaNext
    fillBufBuilderOne Int
    minReq BufferWriter
    writer)

    fillBufBuilderTwo :: ByteString -> B.BufferWriter -> DynaNext
    fillBufBuilderTwo :: ByteString -> BufferWriter -> DynaNext
    fillBufBuilderTwo ByteString
    bs BufferWriter
    writer Buffer
    buf0 Int
    room
        | ByteString -> Int
    BS.length ByteString
    bs Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    <= Int
    room = do
            buf1 <- Buffer -> ByteString -> IO Buffer
    copy Buffer
    buf0 ByteString
    bs
            let len1 = ByteString -> Int
    BS.length ByteString
    bs
            (len2, signal) <- writer buf1 (room - len1)
            return $ nextForBuilder (len1 + len2) signal
        | Bool
    otherwise = do
            let (ByteString
    bs1, ByteString
    bs2) = Int -> ByteString -> (ByteString, ByteString)
    BS.splitAt Int
    room ByteString
    bs
            IO Buffer -> IO ()
    forall (f :: * -> *) a. Functor f => f a -> f ()
    void (IO Buffer -> IO ()) -> IO Buffer -> IO ()
    forall a b. (a -> b) -> a -> b
    $ Buffer -> ByteString -> IO Buffer
    copy Buffer
    buf0 ByteString
    bs1
            Next -> IO Next
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return (Next -> IO Next) -> Next -> IO Next
    forall a b. (a -> b) -> a -> b
    $ Int -> Next -> Next
    nextForBuilder Int
    room (ByteString -> BufferWriter -> Next
    B.Chunk ByteString
    bs2 BufferWriter
    writer)

    nextForBuilder :: BytesFilled -> B.Next -> Next
    nextForBuilder :: Int -> Next -> Next
    nextForBuilder Int
    len Next
    B.Done =
        Int -> Bool -> Maybe DynaNext -> Next
    Next Int
    len Bool
    True Maybe DynaNext
    forall a. Maybe a
    Nothing -- let's flush
    nextForBuilder Int
    len (B.More Int
    minReq BufferWriter
    writer) =
        Int -> Bool -> Maybe DynaNext -> Next
    Next Int
    len Bool
    False (Maybe DynaNext -> Next) -> Maybe DynaNext -> Next
    forall a b. (a -> b) -> a -> b
    $ DynaNext -> Maybe DynaNext
    forall a. a -> Maybe a
    Just (Int -> BufferWriter -> DynaNext
    fillBufBuilderOne Int
    minReq BufferWriter
    writer)
    nextForBuilder Int
    len (B.Chunk ByteString
    bs BufferWriter
    writer) =
        Int -> Bool -> Maybe DynaNext -> Next
    Next Int
    len Bool
    False (Maybe DynaNext -> Next) -> Maybe DynaNext -> Next
    forall a b. (a -> b) -> a -> b
    $ DynaNext -> Maybe DynaNext
    forall a. a -> Maybe a
    Just (ByteString -> BufferWriter -> DynaNext
    fillBufBuilderTwo ByteString
    bs BufferWriter
    writer)

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

    -- | Like 'DynaNext', but with additional argument indicating total bytes written
    --
    -- Since @http2@ uses @DynaNext@ to construct a /single/ @DATA@ frame, the
    -- \"total number of bytes written\" refers to the current size of the payload
    -- of that @DATA@ frame.
    type NextWithTotal = Int -> DynaNext

    -- | Run the chunk, then continue as specified, unless streaming is finished
    runStreamingChunk :: StreamingChunk -> NextWithTotal -> NextWithTotal
    runStreamingChunk :: StreamingChunk -> NextWithTotal -> NextWithTotal
    runStreamingChunk StreamingChunk
    chunk NextWithTotal
    next =
        case StreamingChunk
    chunk of
            StreamingFinished Maybe (IO ())
    mdec -> Maybe (IO ()) -> NextWithTotal
    finished Maybe (IO ())
    mdec
            StreamingCancelled Maybe SomeException
    mErr -> Maybe SomeException -> NextWithTotal
    cancel Maybe SomeException
    mErr
            StreamingChunk
    StreamingFlush -> NextWithTotal
    flush
            StreamingBuilder Builder
    builder IsEndOfStream
    NotEndOfStream -> Builder -> NextWithTotal -> NextWithTotal
    runStreamingBuilder Builder
    builder NextWithTotal
    next
            StreamingBuilder Builder
    builder (EndOfStream Maybe (IO ())
    mdec) -> Builder -> NextWithTotal -> NextWithTotal
    runStreamingBuilder Builder
    builder (Maybe (IO ()) -> NextWithTotal
    finished Maybe (IO ())
    mdec)
      where
        finished :: Maybe CleanupStream -> NextWithTotal
        finished :: Maybe (IO ()) -> NextWithTotal
    finished Maybe (IO ())
    mdec = \Int
    total Buffer
    _buf Int
    _room -> do
            IO () -> Maybe (IO ()) -> IO ()
    forall a. a -> Maybe a -> a
    fromMaybe (() -> IO ()
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return ()) Maybe (IO ())
    mdec
            Next -> IO Next
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return (Next -> IO Next) -> Next -> IO Next
    forall a b. (a -> b) -> a -> b
    $ Int -> Bool -> Maybe DynaNext -> Next
    Next Int
    total Bool
    True Maybe DynaNext
    forall a. Maybe a
    Nothing

        flush :: NextWithTotal
        flush :: NextWithTotal
    flush = \Int
    total Buffer
    _buf Int
    _room -> do
            Next -> IO Next
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return (Next -> IO Next) -> Next -> IO Next
    forall a b. (a -> b) -> a -> b
    $ Int -> Bool -> Maybe DynaNext -> Next
    Next Int
    total Bool
    True (DynaNext -> Maybe DynaNext
    forall a. a -> Maybe a
    Just (DynaNext -> Maybe DynaNext) -> DynaNext -> Maybe DynaNext
    forall a b. (a -> b) -> a -> b
    $ NextWithTotal
    next Int
    0)

        -- Cancel streaming
        --
        -- The @_total@ number of bytes written refers to the @DATA@ frame currently
        -- under construction, but not yet sent (see discussion at 'DynaNext' and
        -- 'NextWithTotal'). Moreover, the documentation of 'outBodyCancel'
        -- explicitly states that such a partially constructed frame, if one exists,
        -- will be discarded on cancellation. We can therefore simply ignore
        -- @_total@ here.
        cancel :: Maybe SomeException -> NextWithTotal
        cancel :: Maybe SomeException -> NextWithTotal
    cancel Maybe SomeException
    mErr = \Int
    _total Buffer
    _buf Int
    _room -> Next -> IO Next
    forall a. a -> IO a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Next -> IO Next) -> Next -> IO Next
    forall a b. (a -> b) -> a -> b
    $ Maybe SomeException -> Next
    CancelNext Maybe SomeException
    mErr

    -- | Run 'Builder' until completion, then continue as specified
    runStreamingBuilder :: Builder -> NextWithTotal -> NextWithTotal
    runStreamingBuilder :: Builder -> NextWithTotal -> NextWithTotal
    runStreamingBuilder Builder
    builder NextWithTotal
    next = \Int
    total Buffer
    buf Int
    room -> do
        writeResult <- Builder -> BufferWriter
    B.runBuilder Builder
    builder Buffer
    buf Int
    room
        ranWriter writeResult total buf room
      where
        ranWriter :: (Int, B.Next) -> NextWithTotal
        ranWriter :: (Int, Next) -> NextWithTotal
    ranWriter (Int
    len, Next
    signal) = \Int
    total Buffer
    buf Int
    room -> do
            let total' :: Int
    total' = Int
    total Int -> Int -> Int
    forall a. Num a => a -> a -> a
    + Int
    len
            case Next
    signal of
                Next
    B.Done ->
                    NextWithTotal
    next Int
    total' (Buffer
    buf Buffer -> Int -> Buffer
    forall a b. Ptr a -> Int -> Ptr b
    `plusPtr` Int
    len) (Int
    room Int -> Int -> Int
    forall a. Num a => a -> a -> a
    - Int
    len)
                B.More Int
    minReq BufferWriter
    writer ->
                    Next -> IO Next
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return (Next -> IO Next) -> Next -> IO Next
    forall a b. (a -> b) -> a -> b
    $ Int -> Bool -> Maybe DynaNext -> Next
    Next Int
    total' Bool
    False (DynaNext -> Maybe DynaNext
    forall a. a -> Maybe a
    Just (DynaNext -> Maybe DynaNext) -> DynaNext -> Maybe DynaNext
    forall a b. (a -> b) -> a -> b
    $ Maybe Int -> BufferWriter -> NextWithTotal
    goMore (Int -> Maybe Int
    forall a. a -> Maybe a
    Just Int
    minReq) BufferWriter
    writer Int
    0)
                B.Chunk ByteString
    bs BufferWriter
    writer ->
                    Next -> IO Next
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return (Next -> IO Next) -> Next -> IO Next
    forall a b. (a -> b) -> a -> b
    $ Int -> Bool -> Maybe DynaNext -> Next
    Next Int
    total' Bool
    False (DynaNext -> Maybe DynaNext
    forall a. a -> Maybe a
    Just (DynaNext -> Maybe DynaNext) -> DynaNext -> Maybe DynaNext
    forall a b. (a -> b) -> a -> b
    $ ByteString -> BufferWriter -> NextWithTotal
    goChunk ByteString
    bs BufferWriter
    writer Int
    0)

        goMore :: Maybe Int -> B.BufferWriter -> NextWithTotal
        goMore :: Maybe Int -> BufferWriter -> NextWithTotal
    goMore Maybe Int
    mMinReq BufferWriter
    writer = \Int
    total Buffer
    buf Int
    room -> do
            let enoughRoom :: Bool
    enoughRoom = Bool -> (Int -> Bool) -> Maybe Int -> Bool
    forall b a. b -> (a -> b) -> Maybe a -> b
    maybe Bool
    True (Int
    room Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    >=) Maybe Int
    mMinReq
            if Bool
    enoughRoom
                then do
                    writeResult <- BufferWriter
    writer Buffer
    buf Int
    room
                    ranWriter writeResult total buf room
                else do
                    Next -> IO Next
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return (Next -> IO Next) -> Next -> IO Next
    forall a b. (a -> b) -> a -> b
    $ Int -> Bool -> Maybe DynaNext -> Next
    Next Int
    total Bool
    True (DynaNext -> Maybe DynaNext
    forall a. a -> Maybe a
    Just (DynaNext -> Maybe DynaNext) -> DynaNext -> Maybe DynaNext
    forall a b. (a -> b) -> a -> b
    $ Maybe Int -> BufferWriter -> NextWithTotal
    goMore Maybe Int
    mMinReq BufferWriter
    writer Int
    0)

        goChunk :: ByteString -> B.BufferWriter -> NextWithTotal
        goChunk :: ByteString -> BufferWriter -> NextWithTotal
    goChunk ByteString
    bs BufferWriter
    writer = \Int
    total Buffer
    buf Int
    room ->
            if ByteString -> Int
    BS.length ByteString
    bs Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    <= Int
    room
                then do
                    buf' <- Buffer -> ByteString -> IO Buffer
    copy Buffer
    buf ByteString
    bs
                    let len = ByteString -> Int
    BS.length ByteString
    bs
                    goMore Nothing writer (total + len) buf' (room - len)
                else do
                    let (ByteString
    bs1, ByteString
    bs2) = Int -> ByteString -> (ByteString, ByteString)
    BS.splitAt Int
    room ByteString
    bs
                    IO Buffer -> IO ()
    forall (f :: * -> *) a. Functor f => f a -> f ()
    void (IO Buffer -> IO ()) -> IO Buffer -> IO ()
    forall a b. (a -> b) -> a -> b
    $ Buffer -> ByteString -> IO Buffer
    copy Buffer
    buf ByteString
    bs1
                    Next -> IO Next
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return (Next -> IO Next) -> Next -> IO Next
    forall a b. (a -> b) -> a -> b
    $ Int -> Bool -> Maybe DynaNext -> Next
    Next (Int
    total Int -> Int -> Int
    forall a. Num a => a -> a -> a
    + Int
    room) Bool
    False (DynaNext -> Maybe DynaNext
    forall a. a -> Maybe a
    Just (DynaNext -> Maybe DynaNext) -> DynaNext -> Maybe DynaNext
    forall a b. (a -> b) -> a -> b
    $ ByteString -> BufferWriter -> NextWithTotal
    goChunk ByteString
    bs2 BufferWriter
    writer Int
    0)

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

    fillBufFile :: PositionRead -> FileOffset -> ByteCount -> Sentinel -> DynaNext
    fillBufFile :: PositionRead -> Int64 -> Int64 -> Sentinel -> DynaNext
    fillBufFile PositionRead
    pread Int64
    start Int64
    bytes Sentinel
    sentinel Buffer
    buf Int
    room = do
        len <- PositionRead
    pread Int64
    start (Int -> Int64 -> Int64
    mini Int
    room Int64
    bytes) Buffer
    buf
        case sentinel of
            Refresher IO ()
    refresh -> IO ()
    refresh
            Sentinel
    _ -> () -> IO ()
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return ()
        let len' = Int64 -> Int
    forall a b. (Integral a, Num b) => a -> b
    fromIntegral Int64
    len
        nextForFile len' pread (start + len) (bytes - len) sentinel

    nextForFile
        :: BytesFilled -> PositionRead -> FileOffset -> ByteCount -> Sentinel -> IO Next
    nextForFile :: Int -> PositionRead -> Int64 -> Int64 -> Sentinel -> IO Next
    nextForFile Int
    0 PositionRead
    _ Int64
    _ Int64
    _ Sentinel
    _ = Next -> IO Next
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return (Next -> IO Next) -> Next -> IO Next
    forall a b. (a -> b) -> a -> b
    $ Int -> Bool -> Maybe DynaNext -> Next
    Next Int
    0 Bool
    True Maybe DynaNext
    forall a. Maybe a
    Nothing -- let's flush
    nextForFile Int
    len PositionRead
    _ Int64
    _ Int64
    0 Sentinel
    sentinel = do
        case Sentinel
    sentinel of
            Closer IO ()
    close -> IO ()
    close
            Sentinel
    _ -> () -> IO ()
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return ()
        Next -> IO Next
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return (Next -> IO Next) -> Next -> IO Next
    forall a b. (a -> b) -> a -> b
    $ Int -> Bool -> Maybe DynaNext -> Next
    Next Int
    len Bool
    False Maybe DynaNext
    forall a. Maybe a
    Nothing
    nextForFile Int
    len PositionRead
    pread Int64
    start Int64
    bytes Sentinel
    refresh =
        Next -> IO Next
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return (Next -> IO Next) -> Next -> IO Next
    forall a b. (a -> b) -> a -> b
    $ Int -> Bool -> Maybe DynaNext -> Next
    Next Int
    len Bool
    False (Maybe DynaNext -> Next) -> Maybe DynaNext -> Next
    forall a b. (a -> b) -> a -> b
    $ DynaNext -> Maybe DynaNext
    forall a. a -> Maybe a
    Just (DynaNext -> Maybe DynaNext) -> DynaNext -> Maybe DynaNext
    forall a b. (a -> b) -> a -> b
    $ PositionRead -> Int64 -> Int64 -> Sentinel -> DynaNext
    fillBufFile PositionRead
    pread Int64
    start Int64
    bytes Sentinel
    refresh

    {-# INLINE mini #-}
    mini :: Int -> Int64 -> Int64
    mini :: Int -> Int64 -> Int64
    mini Int
    i Int64
    n
        | Int -> Int64
    forall a b. (Integral a, Num b) => a -> b
    fromIntegral Int
    i Int64 -> Int64 -> Bool
    forall a. Ord a => a -> a -> Bool
    < Int64
    n = Int -> Int64
    forall a b. (Integral a, Num b) => a -> b
    fromIntegral Int
    i
        | Bool
    otherwise = Int64
    n