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

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

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 Foreign.Ptr (plusPtr)
import Network.ByteOrder
import Network.HTTP.Semantics.Client

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

-- type DynaNext = Buffer -> BufferSize -> WindowSize -> IO Next
type DynaNext = Buffer -> Int -> IO Next

type BytesFilled = Int

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

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

data StreamingChunk
    = -- | Indicate that the stream is finished
      StreamingFinished CleanupStream
    | -- | 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
      --
      -- The optional 'CleanupStream' argument can be used to ensure that the
      -- final DATA frame in the stream is marked as end-of-stream, as opposed
      -- to using a separate, /empty/, data frame with this flag set.
      StreamingBuilder Builder (Maybe CleanupStream)

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

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

fillBuilderBodyGetNext :: Builder -> DynaNext
fillBuilderBodyGetNext :: Builder -> DynaNext
fillBuilderBodyGetNext Builder
bb Buffer
buf Int
room = do
    (Int
len, Next
signal) <- Builder -> BufferWriter
B.runBuilder Builder
bb Buffer
buf Int
room
    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
len Next
signal

fillFileBodyGetNext
    :: PositionRead -> FileOffset -> ByteCount -> IO () -> DynaNext
fillFileBodyGetNext :: PositionRead -> Int64 -> Int64 -> IO () -> DynaNext
fillFileBodyGetNext PositionRead
pread Int64
start Int64
bytecount IO ()
refresh Buffer
buf Int
room = do
    Int64
len <- PositionRead
pread Int64
start (Int -> Int64 -> Int64
mini Int
room Int64
bytecount) Buffer
buf
    let len' :: Int
len' = Int64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int64
len
    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 -> PositionRead -> Int64 -> Int64 -> IO () -> Next
nextForFile Int
len' PositionRead
pread (Int64
start Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
+ Int64
len) (Int64
bytecount Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- Int64
len) IO ()
refresh

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
        Maybe StreamingChunk
mChunk <- IO (Maybe StreamingChunk)
takeQ
        case Maybe StreamingChunk
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
          (Int
len, Next
signal) <- BufferWriter
writer Buffer
buf0 Int
room
          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
len Next
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
        Buffer
buf1 <- Buffer -> ByteString -> IO Buffer
copy Buffer
buf0 ByteString
bs
        let len1 :: Int
len1 = ByteString -> Int
BS.length ByteString
bs
        (Int
len2, Next
signal) <- BufferWriter
writer Buffer
buf1 (Int
room Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
len1)
        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
len1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len2) Next
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
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 IO ()
dec -> IO () -> NextWithTotal
finished IO ()
dec
        StreamingChunk
StreamingFlush -> NextWithTotal
flush
        StreamingBuilder Builder
builder Maybe (IO ())
Nothing -> Builder -> NextWithTotal -> NextWithTotal
runStreamingBuilder Builder
builder NextWithTotal
next
        StreamingBuilder Builder
builder (Just IO ()
dec) -> Builder -> NextWithTotal -> NextWithTotal
runStreamingBuilder Builder
builder (IO () -> NextWithTotal
finished IO ()
dec)
  where
    finished :: CleanupStream -> NextWithTotal
    finished :: IO () -> NextWithTotal
finished IO ()
dec = \Int
total Buffer
_buf Int
_room -> do
        IO ()
dec
        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)

-- | 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
    (Int, Next)
writeResult <- Builder -> BufferWriter
B.runBuilder Builder
builder Buffer
buf Int
room
    (Int, Next) -> NextWithTotal
ranWriter (Int, Next)
writeResult Int
total Buffer
buf Int
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
              (Int, Next)
writeResult <- BufferWriter
writer Buffer
buf Int
room
              (Int, Next) -> NextWithTotal
ranWriter (Int, Next)
writeResult Int
total Buffer
buf Int
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
                Buffer
buf' <- Buffer -> ByteString -> IO Buffer
copy Buffer
buf ByteString
bs
                let len :: Int
len = ByteString -> Int
BS.length ByteString
bs
                Maybe Int -> BufferWriter -> NextWithTotal
goMore Maybe Int
forall a. Maybe a
Nothing BufferWriter
writer (Int
total Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len) Buffer
buf' (Int
room Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
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 -> IO () -> DynaNext
fillBufFile :: PositionRead -> Int64 -> Int64 -> IO () -> DynaNext
fillBufFile PositionRead
pread Int64
start Int64
bytes IO ()
refresh Buffer
buf Int
room = do
    Int64
len <- PositionRead
pread Int64
start (Int -> Int64 -> Int64
mini Int
room Int64
bytes) Buffer
buf
    IO ()
refresh
    let len' :: Int
len' = Int64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int64
len
    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 -> PositionRead -> Int64 -> Int64 -> IO () -> Next
nextForFile Int
len' PositionRead
pread (Int64
start Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
+ Int64
len) (Int64
bytes Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- Int64
len) IO ()
refresh

nextForFile
    :: BytesFilled -> PositionRead -> FileOffset -> ByteCount -> IO () -> Next
nextForFile :: Int -> PositionRead -> Int64 -> Int64 -> IO () -> Next
nextForFile Int
0 PositionRead
_ Int64
_ Int64
_ IO ()
_ = 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 IO ()
_ = 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 IO ()
refresh =
    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 -> IO () -> DynaNext
fillBufFile PositionRead
pread Int64
start Int64
bytes IO ()
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