{-# LANGUAGE ScopedTypeVariables, CPP, BangPatterns, RankNTypes, TupleSections #-}
{-# LANGUAGE Unsafe #-}
{-# 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
  , 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)

#if !(MIN_VERSION_base(4,11,0))
import           Data.Semigroup (Semigroup((<>)))
#endif

import qualified Data.ByteString               as S
import qualified Data.ByteString.Internal      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

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


-- | 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 (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 strict 'S.ByteString'.
{-# INLINE byteStringFromBuffer #-}
byteStringFromBuffer :: Buffer -> S.ByteString
byteStringFromBuffer :: Buffer -> ByteString
byteStringFromBuffer (Buffer ForeignPtr Word8
fpbuf (BufferRange Ptr Word8
op Ptr Word8
_)) =
    ForeignPtr Word8 -> Int -> ByteString
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 lazy 'L.ByteString'
-- trimming it if necessary.
{-# INLINE trimmedChunkFromBuffer #-}
trimmedChunkFromBuffer :: AllocationStrategy -> Buffer
                       -> L.ByteString -> L.ByteString
trimmedChunkFromBuffer :: AllocationStrategy -> Buffer -> ByteString -> ByteString
trimmedChunkFromBuffer (AllocationStrategy Maybe (Buffer, Int) -> IO Buffer
_ Int
_ Int -> Int -> Bool
trim) Buffer
buf ByteString
k
  | ByteString -> Bool
S.null ByteString
bs                           = ByteString
k
  | Int -> Int -> Bool
trim (ByteString -> Int
S.length ByteString
bs) (Buffer -> Int
bufferSize Buffer
buf) = ByteString -> ByteString -> ByteString
L.Chunk (ByteString -> ByteString
S.copy ByteString
bs) ByteString
k
  | Bool
otherwise                           = ByteString -> ByteString -> ByteString
L.Chunk ByteString
bs          ByteString
k
  where
    bs :: ByteString
bs = Buffer -> ByteString
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.ByteString (IO (ChunkIOStream a))
       -- ^ Yield a /non-empty/ strict 'S.ByteString'.

-- | A smart constructor for yielding one chunk that ignores the chunk if
-- it is empty.
{-# INLINE yield1 #-}
yield1 :: S.ByteString -> IO (ChunkIOStream a) -> IO (ChunkIOStream a)
yield1 :: ByteString -> IO (ChunkIOStream a) -> IO (ChunkIOStream a)
yield1 ByteString
bs IO (ChunkIOStream a)
cios | ByteString -> Bool
S.null ByteString
bs = IO (ChunkIOStream a)
cios
               | Bool
otherwise = ChunkIOStream a -> IO (ChunkIOStream 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
$ ByteString -> IO (ChunkIOStream a) -> ChunkIOStream a
forall a. ByteString -> IO (ChunkIOStream a) -> ChunkIOStream a
Yield1 ByteString
bs IO (ChunkIOStream a)
cios

-- | Convert a @'ChunkIOStream' ()@ to a lazy 'L.ByteString' using
-- 'unsafeDupablePerformIO'.
{-# INLINE ciosUnitToLazyByteString #-}
ciosUnitToLazyByteString :: AllocationStrategy
                         -> L.ByteString -> ChunkIOStream () -> L.ByteString
ciosUnitToLazyByteString :: AllocationStrategy -> ByteString -> ChunkIOStream () -> ByteString
ciosUnitToLazyByteString AllocationStrategy
strategy ByteString
k = ChunkIOStream () -> ByteString
forall a. ChunkIOStream a -> ByteString
go
  where
    go :: ChunkIOStream a -> ByteString
go (Finished Buffer
buf a
_) = AllocationStrategy -> Buffer -> ByteString -> ByteString
trimmedChunkFromBuffer AllocationStrategy
strategy Buffer
buf ByteString
k
    go (Yield1 ByteString
bs IO (ChunkIOStream a)
io)   = ByteString -> ByteString -> ByteString
L.Chunk ByteString
bs (ByteString -> ByteString) -> ByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ IO ByteString -> ByteString
forall a. IO a -> a
unsafeDupablePerformIO (ChunkIOStream a -> ByteString
go (ChunkIOStream a -> ByteString)
-> IO (ChunkIOStream a) -> IO ByteString
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.ByteString' using 'unsafeDupablePerformIO'.
{-# INLINE ciosToLazyByteString #-}
ciosToLazyByteString :: AllocationStrategy
                     -> (a -> (b, L.ByteString))
                     -> ChunkIOStream a
                     -> (b, L.ByteString)
ciosToLazyByteString :: AllocationStrategy
-> (a -> (b, ByteString)) -> ChunkIOStream a -> (b, ByteString)
ciosToLazyByteString AllocationStrategy
strategy a -> (b, ByteString)
k =
    ChunkIOStream a -> (b, ByteString)
go
  where
    go :: ChunkIOStream a -> (b, ByteString)
go (Finished Buffer
buf a
x) =
        (ByteString -> ByteString) -> (b, ByteString) -> (b, ByteString)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (AllocationStrategy -> Buffer -> ByteString -> ByteString
trimmedChunkFromBuffer AllocationStrategy
strategy Buffer
buf) ((b, ByteString) -> (b, ByteString))
-> (b, ByteString) -> (b, ByteString)
forall a b. (a -> b) -> a -> b
$ a -> (b, ByteString)
k a
x
    go (Yield1 ByteString
bs IO (ChunkIOStream a)
io)   = (ByteString -> ByteString) -> (b, ByteString) -> (b, ByteString)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (ByteString -> ByteString -> ByteString
L.Chunk ByteString
bs) ((b, ByteString) -> (b, ByteString))
-> (b, ByteString) -> (b, ByteString)
forall a b. (a -> b) -> a -> b
$ IO (b, ByteString) -> (b, ByteString)
forall a. IO a -> a
unsafeDupablePerformIO (ChunkIOStream a -> (b, ByteString)
go (ChunkIOStream a -> (b, ByteString))
-> IO (ChunkIOStream a) -> IO (b, ByteString)
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.ByteString
                     (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 :: 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 :: 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.ByteString' chunk should be inserted directly.
{-# INLINE insertChunk #-}
insertChunk :: Ptr Word8
            -- ^ Next free byte in current 'BufferRange'
            -> S.ByteString
            -- ^ Chunk to insert.
            -> BuildStep a
            -- ^ 'BuildStep' to run on next 'BufferRange'
            -> BuildSignal a
insertChunk :: Ptr Word8 -> ByteString -> BuildStep a -> BuildSignal a
insertChunk = Ptr Word8 -> ByteString -> BuildStep a -> BuildSignal a
forall a. Ptr Word8 -> ByteString -> 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.ByteString -> BuildStep a -> IO b)
    -- ^ Handling the 'insertChunk' signal
    -> BufferRange
    -- ^ Buffer range to fill.
    -> IO b
    -- ^ Value computed while filling this 'BufferRange'.
fillWithBuildStep :: BuildStep a
-> (Ptr Word8 -> a -> IO b)
-> (Ptr Word8 -> Int -> BuildStep a -> IO b)
-> (Ptr Word8 -> ByteString -> 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 -> ByteString -> 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 ByteString
bs BuildStep a
nextStep     -> Ptr Word8 -> ByteString -> BuildStep a -> IO b
fChunk Ptr Word8
op ByteString
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.
        -- Moroever, 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 (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 :: 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 forall r. BuildStep r -> BuildStep r
forall a b. (a -> b) -> a -> b
($)
-- 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.

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

instance Semigroup Builder where
  {-# INLINE (<>) #-}
  <> :: Builder -> Builder -> Builder
(<>) = Builder -> Builder -> Builder
append

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

-- | Flush the current buffer. This introduces a chunk boundary.
{-# INLINE flush #-}
flush :: Builder
flush :: Builder
flush = (forall r. BuildStep r -> BuildStep r) -> Builder
builder 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 (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 -> ByteString -> BuildStep a -> BuildSignal a
forall a. Ptr Word8 -> ByteString -> BuildStep a -> BuildSignal a
insertChunk Ptr Word8
op ByteString
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 reportng 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 { 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. Moroever, 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 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 :: 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 (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 :: (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 -> (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 :: 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 :: 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 :: 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 <*> :: 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 (<*) #-}
  <* :: Put a -> Put b -> Put a
(<*) = Put a -> Put b -> Put a
forall a b. Put a -> Put b -> Put a
ap_l
  {-# INLINE (*>) #-}
  *> :: 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 :: a -> Put a
return = a -> Put a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
  {-# INLINE (>>=) #-}
  Put forall r. (a -> BuildStep r) -> BuildStep r
m >>= :: 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 -> (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 (>>) #-}
  >> :: Put a -> Put b -> Put 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 :: 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 <- String -> Handle -> (Handle__ -> IO (IO a)) -> IO (IO a)
forall a. String -> Handle -> (Handle__ -> IO a) -> IO a
wantWritableHandle String
"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 interuptible operation that may raise an async. exception
        --      as long as we are inside 'wantWritableHandle' and the like.
        --      We possibly run the interuptible '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 (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 (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 =
                  String -> IO (IO a)
forall a. HasCallStack => String -> a
error (String -> IO (IO a)) -> String -> IO (IO a)
forall a b. (a -> b) -> a -> b
$ [String] -> String
unlines
                    [ String
"Data.ByteString.Builder.Internal.hPut: internal error."
                    , String
"  Not enough space after flush."
                    , String
"    required: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
minFree
                    , String
"    free: "     String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
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 -> ByteString -> 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 -> ByteString -> BuildStep a -> IO b)
-> BufferRange
-> IO b
fillWithBuildStep BuildStep a
step Ptr Word8 -> a -> IO (IO a)
forall a b. Ptr a -> b -> IO (IO b)
doneH Ptr Word8 -> Int -> BuildStep a -> IO (IO a)
forall a. Ptr a -> Int -> BuildStep a -> IO (IO a)
fullH Ptr Word8 -> ByteString -> BuildStep a -> IO (IO a)
forall a. Ptr a -> ByteString -> 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 (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 {bufR :: Int
IO.bufR = Int
off'}
                    IORef (Buffer Word8) -> Buffer Word8 -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (Buffer Word8)
refBuf Buffer Word8
buf'

                doneH :: Ptr a -> b -> IO (IO b)
doneH Ptr a
op' b
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 b -> IO (IO b)
forall (m :: * -> *) a. Monad m => a -> m a
return (IO b -> IO (IO b)) -> IO b -> IO (IO b)
forall a b. (a -> b) -> a -> b
$ b -> IO b
forall (m :: * -> *) a. Monad m => a -> m a
return b
x
                        BufferMode
_line_or_no_buffering -> IO b -> IO (IO b)
forall (m :: * -> *) a. Monad m => a -> m a
return (IO b -> IO (IO b)) -> IO b -> IO (IO b)
forall a b. (a -> b) -> a -> b
$ Handle -> IO ()
hFlush Handle
h IO () -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> b -> IO b
forall (m :: * -> *) a. Monad m => a -> m a
return b
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 (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 -> ByteString -> BuildStep a -> IO (IO a)
insertChunkH Ptr a
op' ByteString
bs BuildStep a
nextStep = do
                    Ptr a -> IO ()
forall a. Ptr a -> IO ()
updateBufR Ptr a
op'
                    IO a -> IO (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 -> ByteString -> IO ()
S.hPut Handle
h ByteString
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 lazy 'L.ByteString'.
--
-- 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.ByteString' -> DecodingState -> 'Put' (Maybe DecodingState)
--decodeBase64 = ...
-- @
--
-- The above function takes a strict 'S.ByteString' 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.ByteString' was
-- successful. The checking if the strict 'S.ByteString' 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.ByteString)  -- ^ Result and lazy 'L.ByteString'
                          -- written as its side-effect
putToLazyByteString :: Put a -> (a, ByteString)
putToLazyByteString = AllocationStrategy
-> (a -> (a, ByteString)) -> Put a -> (a, ByteString)
forall a b.
AllocationStrategy
-> (a -> (b, ByteString)) -> Put a -> (b, ByteString)
putToLazyByteStringWith
    (Int -> Int -> AllocationStrategy
safeStrategy Int
L.smallChunkSize Int
L.defaultChunkSize) (, ByteString
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.ByteString))
       -- ^ 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.ByteString)
       -- ^ Resulting lazy 'L.ByteString'
putToLazyByteStringWith :: AllocationStrategy
-> (a -> (b, ByteString)) -> Put a -> (b, ByteString)
putToLazyByteStringWith AllocationStrategy
strategy a -> (b, ByteString)
k Put a
p =
    AllocationStrategy
-> (a -> (b, ByteString)) -> ChunkIOStream a -> (b, ByteString)
forall a b.
AllocationStrategy
-> (a -> (b, ByteString)) -> ChunkIOStream a -> (b, ByteString)
ciosToLazyByteString AllocationStrategy
strategy a -> (b, ByteString)
k (ChunkIOStream a -> (b, ByteString))
-> ChunkIOStream a -> (b, ByteString)
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 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 (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 'BufferRange' into the output stream.
wrappedBytesCopyStep :: BufferRange  -- ^ Input 'BufferRange'.
                     -> BuildStep a -> BuildStep a
wrappedBytesCopyStep :: BufferRange -> BuildStep a -> BuildStep a
wrappedBytesCopyStep (BufferRange Ptr Word8
ip0 Ptr Word8
ipe) BuildStep a
k =
    Ptr Word8 -> BuildStep a
go Ptr Word8
ip0
  where
    go :: Ptr Word8 -> BuildStep a
go !Ptr Word8
ip (BufferRange Ptr Word8
op Ptr Word8
ope)
      | Int
inpRemaining Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
outRemaining = do
          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
          Ptr Word8 -> Ptr Word8 -> Int -> IO ()
forall a. Ptr a -> Ptr a -> Int -> IO ()
copyBytes Ptr Word8
op Ptr Word8
ip Int
outRemaining
          let !ip' :: Ptr b
ip' = Ptr Word8
ip Ptr Word8 -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
outRemaining
          BuildSignal a -> IO (BuildSignal 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 (Ptr Word8 -> BuildStep a
go Ptr Word8
forall b. Ptr b
ip')
      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 = Ptr Word8
ipe Ptr Word8 -> Ptr Word8 -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr Word8
ip


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


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

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

{-# INLINE byteStringCopyStep #-}
byteStringCopyStep :: S.ByteString -> BuildStep a -> BuildStep a
byteStringCopyStep :: ByteString -> BuildStep a -> BuildStep a
byteStringCopyStep (S.BS ForeignPtr Word8
ifp Int
isize) !BuildStep a
k0 br0 :: BufferRange
br0@(BufferRange Ptr Word8
op Ptr Word8
ope)
    -- Ensure that the common case is not recursive and therefore yields
    -- better code.
    | Ptr Word8
forall b. Ptr b
op' Ptr Word8 -> Ptr Word8 -> Bool
forall a. Ord a => a -> a -> Bool
<= Ptr Word8
ope = do Ptr Word8 -> Ptr Word8 -> Int -> IO ()
forall a. Ptr a -> Ptr a -> Int -> IO ()
copyBytes Ptr Word8
op Ptr Word8
ip Int
isize
                      ForeignPtr Word8 -> IO ()
forall a. ForeignPtr a -> IO ()
touchForeignPtr ForeignPtr Word8
ifp
                      BuildStep a
k0 (Ptr Word8 -> Ptr Word8 -> BufferRange
BufferRange Ptr Word8
forall b. Ptr b
op' Ptr Word8
ope)
    | Bool
otherwise  = BufferRange -> BuildStep a -> BuildStep a
forall a. BufferRange -> BuildStep a -> BuildStep a
wrappedBytesCopyStep (Ptr Word8 -> Ptr Word8 -> BufferRange
BufferRange Ptr Word8
ip Ptr Word8
forall b. Ptr b
ipe) BuildStep a
k BufferRange
br0
  where
    op' :: Ptr b
op'  = Ptr Word8
op Ptr Word8 -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
isize
    ip :: Ptr Word8
ip   = ForeignPtr Word8 -> Ptr Word8
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr Word8
ifp
    ipe :: Ptr b
ipe  = Ptr Word8
ip Ptr Word8 -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
isize
    k :: BuildStep a
k BufferRange
br = do ForeignPtr Word8 -> IO ()
forall a. ForeignPtr a -> IO ()
touchForeignPtr ForeignPtr Word8
ifp  -- input consumed: OK to release here
              BuildStep a
k0 BufferRange
br

-- | Construct a 'Builder' that always inserts the strict 'S.ByteString'
-- 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.ByteString's. Otherwise, the generated chunks are too
-- fragmented to be processed efficiently afterwards.
--
{-# INLINE byteStringInsert #-}
byteStringInsert :: S.ByteString -> Builder
byteStringInsert :: ByteString -> Builder
byteStringInsert =
    \ByteString
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 (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 -> ByteString -> BuildStep r -> BuildSignal r
forall a. Ptr Word8 -> ByteString -> BuildStep a -> BuildSignal a
insertChunk Ptr Word8
op ByteString
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 :: 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 (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 lazy 'L.ByteString'.
--
{-# INLINE lazyByteStringThreshold #-}
lazyByteStringThreshold :: Int -> L.ByteString -> Builder
lazyByteStringThreshold :: Int -> ByteString -> Builder
lazyByteStringThreshold Int
maxCopySize =
    (ByteString -> Builder -> Builder)
-> Builder -> ByteString -> Builder
forall a. (ByteString -> a -> a) -> a -> ByteString -> a
L.foldrChunks (\ByteString
bs Builder
b -> Int -> ByteString -> Builder
byteStringThreshold Int
maxCopySize ByteString
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 lazy 'L.ByteString'.
--
{-# INLINE lazyByteStringCopy #-}
lazyByteStringCopy :: L.ByteString -> Builder
lazyByteStringCopy :: ByteString -> Builder
lazyByteStringCopy =
    (ByteString -> Builder -> Builder)
-> Builder -> ByteString -> Builder
forall a. (ByteString -> a -> a) -> a -> ByteString -> a
L.foldrChunks (\ByteString
bs Builder
b -> ByteString -> Builder
byteStringCopy ByteString
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 lazy 'L.ByteString'
-- directly.
--
{-# INLINE lazyByteStringInsert #-}
lazyByteStringInsert :: L.ByteString -> Builder
lazyByteStringInsert :: ByteString -> Builder
lazyByteStringInsert =
    (ByteString -> Builder -> Builder)
-> Builder -> ByteString -> Builder
forall a. (ByteString -> a -> a) -> a -> ByteString -> a
L.foldrChunks (\ByteString
bs Builder
b -> ByteString -> Builder
byteStringInsert ByteString
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 strict
-- 'S.ByteString'.
-- The 'Builder' inserts large 'S.ByteString's directly, but copies small ones
-- to ensure that the generated chunks are large on average.
--
{-# INLINE byteString #-}
byteString :: S.ByteString -> Builder
byteString :: ByteString -> Builder
byteString = Int -> ByteString -> Builder
byteStringThreshold Int
maximalCopySize

-- | Create a 'Builder' denoting the same sequence of bytes as a lazy
-- 'L.ByteString'.
-- The 'Builder' inserts large chunks of the lazy 'L.ByteString' directly,
-- but copies small ones to ensure that the generated chunks are large on
-- average.
--
{-# INLINE lazyByteString #-}
lazyByteString :: L.ByteString -> Builder
lazyByteString :: ByteString -> Builder
lazyByteString = Int -> ByteString -> 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.ByteString' 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.

-- The strategy
--
-- > 'AllocationStrategy' firstBufSize bufSize trim
--
-- states that the first buffer is of size @firstBufSize@, all following buffers
-- are of size @bufSize@, and a buffer of size @n@ filled with @k@ bytes should
-- be trimmed iff @trim k n@ is 'True'.
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 the @oldBuf@, if it can guarantee that this
     -- referentially transparent and @oldBuf@ is large enough.
  -> 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 lazy 'L.ByteString'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 lazy 'L.ByteString'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

-- | /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.toLazyByteString' is defined as follows.
--
-- @
-- {-\# NOINLINE toLazyByteString \#-}
-- toLazyByteString =
--   toLazyByteStringWith ('safeStrategy' 'L.smallChunkSize' 'L.defaultChunkSize') L.empty
-- @
--
-- where @L.empty@ is the zero-length lazy 'L.ByteString'.
--
-- 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.ByteString'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.ByteString
       -- ^ Lazy 'L.ByteString' to use as the tail of the generated lazy
       -- 'L.ByteString'
    -> Builder
       -- ^ 'Builder' to execute
    -> L.ByteString
       -- ^ Resulting lazy 'L.ByteString'
toLazyByteStringWith :: AllocationStrategy -> ByteString -> Builder -> ByteString
toLazyByteStringWith AllocationStrategy
strategy ByteString
k Builder
b =
    AllocationStrategy -> ByteString -> ChunkIOStream () -> ByteString
ciosUnitToLazyByteString AllocationStrategy
strategy ByteString
k (ChunkIOStream () -> ByteString) -> ChunkIOStream () -> ByteString
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
    :: AllocationStrategy          -- ^ Buffer allocation strategy to use
    -> BuildStep a                 -- ^ 'BuildStep' to execute
    -> IO (ChunkIOStream a)
buildStepToCIOS :: 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 (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= BuildStep a -> Buffer -> IO (ChunkIOStream a)
forall a. BuildStep a -> Buffer -> IO (ChunkIOStream a)
fill BuildStep a
step
  where
    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 -> ByteString -> 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 -> ByteString -> BuildStep a -> IO b)
-> BufferRange
-> IO b
fillWithBuildStep BuildStep a
step Ptr Word8 -> a -> IO (ChunkIOStream a)
forall (m :: * -> *) a.
Monad m =>
Ptr Word8 -> a -> m (ChunkIOStream a)
doneH Ptr Word8 -> Int -> BuildStep a -> IO (ChunkIOStream a)
forall a. Ptr a -> Int -> BuildStep a -> IO (ChunkIOStream a)
fullH Ptr Word8 -> ByteString -> BuildStep a -> IO (ChunkIOStream a)
forall a.
Ptr a -> ByteString -> 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 (m :: * -> *) a. Monad m => a -> m a
return ChunkIOStream a
res
      where
        pbuf :: Ptr Word8
pbuf = ForeignPtr Word8 -> Ptr Word8
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr Word8
fpbuf

        doneH :: Ptr Word8 -> a -> m (ChunkIOStream a)
doneH Ptr Word8
op' a
x = ChunkIOStream a -> m (ChunkIOStream a)
forall (m :: * -> *) a. Monad m => a -> m a
return (ChunkIOStream a -> m (ChunkIOStream a))
-> ChunkIOStream a -> m (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 a -> Int -> BuildStep a -> IO (ChunkIOStream a)
fullH Ptr a
op' Int
minSize BuildStep a
nextStep =
            Ptr a -> (Bool -> IO (ChunkIOStream a)) -> IO (ChunkIOStream a)
forall a a.
Ptr a -> (Bool -> IO (ChunkIOStream a)) -> IO (ChunkIOStream a)
wrapChunk Ptr a
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 (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= BuildStep a -> Buffer -> IO (ChunkIOStream a)
fill BuildStep a
nextStep

        insertChunkH :: Ptr a -> ByteString -> BuildStep a -> IO (ChunkIOStream a)
insertChunkH Ptr a
op' ByteString
bs BuildStep a
nextStep =
            Ptr a -> (Bool -> IO (ChunkIOStream a)) -> IO (ChunkIOStream a)
forall a a.
Ptr a -> (Bool -> IO (ChunkIOStream a)) -> IO (ChunkIOStream a)
wrapChunk Ptr a
op' ((Bool -> IO (ChunkIOStream a)) -> IO (ChunkIOStream a))
-> (Bool -> IO (ChunkIOStream a)) -> IO (ChunkIOStream a)
forall a b. (a -> b) -> a -> b
$ \Bool
isEmpty -> ByteString -> IO (ChunkIOStream a) -> IO (ChunkIOStream a)
forall a.
ByteString -> IO (ChunkIOStream a) -> IO (ChunkIOStream a)
yield1 ByteString
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 a -> (Bool -> IO (ChunkIOStream a)) -> IO (ChunkIOStream a)
wrapChunk !Ptr a
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
              ByteString
bs <- Int -> (Ptr Word8 -> IO ()) -> IO ByteString
S.create Int
chunkSize ((Ptr Word8 -> IO ()) -> IO ByteString)
-> (Ptr Word8 -> IO ()) -> IO ByteString
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
pbuf' ->
                        Ptr Word8 -> Ptr Word8 -> Int -> IO ()
forall a. Ptr a -> Ptr a -> Int -> IO ()
copyBytes Ptr Word8
pbuf' Ptr Word8
pbuf Int
chunkSize
              -- FIXME: We could reuse the trimmed buffer here.
              ChunkIOStream a -> IO (ChunkIOStream 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
$ ByteString -> IO (ChunkIOStream a) -> ChunkIOStream a
forall a. ByteString -> IO (ChunkIOStream a) -> ChunkIOStream a
Yield1 ByteString
bs (Bool -> IO (ChunkIOStream a)
mkCIOS Bool
False)
          | Bool
otherwise            =
              ChunkIOStream a -> IO (ChunkIOStream 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
$ ByteString -> IO (ChunkIOStream a) -> ChunkIOStream a
forall a. ByteString -> IO (ChunkIOStream a) -> ChunkIOStream a
Yield1 (ForeignPtr Word8 -> Int -> ByteString
S.BS ForeignPtr Word8
fpbuf Int
chunkSize) (Bool -> IO (ChunkIOStream a)
mkCIOS Bool
False)
          where
            chunkSize :: Int
chunkSize = Ptr a
op' Ptr a -> 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