bytestring-0.11.5.0: Data/ByteString/Builder/Internal.hs
{-# 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.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
-- | 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 fpbuf (BufferRange _ ope)) =
ope `minusPtr` unsafeForeignPtrToPtr fpbuf
-- | Allocate a new buffer of the given size.
{-# INLINE newBuffer #-}
newBuffer :: Int -> IO Buffer
newBuffer size = do
fpbuf <- S.mallocByteString size
let pbuf = unsafeForeignPtrToPtr fpbuf
return $! Buffer fpbuf (BufferRange pbuf (pbuf `plusPtr` size))
-- | Convert the filled part of a 'Buffer' to a strict 'S.ByteString'.
{-# INLINE byteStringFromBuffer #-}
byteStringFromBuffer :: Buffer -> S.ByteString
byteStringFromBuffer (Buffer fpbuf (BufferRange op _)) =
S.BS fpbuf (op `minusPtr` unsafeForeignPtrToPtr 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 _ _ trim) buf k
| S.null bs = k
| trim (S.length bs) (bufferSize buf) = L.Chunk (S.copy bs) k
| otherwise = L.Chunk bs k
where
bs = byteStringFromBuffer 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 bs cios | S.null bs = cios
| otherwise = return $ Yield1 bs cios
-- | Convert a @'ChunkIOStream' ()@ to a lazy 'L.ByteString' using
-- 'unsafeDupablePerformIO'.
{-# INLINE ciosUnitToLazyByteString #-}
ciosUnitToLazyByteString :: AllocationStrategy
-> L.ByteString -> ChunkIOStream () -> L.ByteString
ciosUnitToLazyByteString strategy k = go
where
go (Finished buf _) = trimmedChunkFromBuffer strategy buf k
go (Yield1 bs io) = L.Chunk bs $ unsafeDupablePerformIO (go <$> 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 strategy k =
go
where
go (Finished buf x) =
second (trimmedChunkFromBuffer strategy buf) $ k x
go (Yield1 bs io) = second (L.Chunk bs) $ unsafeDupablePerformIO (go <$> 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 = 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 = 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 = 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 step fDone fFull fChunk !br = do
signal <- step br
case signal of
Done op x -> fDone op x
BufferFull minSize op nextStep -> fFull op minSize nextStep
InsertChunk op bs nextStep -> fChunk op bs 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 = Builder
-- | The final build step that returns the 'done' signal.
finalBuildStep :: BuildStep ()
finalBuildStep (BufferRange op _) = return $ Done 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 b = runBuilderWith b finalBuildStep
-- | Run a 'Builder'.
{-# INLINE runBuilderWith #-}
runBuilderWith :: Builder -- ^ 'Builder' to run
-> BuildStep a -- ^ Continuation 'BuildStep'
-> BuildStep a
runBuilderWith (Builder b) = 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 ($)
-- 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 b1) (Builder b2) = Builder $ b1 . b2
instance Semigroup Builder where
{-# INLINE (<>) #-}
(<>) = append
instance Monoid Builder where
{-# INLINE mempty #-}
mempty = empty
{-# INLINE mappend #-}
mappend = (<>)
{-# INLINE mconcat #-}
mconcat = foldr mappend mempty
-- | Flush the current buffer. This introduces a chunk boundary.
{-# INLINE flush #-}
flush :: Builder
flush = builder step
where
step k (BufferRange op _) = return $ insertChunk op S.empty 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 { 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 = 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 p) = p $ \x (BufferRange op _) -> return $ Done op x
instance Functor Put where
fmap f p = Put $ \k -> unPut p (k . 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 $ \k -> a (\a' -> b (\_ -> k 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 $ \k -> a (\_ -> b k)
instance Applicative Put where
{-# INLINE pure #-}
pure x = Put $ \k -> k x
{-# INLINE (<*>) #-}
Put f <*> Put a = Put $ \k -> f (\f' -> a (k . f'))
{-# INLINE (<*) #-}
(<*) = ap_l
{-# INLINE (*>) #-}
(*>) = ap_r
instance Monad Put where
{-# INLINE return #-}
return = pure
{-# INLINE (>>=) #-}
Put m >>= f = Put $ \k -> m (\m' -> unPut (f m') k)
{-# INLINE (>>) #-}
(>>) = (*>)
-- Conversion between Put and Builder
-------------------------------------
-- | Run a 'Builder' as a side-effect of a @'Put' ()@ action.
{-# INLINE[1] putBuilder #-}
putBuilder :: Builder -> Put ()
putBuilder (Builder b) = Put $ \k -> b (k ())
-- | Convert a @'Put' ()@ action to a 'Builder'.
{-# INLINE fromPut #-}
fromPut :: Put () -> Builder
fromPut (Put p) = Builder $ \k -> p (const 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 h p = do
fillHandle 1 (runPut p)
where
fillHandle :: Int -> BuildStep a -> IO a
fillHandle !minFree step = do
next <- wantWritableHandle "hPut" h fillHandle_
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_ h_ = do
makeSpace =<< readIORef refBuf
fillBuffer =<< readIORef refBuf
where
refBuf = haByteBuffer h_
freeSpace buf = IO.bufSize buf - IO.bufR buf
makeSpace buf
| IO.bufSize buf < minFree = do
flushWriteBuffer h_
s <- IO.bufState <$> readIORef refBuf
IO.newByteBuffer minFree s >>= writeIORef refBuf
| freeSpace buf < minFree = flushWriteBuffer h_
| otherwise =
return ()
fillBuffer buf
| freeSpace buf < minFree =
error $ unlines
[ "Data.ByteString.Builder.Internal.hPut: internal error."
, " Not enough space after flush."
, " required: " ++ show minFree
, " free: " ++ show (freeSpace buf)
]
| otherwise = do
let !br = BufferRange op (pBuf `plusPtr` IO.bufSize buf)
res <- fillWithBuildStep step doneH fullH insertChunkH br
touchForeignPtr fpBuf
return res
where
fpBuf = IO.bufRaw buf
pBuf = unsafeForeignPtrToPtr fpBuf
op = pBuf `plusPtr` IO.bufR buf
{-# INLINE updateBufR #-}
updateBufR op' = do
let !off' = op' `minusPtr` pBuf
!buf' = buf {IO.bufR = off'}
writeIORef refBuf buf'
doneH op' x = do
updateBufR 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 haBufferMode h_ of
BlockBuffering _ -> return $ return x
_line_or_no_buffering -> return $ hFlush h >> return x
fullH op' minSize nextStep = do
updateBufR op'
return $ fillHandle minSize nextStep
-- 'fillHandle' will flush the buffer (provided there is
-- really less than @minSize@ space left) before executing
-- the 'nextStep'.
insertChunkH op' bs nextStep = do
updateBufR op'
return $ do
S.hPut h bs
fillHandle 1 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 = putToLazyByteStringWith
(safeStrategy L.smallChunkSize L.defaultChunkSize) (, 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 strategy k p =
ciosToLazyByteString strategy k $ unsafeDupablePerformIO $
buildStepToCIOS strategy (runPut 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 minFree =
builder step
where
step k br@(BufferRange op ope)
| ope `minusPtr` op < minFree = return $ bufferFull minFree op k
| otherwise = k br
-- | Copy the bytes from a 'BufferRange' into the output stream.
wrappedBytesCopyStep :: BufferRange -- ^ Input 'BufferRange'.
-> BuildStep a -> BuildStep a
wrappedBytesCopyStep (BufferRange ip0 ipe) k =
go ip0
where
go !ip (BufferRange op ope)
| inpRemaining <= outRemaining = do
copyBytes op ip inpRemaining
let !br' = BufferRange (op `plusPtr` inpRemaining) ope
k br'
| otherwise = do
copyBytes op ip outRemaining
let !ip' = ip `plusPtr` outRemaining
return $ bufferFull 1 ope (go ip')
where
outRemaining = ope `minusPtr` op
inpRemaining = ipe `minusPtr` 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 maxCopySize =
\bs -> builder $ step bs
where
step bs@(S.BS _ len) !k br@(BufferRange !op _)
| len <= maxCopySize = byteStringCopyStep bs k br
| otherwise = return $ insertChunk op bs 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 = \bs -> builder $ byteStringCopyStep bs
{-# INLINE byteStringCopyStep #-}
byteStringCopyStep :: S.ByteString -> BuildStep a -> BuildStep a
byteStringCopyStep (S.BS ifp isize) !k0 br0@(BufferRange op ope)
-- Ensure that the common case is not recursive and therefore yields
-- better code.
| op' <= ope = do copyBytes op ip isize
touchForeignPtr ifp
k0 (BufferRange op' ope)
| otherwise = wrappedBytesCopyStep (BufferRange ip ipe) k br0
where
op' = op `plusPtr` isize
ip = unsafeForeignPtrToPtr ifp
ipe = ip `plusPtr` isize
k br = do touchForeignPtr ifp -- input consumed: OK to release here
k0 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 =
\bs -> builder $ \k (BufferRange op _) -> return $ insertChunk op bs k
-- Short bytestrings
------------------------------------------------------------------------------
-- | Construct a 'Builder' that copies the 'SH.ShortByteString'.
--
{-# INLINE shortByteString #-}
shortByteString :: Sh.ShortByteString -> Builder
shortByteString = \sbs -> builder $ shortByteStringCopyStep 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 !sbs k =
go 0 (Sh.length sbs)
where
go !ip !ipe (BufferRange op ope)
| inpRemaining <= outRemaining = do
Sh.copyToPtr sbs ip op inpRemaining
let !br' = BufferRange (op `plusPtr` inpRemaining) ope
k br'
| otherwise = do
Sh.copyToPtr sbs ip op outRemaining
let !ip' = ip + outRemaining
return $ bufferFull 1 ope (go ip' ipe)
where
outRemaining = ope `minusPtr` op
inpRemaining = ipe - 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 maxCopySize =
L.foldrChunks (\bs b -> byteStringThreshold maxCopySize bs `mappend` b) 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 =
L.foldrChunks (\bs b -> byteStringCopy bs `mappend` b) mempty
-- | Construct a 'Builder' that inserts all chunks of the lazy 'L.ByteString'
-- directly.
--
{-# INLINE lazyByteStringInsert #-}
lazyByteStringInsert :: L.ByteString -> Builder
lazyByteStringInsert =
L.foldrChunks (\bs b -> byteStringInsert bs `mappend` b) 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 = byteStringThreshold 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 = lazyByteStringThreshold 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 = 2 * 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 = AllocationStrategy
-- | Sanitize a buffer size; i.e., make it at least the size of an 'Int'.
{-# INLINE sanitize #-}
sanitize :: Int -> Int
sanitize = max (sizeOf (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 firstSize bufSize =
AllocationStrategy nextBuffer (sanitize bufSize) (\_ _ -> False)
where
{-# INLINE nextBuffer #-}
nextBuffer Nothing = newBuffer $ sanitize firstSize
nextBuffer (Just (_, minSize)) = newBuffer 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 firstSize bufSize =
AllocationStrategy nextBuffer (sanitize bufSize) trim
where
trim used size = 2 * used < size
{-# INLINE nextBuffer #-}
nextBuffer Nothing = newBuffer $ sanitize firstSize
nextBuffer (Just (_, minSize)) = newBuffer 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 strategy k b =
ciosUnitToLazyByteString strategy k $ unsafeDupablePerformIO $
buildStepToCIOS strategy (runBuilder 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 (AllocationStrategy nextBuffer bufSize trim) =
\step -> nextBuffer Nothing >>= fill step
where
fill :: BuildStep a -> Buffer -> IO (ChunkIOStream a)
fill !step buf@(Buffer fpbuf br@(BufferRange _ pe)) = do
res <- fillWithBuildStep step doneH fullH insertChunkH br
touchForeignPtr fpbuf
return res
where
pbuf :: Ptr Word8
pbuf = unsafeForeignPtrToPtr fpbuf
doneH :: Ptr Word8 -> a -> IO (ChunkIOStream a)
doneH op' x = return $
Finished (Buffer fpbuf (BufferRange op' pe)) x
fullH :: Ptr Word8 -> Int -> BuildStep a -> IO (ChunkIOStream a)
fullH op' minSize nextStep =
wrapChunk op' $ const $
nextBuffer (Just (buf, max minSize bufSize)) >>= fill nextStep
insertChunkH :: Ptr Word8 -> S.ByteString -> BuildStep a -> IO (ChunkIOStream a)
insertChunkH op' bs nextStep =
wrapChunk op' $ \isEmpty -> yield1 bs $
-- Checking for empty case avoids allocating 'n-1' empty
-- buffers for 'n' insertChunkH right after each other.
if isEmpty
then fill nextStep (Buffer fpbuf (BufferRange pbuf pe))
else do buf' <- nextBuffer (Just (buf, bufSize))
fill nextStep buf'
-- Wrap and yield a chunk, trimming it if necesary
{-# INLINE wrapChunk #-}
wrapChunk :: Ptr Word8 -> (Bool -> IO (ChunkIOStream a)) -> IO (ChunkIOStream a)
wrapChunk !op' mkCIOS
| chunkSize == 0 = mkCIOS True
| trim chunkSize size = do
bs <- S.createFp chunkSize $ \fpbuf' ->
S.memcpyFp fpbuf' fpbuf chunkSize
-- Instead of allocating a new buffer after trimming,
-- we re-use the old buffer and consider it empty.
return $ Yield1 bs (mkCIOS True)
| otherwise =
return $ Yield1 (S.BS fpbuf chunkSize) (mkCIOS False)
where
chunkSize = op' `minusPtr` pbuf
size = pe `minusPtr` pbuf