{-# LANGUAGE BangPatterns      #-}
{-# LANGUAGE CPP               #-}
{-# LANGUAGE OverloadedStrings #-}

-- | Buffering for output streams based on bytestring builders.
--
-- Buffering an output stream can often improve throughput by reducing the
-- number of system calls made through the file descriptor. The @bytestring@
-- package provides an efficient monoidal datatype used for serializing values
-- directly to an output buffer, called a 'Builder', originally implemented in
-- the @blaze-builder@ package by Simon Meier. When compiling with @bytestring@
-- versions older than 0.10.4, (i.e. GHC <= 7.6) users must depend on the
-- @bytestring-builder@ library to get the new builder implementation. Since we
-- try to maintain compatibility with the last three GHC versions, the
-- dependency on @bytestring-builder@ can be dropped after the release of GHC
-- 7.12.
--
--
-- /Using this module/
--
-- Given an 'OutputStream' taking 'ByteString':
--
-- > someOutputStream :: OutputStream ByteString
--
-- You create a new output stream wrapping the original one that accepts
-- 'Builder' values:
--
--
-- @
-- do
--     newStream <- Streams.'builderStream' someOutputStream
--     Streams.'write' ('Just' $ 'Data.ByteString.Builder.byteString' \"hello\") newStream
--     ....
-- @
--
--
-- You can flush the output buffer using 'Data.ByteString.Builder.Extra.flush':
--
-- @
--     ....
--     Streams.'write' ('Just' 'Data.ByteString.Builder.Extra.flush') newStream
--     ....
-- @
--
-- As a convention, 'builderStream' will write the empty string to the wrapped
-- 'OutputStream' upon a builder buffer flush. Output streams which receive
-- 'ByteString' should either ignore the empty string or interpret it as a
-- signal to flush their own buffers, as the @handleToOutputStream@ and
-- "System.IO.Streams.Zlib" functions do.
--
-- /Example/
--
-- @
-- example :: IO [ByteString]
-- example = do
--     let l1 = 'Data.List.intersperse' \" \" [\"the\", \"quick\", \"brown\", \"fox\"]
--     let l2 = 'Data.List.intersperse' \" \" [\"jumped\", \"over\", \"the\"]
--     let l  = map 'Data.ByteString.Builder.byteString' l1 ++ ['Data.ByteString.Builder.Extra.flush'] ++ map 'Data.ByteString.Builder.byteString' l2
--     is          \<- Streams.'System.IO.Streams.fromList' l
--     (os0, grab) \<- Streams.'System.IO.Streams.listOutputStream'
--     os          \<- Streams.'builderStream' os0
--     Streams.'System.IO.Streams.connect' is os >> grab
--
-- ghci> example
-- [\"the quick brown fox\",\"\",\"jumped over the\"]
-- @
--
module System.IO.Streams.Builder
 ( -- * Blaze builder conversion
   builderStream
 , builderStreamWithBufferSize
 , unsafeBuilderStream
 ) where

------------------------------------------------------------------------------
import           Control.Monad                    (when)
import           Data.ByteString.Builder.Internal (Buffer (..), BufferRange (..), Builder, byteStringFromBuffer, defaultChunkSize, fillWithBuildStep, newBuffer, runBuilder)
import           Data.ByteString.Char8            (ByteString)
import qualified Data.ByteString.Char8            as S
import           Data.IORef                       (newIORef, readIORef, writeIORef)

------------------------------------------------------------------------------
import           System.IO.Streams.Internal       (OutputStream, makeOutputStream, write, writeTo)


------------------------------------------------------------------------------
builderStreamWithBufferFunc :: IO Buffer
                            -> OutputStream ByteString
                            -> IO (OutputStream Builder)
builderStreamWithBufferFunc :: IO Buffer -> OutputStream ByteString -> IO (OutputStream Builder)
builderStreamWithBufferFunc IO Buffer
mkNewBuf OutputStream ByteString
os = do
    IORef (Maybe Buffer)
ref <- Maybe Buffer -> IO (IORef (Maybe Buffer))
forall a. a -> IO (IORef a)
newIORef Maybe Buffer
forall a. Maybe a
Nothing
    (Maybe Builder -> IO ()) -> IO (OutputStream Builder)
forall a. (Maybe a -> IO ()) -> IO (OutputStream a)
makeOutputStream ((Maybe Builder -> IO ()) -> IO (OutputStream Builder))
-> (Maybe Builder -> IO ()) -> IO (OutputStream Builder)
forall a b. (a -> b) -> a -> b
$ IORef (Maybe Buffer) -> Maybe Builder -> IO ()
chunk IORef (Maybe Buffer)
ref
  where
    chunk :: IORef (Maybe Buffer) -> Maybe Builder -> IO ()
chunk IORef (Maybe Buffer)
ref Maybe Builder
Nothing = do
        Maybe Buffer
mbuf <- IORef (Maybe Buffer) -> IO (Maybe Buffer)
forall a. IORef a -> IO a
readIORef IORef (Maybe Buffer)
ref
        case Maybe Buffer
mbuf of
          -- If we existing buffer leftovers, write them to the output.
          Maybe Buffer
Nothing  -> () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return (() -> IO ()) -> () -> IO ()
forall a b. (a -> b) -> a -> b
$! ()
          Just Buffer
buf -> Buffer -> IO ()
writeBuf Buffer
buf
        Maybe ByteString -> OutputStream ByteString -> IO ()
forall a. Maybe a -> OutputStream a -> IO ()
write Maybe ByteString
forall a. Maybe a
Nothing OutputStream ByteString
os
    chunk IORef (Maybe Buffer)
ref (Just Builder
builder) = IORef (Maybe Buffer) -> BuildStep () -> IO ()
runStep IORef (Maybe Buffer)
ref (BuildStep () -> IO ()) -> BuildStep () -> IO ()
forall a b. (a -> b) -> a -> b
$ Builder -> BuildStep ()
runBuilder Builder
builder

    getBuf :: IORef (Maybe Buffer) -> IO Buffer
getBuf IORef (Maybe Buffer)
ref = IORef (Maybe Buffer) -> IO (Maybe Buffer)
forall a. IORef a -> IO a
readIORef IORef (Maybe Buffer)
ref IO (Maybe Buffer) -> (Maybe Buffer -> IO Buffer) -> IO Buffer
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= IO Buffer -> (Buffer -> IO Buffer) -> Maybe Buffer -> IO Buffer
forall b a. b -> (a -> b) -> Maybe a -> b
maybe IO Buffer
mkNewBuf Buffer -> IO Buffer
forall (m :: * -> *) a. Monad m => a -> m a
return

    bumpBuf :: Buffer -> Ptr Word8 -> Buffer
bumpBuf (Buffer ForeignPtr Word8
fp (BufferRange !Ptr Word8
_ Ptr Word8
endBuf)) Ptr Word8
endPtr =
        ForeignPtr Word8 -> BufferRange -> Buffer
Buffer ForeignPtr Word8
fp (Ptr Word8 -> Ptr Word8 -> BufferRange
BufferRange Ptr Word8
endPtr Ptr Word8
endBuf)

    updateBuf :: IORef (Maybe Buffer) -> Buffer -> Ptr Word8 -> IO ()
updateBuf IORef (Maybe Buffer)
ref Buffer
buf Ptr Word8
endPtr = IORef (Maybe Buffer) -> Maybe Buffer -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (Maybe Buffer)
ref (Maybe Buffer -> IO ()) -> Maybe Buffer -> IO ()
forall a b. (a -> b) -> a -> b
$! Buffer -> Maybe Buffer
forall a. a -> Maybe a
Just (Buffer -> Maybe Buffer) -> Buffer -> Maybe Buffer
forall a b. (a -> b) -> a -> b
$! Buffer -> Ptr Word8 -> Buffer
bumpBuf Buffer
buf Ptr Word8
endPtr

    writeBuf :: Buffer -> IO ()
writeBuf Buffer
buf = do
        let bs :: ByteString
bs = Buffer -> ByteString
byteStringFromBuffer Buffer
buf
        Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Bool -> Bool
not (Bool -> Bool) -> (ByteString -> Bool) -> ByteString -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Bool
S.null (ByteString -> Bool) -> ByteString -> Bool
forall a b. (a -> b) -> a -> b
$ ByteString
bs) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ OutputStream ByteString -> Maybe ByteString -> IO ()
forall a. OutputStream a -> Maybe a -> IO ()
writeTo OutputStream ByteString
os (Maybe ByteString -> IO ()) -> Maybe ByteString -> IO ()
forall a b. (a -> b) -> a -> b
$! ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just ByteString
bs

    bufRange :: Buffer -> BufferRange
bufRange (Buffer ForeignPtr Word8
_ BufferRange
rng) = BufferRange
rng

    runStep :: IORef (Maybe Buffer) -> BuildStep () -> IO ()
runStep IORef (Maybe Buffer)
ref BuildStep ()
step = do
        Buffer
buf <- IORef (Maybe Buffer) -> IO Buffer
getBuf IORef (Maybe Buffer)
ref
        BuildStep ()
-> (Ptr Word8 -> () -> IO ())
-> (Ptr Word8 -> Int -> BuildStep () -> IO ())
-> (Ptr Word8 -> ByteString -> BuildStep () -> IO ())
-> BufferRange
-> IO ()
forall a b.
BuildStep a
-> (Ptr Word8 -> a -> IO b)
-> (Ptr Word8 -> Int -> BuildStep a -> IO b)
-> (Ptr Word8 -> ByteString -> BuildStep a -> IO b)
-> BufferRange
-> IO b
fillWithBuildStep BuildStep ()
step (Buffer -> Ptr Word8 -> () -> IO ()
cDone Buffer
buf) (Buffer -> Ptr Word8 -> Int -> BuildStep () -> IO ()
forall p. Buffer -> Ptr Word8 -> p -> BuildStep () -> IO ()
cFull Buffer
buf) (Buffer -> Ptr Word8 -> ByteString -> BuildStep () -> IO ()
cInsert Buffer
buf)
                          (Buffer -> BufferRange
bufRange Buffer
buf)
      where
        cDone :: Buffer -> Ptr Word8 -> () -> IO ()
cDone Buffer
buf Ptr Word8
endPtr !() = IORef (Maybe Buffer) -> Buffer -> Ptr Word8 -> IO ()
updateBuf IORef (Maybe Buffer)
ref Buffer
buf Ptr Word8
endPtr
        cFull :: Buffer -> Ptr Word8 -> p -> BuildStep () -> IO ()
cFull Buffer
buf !Ptr Word8
endPtr !p
_ BuildStep ()
newStep = do
            Buffer -> IO ()
writeBuf (Buffer -> IO ()) -> Buffer -> IO ()
forall a b. (a -> b) -> a -> b
$! Buffer -> Ptr Word8 -> Buffer
bumpBuf Buffer
buf Ptr Word8
endPtr
            IORef (Maybe Buffer) -> Maybe Buffer -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (Maybe Buffer)
ref Maybe Buffer
forall a. Maybe a
Nothing
            IORef (Maybe Buffer) -> BuildStep () -> IO ()
runStep IORef (Maybe Buffer)
ref BuildStep ()
newStep
        cInsert :: Buffer -> Ptr Word8 -> ByteString -> BuildStep () -> IO ()
cInsert Buffer
buf !Ptr Word8
endPtr !ByteString
bs BuildStep ()
newStep = do
            Buffer -> IO ()
writeBuf (Buffer -> IO ()) -> Buffer -> IO ()
forall a b. (a -> b) -> a -> b
$! Buffer -> Ptr Word8 -> Buffer
bumpBuf Buffer
buf Ptr Word8
endPtr
            IORef (Maybe Buffer) -> Maybe Buffer -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (Maybe Buffer)
ref Maybe Buffer
forall a. Maybe a
Nothing
            OutputStream ByteString -> Maybe ByteString -> IO ()
forall a. OutputStream a -> Maybe a -> IO ()
writeTo OutputStream ByteString
os (Maybe ByteString -> IO ()) -> Maybe ByteString -> IO ()
forall a b. (a -> b) -> a -> b
$! ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just ByteString
bs
            IORef (Maybe Buffer) -> BuildStep () -> IO ()
runStep IORef (Maybe Buffer)
ref BuildStep ()
newStep


------------------------------------------------------------------------------
-- | Converts a 'ByteString' sink into a 'Builder' sink, using the supplied
-- buffer size.
--
-- Note that if the generated builder receives a
-- 'Blaze.ByteString.Builder.flush', by convention it will send an empty string
-- to the supplied @'OutputStream' 'ByteString'@ to indicate that any output
-- buffers are to be flushed.
--
-- /Since: 1.3.0.0./
builderStreamWithBufferSize :: Int -> OutputStream ByteString -> IO (OutputStream Builder)
builderStreamWithBufferSize :: Int -> OutputStream ByteString -> IO (OutputStream Builder)
builderStreamWithBufferSize Int
bufsiz = IO Buffer -> OutputStream ByteString -> IO (OutputStream Builder)
builderStreamWithBufferFunc (Int -> IO Buffer
newBuffer Int
bufsiz)


------------------------------------------------------------------------------
-- | Converts a 'ByteString' sink into a 'Builder' sink.
--
-- Note that if the generated builder receives a
-- 'Blaze.ByteString.Builder.flush', by convention it will send an empty string
-- to the supplied @'OutputStream' 'ByteString'@ to indicate that any output
-- buffers are to be flushed.
--
builderStream :: OutputStream ByteString -> IO (OutputStream Builder)
builderStream :: OutputStream ByteString -> IO (OutputStream Builder)
builderStream = Int -> OutputStream ByteString -> IO (OutputStream Builder)
builderStreamWithBufferSize Int
defaultChunkSize


------------------------------------------------------------------------------
-- | Unsafe variation on 'builderStream' that reuses an existing buffer for
-- efficiency.
--
-- /NOTE/: because the buffer is reused, subsequent 'ByteString' values written
-- to the wrapped 'OutputString' will cause previous yielded strings to change.
-- Do not retain references to these 'ByteString' values inside the
-- 'OutputStream' you pass to this function, or you will violate referential
-- transparency.
--
-- If you /must/ retain copies of these values, then please use
-- 'Data.ByteString.copy' to ensure that you have a fresh copy of the
-- underlying string.
--
-- You can create a Buffer with 'Data.ByteString.Builder.Internal.newBuffer'.
--
unsafeBuilderStream :: IO Buffer
                    -> OutputStream ByteString
                    -> IO (OutputStream Builder)
unsafeBuilderStream :: IO Buffer -> OutputStream ByteString -> IO (OutputStream Builder)
unsafeBuilderStream IO Buffer
mkBuf OutputStream ByteString
os = do
    Buffer
buf <- IO Buffer
mkBuf
    IO Buffer -> OutputStream ByteString -> IO (OutputStream Builder)
builderStreamWithBufferFunc (Buffer -> IO Buffer
forall (m :: * -> *) a. Monad m => a -> m a
return Buffer
buf) OutputStream ByteString
os