{-# LANGUAGE Unsafe #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE NoMonoLocalBinds #-}

{-# OPTIONS_HADDOCK not-home #-}

-- | Copyright : (c) 2010 - 2011 Simon Meier
-- License     : BSD3-style (see LICENSE)
--
-- Maintainer  : Simon Meier <iridcode@gmail.com>
-- Stability   : unstable, private
-- Portability : GHC
--
-- *Warning:* this module is internal. If you find that you need it then please
-- contact the maintainers and explain what you are trying to do and discuss
-- what you would need in the public API. It is important that you do this as
-- the module may not be exposed at all in future releases.
--
-- Core types and functions for the 'Builder' monoid and its generalization,
-- the 'Put' monad.
--
-- The design of the 'Builder' monoid is optimized such that
--
--   1. buffers of arbitrary size can be filled as efficiently as possible and
--
--   2. sequencing of 'Builder's is as cheap as possible.
--
-- We achieve (1) by completely handing over control over writing to the buffer
-- to the 'BuildStep' implementing the 'Builder'. This 'BuildStep' is just told
-- the start and the end of the buffer (represented as a 'BufferRange'). Then,
-- the 'BuildStep' can write to as big a prefix of this 'BufferRange' in any
-- way it desires. If the 'BuildStep' is done, the 'BufferRange' is full, or a
-- long sequence of bytes should be inserted directly, then the 'BuildStep'
-- signals this to its caller using a 'BuildSignal'.
--
-- We achieve (2) by requiring that every 'Builder' is implemented by a
-- 'BuildStep' that takes a continuation 'BuildStep', which it calls with the
-- updated 'BufferRange' after it is done. Therefore, only two pointers have
-- to be passed in a function call to implement concatenation of 'Builder's.
-- Moreover, many 'Builder's are completely inlined, which enables the compiler
-- to sequence them without a function call and with no boxing at all.
--
-- This design gives the implementation of a 'Builder' full access to the 'IO'
-- monad. Therefore, utmost care has to be taken to not overwrite anything
-- outside the given 'BufferRange's. Moreover, further care has to be taken to
-- ensure that 'Builder's and 'Put's are referentially transparent. See the
-- comments of the 'builder' and 'put' functions for further information.
-- Note that there are /no safety belts/ at all, when implementing a 'Builder'
-- using an 'IO' action: you are writing code that might enable the next
-- buffer-overflow attack on a Haskell server!
--
module Data.ByteString.Builder.Internal (
  -- * Buffer management
    Buffer(..)
  , BufferRange(..)
  , newBuffer
  , bufferSize
  , byteStringFromBuffer

  , ChunkIOStream(..)
  , buildStepToCIOS
  , ciosUnitToLazyByteString
  , ciosToLazyByteString

  -- * Build signals and steps
  , BuildSignal
  , BuildStep
  , finalBuildStep

  , done
  , bufferFull
  , insertChunk

  , fillWithBuildStep

  -- * The Builder monoid
  , Builder
  , builder
  , runBuilder
  , runBuilderWith

  -- ** Primitive combinators
  , empty
  , append
  , flush
  , ensureFree
  -- , sizedChunksInsert

  , byteStringCopy
  , byteStringInsert
  , byteStringThreshold

  , lazyByteStringCopy
  , lazyByteStringInsert
  , lazyByteStringThreshold

  , shortByteString

  , maximalCopySize
  , byteString
  , lazyByteString

  -- ** Execution
  , toLazyByteString
  , toLazyByteStringWith
  , AllocationStrategy
  , safeStrategy
  , untrimmedStrategy
  , customStrategy
  , L.smallChunkSize
  , L.defaultChunkSize
  , L.chunkOverhead

  -- * The Put monad
  , Put
  , put
  , runPut

  -- ** Execution
  , putToLazyByteString
  , putToLazyByteStringWith
  , hPut

  -- ** Conversion to and from Builders
  , putBuilder
  , fromPut

  -- -- ** Lifting IO actions
  -- , putLiftIO

) where

import           Control.Arrow (second)
import           Control.DeepSeq (NFData(..))
import           GHC.Exts (IsList(..))

import           Data.Semigroup (Semigroup(..))
import           Data.List.NonEmpty (NonEmpty(..))

import qualified Data.ByteString               as S
import qualified Data.ByteString.Unsafe        as S
import qualified Data.ByteString.Internal.Type as S
import qualified Data.ByteString.Lazy.Internal as L
import qualified Data.ByteString.Short.Internal as Sh

import qualified GHC.IO.Buffer as IO (Buffer(..), newByteBuffer)
import           GHC.IO.Handle.Internals (wantWritableHandle, flushWriteBuffer)
import           GHC.IO.Handle.Types (Handle__, haByteBuffer, haBufferMode)
import           System.IO (hFlush, BufferMode(..), Handle)
import           Data.IORef

import           Foreign
import           Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
import           System.IO.Unsafe (unsafeDupablePerformIO)

------------------------------------------------------------------------------
-- Buffers
------------------------------------------------------------------------------
-- | A range of bytes in a buffer represented by the pointer to the first byte
-- of the range and the pointer to the first byte /after/ the range.
data BufferRange = BufferRange {-# UNPACK #-} !(Ptr Word8)  -- First byte of range
                               {-# UNPACK #-} !(Ptr Word8)  -- First byte /after/ range

-- | @since 0.12.2.0
instance NFData BufferRange where
  rnf :: BufferRange -> ()
rnf !BufferRange
_ = ()

-- | A 'Buffer' together with the 'BufferRange' of free bytes. The filled
-- space starts at offset 0 and ends at the first free byte.
data Buffer = Buffer {-# UNPACK #-} !(ForeignPtr Word8)
                     {-# UNPACK #-} !BufferRange

-- | Like the @NFData@ instance for @StrictByteString@,
-- this does not force the @ForeignPtrContents@ field
-- of the underlying @ForeignPtr@.
--
-- @since 0.12.2.0
instance NFData Buffer where
  rnf :: Buffer -> ()
rnf !Buffer
_ = ()

-- | Combined size of the filled and free space in the buffer.
{-# INLINE bufferSize #-}
bufferSize :: Buffer -> Int
bufferSize :: Buffer -> Int
bufferSize (Buffer ForeignPtr Word8
fpbuf (BufferRange Ptr Word8
_ Ptr Word8
ope)) =
    Ptr Word8
ope Ptr Word8 -> Ptr Word8 -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` ForeignPtr Word8 -> Ptr Word8
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr Word8
fpbuf

-- | Allocate a new buffer of the given size.
{-# INLINE newBuffer #-}
newBuffer :: Int -> IO Buffer
newBuffer :: Int -> IO Buffer
newBuffer Int
size = do
    ForeignPtr Word8
fpbuf <- Int -> IO (ForeignPtr Word8)
forall a. Int -> IO (ForeignPtr a)
S.mallocByteString Int
size
    let pbuf :: Ptr Word8
pbuf = ForeignPtr Word8 -> Ptr Word8
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr Word8
fpbuf
    Buffer -> IO Buffer
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Buffer -> IO Buffer) -> Buffer -> IO Buffer
forall a b. (a -> b) -> a -> b
$! ForeignPtr Word8 -> BufferRange -> Buffer
Buffer ForeignPtr Word8
fpbuf (Ptr Word8 -> Ptr Word8 -> BufferRange
BufferRange Ptr Word8
pbuf (Ptr Word8
pbuf Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
size))

-- | Convert the filled part of a 'Buffer' to a 'S.StrictByteString'.
{-# INLINE byteStringFromBuffer #-}
byteStringFromBuffer :: Buffer -> S.StrictByteString
byteStringFromBuffer :: Buffer -> StrictByteString
byteStringFromBuffer (Buffer ForeignPtr Word8
fpbuf (BufferRange Ptr Word8
op Ptr Word8
_)) =
    ForeignPtr Word8 -> Int -> StrictByteString
S.BS ForeignPtr Word8
fpbuf (Ptr Word8
op Ptr Word8 -> Ptr Word8 -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` ForeignPtr Word8 -> Ptr Word8
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr Word8
fpbuf)

-- | Prepend the filled part of a 'Buffer' to a 'L.LazyByteString'
-- trimming it if necessary.
{-# INLINE trimmedChunkFromBuffer #-}
trimmedChunkFromBuffer :: AllocationStrategy -> Buffer
                       -> L.LazyByteString -> L.LazyByteString
trimmedChunkFromBuffer :: AllocationStrategy -> Buffer -> LazyByteString -> LazyByteString
trimmedChunkFromBuffer (AllocationStrategy Maybe (Buffer, Int) -> IO Buffer
_ Int
_ Int -> Int -> Bool
trim) Buffer
buf LazyByteString
k
  | StrictByteString -> Bool
S.null StrictByteString
bs                           = LazyByteString
k
  | Int -> Int -> Bool
trim (StrictByteString -> Int
S.length StrictByteString
bs) (Buffer -> Int
bufferSize Buffer
buf) = StrictByteString -> LazyByteString -> LazyByteString
L.Chunk (StrictByteString -> StrictByteString
S.copy StrictByteString
bs) LazyByteString
k
  | Bool
otherwise                           = StrictByteString -> LazyByteString -> LazyByteString
L.Chunk StrictByteString
bs          LazyByteString
k
  where
    bs :: StrictByteString
bs = Buffer -> StrictByteString
byteStringFromBuffer Buffer
buf

------------------------------------------------------------------------------
-- Chunked IO Stream
------------------------------------------------------------------------------

-- | A stream of chunks that are constructed in the 'IO' monad.
--
-- This datatype serves as the common interface for the buffer-by-buffer
-- execution of a 'BuildStep' by 'buildStepToCIOS'. Typical users of this
-- interface are 'ciosToLazyByteString' or iteratee-style libraries like
-- @enumerator@.
data ChunkIOStream a =
       Finished Buffer a
       -- ^ The partially filled last buffer together with the result.
     | Yield1 S.StrictByteString (IO (ChunkIOStream a))
       -- ^ Yield a /non-empty/ 'S.StrictByteString'.

-- | A smart constructor for yielding one chunk that ignores the chunk if
-- it is empty.
{-# INLINE yield1 #-}
yield1 :: S.StrictByteString -> IO (ChunkIOStream a) -> IO (ChunkIOStream a)
yield1 :: forall a.
StrictByteString -> IO (ChunkIOStream a) -> IO (ChunkIOStream a)
yield1 StrictByteString
bs IO (ChunkIOStream a)
cios | StrictByteString -> Bool
S.null StrictByteString
bs = IO (ChunkIOStream a)
cios
               | Bool
otherwise = ChunkIOStream a -> IO (ChunkIOStream a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (ChunkIOStream a -> IO (ChunkIOStream a))
-> ChunkIOStream a -> IO (ChunkIOStream a)
forall a b. (a -> b) -> a -> b
$ StrictByteString -> IO (ChunkIOStream a) -> ChunkIOStream a
forall a.
StrictByteString -> IO (ChunkIOStream a) -> ChunkIOStream a
Yield1 StrictByteString
bs IO (ChunkIOStream a)
cios

-- | Convert a @'ChunkIOStream' ()@ to a 'L.LazyByteString' using
-- 'unsafeDupablePerformIO'.
{-# INLINE ciosUnitToLazyByteString #-}
ciosUnitToLazyByteString :: AllocationStrategy
                         -> L.LazyByteString -> ChunkIOStream () -> L.LazyByteString
ciosUnitToLazyByteString :: AllocationStrategy
-> LazyByteString -> ChunkIOStream () -> LazyByteString
ciosUnitToLazyByteString AllocationStrategy
strategy LazyByteString
k = ChunkIOStream () -> LazyByteString
forall {a}. ChunkIOStream a -> LazyByteString
go
  where
    go :: ChunkIOStream a -> LazyByteString
go (Finished Buffer
buf a
_) = AllocationStrategy -> Buffer -> LazyByteString -> LazyByteString
trimmedChunkFromBuffer AllocationStrategy
strategy Buffer
buf LazyByteString
k
    go (Yield1 StrictByteString
bs IO (ChunkIOStream a)
io)   = StrictByteString -> LazyByteString -> LazyByteString
L.Chunk StrictByteString
bs (LazyByteString -> LazyByteString)
-> LazyByteString -> LazyByteString
forall a b. (a -> b) -> a -> b
$ IO LazyByteString -> LazyByteString
forall a. IO a -> a
unsafeDupablePerformIO (ChunkIOStream a -> LazyByteString
go (ChunkIOStream a -> LazyByteString)
-> IO (ChunkIOStream a) -> IO LazyByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO (ChunkIOStream a)
io)

-- | Convert a 'ChunkIOStream' to a lazy tuple of the result and the written
-- 'L.LazyByteString' using 'unsafeDupablePerformIO'.
{-# INLINE ciosToLazyByteString #-}
ciosToLazyByteString :: AllocationStrategy
                     -> (a -> (b, L.LazyByteString))
                     -> ChunkIOStream a
                     -> (b, L.LazyByteString)
ciosToLazyByteString :: forall a b.
AllocationStrategy
-> (a -> (b, LazyByteString))
-> ChunkIOStream a
-> (b, LazyByteString)
ciosToLazyByteString AllocationStrategy
strategy a -> (b, LazyByteString)
k =
    ChunkIOStream a -> (b, LazyByteString)
go
  where
    go :: ChunkIOStream a -> (b, LazyByteString)
go (Finished Buffer
buf a
x) =
        (LazyByteString -> LazyByteString)
-> (b, LazyByteString) -> (b, LazyByteString)
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (AllocationStrategy -> Buffer -> LazyByteString -> LazyByteString
trimmedChunkFromBuffer AllocationStrategy
strategy Buffer
buf) ((b, LazyByteString) -> (b, LazyByteString))
-> (b, LazyByteString) -> (b, LazyByteString)
forall a b. (a -> b) -> a -> b
$ a -> (b, LazyByteString)
k a
x
    go (Yield1 StrictByteString
bs IO (ChunkIOStream a)
io)   = (LazyByteString -> LazyByteString)
-> (b, LazyByteString) -> (b, LazyByteString)
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (StrictByteString -> LazyByteString -> LazyByteString
L.Chunk StrictByteString
bs) ((b, LazyByteString) -> (b, LazyByteString))
-> (b, LazyByteString) -> (b, LazyByteString)
forall a b. (a -> b) -> a -> b
$ IO (b, LazyByteString) -> (b, LazyByteString)
forall a. IO a -> a
unsafeDupablePerformIO (ChunkIOStream a -> (b, LazyByteString)
go (ChunkIOStream a -> (b, LazyByteString))
-> IO (ChunkIOStream a) -> IO (b, LazyByteString)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO (ChunkIOStream a)
io)

------------------------------------------------------------------------------
-- Build signals
------------------------------------------------------------------------------

-- | 'BuildStep's may be called *multiple times* and they must not rise an
-- async. exception.
type BuildStep a = BufferRange -> IO (BuildSignal a)

-- | 'BuildSignal's abstract signals to the caller of a 'BuildStep'. There are
-- three signals: 'done', 'bufferFull', or 'insertChunks signals
data BuildSignal a =
    Done {-# UNPACK #-} !(Ptr Word8) a
  | BufferFull
      {-# UNPACK #-} !Int
      {-# UNPACK #-} !(Ptr Word8)
                     (BuildStep a)
  | InsertChunk
      {-# UNPACK #-} !(Ptr Word8)
                     S.StrictByteString
                     (BuildStep a)

-- | Signal that the current 'BuildStep' is done and has computed a value.
{-# INLINE done #-}
done :: Ptr Word8      -- ^ Next free byte in current 'BufferRange'
     -> a              -- ^ Computed value
     -> BuildSignal a
done :: forall a. Ptr Word8 -> a -> BuildSignal a
done = Ptr Word8 -> a -> BuildSignal a
forall a. Ptr Word8 -> a -> BuildSignal a
Done

-- | Signal that the current buffer is full.
{-# INLINE bufferFull #-}
bufferFull :: Int
           -- ^ Minimal size of next 'BufferRange'.
           -> Ptr Word8
           -- ^ Next free byte in current 'BufferRange'.
           -> BuildStep a
           -- ^ 'BuildStep' to run on the next 'BufferRange'. This 'BuildStep'
           -- may assume that it is called with a 'BufferRange' of at least the
           -- required minimal size; i.e., the caller of this 'BuildStep' must
           -- guarantee this.
           -> BuildSignal a
bufferFull :: forall a. Int -> Ptr Word8 -> BuildStep a -> BuildSignal a
bufferFull = Int -> Ptr Word8 -> BuildStep a -> BuildSignal a
forall a. Int -> Ptr Word8 -> BuildStep a -> BuildSignal a
BufferFull


-- | Signal that a 'S.StrictByteString' chunk should be inserted directly.
{-# INLINE insertChunk #-}
insertChunk :: Ptr Word8
            -- ^ Next free byte in current 'BufferRange'
            -> S.StrictByteString
            -- ^ Chunk to insert.
            -> BuildStep a
            -- ^ 'BuildStep' to run on next 'BufferRange'
            -> BuildSignal a
insertChunk :: forall a.
Ptr Word8 -> StrictByteString -> BuildStep a -> BuildSignal a
insertChunk = Ptr Word8 -> StrictByteString -> BuildStep a -> BuildSignal a
forall a.
Ptr Word8 -> StrictByteString -> BuildStep a -> BuildSignal a
InsertChunk


-- | Fill a 'BufferRange' using a 'BuildStep'.
{-# INLINE fillWithBuildStep #-}
fillWithBuildStep
    :: BuildStep a
    -- ^ Build step to use for filling the 'BufferRange'.
    -> (Ptr Word8 -> a -> IO b)
    -- ^ Handling the 'done' signal
    -> (Ptr Word8 -> Int -> BuildStep a -> IO b)
    -- ^ Handling the 'bufferFull' signal
    -> (Ptr Word8 -> S.StrictByteString -> BuildStep a -> IO b)
    -- ^ Handling the 'insertChunk' signal
    -> BufferRange
    -- ^ Buffer range to fill.
    -> IO b
    -- ^ Value computed while filling this 'BufferRange'.
fillWithBuildStep :: forall a b.
BuildStep a
-> (Ptr Word8 -> a -> IO b)
-> (Ptr Word8 -> Int -> BuildStep a -> IO b)
-> (Ptr Word8 -> StrictByteString -> BuildStep a -> IO b)
-> BufferRange
-> IO b
fillWithBuildStep BuildStep a
step Ptr Word8 -> a -> IO b
fDone Ptr Word8 -> Int -> BuildStep a -> IO b
fFull Ptr Word8 -> StrictByteString -> BuildStep a -> IO b
fChunk !BufferRange
br = do
    BuildSignal a
signal <- BuildStep a
step BufferRange
br
    case BuildSignal a
signal of
        Done Ptr Word8
op a
x                      -> Ptr Word8 -> a -> IO b
fDone Ptr Word8
op a
x
        BufferFull Int
minSize Ptr Word8
op BuildStep a
nextStep -> Ptr Word8 -> Int -> BuildStep a -> IO b
fFull Ptr Word8
op Int
minSize BuildStep a
nextStep
        InsertChunk Ptr Word8
op StrictByteString
bs BuildStep a
nextStep     -> Ptr Word8 -> StrictByteString -> BuildStep a -> IO b
fChunk Ptr Word8
op StrictByteString
bs BuildStep a
nextStep


------------------------------------------------------------------------------
-- The 'Builder' monoid
------------------------------------------------------------------------------

-- | 'Builder's denote sequences of bytes.
-- They are 'Monoid's where
--   'mempty' is the zero-length sequence and
--   'mappend' is concatenation, which runs in /O(1)/.
newtype Builder = Builder (forall r. BuildStep r -> BuildStep r)

-- | Construct a 'Builder'. In contrast to 'BuildStep's, 'Builder's are
-- referentially transparent.
{-# INLINE builder #-}
builder :: (forall r. BuildStep r -> BuildStep r)
        -- ^ A function that fills a 'BufferRange', calls the continuation with
        -- the updated 'BufferRange' once its done, and signals its caller how
        -- to proceed using 'done', 'bufferFull', or 'insertChunk'.
        --
        -- This function must be referentially transparent; i.e., calling it
        -- multiple times with equally sized 'BufferRange's must result in the
        -- same sequence of bytes being written. If you need mutable state,
        -- then you must allocate it anew upon each call of this function.
        -- Moreover, this function must call the continuation once its done.
        -- Otherwise, concatenation of 'Builder's does not work. Finally, this
        -- function must write to all bytes that it claims it has written.
        -- Otherwise, the resulting 'Builder' is not guaranteed to be
        -- referentially transparent and sensitive data might leak.
        -> Builder
builder :: (forall r. BuildStep r -> BuildStep r) -> Builder
builder = (forall r. BuildStep r -> BuildStep r) -> Builder
Builder

-- | The final build step that returns the 'done' signal.
finalBuildStep :: BuildStep ()
finalBuildStep :: BuildStep ()
finalBuildStep (BufferRange Ptr Word8
op Ptr Word8
_) = BuildSignal () -> IO (BuildSignal ())
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (BuildSignal () -> IO (BuildSignal ()))
-> BuildSignal () -> IO (BuildSignal ())
forall a b. (a -> b) -> a -> b
$ Ptr Word8 -> () -> BuildSignal ()
forall a. Ptr Word8 -> a -> BuildSignal a
Done Ptr Word8
op ()

-- | Run a 'Builder' with the 'finalBuildStep'.
{-# INLINE runBuilder #-}
runBuilder :: Builder      -- ^ 'Builder' to run
           -> BuildStep () -- ^ 'BuildStep' that writes the byte stream of this
                           -- 'Builder' and signals 'done' upon completion.
runBuilder :: Builder -> BuildStep ()
runBuilder Builder
b = Builder -> BuildStep () -> BuildStep ()
forall a. Builder -> BuildStep a -> BuildStep a
runBuilderWith Builder
b BuildStep ()
finalBuildStep

-- | Run a 'Builder'.
{-# INLINE runBuilderWith #-}
runBuilderWith :: Builder      -- ^ 'Builder' to run
               -> BuildStep a -- ^ Continuation 'BuildStep'
               -> BuildStep a
runBuilderWith :: forall a. Builder -> BuildStep a -> BuildStep a
runBuilderWith (Builder forall r. BuildStep r -> BuildStep r
b) = BuildStep a -> BuildStep a
forall r. BuildStep r -> BuildStep r
b

-- | The 'Builder' denoting a zero-length sequence of bytes. This function is
-- only exported for use in rewriting rules. Use 'mempty' otherwise.
{-# INLINE[1] empty #-}
empty :: Builder
empty :: Builder
empty = (forall r. BuildStep r -> BuildStep r) -> Builder
Builder (\BuildStep r
k BufferRange
br -> BuildStep r
k BufferRange
br)
-- This eta expansion (hopefully) allows GHC to worker-wrapper the
-- 'BufferRange' in the 'empty' base case of loops (since
-- worker-wrapper requires (TODO: verify this) that all paths match
-- against the wrapped argument.
--
-- Do not use ($), which has arity 1 since base-4.19.
-- See also https://gitlab.haskell.org/ghc/ghc/-/issues/23822

-- | Concatenate two 'Builder's. This function is only exported for use in rewriting
-- rules. Use 'mappend' otherwise.
{-# INLINE[1] append #-}
append :: Builder -> Builder -> Builder
append :: Builder -> Builder -> Builder
append (Builder forall r. BuildStep r -> BuildStep r
b1) (Builder forall r. BuildStep r -> BuildStep r
b2) = (forall r. BuildStep r -> BuildStep r) -> Builder
Builder ((forall r. BuildStep r -> BuildStep r) -> Builder)
-> (forall r. BuildStep r -> BuildStep r) -> Builder
forall a b. (a -> b) -> a -> b
$ BuildStep r -> BuildStep r
forall r. BuildStep r -> BuildStep r
b1 (BuildStep r -> BuildStep r)
-> (BuildStep r -> BuildStep r) -> BuildStep r -> BuildStep r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BuildStep r -> BuildStep r
forall r. BuildStep r -> BuildStep r
b2

stimesBuilder :: Integral t => t -> Builder -> Builder
{-# INLINABLE stimesBuilder #-}
stimesBuilder :: forall t. Integral t => t -> Builder -> Builder
stimesBuilder t
n Builder
b
  | t
n t -> t -> Bool
forall a. Ord a => a -> a -> Bool
>= t
0 = t -> Builder
forall {t}. (Eq t, Num t) => t -> Builder
go t
n
  | Bool
otherwise = Builder
stimesNegativeErr
  where go :: t -> Builder
go t
0 = Builder
empty
        go t
k = Builder
b Builder -> Builder -> Builder
`append` t -> Builder
go (t
k t -> t -> t
forall a. Num a => a -> a -> a
- t
1)

stimesNegativeErr :: Builder
-- See Note [Float error calls out of INLINABLE things]
-- in Data.ByteString.Internal.Type
stimesNegativeErr :: Builder
stimesNegativeErr
  = [Char] -> Builder
forall a. [Char] -> a
errorWithoutStackTrace [Char]
"stimes @Builder: non-negative multiplier expected"

instance Semigroup Builder where
  {-# INLINE (<>) #-}
  <> :: Builder -> Builder -> Builder
(<>) = Builder -> Builder -> Builder
append
  sconcat :: NonEmpty Builder -> Builder
sconcat (Builder
b:|[Builder]
bs) = Builder
b Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> (Builder -> Builder -> Builder) -> Builder -> [Builder] -> Builder
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Builder -> Builder -> Builder
forall a. Monoid a => a -> a -> a
mappend Builder
forall a. Monoid a => a
mempty [Builder]
bs
  {-# INLINE stimes #-}
  stimes :: forall t. Integral t => t -> Builder -> Builder
stimes = b -> Builder -> Builder
forall t. Integral t => t -> Builder -> Builder
stimesBuilder

instance Monoid Builder where
  {-# INLINE mempty #-}
  mempty :: Builder
mempty = Builder
empty
  {-# INLINE mappend #-}
  mappend :: Builder -> Builder -> Builder
mappend = Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
(<>)
  {-# INLINE mconcat #-}
  mconcat :: [Builder] -> Builder
mconcat = (Builder -> Builder -> Builder) -> Builder -> [Builder] -> Builder
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Builder -> Builder -> Builder
forall a. Monoid a => a -> a -> a
mappend Builder
forall a. Monoid a => a
mempty

-- | For long or infinite lists use 'fromList' because it uses 'LazyByteString' otherwise use 'fromListN' which uses 'StrictByteString'.
instance IsList Builder where
  type Item Builder = Word8
  fromList :: [Item Builder] -> Builder
fromList = LazyByteString -> Builder
lazyByteString (LazyByteString -> Builder)
-> ([Word8] -> LazyByteString) -> [Word8] -> Builder
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Word8] -> LazyByteString
[Item LazyByteString] -> LazyByteString
forall l. IsList l => [Item l] -> l
fromList
  fromListN :: Int -> [Item Builder] -> Builder
fromListN Int
n = StrictByteString -> Builder
byteString (StrictByteString -> Builder)
-> ([Word8] -> StrictByteString) -> [Word8] -> Builder
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [Item StrictByteString] -> StrictByteString
forall l. IsList l => Int -> [Item l] -> l
fromListN Int
n
  toList :: Builder -> [Item Builder]
toList = LazyByteString -> [Word8]
LazyByteString -> [Item LazyByteString]
forall l. IsList l => l -> [Item l]
toList (LazyByteString -> [Word8])
-> (Builder -> LazyByteString) -> Builder -> [Word8]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Builder -> LazyByteString
toLazyByteString

-- | Flush the current buffer. This introduces a chunk boundary.
{-# INLINE flush #-}
flush :: Builder
flush :: Builder
flush = (forall r. BuildStep r -> BuildStep r) -> Builder
builder BuildStep r -> BuildStep r
forall r. BuildStep r -> BuildStep r
forall {m :: * -> *} {a}.
Monad m =>
BuildStep a -> BufferRange -> m (BuildSignal a)
step
  where
    step :: BuildStep a -> BufferRange -> m (BuildSignal a)
step BuildStep a
k (BufferRange Ptr Word8
op Ptr Word8
_) = BuildSignal a -> m (BuildSignal a)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (BuildSignal a -> m (BuildSignal a))
-> BuildSignal a -> m (BuildSignal a)
forall a b. (a -> b) -> a -> b
$ Ptr Word8 -> StrictByteString -> BuildStep a -> BuildSignal a
forall a.
Ptr Word8 -> StrictByteString -> BuildStep a -> BuildSignal a
insertChunk Ptr Word8
op StrictByteString
S.empty BuildStep a
k


------------------------------------------------------------------------------
-- Put
------------------------------------------------------------------------------

-- | A 'Put' action denotes a computation of a value that writes a stream of
-- bytes as a side-effect. 'Put's are strict in their side-effect; i.e., the
-- stream of bytes will always be written before the computed value is
-- returned.
--
-- 'Put's are a generalization of 'Builder's. The typical use case is the
-- implementation of an encoding that might fail (e.g., an interface to the
-- <https://hackage.haskell.org/package/zlib zlib>
-- compression library or the conversion from Base64 encoded data to
-- 8-bit data). For a 'Builder', the only way to handle and report such a
-- failure is ignore it or call 'error'.  In contrast, 'Put' actions are
-- expressive enough to allow reporting and handling such a failure in a pure
-- fashion.
--
-- @'Put' ()@ actions are isomorphic to 'Builder's. The functions 'putBuilder'
-- and 'fromPut' convert between these two types. Where possible, you should
-- use 'Builder's, as sequencing them is slightly cheaper than sequencing
-- 'Put's because they do not carry around a computed value.
newtype Put a = Put { forall a. Put a -> forall r. (a -> BuildStep r) -> BuildStep r
unPut :: forall r. (a -> BuildStep r) -> BuildStep r }

-- | Construct a 'Put' action. In contrast to 'BuildStep's, 'Put's are
-- referentially transparent in the sense that sequencing the same 'Put'
-- multiple times yields every time the same value with the same side-effect.
{-# INLINE put #-}
put :: (forall r. (a -> BuildStep r) -> BuildStep r)
       -- ^ A function that fills a 'BufferRange', calls the continuation with
       -- the updated 'BufferRange' and its computed value once its done, and
       -- signals its caller how to proceed using 'done', 'bufferFull', or
       -- 'insertChunk' signals.
       --
    -- This function must be referentially transparent; i.e., calling it
    -- multiple times with equally sized 'BufferRange's must result in the
    -- same sequence of bytes being written and the same value being
    -- computed. If you need mutable state, then you must allocate it anew
    -- upon each call of this function. Moreover, this function must call
    -- the continuation once its done. Otherwise, monadic sequencing of
    -- 'Put's does not work. Finally, this function must write to all bytes
    -- that it claims it has written. Otherwise, the resulting 'Put' is
    -- not guaranteed to be referentially transparent and sensitive data
    -- might leak.
       -> Put a
put :: forall a. (forall r. (a -> BuildStep r) -> BuildStep r) -> Put a
put = (forall r. (a -> BuildStep r) -> BuildStep r) -> Put a
forall a. (forall r. (a -> BuildStep r) -> BuildStep r) -> Put a
Put

-- | Run a 'Put'.
{-# INLINE runPut #-}
runPut :: Put a       -- ^ Put to run
       -> BuildStep a -- ^ 'BuildStep' that first writes the byte stream of
                      -- this 'Put' and then yields the computed value using
                      -- the 'done' signal.
runPut :: forall a. Put a -> BuildStep a
runPut (Put forall r. (a -> BuildStep r) -> BuildStep r
p) = (a -> BuildStep a) -> BuildStep a
forall r. (a -> BuildStep r) -> BuildStep r
p ((a -> BuildStep a) -> BuildStep a)
-> (a -> BuildStep a) -> BuildStep a
forall a b. (a -> b) -> a -> b
$ \a
x (BufferRange Ptr Word8
op Ptr Word8
_) -> BuildSignal a -> IO (BuildSignal a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (BuildSignal a -> IO (BuildSignal a))
-> BuildSignal a -> IO (BuildSignal a)
forall a b. (a -> b) -> a -> b
$ Ptr Word8 -> a -> BuildSignal a
forall a. Ptr Word8 -> a -> BuildSignal a
Done Ptr Word8
op a
x

instance Functor Put where
  fmap :: forall a b. (a -> b) -> Put a -> Put b
fmap a -> b
f Put a
p = (forall r. (b -> BuildStep r) -> BuildStep r) -> Put b
forall a. (forall r. (a -> BuildStep r) -> BuildStep r) -> Put a
Put ((forall r. (b -> BuildStep r) -> BuildStep r) -> Put b)
-> (forall r. (b -> BuildStep r) -> BuildStep r) -> Put b
forall a b. (a -> b) -> a -> b
$ \b -> BuildStep r
k -> Put a -> forall r. (a -> BuildStep r) -> BuildStep r
forall a. Put a -> forall r. (a -> BuildStep r) -> BuildStep r
unPut Put a
p (b -> BuildStep r
k (b -> BuildStep r) -> (a -> b) -> a -> BuildStep r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> b
f)
  {-# INLINE fmap #-}

-- | Synonym for '<*' from 'Applicative'; used in rewriting rules.
{-# INLINE[1] ap_l #-}
ap_l :: Put a -> Put b -> Put a
ap_l :: forall a b. Put a -> Put b -> Put a
ap_l (Put forall r. (a -> BuildStep r) -> BuildStep r
a) (Put forall r. (b -> BuildStep r) -> BuildStep r
b) = (forall r. (a -> BuildStep r) -> BuildStep r) -> Put a
forall a. (forall r. (a -> BuildStep r) -> BuildStep r) -> Put a
Put ((forall r. (a -> BuildStep r) -> BuildStep r) -> Put a)
-> (forall r. (a -> BuildStep r) -> BuildStep r) -> Put a
forall a b. (a -> b) -> a -> b
$ \a -> BuildStep r
k -> (a -> BuildStep r) -> BuildStep r
forall r. (a -> BuildStep r) -> BuildStep r
a (\a
a' -> (b -> BuildStep r) -> BuildStep r
forall r. (b -> BuildStep r) -> BuildStep r
b (\b
_ -> a -> BuildStep r
k a
a'))

-- | Synonym for '*>' from 'Applicative' and '>>' from 'Monad'; used in
-- rewriting rules.
{-# INLINE[1] ap_r #-}
ap_r :: Put a -> Put b -> Put b
ap_r :: forall a b. Put a -> Put b -> Put b
ap_r (Put forall r. (a -> BuildStep r) -> BuildStep r
a) (Put forall r. (b -> BuildStep r) -> BuildStep r
b) = (forall r. (b -> BuildStep r) -> BuildStep r) -> Put b
forall a. (forall r. (a -> BuildStep r) -> BuildStep r) -> Put a
Put ((forall r. (b -> BuildStep r) -> BuildStep r) -> Put b)
-> (forall r. (b -> BuildStep r) -> BuildStep r) -> Put b
forall a b. (a -> b) -> a -> b
$ \b -> BuildStep r
k -> (a -> BuildStep r) -> BuildStep r
forall r. (a -> BuildStep r) -> BuildStep r
a (\a
_ -> (b -> BuildStep r) -> BuildStep r
forall r. (b -> BuildStep r) -> BuildStep r
b b -> BuildStep r
k)

instance Applicative Put where
  {-# INLINE pure #-}
  pure :: forall a. a -> Put a
pure a
x = (forall r. (a -> BuildStep r) -> BuildStep r) -> Put a
forall a. (forall r. (a -> BuildStep r) -> BuildStep r) -> Put a
Put ((forall r. (a -> BuildStep r) -> BuildStep r) -> Put a)
-> (forall r. (a -> BuildStep r) -> BuildStep r) -> Put a
forall a b. (a -> b) -> a -> b
$ \a -> BuildStep r
k -> a -> BuildStep r
k a
x
  {-# INLINE (<*>) #-}
  Put forall r. ((a -> b) -> BuildStep r) -> BuildStep r
f <*> :: forall a b. Put (a -> b) -> Put a -> Put b
<*> Put forall r. (a -> BuildStep r) -> BuildStep r
a = (forall r. (b -> BuildStep r) -> BuildStep r) -> Put b
forall a. (forall r. (a -> BuildStep r) -> BuildStep r) -> Put a
Put ((forall r. (b -> BuildStep r) -> BuildStep r) -> Put b)
-> (forall r. (b -> BuildStep r) -> BuildStep r) -> Put b
forall a b. (a -> b) -> a -> b
$ \b -> BuildStep r
k -> ((a -> b) -> BuildStep r) -> BuildStep r
forall r. ((a -> b) -> BuildStep r) -> BuildStep r
f (\a -> b
f' -> (a -> BuildStep r) -> BuildStep r
forall r. (a -> BuildStep r) -> BuildStep r
a (b -> BuildStep r
k (b -> BuildStep r) -> (a -> b) -> a -> BuildStep r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> b
f'))
  {-# INLINE (<*) #-}
  <* :: forall a b. Put a -> Put b -> Put a
(<*) = Put a -> Put b -> Put a
forall a b. Put a -> Put b -> Put a
ap_l
  {-# INLINE (*>) #-}
  *> :: forall a b. Put a -> Put b -> Put b
(*>) = Put a -> Put b -> Put b
forall a b. Put a -> Put b -> Put b
ap_r

instance Monad Put where
  {-# INLINE return #-}
  return :: forall a. a -> Put a
return = a -> Put a
forall a. a -> Put a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
  {-# INLINE (>>=) #-}
  Put forall r. (a -> BuildStep r) -> BuildStep r
m >>= :: forall a b. Put a -> (a -> Put b) -> Put b
>>= a -> Put b
f = (forall r. (b -> BuildStep r) -> BuildStep r) -> Put b
forall a. (forall r. (a -> BuildStep r) -> BuildStep r) -> Put a
Put ((forall r. (b -> BuildStep r) -> BuildStep r) -> Put b)
-> (forall r. (b -> BuildStep r) -> BuildStep r) -> Put b
forall a b. (a -> b) -> a -> b
$ \b -> BuildStep r
k -> (a -> BuildStep r) -> BuildStep r
forall r. (a -> BuildStep r) -> BuildStep r
m (\a
m' -> Put b -> forall r. (b -> BuildStep r) -> BuildStep r
forall a. Put a -> forall r. (a -> BuildStep r) -> BuildStep r
unPut (a -> Put b
f a
m') b -> BuildStep r
k)
  {-# INLINE (>>) #-}
  >> :: forall a b. Put a -> Put b -> Put b
(>>) = Put a -> Put b -> Put b
forall a b. Put a -> Put b -> Put b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
(*>)

-- Conversion between Put and Builder
-------------------------------------

-- | Run a 'Builder' as a side-effect of a @'Put' ()@ action.
{-# INLINE[1] putBuilder #-}
putBuilder :: Builder -> Put ()
putBuilder :: Builder -> Put ()
putBuilder (Builder forall r. BuildStep r -> BuildStep r
b) = (forall r. (() -> BuildStep r) -> BuildStep r) -> Put ()
forall a. (forall r. (a -> BuildStep r) -> BuildStep r) -> Put a
Put ((forall r. (() -> BuildStep r) -> BuildStep r) -> Put ())
-> (forall r. (() -> BuildStep r) -> BuildStep r) -> Put ()
forall a b. (a -> b) -> a -> b
$ \() -> BuildStep r
k -> BuildStep r -> BuildStep r
forall r. BuildStep r -> BuildStep r
b (() -> BuildStep r
k ())

-- | Convert a @'Put' ()@ action to a 'Builder'.
{-# INLINE fromPut #-}
fromPut :: Put () -> Builder
fromPut :: Put () -> Builder
fromPut (Put forall r. (() -> BuildStep r) -> BuildStep r
p) = (forall r. BuildStep r -> BuildStep r) -> Builder
Builder ((forall r. BuildStep r -> BuildStep r) -> Builder)
-> (forall r. BuildStep r -> BuildStep r) -> Builder
forall a b. (a -> b) -> a -> b
$ \BuildStep r
k -> (() -> BuildStep r) -> BuildStep r
forall r. (() -> BuildStep r) -> BuildStep r
p (BuildStep r -> () -> BuildStep r
forall a b. a -> b -> a
const BuildStep r
k)

-- We rewrite consecutive uses of 'putBuilder' such that the append of the
-- involved 'Builder's is used. This can significantly improve performance,
-- when the bound-checks of the concatenated builders are fused.

-- ap_l rules
{-# RULES

"ap_l/putBuilder" forall b1 b2.
       ap_l (putBuilder b1) (putBuilder b2)
     = putBuilder (append b1 b2)

"ap_l/putBuilder/assoc_r" forall b1 b2 (p :: Put a).
       ap_l (putBuilder b1) (ap_l (putBuilder b2) p)
     = ap_l (putBuilder (append b1 b2)) p

"ap_l/putBuilder/assoc_l" forall (p :: Put a) b1 b2.
       ap_l (ap_l p (putBuilder b1)) (putBuilder b2)
     = ap_l p (putBuilder (append b1 b2))
 #-}

-- ap_r rules
{-# RULES

"ap_r/putBuilder" forall b1 b2.
       ap_r (putBuilder b1) (putBuilder b2)
     = putBuilder (append b1 b2)

"ap_r/putBuilder/assoc_r" forall b1 b2 (p :: Put a).
       ap_r (putBuilder b1) (ap_r (putBuilder b2) p)
     = ap_r (putBuilder (append b1 b2)) p

"ap_r/putBuilder/assoc_l" forall (p :: Put a) b1 b2.
       ap_r (ap_r p (putBuilder b1)) (putBuilder b2)
     = ap_r p (putBuilder (append b1 b2))

 #-}

-- combined ap_l/ap_r rules
{-# RULES

"ap_l/ap_r/putBuilder/assoc_r" forall b1 b2 (p :: Put a).
       ap_l (putBuilder b1) (ap_r (putBuilder b2) p)
     = ap_l (putBuilder (append b1 b2)) p

"ap_r/ap_l/putBuilder/assoc_r" forall b1 b2 (p :: Put a).
       ap_r (putBuilder b1) (ap_l (putBuilder b2) p)
     = ap_l (putBuilder (append b1 b2)) p

"ap_l/ap_r/putBuilder/assoc_l" forall (p :: Put a) b1 b2.
       ap_l (ap_r p (putBuilder b1)) (putBuilder b2)
     = ap_r p (putBuilder (append b1 b2))

"ap_r/ap_l/putBuilder/assoc_l" forall (p :: Put a) b1 b2.
       ap_r (ap_l p (putBuilder b1)) (putBuilder b2)
     = ap_r p (putBuilder (append b1 b2))

 #-}


-- Lifting IO actions
---------------------

{-
-- | Lift an 'IO' action to a 'Put' action.
{-# INLINE putLiftIO #-}
putLiftIO :: IO a -> Put a
putLiftIO io = put $ \k br -> io >>= (`k` br)
-}


------------------------------------------------------------------------------
-- Executing a Put directly on a buffered Handle
------------------------------------------------------------------------------

-- | Run a 'Put' action redirecting the produced output to a 'Handle'.
--
-- The output is buffered using the 'Handle's associated buffer. If this
-- buffer is too small to execute one step of the 'Put' action, then
-- it is replaced with a large enough buffer.
hPut :: forall a. Handle -> Put a -> IO a
hPut :: forall a. Handle -> Put a -> IO a
hPut Handle
h Put a
p = do
    Int -> BuildStep a -> IO a
fillHandle Int
1 (Put a -> BuildStep a
forall a. Put a -> BuildStep a
runPut Put a
p)
  where
    fillHandle :: Int -> BuildStep a -> IO a
    fillHandle :: Int -> BuildStep a -> IO a
fillHandle !Int
minFree BuildStep a
step = do
        IO a
next <- [Char] -> Handle -> (Handle__ -> IO (IO a)) -> IO (IO a)
forall a. [Char] -> Handle -> (Handle__ -> IO a) -> IO a
wantWritableHandle [Char]
"hPut" Handle
h Handle__ -> IO (IO a)
fillHandle_
        IO a
next
      where
        -- | We need to return an inner IO action that is executed outside
        -- the lock taken on the Handle for two reasons:
        --
        --   1. GHC.IO.Handle.Internals mentions in "Note [async]" that
        --      we should never do any side-effecting operations before
        --      an interruptible operation that may raise an async. exception
        --      as long as we are inside 'wantWritableHandle' and the like.
        --      We possibly run the interruptible 'flushWriteBuffer' right at
        --      the start of 'fillHandle', hence entering it a second time is
        --      not safe, as it could lead to a 'BuildStep' being run twice.
        --
        --      FIXME (SM): Adapt this function or at least its documentation,
        --      as it is OK to run a 'BuildStep' twice. We dropped this
        --      requirement in favor of being able to use
        --      'unsafeDupablePerformIO' and the speed improvement that it
        --      brings.
        --
        --   2. We use the 'S.hPut' function to also write to the handle.
        --      This function tries to take the same lock taken by
        --      'wantWritableHandle'. Therefore, we cannot call 'S.hPut'
        --      inside 'wantWritableHandle'.
        --
        fillHandle_ :: Handle__ -> IO (IO a)
        fillHandle_ :: Handle__ -> IO (IO a)
fillHandle_ Handle__
h_ = do
            Buffer Word8 -> IO ()
forall {e}. Buffer e -> IO ()
makeSpace  (Buffer Word8 -> IO ()) -> IO (Buffer Word8) -> IO ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< IORef (Buffer Word8) -> IO (Buffer Word8)
forall a. IORef a -> IO a
readIORef IORef (Buffer Word8)
refBuf
            Buffer Word8 -> IO (IO a)
fillBuffer (Buffer Word8 -> IO (IO a)) -> IO (Buffer Word8) -> IO (IO a)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< IORef (Buffer Word8) -> IO (Buffer Word8)
forall a. IORef a -> IO a
readIORef IORef (Buffer Word8)
refBuf
          where
            refBuf :: IORef (Buffer Word8)
refBuf        = Handle__ -> IORef (Buffer Word8)
haByteBuffer Handle__
h_
            freeSpace :: Buffer e -> Int
freeSpace Buffer e
buf = Buffer e -> Int
forall e. Buffer e -> Int
IO.bufSize Buffer e
buf Int -> Int -> Int
forall a. Num a => a -> a -> a
- Buffer e -> Int
forall e. Buffer e -> Int
IO.bufR Buffer e
buf

            makeSpace :: Buffer e -> IO ()
makeSpace Buffer e
buf
              | Buffer e -> Int
forall e. Buffer e -> Int
IO.bufSize Buffer e
buf Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
minFree = do
                  Handle__ -> IO ()
flushWriteBuffer Handle__
h_
                  BufferState
s <- Buffer Word8 -> BufferState
forall e. Buffer e -> BufferState
IO.bufState (Buffer Word8 -> BufferState)
-> IO (Buffer Word8) -> IO BufferState
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IORef (Buffer Word8) -> IO (Buffer Word8)
forall a. IORef a -> IO a
readIORef IORef (Buffer Word8)
refBuf
                  Int -> BufferState -> IO (Buffer Word8)
IO.newByteBuffer Int
minFree BufferState
s IO (Buffer Word8) -> (Buffer Word8 -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= IORef (Buffer Word8) -> Buffer Word8 -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (Buffer Word8)
refBuf

              | Buffer e -> Int
forall e. Buffer e -> Int
freeSpace Buffer e
buf Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
minFree = Handle__ -> IO ()
flushWriteBuffer Handle__
h_
              | Bool
otherwise               =
                                          () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()

            fillBuffer :: Buffer Word8 -> IO (IO a)
fillBuffer Buffer Word8
buf
              | Buffer Word8 -> Int
forall e. Buffer e -> Int
freeSpace Buffer Word8
buf Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
minFree =
                  [Char] -> IO (IO a)
forall a. HasCallStack => [Char] -> a
error ([Char] -> IO (IO a)) -> [Char] -> IO (IO a)
forall a b. (a -> b) -> a -> b
$ [[Char]] -> [Char]
unlines
                    [ [Char]
"Data.ByteString.Builder.Internal.hPut: internal error."
                    , [Char]
"  Not enough space after flush."
                    , [Char]
"    required: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
minFree
                    , [Char]
"    free: "     [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show (Buffer Word8 -> Int
forall e. Buffer e -> Int
freeSpace Buffer Word8
buf)
                    ]
              | Bool
otherwise = do
                  let !br :: BufferRange
br = Ptr Word8 -> Ptr Word8 -> BufferRange
BufferRange Ptr Word8
forall {b}. Ptr b
op (Ptr Word8
pBuf Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Buffer Word8 -> Int
forall e. Buffer e -> Int
IO.bufSize Buffer Word8
buf)
                  IO a
res <- BuildStep a
-> (Ptr Word8 -> a -> IO (IO a))
-> (Ptr Word8 -> Int -> BuildStep a -> IO (IO a))
-> (Ptr Word8 -> StrictByteString -> BuildStep a -> IO (IO a))
-> BufferRange
-> IO (IO a)
forall a b.
BuildStep a
-> (Ptr Word8 -> a -> IO b)
-> (Ptr Word8 -> Int -> BuildStep a -> IO b)
-> (Ptr Word8 -> StrictByteString -> BuildStep a -> IO b)
-> BufferRange
-> IO b
fillWithBuildStep BuildStep a
step Ptr Word8 -> a -> IO (IO a)
forall {a} {a}. Ptr a -> a -> IO (IO a)
doneH Ptr Word8 -> Int -> BuildStep a -> IO (IO a)
forall {a}. Ptr a -> Int -> BuildStep a -> IO (IO a)
fullH Ptr Word8 -> StrictByteString -> BuildStep a -> IO (IO a)
forall {a}. Ptr a -> StrictByteString -> BuildStep a -> IO (IO a)
insertChunkH BufferRange
br
                  ForeignPtr Word8 -> IO ()
forall a. ForeignPtr a -> IO ()
touchForeignPtr ForeignPtr Word8
fpBuf
                  IO a -> IO (IO a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return IO a
res
              where
                fpBuf :: ForeignPtr Word8
fpBuf = Buffer Word8 -> ForeignPtr Word8
forall e. Buffer e -> RawBuffer e
IO.bufRaw Buffer Word8
buf
                pBuf :: Ptr Word8
pBuf  = ForeignPtr Word8 -> Ptr Word8
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr Word8
fpBuf
                op :: Ptr b
op    = Ptr Word8
pBuf Ptr Word8 -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Buffer Word8 -> Int
forall e. Buffer e -> Int
IO.bufR Buffer Word8
buf

                {-# INLINE updateBufR #-}
                updateBufR :: Ptr a -> IO ()
updateBufR Ptr a
op' = do
                    let !off' :: Int
off' = Ptr a
op' Ptr a -> Ptr Word8 -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr Word8
pBuf
                        !buf' :: Buffer Word8
buf' = Buffer Word8
buf {IO.bufR = off'}
                    IORef (Buffer Word8) -> Buffer Word8 -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (Buffer Word8)
refBuf Buffer Word8
buf'

                doneH :: Ptr a -> a -> IO (IO a)
doneH Ptr a
op' a
x = do
                    Ptr a -> IO ()
forall {a}. Ptr a -> IO ()
updateBufR Ptr a
op'
                    -- We must flush if this Handle is set to NoBuffering.
                    -- If it is set to LineBuffering, be conservative and
                    -- flush anyway (we didn't check for newlines in the data).
                    -- Flushing must happen outside this 'wantWriteableHandle'
                    -- due to the possible async. exception.
                    case Handle__ -> BufferMode
haBufferMode Handle__
h_ of
                        BlockBuffering Maybe Int
_      -> IO a -> IO (IO a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (IO a -> IO (IO a)) -> IO a -> IO (IO a)
forall a b. (a -> b) -> a -> b
$ a -> IO a
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return a
x
                        BufferMode
_line_or_no_buffering -> IO a -> IO (IO a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (IO a -> IO (IO a)) -> IO a -> IO (IO a)
forall a b. (a -> b) -> a -> b
$ Handle -> IO ()
hFlush Handle
h IO () -> IO a -> IO a
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> a -> IO a
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return a
x

                fullH :: Ptr a -> Int -> BuildStep a -> IO (IO a)
fullH Ptr a
op' Int
minSize BuildStep a
nextStep = do
                    Ptr a -> IO ()
forall {a}. Ptr a -> IO ()
updateBufR Ptr a
op'
                    IO a -> IO (IO a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (IO a -> IO (IO a)) -> IO a -> IO (IO a)
forall a b. (a -> b) -> a -> b
$ Int -> BuildStep a -> IO a
fillHandle Int
minSize BuildStep a
nextStep
                    -- 'fillHandle' will flush the buffer (provided there is
                    -- really less than @minSize@ space left) before executing
                    -- the 'nextStep'.

                insertChunkH :: Ptr a -> StrictByteString -> BuildStep a -> IO (IO a)
insertChunkH Ptr a
op' StrictByteString
bs BuildStep a
nextStep = do
                    Ptr a -> IO ()
forall {a}. Ptr a -> IO ()
updateBufR Ptr a
op'
                    IO a -> IO (IO a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (IO a -> IO (IO a)) -> IO a -> IO (IO a)
forall a b. (a -> b) -> a -> b
$ do
                        Handle -> StrictByteString -> IO ()
S.hPut Handle
h StrictByteString
bs
                        Int -> BuildStep a -> IO a
fillHandle Int
1 BuildStep a
nextStep

-- | Execute a 'Put' and return the computed result and the bytes
-- written during the computation as a 'L.LazyByteString'.
--
-- This function is strict in the computed result and lazy in the writing of
-- the bytes. For example, given
--
-- @
--infinitePut = sequence_ (repeat (putBuilder (word8 1))) >> return 0
-- @
--
-- evaluating the expression
--
-- @
--fst $ putToLazyByteString infinitePut
-- @
--
-- does not terminate, while evaluating the expression
--
-- @
--L.head $ snd $ putToLazyByteString infinitePut
-- @
--
-- does terminate and yields the value @1 :: Word8@.
--
-- An illustrative example for these strictness properties is the
-- implementation of Base64 decoding (<http://en.wikipedia.org/wiki/Base64>).
--
-- @
--type DecodingState = ...
--
--decodeBase64 :: 'S.StrictByteString' -> DecodingState -> 'Put' (Maybe DecodingState)
--decodeBase64 = ...
-- @
--
-- The above function takes a 'S.StrictByteString' supposed to represent
-- Base64 encoded data and the current decoding state.
-- It writes the decoded bytes as the side-effect of the 'Put' and returns the
-- new decoding state, if the decoding of all data in the 'S.StrictByteString' was
-- successful. The checking if the 'S.StrictByteString' represents Base64
-- encoded data and the actual decoding are fused. This makes the common case,
-- where all data represents Base64 encoded data, more efficient. It also
-- implies that all data must be decoded before the final decoding
-- state can be returned. 'Put's are intended for implementing such fused
-- checking and decoding/encoding, which is reflected in their strictness
-- properties.
{-# NOINLINE putToLazyByteString #-}
putToLazyByteString
    :: Put a              -- ^ 'Put' to execute
    -> (a, L.LazyByteString)  -- ^ Result and 'L.LazyByteString'
                          -- written as its side-effect
putToLazyByteString :: forall a. Put a -> (a, LazyByteString)
putToLazyByteString = AllocationStrategy
-> (a -> (a, LazyByteString)) -> Put a -> (a, LazyByteString)
forall a b.
AllocationStrategy
-> (a -> (b, LazyByteString)) -> Put a -> (b, LazyByteString)
putToLazyByteStringWith
    (Int -> Int -> AllocationStrategy
safeStrategy Int
L.smallChunkSize Int
L.defaultChunkSize) (, LazyByteString
L.Empty)


-- | Execute a 'Put' with a buffer-allocation strategy and a continuation. For
-- example, 'putToLazyByteString' is implemented as follows.
--
-- @
--putToLazyByteString = 'putToLazyByteStringWith'
--    ('safeStrategy' 'L.smallChunkSize' 'L.defaultChunkSize') (\x -> (x, L.empty))
-- @
--
{-# INLINE putToLazyByteStringWith #-}
putToLazyByteStringWith
    :: AllocationStrategy
       -- ^ Buffer allocation strategy to use
    -> (a -> (b, L.LazyByteString))
       -- ^ Continuation to use for computing the final result and the tail of
       -- its side-effect (the written bytes).
    -> Put a
       -- ^ 'Put' to execute
    -> (b, L.LazyByteString)
       -- ^ Resulting 'L.LazyByteString'
putToLazyByteStringWith :: forall a b.
AllocationStrategy
-> (a -> (b, LazyByteString)) -> Put a -> (b, LazyByteString)
putToLazyByteStringWith AllocationStrategy
strategy a -> (b, LazyByteString)
k Put a
p =
    AllocationStrategy
-> (a -> (b, LazyByteString))
-> ChunkIOStream a
-> (b, LazyByteString)
forall a b.
AllocationStrategy
-> (a -> (b, LazyByteString))
-> ChunkIOStream a
-> (b, LazyByteString)
ciosToLazyByteString AllocationStrategy
strategy a -> (b, LazyByteString)
k (ChunkIOStream a -> (b, LazyByteString))
-> ChunkIOStream a -> (b, LazyByteString)
forall a b. (a -> b) -> a -> b
$ IO (ChunkIOStream a) -> ChunkIOStream a
forall a. IO a -> a
unsafeDupablePerformIO (IO (ChunkIOStream a) -> ChunkIOStream a)
-> IO (ChunkIOStream a) -> ChunkIOStream a
forall a b. (a -> b) -> a -> b
$
        AllocationStrategy -> BuildStep a -> IO (ChunkIOStream a)
forall a. AllocationStrategy -> BuildStep a -> IO (ChunkIOStream a)
buildStepToCIOS AllocationStrategy
strategy (Put a -> BuildStep a
forall a. Put a -> BuildStep a
runPut Put a
p)



------------------------------------------------------------------------------
-- ByteString insertion / controlling chunk boundaries
------------------------------------------------------------------------------

-- Raw memory
-------------

-- | @'ensureFree' n@ ensures that there are at least @n@ free bytes
-- for the following 'Builder'.
{-# INLINE ensureFree #-}
ensureFree :: Int -> Builder
ensureFree :: Int -> Builder
ensureFree Int
minFree =
    (forall r. BuildStep r -> BuildStep r) -> Builder
builder BuildStep r -> BuildStep r
forall r. BuildStep r -> BuildStep r
step
  where
    step :: BuildStep a -> BuildStep a
step BuildStep a
k br :: BufferRange
br@(BufferRange Ptr Word8
op Ptr Word8
ope)
      | Ptr Word8
ope Ptr Word8 -> Ptr Word8 -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr Word8
op Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
minFree = BuildSignal a -> IO (BuildSignal a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (BuildSignal a -> IO (BuildSignal a))
-> BuildSignal a -> IO (BuildSignal a)
forall a b. (a -> b) -> a -> b
$ Int -> Ptr Word8 -> BuildStep a -> BuildSignal a
forall a. Int -> Ptr Word8 -> BuildStep a -> BuildSignal a
bufferFull Int
minFree Ptr Word8
op BuildStep a
k
      | Bool
otherwise                   = BuildStep a
k BufferRange
br

-- | Copy the bytes from a 'S.StrictByteString' into the output stream.
wrappedBytesCopyStep :: S.StrictByteString  -- ^ Input 'S.StrictByteString'.
                     -> BuildStep a -> BuildStep a
-- See Note [byteStringCopyStep and wrappedBytesCopyStep]
wrappedBytesCopyStep :: forall a. StrictByteString -> BuildStep a -> BuildStep a
wrappedBytesCopyStep StrictByteString
bs0 BuildStep a
k =
    StrictByteString -> BuildStep a
go StrictByteString
bs0
  where
    go :: StrictByteString -> BuildStep a
go !bs :: StrictByteString
bs@(S.BS ForeignPtr Word8
ifp Int
inpRemaining) (BufferRange Ptr Word8
op Ptr Word8
ope)
      | Int
inpRemaining Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
outRemaining = do
          ForeignPtr Word8 -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
S.unsafeWithForeignPtr ForeignPtr Word8
ifp ((Ptr Word8 -> IO ()) -> IO ()) -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
ip -> Ptr Word8 -> Ptr Word8 -> Int -> IO ()
forall a. Ptr a -> Ptr a -> Int -> IO ()
copyBytes Ptr Word8
op Ptr Word8
ip Int
inpRemaining
          let !br' :: BufferRange
br' = Ptr Word8 -> Ptr Word8 -> BufferRange
BufferRange (Ptr Word8
op Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
inpRemaining) Ptr Word8
ope
          BuildStep a
k BufferRange
br'
      | Bool
otherwise = do
          ForeignPtr Word8 -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
S.unsafeWithForeignPtr ForeignPtr Word8
ifp ((Ptr Word8 -> IO ()) -> IO ()) -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
ip -> Ptr Word8 -> Ptr Word8 -> Int -> IO ()
forall a. Ptr a -> Ptr a -> Int -> IO ()
copyBytes Ptr Word8
op Ptr Word8
ip Int
outRemaining
          let !bs' :: StrictByteString
bs' = Int -> StrictByteString -> StrictByteString
S.unsafeDrop Int
outRemaining StrictByteString
bs
          BuildSignal a -> IO (BuildSignal a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (BuildSignal a -> IO (BuildSignal a))
-> BuildSignal a -> IO (BuildSignal a)
forall a b. (a -> b) -> a -> b
$ Int -> Ptr Word8 -> BuildStep a -> BuildSignal a
forall a. Int -> Ptr Word8 -> BuildStep a -> BuildSignal a
bufferFull Int
1 Ptr Word8
ope (StrictByteString -> BuildStep a
go StrictByteString
bs')
      where
        outRemaining :: Int
outRemaining = Ptr Word8
ope Ptr Word8 -> Ptr Word8 -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr Word8
op


-- Strict ByteStrings
------------------------------------------------------------------------------


-- | Construct a 'Builder' that copies the 'S.StrictByteString's, if it is
-- smaller than the treshold, and inserts it directly otherwise.
--
-- For example, @byteStringThreshold 1024@ copies 'S.StrictByteString's whose size
-- is less or equal to 1kb, and inserts them directly otherwise. This implies
-- that the average chunk-size of the generated 'L.LazyByteString' may be as
-- low as 513 bytes, as there could always be just a single byte between the
-- directly inserted 1025 byte, 'S.StrictByteString's.
--
{-# INLINE byteStringThreshold #-}
byteStringThreshold :: Int -> S.StrictByteString -> Builder
byteStringThreshold :: Int -> StrictByteString -> Builder
byteStringThreshold Int
maxCopySize =
    \StrictByteString
bs -> (forall r. BuildStep r -> BuildStep r) -> Builder
builder ((forall r. BuildStep r -> BuildStep r) -> Builder)
-> (forall r. BuildStep r -> BuildStep r) -> Builder
forall a b. (a -> b) -> a -> b
$ StrictByteString -> BuildStep r -> BuildStep r
forall a. StrictByteString -> BuildStep a -> BuildStep a
step StrictByteString
bs
  where
    step :: StrictByteString -> BuildStep a -> BuildStep a
step bs :: StrictByteString
bs@(S.BS ForeignPtr Word8
_ Int
len) BuildStep a
k br :: BufferRange
br@(BufferRange !Ptr Word8
op Ptr Word8
_)
      | Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
maxCopySize = StrictByteString -> BuildStep a -> BuildStep a
forall a. StrictByteString -> BuildStep a -> BuildStep a
byteStringCopyStep StrictByteString
bs BuildStep a
k BufferRange
br
      | Bool
otherwise          = BuildSignal a -> IO (BuildSignal a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (BuildSignal a -> IO (BuildSignal a))
-> BuildSignal a -> IO (BuildSignal a)
forall a b. (a -> b) -> a -> b
$ Ptr Word8 -> StrictByteString -> BuildStep a -> BuildSignal a
forall a.
Ptr Word8 -> StrictByteString -> BuildStep a -> BuildSignal a
insertChunk Ptr Word8
op StrictByteString
bs BuildStep a
k

-- | Construct a 'Builder' that copies the 'S.StrictByteString'.
--
-- Use this function to create 'Builder's from smallish (@<= 4kb@)
-- 'S.StrictByteString's or if you need to guarantee that the 'S.StrictByteString' is not
-- shared with the chunks generated by the 'Builder'.
--
{-# INLINE byteStringCopy #-}
byteStringCopy :: S.StrictByteString -> Builder
byteStringCopy :: StrictByteString -> Builder
byteStringCopy = \StrictByteString
bs -> (forall r. BuildStep r -> BuildStep r) -> Builder
builder ((forall r. BuildStep r -> BuildStep r) -> Builder)
-> (forall r. BuildStep r -> BuildStep r) -> Builder
forall a b. (a -> b) -> a -> b
$ StrictByteString -> BuildStep r -> BuildStep r
forall a. StrictByteString -> BuildStep a -> BuildStep a
byteStringCopyStep StrictByteString
bs

{-
Note [byteStringCopyStep and wrappedBytesCopyStep]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
A Builder that copies the contents of an arbitrary ByteString needs a
recursive loop, since the bytes to be copied might not fit into the
first few chunk buffers provided by the driver.  That loop is
implemented in 'wrappedBytesCopyStep'.  But we also have a
non-recursive wrapper, 'byteStringCopyStep', which performs exactly
the first iteration of that loop, falling back to 'wrappedBytesCopyStep'
if a chunk boundary is reached before the entire ByteString is copied.

This is very strange!  Why do we do this?  Perhaps mostly for
historical reasons.  But sadly, changing this to use a single
recursive loop regresses the benchmark 'foldMap byteStringCopy' by
about 30% as of 2024, in one of two ways:

 1. If the continuation 'k' is taken as an argument of the
    inner copying loop, it remains an unknown function call.
    So for each bytestring copied, that continuation must be
    entered later via a gen-apply function, which incurs dozens
    of cycles of extra overhead.
 2. If the continuation 'k' is lifted out of the inner copying
    loop, it becomes a free variable.  And after a bit of
    inlining, there will be no unknown function call.  But, if
    the continuation function has any free variables, these
    become free variables of the inner copying loop, which
    prevent the loop from floating out.  (In the actual
    benchmark, the tail of the list of bytestrings to copy is
    such a free variable of the continuation.)  As a result,
    the inner copying loop becomes a function closure object
    rather than a top-level function.  And that means a new
    inner-copying-loop function-closure-object must be
    allocated on the heap for every bytestring copied, which
    is expensive.

    In theory, GHC's late-lambda-lifting pass can clean this up by
    abstracting over the problematic free variables.  But for some
    unknown reason (perhaps a bug in ghc-9.10.1) this optimization
    does not fire on the relevant benchmark code, even with a
    sufficiently high value of -fstg-lift-lams-rec-args.



Alternatively, it is possible to avoid recursion altogether by
requesting that the next chunk be large enough to accommodate the
entire remainder of the input when a chunk boundary is reached.
But:
 * For very large ByteStrings, this may incur unwanted latency.
 * Large next-chunk-size requests have caused breakage downstream
   in the past.  See also https://github.com/yesodweb/wai/issues/894
-}

{-# INLINE byteStringCopyStep #-}
byteStringCopyStep :: S.StrictByteString -> BuildStep a -> BuildStep a
-- See Note [byteStringCopyStep and wrappedBytesCopyStep]
byteStringCopyStep :: forall a. StrictByteString -> BuildStep a -> BuildStep a
byteStringCopyStep bs :: StrictByteString
bs@(S.BS ForeignPtr Word8
ifp Int
isize) BuildStep a
k br :: BufferRange
br@(BufferRange Ptr Word8
op Ptr Word8
ope)
    | Int
isize Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
osize = do
        ForeignPtr Word8 -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
S.unsafeWithForeignPtr ForeignPtr Word8
ifp ((Ptr Word8 -> IO ()) -> IO ()) -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
ip -> Ptr Word8 -> Ptr Word8 -> Int -> IO ()
forall a. Ptr a -> Ptr a -> Int -> IO ()
copyBytes Ptr Word8
op Ptr Word8
ip Int
isize
        BuildStep a
k (Ptr Word8 -> Ptr Word8 -> BufferRange
BufferRange Ptr Word8
forall {b}. Ptr b
op' Ptr Word8
ope)
    | Bool
otherwise  = StrictByteString -> BuildStep a -> BuildStep a
forall a. StrictByteString -> BuildStep a -> BuildStep a
wrappedBytesCopyStep StrictByteString
bs BuildStep a
k BufferRange
br
  where
    osize :: Int
osize = Ptr Word8
ope Ptr Word8 -> Ptr Word8 -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr Word8
op
    op' :: Ptr b
op'  = Ptr Word8
op Ptr Word8 -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
isize

-- | Construct a 'Builder' that always inserts the 'S.StrictByteString'
-- directly as a chunk.
--
-- This implies flushing the output buffer, even if it contains just
-- a single byte. You should therefore use 'byteStringInsert' only for large
-- (@> 8kb@) 'S.StrictByteString's. Otherwise, the generated chunks are too
-- fragmented to be processed efficiently afterwards.
--
{-# INLINE byteStringInsert #-}
byteStringInsert :: S.StrictByteString -> Builder
byteStringInsert :: StrictByteString -> Builder
byteStringInsert =
    \StrictByteString
bs -> (forall r. BuildStep r -> BuildStep r) -> Builder
builder ((forall r. BuildStep r -> BuildStep r) -> Builder)
-> (forall r. BuildStep r -> BuildStep r) -> Builder
forall a b. (a -> b) -> a -> b
$ \BuildStep r
k (BufferRange Ptr Word8
op Ptr Word8
_) -> BuildSignal r -> IO (BuildSignal r)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (BuildSignal r -> IO (BuildSignal r))
-> BuildSignal r -> IO (BuildSignal r)
forall a b. (a -> b) -> a -> b
$ Ptr Word8 -> StrictByteString -> BuildStep r -> BuildSignal r
forall a.
Ptr Word8 -> StrictByteString -> BuildStep a -> BuildSignal a
insertChunk Ptr Word8
op StrictByteString
bs BuildStep r
k

-- Short bytestrings
------------------------------------------------------------------------------

-- | Construct a 'Builder' that copies the 'SH.ShortByteString'.
--
{-# INLINE shortByteString #-}
shortByteString :: Sh.ShortByteString -> Builder
shortByteString :: ShortByteString -> Builder
shortByteString = \ShortByteString
sbs -> (forall r. BuildStep r -> BuildStep r) -> Builder
builder ((forall r. BuildStep r -> BuildStep r) -> Builder)
-> (forall r. BuildStep r -> BuildStep r) -> Builder
forall a b. (a -> b) -> a -> b
$ ShortByteString -> BuildStep r -> BuildStep r
forall a. ShortByteString -> BuildStep a -> BuildStep a
shortByteStringCopyStep ShortByteString
sbs

-- | Copy the bytes from a 'SH.ShortByteString' into the output stream.
{-# INLINE shortByteStringCopyStep #-}
shortByteStringCopyStep :: Sh.ShortByteString  -- ^ Input 'SH.ShortByteString'.
                        -> BuildStep a -> BuildStep a
shortByteStringCopyStep :: forall a. ShortByteString -> BuildStep a -> BuildStep a
shortByteStringCopyStep !ShortByteString
sbs BuildStep a
k =
    Int -> Int -> BuildStep a
go Int
0 (ShortByteString -> Int
Sh.length ShortByteString
sbs)
  where
    go :: Int -> Int -> BuildStep a
go !Int
ip !Int
ipe (BufferRange Ptr Word8
op Ptr Word8
ope)
      | Int
inpRemaining Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
outRemaining = do
          ShortByteString -> Int -> Ptr Word8 -> Int -> IO ()
forall a. ShortByteString -> Int -> Ptr a -> Int -> IO ()
Sh.copyToPtr ShortByteString
sbs Int
ip Ptr Word8
op Int
inpRemaining
          let !br' :: BufferRange
br' = Ptr Word8 -> Ptr Word8 -> BufferRange
BufferRange (Ptr Word8
op Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
inpRemaining) Ptr Word8
ope
          BuildStep a
k BufferRange
br'
      | Bool
otherwise = do
          ShortByteString -> Int -> Ptr Word8 -> Int -> IO ()
forall a. ShortByteString -> Int -> Ptr a -> Int -> IO ()
Sh.copyToPtr ShortByteString
sbs Int
ip Ptr Word8
op Int
outRemaining
          let !ip' :: Int
ip' = Int
ip Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
outRemaining
          BuildSignal a -> IO (BuildSignal a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (BuildSignal a -> IO (BuildSignal a))
-> BuildSignal a -> IO (BuildSignal a)
forall a b. (a -> b) -> a -> b
$ Int -> Ptr Word8 -> BuildStep a -> BuildSignal a
forall a. Int -> Ptr Word8 -> BuildStep a -> BuildSignal a
bufferFull Int
1 Ptr Word8
ope (Int -> Int -> BuildStep a
go Int
ip' Int
ipe)
      where
        outRemaining :: Int
outRemaining = Ptr Word8
ope Ptr Word8 -> Ptr Word8 -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr Word8
op
        inpRemaining :: Int
inpRemaining = Int
ipe Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
ip


-- Lazy bytestrings
------------------------------------------------------------------------------

-- | Construct a 'Builder' that uses the thresholding strategy of 'byteStringThreshold'
-- for each chunk of the 'L.LazyByteString'.
--
{-# INLINE lazyByteStringThreshold #-}
lazyByteStringThreshold :: Int -> L.LazyByteString -> Builder
lazyByteStringThreshold :: Int -> LazyByteString -> Builder
lazyByteStringThreshold Int
maxCopySize =
    (StrictByteString -> Builder -> Builder)
-> Builder -> LazyByteString -> Builder
forall a. (StrictByteString -> a -> a) -> a -> LazyByteString -> a
L.foldrChunks (\StrictByteString
bs Builder
b -> Int -> StrictByteString -> Builder
byteStringThreshold Int
maxCopySize StrictByteString
bs Builder -> Builder -> Builder
forall a. Monoid a => a -> a -> a
`mappend` Builder
b) Builder
forall a. Monoid a => a
mempty
    -- TODO: We could do better here. Currently, Large, Small, Large, leads to
    -- an unnecessary copy of the 'Small' chunk.

-- | Construct a 'Builder' that copies the 'L.LazyByteString'.
--
{-# INLINE lazyByteStringCopy #-}
lazyByteStringCopy :: L.LazyByteString -> Builder
lazyByteStringCopy :: LazyByteString -> Builder
lazyByteStringCopy =
    (StrictByteString -> Builder -> Builder)
-> Builder -> LazyByteString -> Builder
forall a. (StrictByteString -> a -> a) -> a -> LazyByteString -> a
L.foldrChunks (\StrictByteString
bs Builder
b -> StrictByteString -> Builder
byteStringCopy StrictByteString
bs Builder -> Builder -> Builder
forall a. Monoid a => a -> a -> a
`mappend` Builder
b) Builder
forall a. Monoid a => a
mempty

-- | Construct a 'Builder' that inserts all chunks of the 'L.LazyByteString'
-- directly.
--
{-# INLINE lazyByteStringInsert #-}
lazyByteStringInsert :: L.LazyByteString -> Builder
lazyByteStringInsert :: LazyByteString -> Builder
lazyByteStringInsert =
    (StrictByteString -> Builder -> Builder)
-> Builder -> LazyByteString -> Builder
forall a. (StrictByteString -> a -> a) -> a -> LazyByteString -> a
L.foldrChunks (\StrictByteString
bs Builder
b -> StrictByteString -> Builder
byteStringInsert StrictByteString
bs Builder -> Builder -> Builder
forall a. Monoid a => a -> a -> a
`mappend` Builder
b) Builder
forall a. Monoid a => a
mempty

-- | Create a 'Builder' denoting the same sequence of bytes as a
-- 'S.StrictByteString'.
-- The 'Builder' inserts large 'S.StrictByteString's directly, but copies small ones
-- to ensure that the generated chunks are large on average.
--
{-# INLINE byteString #-}
byteString :: S.StrictByteString -> Builder
byteString :: StrictByteString -> Builder
byteString = Int -> StrictByteString -> Builder
byteStringThreshold Int
maximalCopySize

-- | Create a 'Builder' denoting the same sequence of bytes as a lazy
-- 'L.LazyByteString'.
-- The 'Builder' inserts large chunks of the 'L.LazyByteString' directly,
-- but copies small ones to ensure that the generated chunks are large on
-- average.
--
{-# INLINE lazyByteString #-}
lazyByteString :: L.LazyByteString -> Builder
lazyByteString :: LazyByteString -> Builder
lazyByteString = Int -> LazyByteString -> Builder
lazyByteStringThreshold Int
maximalCopySize
-- FIXME: also insert the small chunk for [large,small,large] directly.
-- Perhaps it makes even sense to concatenate the small chunks in
-- [large,small,small,small,large] and insert them directly afterwards to avoid
-- unnecessary buffer spilling. Hmm, but that uncontrollably increases latency
-- => no good!

-- | The maximal size of a 'S.StrictByteString' that is copied.
-- @2 * 'L.smallChunkSize'@ to guarantee that on average a chunk is of
-- 'L.smallChunkSize'.
maximalCopySize :: Int
maximalCopySize :: Int
maximalCopySize = Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
L.smallChunkSize

------------------------------------------------------------------------------
-- Builder execution
------------------------------------------------------------------------------

-- | A buffer allocation strategy for executing 'Builder's.
data AllocationStrategy = AllocationStrategy
         (Maybe (Buffer, Int) -> IO Buffer)
         {-# UNPACK #-} !Int
         (Int -> Int -> Bool)

-- | Create a custom allocation strategy. See the code for 'safeStrategy' and
-- 'untrimmedStrategy' for examples.
{-# INLINE customStrategy #-}
customStrategy
  :: (Maybe (Buffer, Int) -> IO Buffer)
     -- ^ Buffer allocation function.
     --
     -- * If 'Nothing' is given, then a new first buffer should be allocated.
     --
     -- * If @'Just' (oldBuf, minSize)@ is given, then a buffer with minimal
     -- size @minSize@ must be returned. The strategy may reuse @oldBuf@ only if
     -- @oldBuf@ is large enough and the consumer can guarantee that this will
     -- not result in a violation of referential transparency.
     --
     -- /Warning:/ for multithreaded programs, it is generally unsafe to reuse
     -- buffers when using the consumers of 'Builder' in this package. For
     -- example, if 'toLazyByteStringWith' is called with an
     --  'AllocationStrategy' that reuses buffers, evaluating the result by
     -- multiple threads simultaneously may lead to corrupted output.
  -> Int
     -- ^ Default buffer size.
  -> (Int -> Int -> Bool)
     -- ^ A predicate @trim used allocated@ returning 'True', if the buffer
     -- should be trimmed before it is returned.
  -> AllocationStrategy
customStrategy :: (Maybe (Buffer, Int) -> IO Buffer)
-> Int -> (Int -> Int -> Bool) -> AllocationStrategy
customStrategy = (Maybe (Buffer, Int) -> IO Buffer)
-> Int -> (Int -> Int -> Bool) -> AllocationStrategy
AllocationStrategy

-- | Sanitize a buffer size; i.e., make it at least the size of an 'Int'.
{-# INLINE sanitize #-}
sanitize :: Int -> Int
sanitize :: Int -> Int
sanitize = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max (Int -> Int
forall a. Storable a => a -> Int
sizeOf (Int
forall a. HasCallStack => a
undefined :: Int))

-- | Use this strategy for generating 'L.LazyByteString's whose chunks are
-- discarded right after they are generated. For example, if you just generate
-- them to write them to a network socket.
{-# INLINE untrimmedStrategy #-}
untrimmedStrategy :: Int -- ^ Size of the first buffer
                  -> Int -- ^ Size of successive buffers
                  -> AllocationStrategy
                  -- ^ An allocation strategy that does not trim any of the
                  -- filled buffers before converting it to a chunk
untrimmedStrategy :: Int -> Int -> AllocationStrategy
untrimmedStrategy Int
firstSize Int
bufSize =
    (Maybe (Buffer, Int) -> IO Buffer)
-> Int -> (Int -> Int -> Bool) -> AllocationStrategy
AllocationStrategy Maybe (Buffer, Int) -> IO Buffer
forall {a}. Maybe (a, Int) -> IO Buffer
nextBuffer (Int -> Int
sanitize Int
bufSize) (\Int
_ Int
_ -> Bool
False)
  where
    {-# INLINE nextBuffer #-}
    nextBuffer :: Maybe (a, Int) -> IO Buffer
nextBuffer Maybe (a, Int)
Nothing             = Int -> IO Buffer
newBuffer (Int -> IO Buffer) -> Int -> IO Buffer
forall a b. (a -> b) -> a -> b
$ Int -> Int
sanitize Int
firstSize
    nextBuffer (Just (a
_, Int
minSize)) = Int -> IO Buffer
newBuffer Int
minSize


-- | Use this strategy for generating 'L.LazyByteString's whose chunks are
-- likely to survive one garbage collection. This strategy trims buffers
-- that are filled less than half in order to avoid spilling too much memory.
{-# INLINE safeStrategy #-}
safeStrategy :: Int  -- ^ Size of first buffer
             -> Int  -- ^ Size of successive buffers
             -> AllocationStrategy
             -- ^ An allocation strategy that guarantees that at least half
             -- of the allocated memory is used for live data
safeStrategy :: Int -> Int -> AllocationStrategy
safeStrategy Int
firstSize Int
bufSize =
    (Maybe (Buffer, Int) -> IO Buffer)
-> Int -> (Int -> Int -> Bool) -> AllocationStrategy
AllocationStrategy Maybe (Buffer, Int) -> IO Buffer
forall {a}. Maybe (a, Int) -> IO Buffer
nextBuffer (Int -> Int
sanitize Int
bufSize) Int -> Int -> Bool
forall {a}. (Ord a, Num a) => a -> a -> Bool
trim
  where
    trim :: a -> a -> Bool
trim a
used a
size                 = a
2 a -> a -> a
forall a. Num a => a -> a -> a
* a
used a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
size
    {-# INLINE nextBuffer #-}
    nextBuffer :: Maybe (a, Int) -> IO Buffer
nextBuffer Maybe (a, Int)
Nothing             = Int -> IO Buffer
newBuffer (Int -> IO Buffer) -> Int -> IO Buffer
forall a b. (a -> b) -> a -> b
$ Int -> Int
sanitize Int
firstSize
    nextBuffer (Just (a
_, Int
minSize)) = Int -> IO Buffer
newBuffer Int
minSize

-- | Execute a 'Builder' and return the generated chunks as a 'L.LazyByteString'.
-- The work is performed lazy, i.e., only when a chunk of the 'L.LazyByteString'
-- is forced.
{-# NOINLINE toLazyByteString #-} -- ensure code is shared
toLazyByteString :: Builder -> L.LazyByteString
toLazyByteString :: Builder -> LazyByteString
toLazyByteString = AllocationStrategy -> LazyByteString -> Builder -> LazyByteString
toLazyByteStringWith
    (Int -> Int -> AllocationStrategy
safeStrategy Int
L.smallChunkSize Int
L.defaultChunkSize) LazyByteString
L.Empty

-- | /Heavy inlining./ Execute a 'Builder' with custom execution parameters.
--
-- This function is inlined despite its heavy code-size to allow fusing with
-- the allocation strategy. For example, the default 'Builder' execution
-- function 'Data.ByteString.Builder.Internal.toLazyByteString' is defined as follows.
--
-- @
-- {-\# NOINLINE toLazyByteString \#-}
-- toLazyByteString =
--   toLazyByteStringWith ('safeStrategy' 'L.smallChunkSize' 'L.defaultChunkSize') L.Empty
-- @
--
-- where @L.Empty@ is the zero-length 'L.LazyByteString'.
--
-- In most cases, the parameters used by 'Data.ByteString.Builder.toLazyByteString' give good
-- performance. A sub-performing case of 'Data.ByteString.Builder.toLazyByteString' is executing short
-- (<128 bytes) 'Builder's. In this case, the allocation overhead for the first
-- 4kb buffer and the trimming cost dominate the cost of executing the
-- 'Builder'. You can avoid this problem using
--
-- >toLazyByteStringWith (safeStrategy 128 smallChunkSize) L.Empty
--
-- This reduces the allocation and trimming overhead, as all generated
-- 'L.LazyByteString's fit into the first buffer and there is no trimming
-- required, if more than 64 bytes and less than 128 bytes are written.
--
{-# INLINE toLazyByteStringWith #-}
toLazyByteStringWith
    :: AllocationStrategy
       -- ^ Buffer allocation strategy to use
    -> L.LazyByteString
       -- ^ 'L.LazyByteString' to use as the tail of the generated lazy
       -- 'L.LazyByteString'
    -> Builder
       -- ^ 'Builder' to execute
    -> L.LazyByteString
       -- ^ Resulting 'L.LazyByteString'
toLazyByteStringWith :: AllocationStrategy -> LazyByteString -> Builder -> LazyByteString
toLazyByteStringWith AllocationStrategy
strategy LazyByteString
k Builder
b =
    AllocationStrategy
-> LazyByteString -> ChunkIOStream () -> LazyByteString
ciosUnitToLazyByteString AllocationStrategy
strategy LazyByteString
k (ChunkIOStream () -> LazyByteString)
-> ChunkIOStream () -> LazyByteString
forall a b. (a -> b) -> a -> b
$ IO (ChunkIOStream ()) -> ChunkIOStream ()
forall a. IO a -> a
unsafeDupablePerformIO (IO (ChunkIOStream ()) -> ChunkIOStream ())
-> IO (ChunkIOStream ()) -> ChunkIOStream ()
forall a b. (a -> b) -> a -> b
$
        AllocationStrategy -> BuildStep () -> IO (ChunkIOStream ())
forall a. AllocationStrategy -> BuildStep a -> IO (ChunkIOStream a)
buildStepToCIOS AllocationStrategy
strategy (Builder -> BuildStep ()
runBuilder Builder
b)

-- | Convert a 'BuildStep' to a 'ChunkIOStream' stream by executing it on
-- 'Buffer's allocated according to the given 'AllocationStrategy'.
{-# INLINE buildStepToCIOS #-}
buildStepToCIOS
    :: forall a.
       AllocationStrategy          -- ^ Buffer allocation strategy to use
    -> BuildStep a                 -- ^ 'BuildStep' to execute
    -> IO (ChunkIOStream a)
buildStepToCIOS :: forall a. AllocationStrategy -> BuildStep a -> IO (ChunkIOStream a)
buildStepToCIOS (AllocationStrategy Maybe (Buffer, Int) -> IO Buffer
nextBuffer Int
bufSize Int -> Int -> Bool
trim) =
    \BuildStep a
step -> Maybe (Buffer, Int) -> IO Buffer
nextBuffer Maybe (Buffer, Int)
forall a. Maybe a
Nothing IO Buffer
-> (Buffer -> IO (ChunkIOStream a)) -> IO (ChunkIOStream a)
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= BuildStep a -> Buffer -> IO (ChunkIOStream a)
fill BuildStep a
step
  where
    fill :: BuildStep a -> Buffer -> IO (ChunkIOStream a)
    fill :: BuildStep a -> Buffer -> IO (ChunkIOStream a)
fill !BuildStep a
step buf :: Buffer
buf@(Buffer ForeignPtr Word8
fpbuf br :: BufferRange
br@(BufferRange Ptr Word8
_ Ptr Word8
pe)) = do
        ChunkIOStream a
res <- BuildStep a
-> (Ptr Word8 -> a -> IO (ChunkIOStream a))
-> (Ptr Word8 -> Int -> BuildStep a -> IO (ChunkIOStream a))
-> (Ptr Word8
    -> StrictByteString -> BuildStep a -> IO (ChunkIOStream a))
-> BufferRange
-> IO (ChunkIOStream a)
forall a b.
BuildStep a
-> (Ptr Word8 -> a -> IO b)
-> (Ptr Word8 -> Int -> BuildStep a -> IO b)
-> (Ptr Word8 -> StrictByteString -> BuildStep a -> IO b)
-> BufferRange
-> IO b
fillWithBuildStep BuildStep a
step Ptr Word8 -> a -> IO (ChunkIOStream a)
doneH Ptr Word8 -> Int -> BuildStep a -> IO (ChunkIOStream a)
fullH Ptr Word8
-> StrictByteString -> BuildStep a -> IO (ChunkIOStream a)
insertChunkH BufferRange
br
        ForeignPtr Word8 -> IO ()
forall a. ForeignPtr a -> IO ()
touchForeignPtr ForeignPtr Word8
fpbuf
        ChunkIOStream a -> IO (ChunkIOStream a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ChunkIOStream a
res
      where
        pbuf :: Ptr Word8
        pbuf :: Ptr Word8
pbuf = ForeignPtr Word8 -> Ptr Word8
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr Word8
fpbuf

        doneH :: Ptr Word8 -> a -> IO (ChunkIOStream a)
        doneH :: Ptr Word8 -> a -> IO (ChunkIOStream a)
doneH Ptr Word8
op' a
x = ChunkIOStream a -> IO (ChunkIOStream a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (ChunkIOStream a -> IO (ChunkIOStream a))
-> ChunkIOStream a -> IO (ChunkIOStream a)
forall a b. (a -> b) -> a -> b
$
            Buffer -> a -> ChunkIOStream a
forall a. Buffer -> a -> ChunkIOStream a
Finished (ForeignPtr Word8 -> BufferRange -> Buffer
Buffer ForeignPtr Word8
fpbuf (Ptr Word8 -> Ptr Word8 -> BufferRange
BufferRange Ptr Word8
op' Ptr Word8
pe)) a
x

        fullH :: Ptr Word8 -> Int -> BuildStep a -> IO (ChunkIOStream a)
        fullH :: Ptr Word8 -> Int -> BuildStep a -> IO (ChunkIOStream a)
fullH Ptr Word8
op' Int
minSize BuildStep a
nextStep =
            Ptr Word8 -> (Bool -> IO (ChunkIOStream a)) -> IO (ChunkIOStream a)
wrapChunk Ptr Word8
op' ((Bool -> IO (ChunkIOStream a)) -> IO (ChunkIOStream a))
-> (Bool -> IO (ChunkIOStream a)) -> IO (ChunkIOStream a)
forall a b. (a -> b) -> a -> b
$ IO (ChunkIOStream a) -> Bool -> IO (ChunkIOStream a)
forall a b. a -> b -> a
const (IO (ChunkIOStream a) -> Bool -> IO (ChunkIOStream a))
-> IO (ChunkIOStream a) -> Bool -> IO (ChunkIOStream a)
forall a b. (a -> b) -> a -> b
$
                Maybe (Buffer, Int) -> IO Buffer
nextBuffer ((Buffer, Int) -> Maybe (Buffer, Int)
forall a. a -> Maybe a
Just (Buffer
buf, Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
minSize Int
bufSize)) IO Buffer
-> (Buffer -> IO (ChunkIOStream a)) -> IO (ChunkIOStream a)
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= BuildStep a -> Buffer -> IO (ChunkIOStream a)
fill BuildStep a
nextStep

        insertChunkH :: Ptr Word8 -> S.StrictByteString -> BuildStep a -> IO (ChunkIOStream a)
        insertChunkH :: Ptr Word8
-> StrictByteString -> BuildStep a -> IO (ChunkIOStream a)
insertChunkH Ptr Word8
op' StrictByteString
bs BuildStep a
nextStep =
            Ptr Word8 -> (Bool -> IO (ChunkIOStream a)) -> IO (ChunkIOStream a)
wrapChunk Ptr Word8
op' ((Bool -> IO (ChunkIOStream a)) -> IO (ChunkIOStream a))
-> (Bool -> IO (ChunkIOStream a)) -> IO (ChunkIOStream a)
forall a b. (a -> b) -> a -> b
$ \Bool
isEmpty -> StrictByteString -> IO (ChunkIOStream a) -> IO (ChunkIOStream a)
forall a.
StrictByteString -> IO (ChunkIOStream a) -> IO (ChunkIOStream a)
yield1 StrictByteString
bs (IO (ChunkIOStream a) -> IO (ChunkIOStream a))
-> IO (ChunkIOStream a) -> IO (ChunkIOStream a)
forall a b. (a -> b) -> a -> b
$
                -- Checking for empty case avoids allocating 'n-1' empty
                -- buffers for 'n' insertChunkH right after each other.
                if Bool
isEmpty
                  then BuildStep a -> Buffer -> IO (ChunkIOStream a)
fill BuildStep a
nextStep Buffer
buf
                  else do Buffer
buf' <- Maybe (Buffer, Int) -> IO Buffer
nextBuffer ((Buffer, Int) -> Maybe (Buffer, Int)
forall a. a -> Maybe a
Just (Buffer
buf, Int
bufSize))
                          BuildStep a -> Buffer -> IO (ChunkIOStream a)
fill BuildStep a
nextStep Buffer
buf'

        -- Wrap and yield a chunk, trimming it if necesary
        {-# INLINE wrapChunk #-}
        wrapChunk :: Ptr Word8 -> (Bool -> IO (ChunkIOStream a)) -> IO (ChunkIOStream a)
        wrapChunk :: Ptr Word8 -> (Bool -> IO (ChunkIOStream a)) -> IO (ChunkIOStream a)
wrapChunk !Ptr Word8
op' Bool -> IO (ChunkIOStream a)
mkCIOS
          | Int
chunkSize Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0      = Bool -> IO (ChunkIOStream a)
mkCIOS Bool
True
          | Int -> Int -> Bool
trim Int
chunkSize Int
size = do
              StrictByteString
bs <- Int -> (ForeignPtr Word8 -> IO ()) -> IO StrictByteString
S.createFp Int
chunkSize ((ForeignPtr Word8 -> IO ()) -> IO StrictByteString)
-> (ForeignPtr Word8 -> IO ()) -> IO StrictByteString
forall a b. (a -> b) -> a -> b
$ \ForeignPtr Word8
fpbuf' ->
                        ForeignPtr Word8 -> ForeignPtr Word8 -> Int -> IO ()
S.memcpyFp ForeignPtr Word8
fpbuf' ForeignPtr Word8
fpbuf Int
chunkSize
              -- It is not safe to re-use the old buffer (see #690),
              -- so we allocate a new buffer after trimming.
              ChunkIOStream a -> IO (ChunkIOStream a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (ChunkIOStream a -> IO (ChunkIOStream a))
-> ChunkIOStream a -> IO (ChunkIOStream a)
forall a b. (a -> b) -> a -> b
$ StrictByteString -> IO (ChunkIOStream a) -> ChunkIOStream a
forall a.
StrictByteString -> IO (ChunkIOStream a) -> ChunkIOStream a
Yield1 StrictByteString
bs (Bool -> IO (ChunkIOStream a)
mkCIOS Bool
False)
          | Bool
otherwise            =
              ChunkIOStream a -> IO (ChunkIOStream a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (ChunkIOStream a -> IO (ChunkIOStream a))
-> ChunkIOStream a -> IO (ChunkIOStream a)
forall a b. (a -> b) -> a -> b
$ StrictByteString -> IO (ChunkIOStream a) -> ChunkIOStream a
forall a.
StrictByteString -> IO (ChunkIOStream a) -> ChunkIOStream a
Yield1 (ForeignPtr Word8 -> Int -> StrictByteString
S.BS ForeignPtr Word8
fpbuf Int
chunkSize) (Bool -> IO (ChunkIOStream a)
mkCIOS Bool
False)
          where
            chunkSize :: Int
chunkSize = Ptr Word8
op' Ptr Word8 -> Ptr Word8 -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr Word8
pbuf
            size :: Int
size      = Ptr Word8
pe  Ptr Word8 -> Ptr Word8 -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr Word8
pbuf