{-# LANGUAGE MagicHash #-}
{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE RankNTypes #-}

-- | This is an internal module; its interface is unstable.
module Data.ByteString.FastBuilder.Internal
  (
  -- * Builder and related types
    Builder(..)
  , BuilderState
  , DataSink(..)
  , DynamicSink(..)
  , Queue(..)
  , Request(..)
  , Response(..)

  -- * Internally used exceptions
  , SuspendBuilderException(..)
  , ChunkOverflowException(..)

  -- * Builder building blocks
  , BuildM(..)
  , mkBuilder
  , useBuilder
  , getSink
  , getCur
  , getEnd
  , setCur
  , setEnd

  -- * Running builders
  , runBuilder
  , toLazyByteString
  , toLazyByteStringWith
  , toStrictByteString
  , hPutBuilder
  , hPutBuilderLen
  , hPutBuilderWith

  -- * Basic builders
  , primBounded
  , primFixed
  , primMapListBounded
  , primMapListFixed
  , byteString
  , byteStringThreshold
  , byteStringCopy
  , byteStringCopyNoCheck
  , byteStringInsert
  , unsafeCString
  , unsafeCStringLen
  , ensureBytes
  , getBytes

  -- * Performance tuning
  , rebuild
  ) where

import Control.Concurrent (forkIOWithUnmask, myThreadId)
import Control.Concurrent.MVar
import qualified Control.Exception as E
import Control.Monad
import qualified Data.ByteString as S
import qualified Data.ByteString.Internal as S
import qualified Data.ByteString.Unsafe as S
import qualified Data.ByteString.Lazy as L
import Data.IORef
import Data.Semigroup as Sem
import Data.String
import Data.Word
import Foreign.C.String
import Foreign.C.Types
import Foreign.ForeignPtr
import Foreign.ForeignPtr.Unsafe
import Foreign.Marshal.Utils
import Foreign.Ptr
import qualified System.IO as IO
import System.IO.Unsafe

import GHC.Exts (Addr#, State#, RealWorld, Ptr(..), Int(..), Int#)
import GHC.Magic (oneShot)
import GHC.IO (IO(..), unIO)
import GHC.CString (unpackCString#)

import qualified Data.ByteString.Builder.Prim as P
import qualified Data.ByteString.Builder.Prim.Internal as PI
import qualified Data.ByteString.Builder.Extra as X

-- | 'Builder' is an auxiliary type for efficiently generating a long
-- 'L.ByteString'. It is isomorphic to lazy 'L.ByteString', but offers
-- constant-time concatanation via '<>'.
--
-- Use 'toLazyByteString' to turn a 'Builder' into a 'L.ByteString'
newtype Builder = Builder
  { Builder -> DataSink -> BuilderState -> BuilderState
unBuilder :: DataSink -> BuilderState -> BuilderState
  }
  -- It takes and returns two pointers, "cur" and "end". "cur" points to
  -- the next location to put bytes to, and "end" points to the end of the
  -- buffer.

-- | The state of a builder. The components are:
--
-- * The "cur" pointer
-- * The "end" pointer
-- * The state token
type BuilderState = (# Addr#, Addr#, State# RealWorld #)

instance Sem.Semigroup Builder where
  <> :: Builder -> Builder -> Builder
(<>) = Builder -> Builder -> Builder
appendBuilder
  {-# INLINE (<>) #-}

appendBuilder :: Builder -> Builder -> Builder
appendBuilder :: Builder -> Builder -> Builder
appendBuilder (Builder DataSink -> BuilderState -> BuilderState
a) (Builder DataSink -> BuilderState -> BuilderState
b)
  = Builder -> Builder
rebuild (Builder -> Builder) -> Builder -> Builder
forall a b. (a -> b) -> a -> b
$ (DataSink -> BuilderState -> BuilderState) -> Builder
Builder ((DataSink -> BuilderState -> BuilderState) -> Builder)
-> (DataSink -> BuilderState -> BuilderState) -> Builder
forall a b. (a -> b) -> a -> b
$ \DataSink
dex BuilderState
bs -> DataSink -> BuilderState -> BuilderState
b DataSink
dex (DataSink -> BuilderState -> BuilderState
a DataSink
dex BuilderState
bs)
{-# INLINE[1] appendBuilder #-}

{-# RULES "appendBuilder/assoc"
  forall x y z.
    appendBuilder (appendBuilder x y) z = appendBuilder x (appendBuilder y z)
  #-}

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

-- | 'fromString' = 'stringUtf8'
instance IsString Builder where
  fromString :: String -> Builder
fromString = String -> Builder
builderFromString
  {-# INLINE fromString #-}

-- | Specifies where bytes generated by a builder go.
data DataSink
  = DynamicSink !(IORef DynamicSink)
    -- ^ The destination of data changes while the builder is running.
  | GrowingBuffer !(IORef (ForeignPtr Word8))
    -- ^ Bytes are accumulated in a contiguous buffer.
  | HandleSink !IO.Handle !Int{-next buffer size-} !(IORef Queue)
    -- ^ Bytes are first accumulated in the 'Queue', then flushed to the
    -- 'IO.Handle'.

-- | Variable-destination cases.
data DynamicSink
  = ThreadedSink !(MVar Request) !(MVar Response)
      -- ^ Bytes are sent to another thread.
  | BoundedGrowingBuffer {-# UNPACK #-} !(ForeignPtr Word8) !Int{-bound-}
      -- ^ Bytes are accumulated in a contiguous buffer until the
      -- size limit is reached. After that, the destination switches
      -- to a 'ThreadedSink'.

-- | A mutable buffer.
data Queue = Queue
  { Queue -> ForeignPtr Word8
queueBuffer :: !(ForeignPtr Word8)
  , Queue -> Int
queueStart :: !Int
    -- ^ Starting position.
  , Queue -> Int
queueTotal :: !Int
    -- ^ Bytes written to the handle so far.
  }
  -- TODO: this is not really needed

-- | A request from the driver thread to the builder thread.
data Request
  = Request {-# UNPACK #-} !(Ptr Word8) {-# UNPACK #-} !(Ptr Word8)

-- | A response from the builder thread to the driver thread.
data Response
  = Error E.SomeException
      -- ^ A synchronous exception was thrown by the builder
  | Done !(Ptr Word8)
      -- ^ The builder thread has completed.
  | MoreBuffer !(Ptr Word8) !Int
      -- ^ The builder thread has finished generating one chunk,
      -- and waits for another request with the specified minimum size.
  | InsertByteString !(Ptr Word8) !S.ByteString
      -- ^ The builder thread has partially filled the current chunk,
      -- and wants to emit the bytestring to be included in the final
      -- output.
  deriving (Int -> Response -> ShowS
[Response] -> ShowS
Response -> String
(Int -> Response -> ShowS)
-> (Response -> String) -> ([Response] -> ShowS) -> Show Response
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Response -> ShowS
showsPrec :: Int -> Response -> ShowS
$cshow :: Response -> String
show :: Response -> String
$cshowList :: [Response] -> ShowS
showList :: [Response] -> ShowS
Show)

----------------------------------------------------------------
-- Internally used exceptions

-- | Used in the implementation of 'toLazyByteString'. This is an exception
-- thrown by the consumer thread to itself when it has finished filling the
-- first chunk of the output. After this, a thread will be forked, and the
-- execution of the builder will be resumed in the new thread, using
-- 'ThreadedSink'.
data ChunkOverflowException
  = ChunkOverflowException
      !S.ByteString !(MVar Request) !(MVar Response) !Int

instance Show ChunkOverflowException where
  show :: ChunkOverflowException -> String
show (ChunkOverflowException ByteString
buf MVar Request
_ MVar Response
_ Int
req) =
    String
"ChunkOverflowException " String -> ShowS
forall a. [a] -> [a] -> [a]
++ ByteString -> String
forall a. Show a => a -> String
show ByteString
buf String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" _ _ " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
req

instance E.Exception ChunkOverflowException

-- | Used in the implementation of 'toLazyByteString'. This is a message sent
-- from the consumer thread to the builder thread, requesting the builder
-- thread to temporarily pause execution. Later, the consumer thread may
-- request resumption by filling the 'MVar'.
data SuspendBuilderException = SuspendBuilderException !(MVar ())

instance Show SuspendBuilderException where
  show :: SuspendBuilderException -> String
show SuspendBuilderException
_ = String
"SuspendBuilderException"

instance E.Exception SuspendBuilderException

----------------------------------------------------------------
-- Builder building blocks

-- | An internal type for making it easier to define builders. A value of
-- @'BuildM' a@ can do everything a 'Builder' can do, and in addition,
-- returns a value of type @a@ upon completion.
newtype BuildM a = BuildM { forall a. BuildM a -> (a -> Builder) -> Builder
runBuildM :: (a -> Builder) -> Builder }
  deriving ((forall a b. (a -> b) -> BuildM a -> BuildM b)
-> (forall a b. a -> BuildM b -> BuildM a) -> Functor BuildM
forall a b. a -> BuildM b -> BuildM a
forall a b. (a -> b) -> BuildM a -> BuildM b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall a b. (a -> b) -> BuildM a -> BuildM b
fmap :: forall a b. (a -> b) -> BuildM a -> BuildM b
$c<$ :: forall a b. a -> BuildM b -> BuildM a
<$ :: forall a b. a -> BuildM b -> BuildM a
Functor)

instance Applicative BuildM where
  pure :: forall a. a -> BuildM a
pure = a -> BuildM a
forall a. a -> BuildM a
forall (m :: * -> *) a. Monad m => a -> m a
return
  <*> :: forall a b. BuildM (a -> b) -> BuildM a -> BuildM b
(<*>) = BuildM (a -> b) -> BuildM a -> BuildM b
forall (m :: * -> *) a b. Monad m => m (a -> b) -> m a -> m b
ap

instance Monad BuildM where
  return :: forall a. a -> BuildM a
return a
x = ((a -> Builder) -> Builder) -> BuildM a
forall a. ((a -> Builder) -> Builder) -> BuildM a
BuildM (((a -> Builder) -> Builder) -> BuildM a)
-> ((a -> Builder) -> Builder) -> BuildM a
forall a b. (a -> b) -> a -> b
$ \a -> Builder
k -> a -> Builder
k a
x
  {-# INLINE return #-}
  BuildM (a -> Builder) -> Builder
b >>= :: forall a b. BuildM a -> (a -> BuildM b) -> BuildM b
>>= a -> BuildM b
f = ((b -> Builder) -> Builder) -> BuildM b
forall a. ((a -> Builder) -> Builder) -> BuildM a
BuildM (((b -> Builder) -> Builder) -> BuildM b)
-> ((b -> Builder) -> Builder) -> BuildM b
forall a b. (a -> b) -> a -> b
$ \b -> Builder
k -> (a -> Builder) -> Builder
b ((a -> Builder) -> Builder) -> (a -> Builder) -> Builder
forall a b. (a -> b) -> a -> b
$ \a
r -> BuildM b -> (b -> Builder) -> Builder
forall a. BuildM a -> (a -> Builder) -> Builder
runBuildM (a -> BuildM b
f a
r) b -> Builder
k
  {-# INLINE (>>=) #-}

-- | Create a builder from a BuildM.
mkBuilder :: BuildM () -> Builder
mkBuilder :: BuildM () -> Builder
mkBuilder (BuildM (() -> Builder) -> Builder
bb) = (() -> Builder) -> Builder
bb ((() -> Builder) -> Builder) -> (() -> Builder) -> Builder
forall a b. (a -> b) -> a -> b
$ \()
_ -> Builder
forall a. Monoid a => a
mempty
{-# INLINE mkBuilder #-}

-- | Embed a builder in the BuildM context.
useBuilder :: Builder -> BuildM ()
useBuilder :: Builder -> BuildM ()
useBuilder Builder
b = ((() -> Builder) -> Builder) -> BuildM ()
forall a. ((a -> Builder) -> Builder) -> BuildM a
BuildM (((() -> Builder) -> Builder) -> BuildM ())
-> ((() -> Builder) -> Builder) -> BuildM ()
forall a b. (a -> b) -> a -> b
$ \() -> Builder
k -> Builder
b Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> () -> Builder
k ()
{-# INLINE useBuilder #-}

-- | Get the 'DataSink'.
getSink :: BuildM DataSink
getSink :: BuildM DataSink
getSink = ((DataSink -> Builder) -> Builder) -> BuildM DataSink
forall a. ((a -> Builder) -> Builder) -> BuildM a
BuildM (((DataSink -> Builder) -> Builder) -> BuildM DataSink)
-> ((DataSink -> Builder) -> Builder) -> BuildM DataSink
forall a b. (a -> b) -> a -> b
$ \DataSink -> Builder
k -> (DataSink -> BuilderState -> BuilderState) -> Builder
Builder ((DataSink -> BuilderState -> BuilderState) -> Builder)
-> (DataSink -> BuilderState -> BuilderState) -> Builder
forall a b. (a -> b) -> a -> b
$ \DataSink
dex (# Addr#
cur, Addr#
end, State# RealWorld
s #) ->
  Builder -> DataSink -> BuilderState -> BuilderState
unBuilder (DataSink -> Builder
k DataSink
dex) DataSink
dex (# Addr#
cur, Addr#
end, State# RealWorld
s #)

-- | Get the current pointer.
getCur :: BuildM (Ptr Word8)
getCur :: BuildM (Ptr Word8)
getCur = ((Ptr Word8 -> Builder) -> Builder) -> BuildM (Ptr Word8)
forall a. ((a -> Builder) -> Builder) -> BuildM a
BuildM (((Ptr Word8 -> Builder) -> Builder) -> BuildM (Ptr Word8))
-> ((Ptr Word8 -> Builder) -> Builder) -> BuildM (Ptr Word8)
forall a b. (a -> b) -> a -> b
$ \Ptr Word8 -> Builder
k -> (DataSink -> BuilderState -> BuilderState) -> Builder
Builder ((DataSink -> BuilderState -> BuilderState) -> Builder)
-> (DataSink -> BuilderState -> BuilderState) -> Builder
forall a b. (a -> b) -> a -> b
$ \DataSink
dex (# Addr#
cur, Addr#
end, State# RealWorld
s #) ->
  Builder -> DataSink -> BuilderState -> BuilderState
unBuilder (Ptr Word8 -> Builder
k (Addr# -> Ptr Word8
forall a. Addr# -> Ptr a
Ptr Addr#
cur)) DataSink
dex (# Addr#
cur, Addr#
end, State# RealWorld
s #)

-- | Get the end-of-buffer pointer.
getEnd :: BuildM (Ptr Word8)
getEnd :: BuildM (Ptr Word8)
getEnd = ((Ptr Word8 -> Builder) -> Builder) -> BuildM (Ptr Word8)
forall a. ((a -> Builder) -> Builder) -> BuildM a
BuildM (((Ptr Word8 -> Builder) -> Builder) -> BuildM (Ptr Word8))
-> ((Ptr Word8 -> Builder) -> Builder) -> BuildM (Ptr Word8)
forall a b. (a -> b) -> a -> b
$ \Ptr Word8 -> Builder
k -> (DataSink -> BuilderState -> BuilderState) -> Builder
Builder ((DataSink -> BuilderState -> BuilderState) -> Builder)
-> (DataSink -> BuilderState -> BuilderState) -> Builder
forall a b. (a -> b) -> a -> b
$ \DataSink
dex (# Addr#
cur, Addr#
end, State# RealWorld
s #) ->
  Builder -> DataSink -> BuilderState -> BuilderState
unBuilder (Ptr Word8 -> Builder
k (Addr# -> Ptr Word8
forall a. Addr# -> Ptr a
Ptr Addr#
end)) DataSink
dex (# Addr#
cur, Addr#
end, State# RealWorld
s #)

-- | Set the current pointer.
setCur :: Ptr Word8 -> BuildM ()
setCur :: Ptr Word8 -> BuildM ()
setCur (Ptr Addr#
p) = ((() -> Builder) -> Builder) -> BuildM ()
forall a. ((a -> Builder) -> Builder) -> BuildM a
BuildM (((() -> Builder) -> Builder) -> BuildM ())
-> ((() -> Builder) -> Builder) -> BuildM ()
forall a b. (a -> b) -> a -> b
$ \() -> Builder
k -> (DataSink -> BuilderState -> BuilderState) -> Builder
Builder ((DataSink -> BuilderState -> BuilderState) -> Builder)
-> (DataSink -> BuilderState -> BuilderState) -> Builder
forall a b. (a -> b) -> a -> b
$ \DataSink
dex (# Addr#
_, Addr#
end, State# RealWorld
s #) ->
  Builder -> DataSink -> BuilderState -> BuilderState
unBuilder (() -> Builder
k ()) DataSink
dex (# Addr#
p, Addr#
end, State# RealWorld
s #)

-- | Set the end-of-buffer pointer.
setEnd :: Ptr Word8 -> BuildM ()
setEnd :: Ptr Word8 -> BuildM ()
setEnd (Ptr Addr#
p) = ((() -> Builder) -> Builder) -> BuildM ()
forall a. ((a -> Builder) -> Builder) -> BuildM a
BuildM (((() -> Builder) -> Builder) -> BuildM ())
-> ((() -> Builder) -> Builder) -> BuildM ()
forall a b. (a -> b) -> a -> b
$ \() -> Builder
k -> (DataSink -> BuilderState -> BuilderState) -> Builder
Builder ((DataSink -> BuilderState -> BuilderState) -> Builder)
-> (DataSink -> BuilderState -> BuilderState) -> Builder
forall a b. (a -> b) -> a -> b
$ \DataSink
dex (# Addr#
cur, Addr#
_, State# RealWorld
s #) ->
  Builder -> DataSink -> BuilderState -> BuilderState
unBuilder (() -> Builder
k ()) DataSink
dex (# Addr#
cur, Addr#
p, State# RealWorld
s #)

-- | Perform IO.
io :: IO a -> BuildM a
io :: forall a. IO a -> BuildM a
io (IO State# RealWorld -> (# State# RealWorld, a #)
x) = ((a -> Builder) -> Builder) -> BuildM a
forall a. ((a -> Builder) -> Builder) -> BuildM a
BuildM (((a -> Builder) -> Builder) -> BuildM a)
-> ((a -> Builder) -> Builder) -> BuildM a
forall a b. (a -> b) -> a -> b
$ \a -> Builder
k -> (DataSink -> BuilderState -> BuilderState) -> Builder
Builder ((DataSink -> BuilderState -> BuilderState) -> Builder)
-> (DataSink -> BuilderState -> BuilderState) -> Builder
forall a b. (a -> b) -> a -> b
$ \DataSink
dex (# Addr#
cur, Addr#
end, State# RealWorld
s #) -> case State# RealWorld -> (# State# RealWorld, a #)
x State# RealWorld
s of
  (# State# RealWorld
s', a
val #) -> Builder -> DataSink -> BuilderState -> BuilderState
unBuilder (a -> Builder
k a
val) DataSink
dex (# Addr#
cur, Addr#
end, State# RealWorld
s' #)

-- | Embed a 'BuilderState' transformer into `BuildM`.
updateState :: (BuilderState -> BuilderState) -> BuildM ()
updateState :: (BuilderState -> BuilderState) -> BuildM ()
updateState BuilderState -> BuilderState
f = ((() -> Builder) -> Builder) -> BuildM ()
forall a. ((a -> Builder) -> Builder) -> BuildM a
BuildM (((() -> Builder) -> Builder) -> BuildM ())
-> ((() -> Builder) -> Builder) -> BuildM ()
forall a b. (a -> b) -> a -> b
$ \() -> Builder
k -> (DataSink -> BuilderState -> BuilderState) -> Builder
Builder ((DataSink -> BuilderState -> BuilderState) -> Builder)
-> (DataSink -> BuilderState -> BuilderState) -> Builder
forall a b. (a -> b) -> a -> b
$ \DataSink
sink BuilderState
bs ->
  Builder -> DataSink -> BuilderState -> BuilderState
unBuilder (() -> Builder
k ()) DataSink
sink (BuilderState -> BuilderState
f BuilderState
bs)

-- | A 'Write' is like a 'Builder', but an upper bound of its size is known
-- before it actually starts filling buffers. It means just one overflow check
-- is sufficient for each 'Write'.
data Write = Write !Int (BuilderState -> BuilderState)

instance Sem.Semigroup Write where
  Write Int
s0 BuilderState -> BuilderState
w0 <> :: Write -> Write -> Write
<> Write Int
s1 BuilderState -> BuilderState
w1 = Int -> (BuilderState -> BuilderState) -> Write
Write (Int
s0 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
s1) (\BuilderState
s -> BuilderState -> BuilderState
w1 (BuilderState -> BuilderState
w0 BuilderState
s))

instance Monoid Write where
  mempty :: Write
mempty = Int -> (BuilderState -> BuilderState) -> Write
Write Int
0 (\BuilderState
s -> BuilderState
s)
  mappend :: Write -> Write -> Write
mappend = Write -> Write -> Write
forall a. Semigroup a => a -> a -> a
(<>)
  {-# INLINE mappend #-}

-- | Turn a 'PI.BoundedPrim' into a 'Write'.
writeBoundedPrim :: PI.BoundedPrim a -> a -> Write
writeBoundedPrim :: forall a. BoundedPrim a -> a -> Write
writeBoundedPrim BoundedPrim a
prim a
x =
  Int -> (BuilderState -> BuilderState) -> Write
Write (BoundedPrim a -> Int
forall a. BoundedPrim a -> Int
PI.sizeBound BoundedPrim a
prim) ((BuilderState -> BuilderState) -> Write)
-> (BuilderState -> BuilderState) -> Write
forall a b. (a -> b) -> a -> b
$ \(# Addr#
cur, Addr#
end, State# RealWorld
s #) ->
    case IO (Ptr Word8)
-> State# RealWorld -> (# State# RealWorld, Ptr Word8 #)
forall a. IO a -> State# RealWorld -> (# State# RealWorld, a #)
unIO (BoundedPrim a -> a -> Ptr Word8 -> IO (Ptr Word8)
forall a. BoundedPrim a -> a -> Ptr Word8 -> IO (Ptr Word8)
PI.runB BoundedPrim a
prim a
x (Addr# -> Ptr Word8
forall a. Addr# -> Ptr a
Ptr Addr#
cur)) State# RealWorld
s of
      (# State# RealWorld
s', Ptr Addr#
cur' #) -> (# Addr#
cur', Addr#
end, State# RealWorld
s' #)

----------------------------------------------------------------
--
-- Running builders.

-- | Run a builder.
runBuilder :: Builder -> DataSink -> Ptr Word8 -> Ptr Word8 -> IO (Ptr Word8)
runBuilder :: Builder -> DataSink -> Ptr Word8 -> Ptr Word8 -> IO (Ptr Word8)
runBuilder (Builder DataSink -> BuilderState -> BuilderState
f) DataSink
sink (Ptr Addr#
cur) (Ptr Addr#
end) = (State# RealWorld -> (# State# RealWorld, Ptr Word8 #))
-> IO (Ptr Word8)
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO ((State# RealWorld -> (# State# RealWorld, Ptr Word8 #))
 -> IO (Ptr Word8))
-> (State# RealWorld -> (# State# RealWorld, Ptr Word8 #))
-> IO (Ptr Word8)
forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s ->
  case DataSink -> BuilderState -> BuilderState
f DataSink
sink (# Addr#
cur, Addr#
end, State# RealWorld
s #) of
    (# Addr#
cur', Addr#
_, State# RealWorld
s' #) -> (# State# RealWorld
s', Addr# -> Ptr Word8
forall a. Addr# -> Ptr a
Ptr Addr#
cur' #)

-- | Turn a 'Builder' into a lazy 'L.ByteString'.
--
-- __Performance hint__: when the resulting 'L.ByteString' does not fit
-- in one chunk, this function forks a thread. Due to this, the performance
-- degrades sharply if you use this function from a bound thread. Note in
-- particular that the main thread is a bound thread when you use @ghc
-- -threaded@.
--
-- To avoid this problem, do one of these:
--
-- * Make sure the resulting 'L.ByteString' is consumed in an unbound
--    thread. Consider using 'runInUnboundThread' for this.
-- * Use other function to run the 'Builder' instead. Functions that don't
--    return a lazy 'L.ByteString' do not have this issue.
-- * Link your program without @-threaded@.
toLazyByteString :: Builder -> L.ByteString
toLazyByteString :: Builder -> ByteString
toLazyByteString = Int -> Int -> Builder -> ByteString
toLazyByteStringWith Int
100 Int
32768

-- | Like 'toLazyByteString', but allows the user to specify the initial
-- and the subsequent desired buffer sizes.
toLazyByteStringWith :: Int -> Int -> Builder -> L.ByteString

-- The implementation employs a two-phase strategy to minimize the overhead:
--
-- 0. Fill the first chunk in a single-threaded way. Start from 'initialSize'-
--    sized buffer and double the size whenever the buffer is full. This uses a
--    'BoundedGrowingBuffer' sink.
--
-- 1. If the first chunk is big enough and the builder still hasn't finished,
--    suspend the execution of the builder, fork a new thread and resume
--    execution of the builder in the new thread, using a 'ThreadedSink'.
toLazyByteStringWith :: Int -> Int -> Builder -> ByteString
toLazyByteStringWith !Int
initialSize !Int
maxSize Builder
builder = IO ByteString -> ByteString
forall a. IO a -> a
unsafePerformIO (IO ByteString -> ByteString) -> IO ByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ do
  ForeignPtr Word8
fptr <- Int -> IO (ForeignPtr Word8)
forall a. Int -> IO (ForeignPtr a)
mallocForeignPtrBytes Int
initialSize
  IORef DynamicSink
sink <- DynamicSink -> IO (IORef DynamicSink)
forall a. a -> IO (IORef a)
newIORef (DynamicSink -> IO (IORef DynamicSink))
-> DynamicSink -> IO (IORef DynamicSink)
forall a b. (a -> b) -> a -> b
$ ForeignPtr Word8 -> Int -> DynamicSink
BoundedGrowingBuffer ForeignPtr Word8
fptr Int
maxSize
  let !base :: Ptr Word8
base = ForeignPtr Word8 -> Ptr Word8
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr Word8
fptr
  let
    finalPtr :: Ptr Word8
finalPtr = IO (Ptr Word8) -> Ptr Word8
forall a. IO a -> a
unsafeDupablePerformIO (IO (Ptr Word8) -> Ptr Word8) -> IO (Ptr Word8) -> Ptr Word8
forall a b. (a -> b) -> a -> b
$
      -- The use of unsafeDupablePerformIO is safe here, because at any given
      -- time, at most one thread can be attempting to evaluate this finalPtr
      -- thunk.
      Builder -> DataSink -> Ptr Word8 -> Ptr Word8 -> IO (Ptr Word8)
runBuilder Builder
builder (IORef DynamicSink -> DataSink
DynamicSink IORef DynamicSink
sink) Ptr Word8
base (Ptr Word8
base Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
initialSize)
    {-# NOINLINE finalPtr #-}

    loop :: Ptr Word8 -> IO ByteString
loop Ptr Word8
thunk = do
      -- Pass around `thunk` as an argument, otherwise GHC 7.10.1 inlines it
      -- despite the NOINLINE pragma.
      Either SomeException (Ptr Word8)
r <- IO (Ptr Word8) -> IO (Either SomeException (Ptr Word8))
forall e a. Exception e => IO a -> IO (Either e a)
E.try (IO (Ptr Word8) -> IO (Either SomeException (Ptr Word8)))
-> IO (Ptr Word8) -> IO (Either SomeException (Ptr Word8))
forall a b. (a -> b) -> a -> b
$ Ptr Word8 -> IO (Ptr Word8)
forall a. a -> IO a
E.evaluate Ptr Word8
thunk
      case Either SomeException (Ptr Word8)
r of
        Right Ptr Word8
p -> do
          BoundedGrowingBuffer ForeignPtr Word8
finalFptr Int
_ <- IORef DynamicSink -> IO DynamicSink
forall a. IORef a -> IO a
readIORef IORef DynamicSink
sink
          let !finalBase :: Ptr Word8
finalBase = ForeignPtr Word8 -> Ptr Word8
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr Word8
finalFptr
          ByteString -> IO ByteString
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (ByteString -> IO ByteString) -> ByteString -> IO ByteString
forall a b. (a -> b) -> a -> b
$! ByteString -> ByteString
L.fromStrict (ByteString -> ByteString) -> ByteString -> ByteString
forall a b. (a -> b) -> a -> b
$
            ForeignPtr Word8 -> Int -> Int -> ByteString
S.fromForeignPtr ForeignPtr Word8
finalFptr Int
0 (Ptr Word8
p Ptr Word8 -> Ptr Word8 -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr Word8
finalBase)
        Left SomeException
ex
          | Just (ChunkOverflowException ByteString
chunk MVar Request
reqV MVar Response
respV Int
minSize)
              <- SomeException -> Maybe ChunkOverflowException
forall e. Exception e => SomeException -> Maybe e
E.fromException SomeException
ex
            -> do
              let rest :: [ByteString]
rest = MVar Request
-> MVar Response -> Int -> Int -> Ptr Word8 -> [ByteString]
continueBuilderThreaded MVar Request
reqV MVar Response
respV Int
minSize Int
maxSize Ptr Word8
thunk
              ByteString -> IO ByteString
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (ByteString -> IO ByteString) -> ByteString -> IO ByteString
forall a b. (a -> b) -> a -> b
$ [ByteString] -> ByteString
L.fromChunks ([ByteString] -> ByteString) -> [ByteString] -> ByteString
forall a b. (a -> b) -> a -> b
$
                if ByteString -> Bool
S.null ByteString
chunk then [ByteString]
rest else ByteString
chunk ByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
: [ByteString]
rest
          | Bool
otherwise -> do
              -- Here, there is no way to tell whether 'ex' is an asynchronous
              -- exception or not. We re-throw is as if it were async. This is
              -- a safe assumption, because if it is actually a synchronous
              -- exception, it will be re-thrown when we try to resume
              -- the evaluation of 'thunk'.
              ThreadId
myTid <- IO ThreadId
myThreadId
              ThreadId -> SomeException -> IO ()
forall e. Exception e => ThreadId -> e -> IO ()
E.throwTo ThreadId
myTid SomeException
ex

              Ptr Word8 -> IO ByteString
loop Ptr Word8
thunk

  Ptr Word8 -> IO ByteString
loop Ptr Word8
finalPtr

-- | Continue a suspended builder using threads.
continueBuilderThreaded
  :: MVar Request -> MVar Response -> Int -> Int -> Ptr Word8
  -> [S.ByteString]
continueBuilderThreaded :: MVar Request
-> MVar Response -> Int -> Int -> Ptr Word8 -> [ByteString]
continueBuilderThreaded !MVar Request
reqV !MVar Response
respV !Int
initialSize !Int
maxSize Ptr Word8
thunk =
  Int -> Int -> BufferWriter -> [ByteString]
makeChunks (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
maxSize Int
initialSize) Int
maxSize (BufferWriter -> [ByteString]) -> BufferWriter -> [ByteString]
forall a b. (a -> b) -> a -> b
$ MVar Request -> MVar Response -> Ptr Word8 -> BufferWriter
toBufferWriter MVar Request
reqV MVar Response
respV Ptr Word8
thunk

-- | Run the given suspended builder using a new thread.
toBufferWriter :: MVar Request -> MVar Response -> Ptr Word8 -> X.BufferWriter
toBufferWriter :: MVar Request -> MVar Response -> Ptr Word8 -> BufferWriter
toBufferWriter !MVar Request
reqV !MVar Response
respV Ptr Word8
thunk Ptr Word8
buf0 Int
sz0 = IO (Int, Next) -> IO (Int, Next)
forall a. IO a -> IO a
E.mask_ (IO (Int, Next) -> IO (Int, Next))
-> IO (Int, Next) -> IO (Int, Next)
forall a b. (a -> b) -> a -> b
$
  Maybe ThreadId -> BufferWriter
writer Maybe ThreadId
forall a. Maybe a
Nothing Ptr Word8
buf0 Int
sz0
  where
    writer :: Maybe ThreadId -> BufferWriter
writer !Maybe ThreadId
maybeBuilderTid !Ptr Word8
buf !Int
sz = do
      MVar Request -> Request -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar Request
reqV (Request -> IO ()) -> Request -> IO ()
forall a b. (a -> b) -> a -> b
$ Ptr Word8 -> Ptr Word8 -> Request
Request Ptr Word8
buf (Ptr Word8
buf Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
sz)
      -- Fork after putMVar, in order to minimize the chance that
      -- the new thread is scheduled on a different CPU.
      ThreadId
builderTid <- case Maybe ThreadId
maybeBuilderTid of
        Just ThreadId
t -> ThreadId -> IO ThreadId
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ThreadId
t
        Maybe ThreadId
Nothing -> ((forall a. IO a -> IO a) -> IO ()) -> IO ThreadId
forkIOWithUnmask (((forall a. IO a -> IO a) -> IO ()) -> IO ThreadId)
-> ((forall a. IO a -> IO a) -> IO ()) -> IO ThreadId
forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
u ->
          (forall a. IO a -> IO a) -> MVar Response -> Ptr Word8 -> IO ()
builderThreadWithUnmask IO a -> IO a
forall a. IO a -> IO a
u MVar Response
respV Ptr Word8
thunk
      Response
resp <- ThreadId -> IO Response
wait ThreadId
builderTid
      let go :: Ptr a -> b -> m (Int, b)
go Ptr a
cur b
next = (Int, b) -> m (Int, b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return(Int
written, b
next)
            where !written :: Int
written = Ptr a
cur Ptr a -> Ptr Word8 -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr Word8
buf
      case Response
resp of
        Error SomeException
ex -> SomeException -> IO (Int, Next)
forall e a. Exception e => e -> IO a
E.throwIO SomeException
ex
        Done Ptr Word8
cur -> Ptr Word8 -> Next -> IO (Int, Next)
forall {m :: * -> *} {a} {b}. Monad m => Ptr a -> b -> m (Int, b)
go Ptr Word8
cur Next
X.Done
        MoreBuffer Ptr Word8
cur Int
k -> Ptr Word8 -> Next -> IO (Int, Next)
forall {m :: * -> *} {a} {b}. Monad m => Ptr a -> b -> m (Int, b)
go Ptr Word8
cur (Next -> IO (Int, Next)) -> Next -> IO (Int, Next)
forall a b. (a -> b) -> a -> b
$ Int -> BufferWriter -> Next
X.More Int
k (BufferWriter -> Next) -> BufferWriter -> Next
forall a b. (a -> b) -> a -> b
$ Maybe ThreadId -> BufferWriter
writer (ThreadId -> Maybe ThreadId
forall a. a -> Maybe a
Just ThreadId
builderTid)
        InsertByteString Ptr Word8
cur ByteString
str -> Ptr Word8 -> Next -> IO (Int, Next)
forall {m :: * -> *} {a} {b}. Monad m => Ptr a -> b -> m (Int, b)
go Ptr Word8
cur (Next -> IO (Int, Next)) -> Next -> IO (Int, Next)
forall a b. (a -> b) -> a -> b
$ ByteString -> BufferWriter -> Next
X.Chunk ByteString
str (BufferWriter -> Next) -> BufferWriter -> Next
forall a b. (a -> b) -> a -> b
$ Maybe ThreadId -> BufferWriter
writer (ThreadId -> Maybe ThreadId
forall a. a -> Maybe a
Just ThreadId
builderTid)

    wait :: ThreadId -> IO Response
wait !ThreadId
builderTid = do
      Either SomeException Response
r <- IO Response -> IO (Either SomeException Response)
forall e a. Exception e => IO a -> IO (Either e a)
E.try (IO Response -> IO (Either SomeException Response))
-> IO Response -> IO (Either SomeException Response)
forall a b. (a -> b) -> a -> b
$ MVar Response -> IO Response
forall a. MVar a -> IO a
takeMVar MVar Response
respV
      case Either SomeException Response
r of
        Right Response
resp -> Response -> IO Response
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Response
resp
        Left SomeException
exn -> do
          -- exn must be an async exception, because takeMVar throws no
          -- synchronous exceptions.
          MVar ()
resumeVar <- IO (MVar ())
forall a. IO (MVar a)
newEmptyMVar
          ThreadId -> SuspendBuilderException -> IO ()
forall e. Exception e => ThreadId -> e -> IO ()
E.throwTo ThreadId
builderTid (SuspendBuilderException -> IO ())
-> SuspendBuilderException -> IO ()
forall a b. (a -> b) -> a -> b
$ MVar () -> SuspendBuilderException
SuspendBuilderException MVar ()
resumeVar
          ThreadId
thisTid <- IO ThreadId
myThreadId
          ThreadId -> SomeException -> IO ()
forall e. Exception e => ThreadId -> e -> IO ()
E.throwTo ThreadId
thisTid (SomeException
exn :: E.SomeException)

          -- A thunk containing this computation has been resumed.
          -- Resume the builder thread, and retry.
          MVar () -> () -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar ()
resumeVar ()
          ThreadId -> IO Response
wait ThreadId
builderTid

-- | The body of the builder thread.
builderThreadWithUnmask
  :: (forall a. IO a -> IO a) -> MVar Response -> Ptr Word8
  -> IO ()
builderThreadWithUnmask :: (forall a. IO a -> IO a) -> MVar Response -> Ptr Word8 -> IO ()
builderThreadWithUnmask forall a. IO a -> IO a
unmask !MVar Response
respV Ptr Word8
thunk = IO ()
loop
  where
    loop :: IO ()
loop = do
      Either SomeException (Ptr Word8)
r <- IO (Ptr Word8) -> IO (Either SomeException (Ptr Word8))
forall e a. Exception e => IO a -> IO (Either e a)
E.try (IO (Ptr Word8) -> IO (Either SomeException (Ptr Word8)))
-> IO (Ptr Word8) -> IO (Either SomeException (Ptr Word8))
forall a b. (a -> b) -> a -> b
$ IO (Ptr Word8) -> IO (Ptr Word8)
forall a. IO a -> IO a
unmask (IO (Ptr Word8) -> IO (Ptr Word8))
-> IO (Ptr Word8) -> IO (Ptr Word8)
forall a b. (a -> b) -> a -> b
$ Ptr Word8 -> IO (Ptr Word8)
forall a. a -> IO a
E.evaluate Ptr Word8
thunk
      case Either SomeException (Ptr Word8)
r of
        Right Ptr Word8
p -> MVar Response -> Response -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar Response
respV (Response -> IO ()) -> Response -> IO ()
forall a b. (a -> b) -> a -> b
$ Ptr Word8 -> Response
Done Ptr Word8
p
        Left SomeException
ex
          | Just (SuspendBuilderException MVar ()
lock) <- SomeException -> Maybe SuspendBuilderException
forall e. Exception e => SomeException -> Maybe e
E.fromException SomeException
ex
            -> do MVar () -> IO ()
forall a. MVar a -> IO a
takeMVar MVar ()
lock; IO ()
loop
          | Bool
otherwise -> MVar Response -> Response -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar Response
respV (Response -> IO ()) -> Response -> IO ()
forall a b. (a -> b) -> a -> b
$ SomeException -> Response
Error SomeException
ex

-- | Run a 'X.BufferWriter'.
makeChunks :: Int -> Int -> X.BufferWriter -> [S.ByteString]
makeChunks :: Int -> Int -> BufferWriter -> [ByteString]
makeChunks !Int
initialBufSize Int
maxBufSize = Int -> BufferWriter -> [ByteString]
go Int
initialBufSize
  where
    go :: Int -> BufferWriter -> [ByteString]
go !Int
bufSize BufferWriter
w = IO [ByteString] -> [ByteString]
forall a. IO a -> a
unsafePerformIO (IO [ByteString] -> [ByteString])
-> IO [ByteString] -> [ByteString]
forall a b. (a -> b) -> a -> b
$ do
      ForeignPtr Word8
fptr <- Int -> IO (ForeignPtr Word8)
forall a. Int -> IO (ForeignPtr a)
S.mallocByteString Int
bufSize
      (Int
written, Next
next) <- ForeignPtr Word8 -> (Ptr Word8 -> IO (Int, Next)) -> IO (Int, Next)
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr Word8
fptr ((Ptr Word8 -> IO (Int, Next)) -> IO (Int, Next))
-> (Ptr Word8 -> IO (Int, Next)) -> IO (Int, Next)
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
buf -> BufferWriter
w Ptr Word8
buf Int
bufSize
      let rest :: [ByteString]
rest = case Next
next of
            Next
X.Done -> []
            X.More Int
reqSize BufferWriter
w' -> Int -> BufferWriter -> [ByteString]
go (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
reqSize Int
maxBufSize) BufferWriter
w'
            X.Chunk ByteString
chunk BufferWriter
w' -> ByteString
chunk ByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
: Int -> BufferWriter -> [ByteString]
go Int
maxBufSize BufferWriter
w'
              -- TODO: don't throw away the remaining part of the buffer
      [ByteString] -> IO [ByteString]
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ([ByteString] -> IO [ByteString])
-> [ByteString] -> IO [ByteString]
forall a b. (a -> b) -> a -> b
$ if Int
written Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
        then [ByteString]
rest
        else ForeignPtr Word8 -> Int -> Int -> ByteString
S.fromForeignPtr ForeignPtr Word8
fptr Int
0 Int
written ByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
: [ByteString]
rest

-- | Turn a 'Builder' into a strict 'S.ByteString'.
toStrictByteString :: Builder -> S.ByteString
toStrictByteString :: Builder -> ByteString
toStrictByteString Builder
builder = IO ByteString -> ByteString
forall a. IO a -> a
unsafePerformIO (IO ByteString -> ByteString) -> IO ByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ do
  let cap :: Int
cap = Int
100
  ForeignPtr Word8
fptr <- Int -> IO (ForeignPtr Word8)
forall a. Int -> IO (ForeignPtr a)
mallocForeignPtrBytes Int
cap
  IORef (ForeignPtr Word8)
bufferRef <- ForeignPtr Word8 -> IO (IORef (ForeignPtr Word8))
forall a. a -> IO (IORef a)
newIORef ForeignPtr Word8
fptr
  let !base :: Ptr Word8
base = ForeignPtr Word8 -> Ptr Word8
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr Word8
fptr
  Ptr Word8
cur <- Builder -> DataSink -> Ptr Word8 -> Ptr Word8 -> IO (Ptr Word8)
runBuilder Builder
builder (IORef (ForeignPtr Word8) -> DataSink
GrowingBuffer IORef (ForeignPtr Word8)
bufferRef) Ptr Word8
base (Ptr Word8
base Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
cap)
  ForeignPtr Word8
endFptr <- IORef (ForeignPtr Word8) -> IO (ForeignPtr Word8)
forall a. IORef a -> IO a
readIORef IORef (ForeignPtr Word8)
bufferRef
  let !written :: Int
written = Ptr Word8
cur 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
endFptr
  ByteString -> IO ByteString
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (ByteString -> IO ByteString) -> ByteString -> IO ByteString
forall a b. (a -> b) -> a -> b
$ ForeignPtr Word8 -> Int -> Int -> ByteString
S.fromForeignPtr ForeignPtr Word8
endFptr Int
0 Int
written

-- | Output a 'Builder' to a 'IO.Handle'.
hPutBuilder :: IO.Handle -> Builder -> IO ()
hPutBuilder :: Handle -> Builder -> IO ()
hPutBuilder !Handle
h Builder
builder = IO Int -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Int -> IO ()) -> IO Int -> IO ()
forall a b. (a -> b) -> a -> b
$ Handle -> Builder -> IO Int
hPutBuilderLen Handle
h Builder
builder
{-# INLINE hPutBuilder #-}

-- | Output a 'Builder' to a 'IO.Handle'. Returns the number of bytes written.
hPutBuilderLen :: IO.Handle -> Builder -> IO Int
hPutBuilderLen :: Handle -> Builder -> IO Int
hPutBuilderLen !Handle
h Builder
builder = Handle -> Int -> Int -> Builder -> IO Int
hPutBuilderWith Handle
h Int
100 Int
4096 Builder
builder

-- | Like 'hPutBuffer', but allows the user to specify the initial
-- and the subsequent desired buffer sizes. This function may be useful for
-- setting large buffer when high throughput I/O is needed.
hPutBuilderWith :: IO.Handle -> Int -> Int -> Builder -> IO Int
hPutBuilderWith :: Handle -> Int -> Int -> Builder -> IO Int
hPutBuilderWith !Handle
h !Int
initialCap !Int
nextCap Builder
builder = do
  ForeignPtr Word8
fptr <- Int -> IO (ForeignPtr Word8)
forall a. Int -> IO (ForeignPtr a)
mallocForeignPtrBytes Int
initialCap
  IORef Queue
qRef <- Queue -> IO (IORef Queue)
forall a. a -> IO (IORef a)
newIORef (Queue -> IO (IORef Queue)) -> Queue -> IO (IORef Queue)
forall a b. (a -> b) -> a -> b
$ Queue
    { queueBuffer :: ForeignPtr Word8
queueBuffer = ForeignPtr Word8
fptr
    , queueStart :: Int
queueStart =  Int
0
    , queueTotal :: Int
queueTotal = Int
0
    }
  let !base :: Ptr Word8
base = ForeignPtr Word8 -> Ptr Word8
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr Word8
fptr
  Ptr Word8
cur <- Builder -> DataSink -> Ptr Word8 -> Ptr Word8 -> IO (Ptr Word8)
runBuilder Builder
builder (Handle -> Int -> IORef Queue -> DataSink
HandleSink Handle
h Int
nextCap IORef Queue
qRef)
    Ptr Word8
base (Ptr Word8
base Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
initialCap)
  Handle -> IORef Queue -> Ptr Word8 -> IO ()
flushQueue Handle
h IORef Queue
qRef Ptr Word8
cur
  Queue{ queueTotal :: Queue -> Int
queueTotal = Int
len } <- IORef Queue -> IO Queue
forall a. IORef a -> IO a
readIORef IORef Queue
qRef
  Int -> IO Int
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Int
len

----------------------------------------------------------------
-- builders

-- | Turn a 'String' into a 'Builder', using UTF-8,
builderFromString :: String -> Builder
builderFromString :: String -> Builder
builderFromString = BoundedPrim Char -> String -> Builder
forall a. BoundedPrim a -> [a] -> Builder
primMapListBounded BoundedPrim Char
P.charUtf8
{-# NOINLINE[0] builderFromString #-}

{-# RULES "FastBuilder: builderFromString/unpackCString#"
  forall addr.
    builderFromString (unpackCString# addr) = unsafeCString (Ptr addr)
  #-}

-- | Turn a value of type @a@ into a 'Builder', using the given 'PI.BoundedPrim'.
primBounded :: PI.BoundedPrim a -> a -> Builder
primBounded :: forall a. BoundedPrim a -> a -> Builder
primBounded BoundedPrim a
prim = Write -> Builder
write (Write -> Builder) -> (a -> Write) -> a -> Builder
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BoundedPrim a -> a -> Write
forall a. BoundedPrim a -> a -> Write
writeBoundedPrim BoundedPrim a
prim
{-# INLINE primBounded #-}

-- | Turn a 'Write' into a 'Builder'.
write :: Write -> Builder
write :: Write -> Builder
write (Write Int
size BuilderState -> BuilderState
w) = Builder -> Builder
rebuild (Builder -> Builder) -> Builder -> Builder
forall a b. (a -> b) -> a -> b
$ BuildM () -> Builder
mkBuilder (BuildM () -> Builder) -> BuildM () -> Builder
forall a b. (a -> b) -> a -> b
$ do
  Builder -> BuildM ()
useBuilder (Builder -> BuildM ()) -> Builder -> BuildM ()
forall a b. (a -> b) -> a -> b
$ Int -> Builder
ensureBytes Int
size
  (BuilderState -> BuilderState) -> BuildM ()
updateState BuilderState -> BuilderState
w
{-# INLINE[1] write #-}

{-# RULES "fast-builder: write/write"
  forall w0 w1.
    appendBuilder (write w0) (write w1) = write (w0 <> w1)
  #-}

{-# RULES "fast-builder: write/write/x"
  forall w0 w1 x.
    appendBuilder (write w0) (appendBuilder (write w1) x)
      = appendBuilder (write (w0 <> w1)) x
  #-}

-- | Turn a value of type @a@ into a 'Builder', using the given 'PI.FixedPrim'.
primFixed :: PI.FixedPrim a -> a -> Builder
primFixed :: forall a. FixedPrim a -> a -> Builder
primFixed FixedPrim a
prim = BoundedPrim a -> a -> Builder
forall a. BoundedPrim a -> a -> Builder
primBounded (FixedPrim a -> BoundedPrim a
forall a. FixedPrim a -> BoundedPrim a
PI.toB FixedPrim a
prim)
{-# INLINE primFixed #-}

-- | Turn a list of values of type @a@ into a 'Builder', using the given
-- 'PI.BoundedPrim'.
primMapListBounded :: PI.BoundedPrim a -> [a] -> Builder
primMapListBounded :: forall a. BoundedPrim a -> [a] -> Builder
primMapListBounded BoundedPrim a
prim = [Builder] -> Builder
forall a. Monoid a => [a] -> a
mconcat ([Builder] -> Builder) -> ([a] -> [Builder]) -> [a] -> Builder
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Builder) -> [a] -> [Builder]
forall a b. (a -> b) -> [a] -> [b]
map (BoundedPrim a -> a -> Builder
forall a. BoundedPrim a -> a -> Builder
primBounded BoundedPrim a
prim)
{-# INLINE primMapListBounded #-}

-- | Turn a list of values of type @a@ into a 'Builder', using the given
-- 'PI.FixedPrim'.
primMapListFixed :: PI.FixedPrim a -> [a] -> Builder
primMapListFixed :: forall a. FixedPrim a -> [a] -> Builder
primMapListFixed FixedPrim a
prim = BoundedPrim a -> [a] -> Builder
forall a. BoundedPrim a -> [a] -> Builder
primMapListBounded (FixedPrim a -> BoundedPrim a
forall a. FixedPrim a -> BoundedPrim a
PI.toB FixedPrim a
prim)
{-# INLINE primMapListFixed #-}

-- | Turn a 'S.ByteString' to a 'Builder'.
byteString :: S.ByteString -> Builder
byteString :: ByteString -> Builder
byteString = Int -> ByteString -> Builder
byteStringThreshold Int
maximalCopySize
{-# INLINE byteString #-}

maximalCopySize :: Int
maximalCopySize :: Int
maximalCopySize = Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
X.smallChunkSize

-- | Turn a 'S.ByteString' to a 'Builder'. If the size of the 'S.ByteString'
-- is larger than the given threshold, avoid copying it as much
-- as possible.
byteStringThreshold :: Int -> S.ByteString -> Builder
byteStringThreshold :: Int -> ByteString -> Builder
byteStringThreshold Int
th ByteString
bstr = Builder -> Builder
rebuild (Builder -> Builder) -> Builder -> Builder
forall a b. (a -> b) -> a -> b
$
  if ByteString -> Int
S.length ByteString
bstr Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
th
    then ByteString -> Builder
byteStringInsert ByteString
bstr
    else ByteString -> Builder
byteStringCopy ByteString
bstr

-- | Turn a 'S.ByteString' to a 'Builder'. The 'S.ByteString' will be copied
-- to the buffer, regardless of the size.
byteStringCopy :: S.ByteString -> Builder
byteStringCopy :: ByteString -> Builder
byteStringCopy !ByteString
bstr =
  -- TODO: this is suboptimal; should keep using the same buffer size.
  Int -> Builder
ensureBytes (ByteString -> Int
S.length ByteString
bstr) Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> ByteString -> Builder
byteStringCopyNoCheck ByteString
bstr

-- | Like 'byteStringCopy', but assumes that the current buffer is large enough.
byteStringCopyNoCheck :: S.ByteString -> Builder
byteStringCopyNoCheck :: ByteString -> Builder
byteStringCopyNoCheck !ByteString
bstr = BuildM () -> Builder
mkBuilder (BuildM () -> Builder) -> BuildM () -> Builder
forall a b. (a -> b) -> a -> b
$ do
  Ptr Word8
cur <- BuildM (Ptr Word8)
getCur
  IO () -> BuildM ()
forall a. IO a -> BuildM a
io (IO () -> BuildM ()) -> IO () -> BuildM ()
forall a b. (a -> b) -> a -> b
$ ByteString -> (CString -> IO ()) -> IO ()
forall a. ByteString -> (CString -> IO a) -> IO a
S.unsafeUseAsCString ByteString
bstr ((CString -> IO ()) -> IO ()) -> (CString -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \CString
ptr ->
    Ptr Word8 -> Ptr Word8 -> Int -> IO ()
forall a. Ptr a -> Ptr a -> Int -> IO ()
copyBytes Ptr Word8
cur (CString -> Ptr Word8
forall a b. Ptr a -> Ptr b
castPtr CString
ptr) Int
len
  Ptr Word8 -> BuildM ()
setCur (Ptr Word8 -> BuildM ()) -> Ptr Word8 -> BuildM ()
forall a b. (a -> b) -> a -> b
$ Ptr Word8
cur Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
len
  where
    !len :: Int
len = ByteString -> Int
S.length ByteString
bstr

-- | Turn a 'S.ByteString' to a 'Builder'. When possible, the given
-- 'S.ByteString' will not be copied, and inserted directly into the output
-- instead.
byteStringInsert :: S.ByteString -> Builder
byteStringInsert :: ByteString -> Builder
byteStringInsert !ByteString
bstr = ByteString -> Builder
byteStringInsert_ ByteString
bstr

-- | The body of the 'byteStringInsert', worker-wrappered manually.
byteStringInsert_ :: S.ByteString -> Builder
byteStringInsert_ :: ByteString -> Builder
byteStringInsert_ ByteString
bstr = BuildM () -> Builder
mkBuilder (BuildM () -> Builder) -> BuildM () -> Builder
forall a b. (a -> b) -> a -> b
$ do
  DataSink
sink <- BuildM DataSink
getSink
  case DataSink
sink of
    DynamicSink IORef DynamicSink
dRef -> do
      DynamicSink
dyn <- IO DynamicSink -> BuildM DynamicSink
forall a. IO a -> BuildM a
io (IO DynamicSink -> BuildM DynamicSink)
-> IO DynamicSink -> BuildM DynamicSink
forall a b. (a -> b) -> a -> b
$ IORef DynamicSink -> IO DynamicSink
forall a. IORef a -> IO a
readIORef IORef DynamicSink
dRef
      case DynamicSink
dyn of
        ThreadedSink MVar Request
reqV MVar Response
respV -> do
          Ptr Word8
cur <- BuildM (Ptr Word8)
getCur
          IO () -> BuildM ()
forall a. IO a -> BuildM a
io (IO () -> BuildM ()) -> IO () -> BuildM ()
forall a b. (a -> b) -> a -> b
$ MVar Response -> Response -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar Response
respV (Response -> IO ()) -> Response -> IO ()
forall a b. (a -> b) -> a -> b
$ Ptr Word8 -> ByteString -> Response
InsertByteString Ptr Word8
cur ByteString
bstr
          MVar Request -> BuildM ()
handleRequest MVar Request
reqV
        BoundedGrowingBuffer ForeignPtr Word8
fptr Int
bound -> do
          Int
r <- BuildM Int
remainingBytes
          Bool -> BuildM () -> BuildM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< ByteString -> Int
S.length ByteString
bstr) (BuildM () -> BuildM ()) -> BuildM () -> BuildM ()
forall a b. (a -> b) -> a -> b
$
            IORef DynamicSink -> ForeignPtr Word8 -> Int -> Int -> BuildM ()
growBufferBounded IORef DynamicSink
dRef ForeignPtr Word8
fptr Int
bound (ByteString -> Int
S.length ByteString
bstr)
          -- TODO: insert rather than copy if the first chunk
          -- is full.
          Builder -> BuildM ()
useBuilder (Builder -> BuildM ()) -> Builder -> BuildM ()
forall a b. (a -> b) -> a -> b
$ ByteString -> Builder
byteStringCopyNoCheck ByteString
bstr
    GrowingBuffer IORef (ForeignPtr Word8)
bufRef -> do
      Int
r <- BuildM Int
remainingBytes
      Bool -> BuildM () -> BuildM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< ByteString -> Int
S.length ByteString
bstr) (BuildM () -> BuildM ()) -> BuildM () -> BuildM ()
forall a b. (a -> b) -> a -> b
$
        IORef (ForeignPtr Word8) -> Int -> BuildM ()
growBuffer IORef (ForeignPtr Word8)
bufRef (ByteString -> Int
S.length ByteString
bstr)
      Builder -> BuildM ()
useBuilder (Builder -> BuildM ()) -> Builder -> BuildM ()
forall a b. (a -> b) -> a -> b
$ ByteString -> Builder
byteStringCopyNoCheck ByteString
bstr
    HandleSink Handle
h Int
_nextCap IORef Queue
queueRef -> do
      Ptr Word8
cur <- BuildM (Ptr Word8)
getCur
      IO () -> BuildM ()
forall a. IO a -> BuildM a
io (IO () -> BuildM ()) -> IO () -> BuildM ()
forall a b. (a -> b) -> a -> b
$ Handle -> IORef Queue -> Ptr Word8 -> IO ()
flushQueue Handle
h IORef Queue
queueRef Ptr Word8
cur
      IO () -> BuildM ()
forall a. IO a -> BuildM a
io (IO () -> BuildM ()) -> IO () -> BuildM ()
forall a b. (a -> b) -> a -> b
$ Handle -> ByteString -> IO ()
S.hPut Handle
h ByteString
bstr
      IO () -> BuildM ()
forall a. IO a -> BuildM a
io (IO () -> BuildM ()) -> IO () -> BuildM ()
forall a b. (a -> b) -> a -> b
$ IORef Queue -> (Queue -> Queue) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' IORef Queue
queueRef
        ((Queue -> Queue) -> IO ()) -> (Queue -> Queue) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Queue
q -> Queue
q { queueTotal = queueTotal q + S.length bstr }
{-# NOINLINE byteStringInsert_ #-}

-- | Turn a C String into a 'Builder'. The behavior is undefined if the given
-- 'CString' does not point to a constant null-terminated string.
unsafeCString :: CString -> Builder
unsafeCString :: CString -> Builder
unsafeCString CString
cstr = Builder -> Builder
rebuild (Builder -> Builder) -> Builder -> Builder
forall a b. (a -> b) -> a -> b
$ let
    !len :: Int
len = CSize -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (CSize -> Int) -> CSize -> Int
forall a b. (a -> b) -> a -> b
$ CString -> CSize
c_pure_strlen CString
cstr
  in CStringLen -> Builder
unsafeCStringLen (CString
cstr, Int
len)

foreign import ccall unsafe "strlen" c_pure_strlen :: CString -> CSize

-- | Turn a 'CStringLen' into a 'Builder'. The behavior is undefined if the
-- given 'CStringLen' does not point to a constant memory block.
unsafeCStringLen :: CStringLen -> Builder
unsafeCStringLen :: CStringLen -> Builder
unsafeCStringLen (CString
ptr, Int
len) = Builder -> Builder -> Builder
forall a. Monoid a => a -> a -> a
mappend (Int -> Builder
ensureBytes Int
len) (Builder -> Builder) -> Builder -> Builder
forall a b. (a -> b) -> a -> b
$ BuildM () -> Builder
mkBuilder (BuildM () -> Builder) -> BuildM () -> Builder
forall a b. (a -> b) -> a -> b
$ do
  Ptr Word8
cur <- BuildM (Ptr Word8)
getCur
  IO () -> BuildM ()
forall a. IO a -> BuildM a
io (IO () -> BuildM ()) -> IO () -> BuildM ()
forall a b. (a -> b) -> a -> b
$ Ptr Word8 -> Ptr Word8 -> Int -> IO ()
forall a. Ptr a -> Ptr a -> Int -> IO ()
copyBytes Ptr Word8
cur (CString -> Ptr Word8
forall a b. Ptr a -> Ptr b
castPtr CString
ptr) Int
len
  Ptr Word8 -> BuildM ()
setCur (Ptr Word8 -> BuildM ()) -> Ptr Word8 -> BuildM ()
forall a b. (a -> b) -> a -> b
$ Ptr Word8
cur Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
len

-- | @'ensureBytes' n@ ensures that at least @n@ bytes of free space is
-- available in the current buffer, by allocating a new buffer when
-- necessary.
ensureBytes :: Int -> Builder
ensureBytes :: Int -> Builder
ensureBytes !Int
n = BuildM () -> Builder
mkBuilder (BuildM () -> Builder) -> BuildM () -> Builder
forall a b. (a -> b) -> a -> b
$ do
  Int
r <- BuildM Int
remainingBytes
  Bool -> BuildM () -> BuildM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n) (BuildM () -> BuildM ()) -> BuildM () -> BuildM ()
forall a b. (a -> b) -> a -> b
$ Builder -> BuildM ()
useBuilder (Builder -> BuildM ()) -> Builder -> BuildM ()
forall a b. (a -> b) -> a -> b
$ Int -> Builder
getBytes Int
n
{-# INLINE ensureBytes #-}

-- | @'getBytes' n@ allocates a new buffer, containing at least @n@ bytes.
getBytes :: Int -> Builder
getBytes :: Int -> Builder
getBytes (I# Int#
n) = Int# -> Builder
getBytes_ Int#
n

-- | The body of the 'getBytes' function, worker-wrappered manually.
getBytes_ :: Int# -> Builder
getBytes_ :: Int# -> Builder
getBytes_ Int#
n = BuildM () -> Builder
mkBuilder (BuildM () -> Builder) -> BuildM () -> Builder
forall a b. (a -> b) -> a -> b
$ do
  DataSink
sink <- BuildM DataSink
getSink
  case DataSink
sink of
    DynamicSink IORef DynamicSink
dRef -> do
      DynamicSink
dyn <- IO DynamicSink -> BuildM DynamicSink
forall a. IO a -> BuildM a
io (IO DynamicSink -> BuildM DynamicSink)
-> IO DynamicSink -> BuildM DynamicSink
forall a b. (a -> b) -> a -> b
$ IORef DynamicSink -> IO DynamicSink
forall a. IORef a -> IO a
readIORef IORef DynamicSink
dRef
      case DynamicSink
dyn of
        ThreadedSink MVar Request
reqV MVar Response
respV -> do
          Ptr Word8
cur <- BuildM (Ptr Word8)
getCur
          IO () -> BuildM ()
forall a. IO a -> BuildM a
io (IO () -> BuildM ()) -> IO () -> BuildM ()
forall a b. (a -> b) -> a -> b
$ MVar Response -> Response -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar Response
respV (Response -> IO ()) -> Response -> IO ()
forall a b. (a -> b) -> a -> b
$ Ptr Word8 -> Int -> Response
MoreBuffer Ptr Word8
cur (Int -> Response) -> Int -> Response
forall a b. (a -> b) -> a -> b
$ Int# -> Int
I# Int#
n
          MVar Request -> BuildM ()
handleRequest MVar Request
reqV
        BoundedGrowingBuffer ForeignPtr Word8
fptr Int
bound ->
          IORef DynamicSink -> ForeignPtr Word8 -> Int -> Int -> BuildM ()
growBufferBounded IORef DynamicSink
dRef ForeignPtr Word8
fptr Int
bound (Int# -> Int
I# Int#
n)
    GrowingBuffer IORef (ForeignPtr Word8)
bufRef -> IORef (ForeignPtr Word8) -> Int -> BuildM ()
growBuffer IORef (ForeignPtr Word8)
bufRef (Int# -> Int
I# Int#
n)
    HandleSink Handle
h Int
nextCap IORef Queue
queueRef -> do
      Ptr Word8
cur <- BuildM (Ptr Word8)
getCur
      IO () -> BuildM ()
forall a. IO a -> BuildM a
io (IO () -> BuildM ()) -> IO () -> BuildM ()
forall a b. (a -> b) -> a -> b
$ Handle -> IORef Queue -> Ptr Word8 -> IO ()
flushQueue Handle
h IORef Queue
queueRef Ptr Word8
cur
      IORef Queue -> Int -> BuildM ()
switchQueue IORef Queue
queueRef (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
nextCap (Int# -> Int
I# Int#
n))
{-# NOINLINE getBytes_ #-}

-- | Return the remaining size of the current buffer, in bytes.
remainingBytes :: BuildM Int
remainingBytes :: BuildM Int
remainingBytes = Ptr Word8 -> Ptr Word8 -> Int
forall a b. Ptr a -> Ptr b -> Int
minusPtr (Ptr Word8 -> Ptr Word8 -> Int)
-> BuildM (Ptr Word8) -> BuildM (Ptr Word8 -> Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> BuildM (Ptr Word8)
getEnd BuildM (Ptr Word8 -> Int) -> BuildM (Ptr Word8) -> BuildM Int
forall a b. BuildM (a -> b) -> BuildM a -> BuildM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> BuildM (Ptr Word8)
getCur
{-# INLINE remainingBytes #-}

----------------------------------------------------------------
-- Performance tuning

-- | @'rebuild' b@ is equivalent to @b@, but it allows GHC to assume
-- that @b@ will be run at most once. This can enable various
-- optimizations that greately improve performance.
--
-- There are two types of typical situations where a use of 'rebuild'
-- is often a win:
--
-- * When constructing a builder using a recursive function. e.g.
--  @rebuild $ foldr ...@.
-- * When constructing a builder using a conditional expression. e.g.
--  @rebuild $ case x of ... @
rebuild :: Builder -> Builder
rebuild :: Builder -> Builder
rebuild (Builder DataSink -> BuilderState -> BuilderState
f) = (DataSink -> BuilderState -> BuilderState) -> Builder
Builder ((DataSink -> BuilderState -> BuilderState) -> Builder)
-> (DataSink -> BuilderState -> BuilderState) -> Builder
forall a b. (a -> b) -> a -> b
$ (DataSink -> BuilderState -> BuilderState)
-> DataSink -> BuilderState -> BuilderState
forall a b. (a -> b) -> a -> b
oneShot ((DataSink -> BuilderState -> BuilderState)
 -> DataSink -> BuilderState -> BuilderState)
-> (DataSink -> BuilderState -> BuilderState)
-> DataSink
-> BuilderState
-> BuilderState
forall a b. (a -> b) -> a -> b
$ \DataSink
dex -> (BuilderState -> BuilderState) -> BuilderState -> BuilderState
forall a b. (a -> b) -> a -> b
oneShot ((BuilderState -> BuilderState) -> BuilderState -> BuilderState)
-> (BuilderState -> BuilderState) -> BuilderState -> BuilderState
forall a b. (a -> b) -> a -> b
$
  \(# Addr#
cur, Addr#
end, State# RealWorld
s #) -> DataSink -> BuilderState -> BuilderState
f DataSink
dex (# Addr#
cur, Addr#
end, State# RealWorld
s #)

----------------------------------------------------------------
-- ThreadedSink

-- | Wait for a request, and switch to a new buffer.
handleRequest :: MVar Request -> BuildM ()
handleRequest :: MVar Request -> BuildM ()
handleRequest MVar Request
reqV = do
  Request Ptr Word8
newCur Ptr Word8
newEnd <- IO Request -> BuildM Request
forall a. IO a -> BuildM a
io (IO Request -> BuildM Request) -> IO Request -> BuildM Request
forall a b. (a -> b) -> a -> b
$ MVar Request -> IO Request
forall a. MVar a -> IO a
takeMVar MVar Request
reqV
  Ptr Word8 -> BuildM ()
setCur Ptr Word8
newCur
  Ptr Word8 -> BuildM ()
setEnd Ptr Word8
newEnd

----------------------------------------------------------------
-- GrowingBuffer

-- | @growBuffer bufRef req@ reallocates the buffer, growing it
-- by at least @req@.
growBuffer :: IORef (ForeignPtr Word8) -> Int -> BuildM ()
growBuffer :: IORef (ForeignPtr Word8) -> Int -> BuildM ()
growBuffer !IORef (ForeignPtr Word8)
bufRef !Int
req = do
  Ptr Word8
cur <- BuildM (Ptr Word8)
getCur
  Ptr Word8
end <- BuildM (Ptr Word8)
getEnd
  ForeignPtr Word8
fptr <- IO (ForeignPtr Word8) -> BuildM (ForeignPtr Word8)
forall a. IO a -> BuildM a
io (IO (ForeignPtr Word8) -> BuildM (ForeignPtr Word8))
-> IO (ForeignPtr Word8) -> BuildM (ForeignPtr Word8)
forall a b. (a -> b) -> a -> b
$ IORef (ForeignPtr Word8) -> IO (ForeignPtr Word8)
forall a. IORef a -> IO a
readIORef IORef (ForeignPtr Word8)
bufRef
  let !base :: Ptr Word8
base = ForeignPtr Word8 -> Ptr Word8
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr Word8
fptr
  let !size :: Int
size = Ptr Word8
cur Ptr Word8 -> Ptr Word8 -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr Word8
base
  let !cap :: Int
cap = Ptr Word8
end Ptr Word8 -> Ptr Word8 -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr Word8
base
  let !newCap :: Int
newCap = Int
cap Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
cap Int
req
  ForeignPtr Word8
newFptr <- IO (ForeignPtr Word8) -> BuildM (ForeignPtr Word8)
forall a. IO a -> BuildM a
io (IO (ForeignPtr Word8) -> BuildM (ForeignPtr Word8))
-> IO (ForeignPtr Word8) -> BuildM (ForeignPtr Word8)
forall a b. (a -> b) -> a -> b
$ Int -> IO (ForeignPtr Word8)
forall a. Int -> IO (ForeignPtr a)
mallocForeignPtrBytes Int
newCap
  let !newBase :: Ptr Word8
newBase = ForeignPtr Word8 -> Ptr Word8
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr Word8
newFptr
  Ptr Word8 -> BuildM ()
setCur (Ptr Word8 -> BuildM ()) -> Ptr Word8 -> BuildM ()
forall a b. (a -> b) -> a -> b
$ Ptr Word8
newBase Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
size
  Ptr Word8 -> BuildM ()
setEnd (Ptr Word8 -> BuildM ()) -> Ptr Word8 -> BuildM ()
forall a b. (a -> b) -> a -> b
$ Ptr Word8
newBase Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
newCap
  IO () -> BuildM ()
forall a. IO a -> BuildM a
io (IO () -> BuildM ()) -> IO () -> BuildM ()
forall a b. (a -> b) -> a -> b
$ do
    Ptr Word8 -> Ptr Word8 -> Int -> IO ()
forall a. Ptr a -> Ptr a -> Int -> IO ()
copyBytes Ptr Word8
newBase Ptr Word8
base Int
size
    ForeignPtr Word8 -> IO ()
forall a. ForeignPtr a -> IO ()
touchForeignPtr ForeignPtr Word8
fptr
    ForeignPtr Word8 -> IO ()
forall a. ForeignPtr a -> IO ()
touchForeignPtr ForeignPtr Word8
newFptr
    IORef (ForeignPtr Word8) -> ForeignPtr Word8 -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (ForeignPtr Word8)
bufRef ForeignPtr Word8
newFptr
{-# INLINE growBuffer #-}

----------------------------------------------------------------
-- HandleSink

-- | Put the content of the 'Queue' to the 'IO.Handle', and empty
-- the 'Queue'.
flushQueue :: IO.Handle -> IORef Queue -> Ptr Word8 -> IO ()
flushQueue :: Handle -> IORef Queue -> Ptr Word8 -> IO ()
flushQueue !Handle
h !IORef Queue
qRef !Ptr Word8
cur = do
  Queue{ queueBuffer :: Queue -> ForeignPtr Word8
queueBuffer = ForeignPtr Word8
fptr, queueStart :: Queue -> Int
queueStart = Int
start, queueTotal :: Queue -> Int
queueTotal = Int
total }
    <- IORef Queue -> IO Queue
forall a. IORef a -> IO a
readIORef IORef Queue
qRef
  let !end :: Int
end = Ptr Word8
cur 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
fptr
  Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
end Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
start) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
    Handle -> ByteString -> IO ()
S.hPut Handle
h (ByteString -> IO ()) -> ByteString -> IO ()
forall a b. (a -> b) -> a -> b
$ ForeignPtr Word8 -> Int -> Int -> ByteString
S.fromForeignPtr ForeignPtr Word8
fptr Int
start (Int
end Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
start)
    IORef Queue -> Queue -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef Queue
qRef Queue
      { queueBuffer :: ForeignPtr Word8
queueBuffer = ForeignPtr Word8
fptr
      , queueStart :: Int
queueStart = Int
end
      , queueTotal :: Int
queueTotal = Int
total Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
end Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
start
      }

-- | @switchQueue qRef minSize adv@ discards the old 'Queue' and sets up
-- a new empty 'Queue' of at least @minSize@ large. If the old 'Queue'
-- is large enough, it is re-used.
switchQueue :: IORef Queue -> Int -> BuildM ()
switchQueue :: IORef Queue -> Int -> BuildM ()
switchQueue !IORef Queue
qRef !Int
minSize = do
  Ptr Word8
end <- BuildM (Ptr Word8)
getCur
  Queue{ queueBuffer :: Queue -> ForeignPtr Word8
queueBuffer = ForeignPtr Word8
fptr, queueTotal :: Queue -> Int
queueTotal = Int
total } <- IO Queue -> BuildM Queue
forall a. IO a -> BuildM a
io (IO Queue -> BuildM Queue) -> IO Queue -> BuildM Queue
forall a b. (a -> b) -> a -> b
$ IORef Queue -> IO Queue
forall a. IORef a -> IO a
readIORef IORef Queue
qRef
  let !base :: Ptr Word8
base = ForeignPtr Word8 -> Ptr Word8
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr Word8
fptr
  let !cap :: Int
cap = Ptr Word8
end Ptr Word8 -> Ptr Word8 -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr Word8
base
  ForeignPtr Word8
newFptr <- if Int
minSize Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
cap
    then ForeignPtr Word8 -> BuildM (ForeignPtr Word8)
forall a. a -> BuildM a
forall (m :: * -> *) a. Monad m => a -> m a
return ForeignPtr Word8
fptr
    else IO (ForeignPtr Word8) -> BuildM (ForeignPtr Word8)
forall a. IO a -> BuildM a
io (IO (ForeignPtr Word8) -> BuildM (ForeignPtr Word8))
-> IO (ForeignPtr Word8) -> BuildM (ForeignPtr Word8)
forall a b. (a -> b) -> a -> b
$ Int -> IO (ForeignPtr Word8)
forall a. Int -> IO (ForeignPtr a)
mallocForeignPtrBytes Int
minSize
  let !newBase :: Ptr Word8
newBase = ForeignPtr Word8 -> Ptr Word8
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr Word8
newFptr
  IO () -> BuildM ()
forall a. IO a -> BuildM a
io (IO () -> BuildM ()) -> IO () -> BuildM ()
forall a b. (a -> b) -> a -> b
$ IORef Queue -> Queue -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef Queue
qRef Queue
    { queueBuffer :: ForeignPtr Word8
queueBuffer = ForeignPtr Word8
newFptr
    , queueStart :: Int
queueStart = Int
0
    , queueTotal :: Int
queueTotal = Int
total
    }
  Ptr Word8 -> BuildM ()
setCur Ptr Word8
newBase
  Ptr Word8 -> BuildM ()
setEnd (Ptr Word8 -> BuildM ()) -> Ptr Word8 -> BuildM ()
forall a b. (a -> b) -> a -> b
$ Ptr Word8
newBase Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
minSize Int
cap

----------------------------------------------------------------
-- BoundedGrowingBuffer

-- | @growBufferBounded dRef fptr bound req@ reallocates the buffer, growing it
-- by at least @req@. If the buffer size would exceed @bound@, it instead
-- interrupts execution by throwing a 'ChunkOverflowException', and switches
-- to a 'ThreadedSink'.
growBufferBounded
  :: IORef DynamicSink -> ForeignPtr Word8 -> Int -> Int -> BuildM ()
growBufferBounded :: IORef DynamicSink -> ForeignPtr Word8 -> Int -> Int -> BuildM ()
growBufferBounded !IORef DynamicSink
dRef !ForeignPtr Word8
fptr !Int
bound !Int
req = do
  Ptr Word8
cur <- BuildM (Ptr Word8)
getCur
  Ptr Word8
end <- BuildM (Ptr Word8)
getCur
  let !base :: Ptr Word8
base = ForeignPtr Word8 -> Ptr Word8
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr Word8
fptr
  let !size :: Int
size = Ptr Word8
cur Ptr Word8 -> Ptr Word8 -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr Word8
base
  let !cap :: Int
cap = Ptr Word8
end Ptr Word8 -> Ptr Word8 -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr Word8
base
  let !newCap :: Int
newCap = Int
cap Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
cap Int
req
  if Int
bound Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
newCap
    then IORef DynamicSink -> Int -> ByteString -> BuildM ()
chunkOverflow IORef DynamicSink
dRef Int
req (ByteString -> BuildM ()) -> ByteString -> BuildM ()
forall a b. (a -> b) -> a -> b
$ ForeignPtr Word8 -> Int -> Int -> ByteString
S.fromForeignPtr ForeignPtr Word8
fptr Int
0 Int
size
    else do
      ForeignPtr Word8
newFptr <- IO (ForeignPtr Word8) -> BuildM (ForeignPtr Word8)
forall a. IO a -> BuildM a
io (IO (ForeignPtr Word8) -> BuildM (ForeignPtr Word8))
-> IO (ForeignPtr Word8) -> BuildM (ForeignPtr Word8)
forall a b. (a -> b) -> a -> b
$ Int -> IO (ForeignPtr Word8)
forall a. Int -> IO (ForeignPtr a)
mallocForeignPtrBytes Int
newCap
      let !newBase :: Ptr Word8
newBase = ForeignPtr Word8 -> Ptr Word8
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr Word8
newFptr
      Ptr Word8 -> BuildM ()
setCur (Ptr Word8 -> BuildM ()) -> Ptr Word8 -> BuildM ()
forall a b. (a -> b) -> a -> b
$ Ptr Word8
newBase Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
size
      Ptr Word8 -> BuildM ()
setEnd (Ptr Word8 -> BuildM ()) -> Ptr Word8 -> BuildM ()
forall a b. (a -> b) -> a -> b
$ Ptr Word8
newBase Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
newCap
      IO () -> BuildM ()
forall a. IO a -> BuildM a
io (IO () -> BuildM ()) -> IO () -> BuildM ()
forall a b. (a -> b) -> a -> b
$ do
        Ptr Word8 -> Ptr Word8 -> Int -> IO ()
forall a. Ptr a -> Ptr a -> Int -> IO ()
copyBytes Ptr Word8
newBase Ptr Word8
base Int
size
        ForeignPtr Word8 -> IO ()
forall a. ForeignPtr a -> IO ()
touchForeignPtr ForeignPtr Word8
fptr
        ForeignPtr Word8 -> IO ()
forall a. ForeignPtr a -> IO ()
touchForeignPtr ForeignPtr Word8
newFptr
        IORef DynamicSink -> DynamicSink -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef DynamicSink
dRef (DynamicSink -> IO ()) -> DynamicSink -> IO ()
forall a b. (a -> b) -> a -> b
$ ForeignPtr Word8 -> Int -> DynamicSink
BoundedGrowingBuffer ForeignPtr Word8
newFptr Int
bound
{-# INLINE growBufferBounded #-}

-- | Throw a 'ChunkOverflowException' and switches to a 'ThreadedSink'.
chunkOverflow :: IORef DynamicSink -> Int -> S.ByteString -> BuildM ()
chunkOverflow :: IORef DynamicSink -> Int -> ByteString -> BuildM ()
chunkOverflow !IORef DynamicSink
dRef !Int
minSize !ByteString
chunk = do
  ThreadId
myTid <- IO ThreadId -> BuildM ThreadId
forall a. IO a -> BuildM a
io IO ThreadId
myThreadId
  MVar Request
reqV <- IO (MVar Request) -> BuildM (MVar Request)
forall a. IO a -> BuildM a
io IO (MVar Request)
forall a. IO (MVar a)
newEmptyMVar
  MVar Response
respV <- IO (MVar Response) -> BuildM (MVar Response)
forall a. IO a -> BuildM a
io IO (MVar Response)
forall a. IO (MVar a)
newEmptyMVar
  IO () -> BuildM ()
forall a. IO a -> BuildM a
io (IO () -> BuildM ()) -> IO () -> BuildM ()
forall a b. (a -> b) -> a -> b
$ ThreadId -> ChunkOverflowException -> IO ()
forall e. Exception e => ThreadId -> e -> IO ()
E.throwTo ThreadId
myTid (ChunkOverflowException -> IO ())
-> ChunkOverflowException -> IO ()
forall a b. (a -> b) -> a -> b
$ ByteString
-> MVar Request -> MVar Response -> Int -> ChunkOverflowException
ChunkOverflowException ByteString
chunk MVar Request
reqV MVar Response
respV Int
minSize
  IO () -> BuildM ()
forall a. IO a -> BuildM a
io (IO () -> BuildM ()) -> IO () -> BuildM ()
forall a b. (a -> b) -> a -> b
$ IORef DynamicSink -> DynamicSink -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef DynamicSink
dRef (DynamicSink -> IO ()) -> DynamicSink -> IO ()
forall a b. (a -> b) -> a -> b
$ MVar Request -> MVar Response -> DynamicSink
ThreadedSink MVar Request
reqV MVar Response
respV
  MVar Request -> BuildM ()
handleRequest MVar Request
reqV