{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE RankNTypes #-}

-- | Convert a stream of bytestring @Builder@s into a stream of @ByteString@s.
--
-- Adapted from blaze-builder-enumerator, written by Michael Snoyman and Simon Meier.
--
-- Note that the functions here can work in any monad built on top of @IO@ or
-- @ST@.
--
-- Also provides @toByteStringIO*@ like "Blaze.ByteString.Builder"s, for
-- "Data.ByteString.Builder".
--
-- Since 0.1.9
--
module Data.Streaming.ByteString.Builder
    ( BuilderRecv
    , BuilderPopper
    , BuilderFinish
    , newBuilderRecv
    , newByteStringBuilderRecv

    -- * toByteStringIO
    , toByteStringIO
    , toByteStringIOWith
    , toByteStringIOWithBuffer

    -- * Buffers
    , Buffer

    -- ** Status information
    , freeSize
    , sliceSize
    , bufferSize

    -- ** Creation and modification
    , allocBuffer
    , reuseBuffer
    , nextSlice

    -- ** Conversion to bytestings
    , unsafeFreezeBuffer
    , unsafeFreezeNonEmptyBuffer

    -- * Buffer allocation strategies
    , BufferAllocStrategy
    , allNewBuffersStrategy
    , reuseBufferStrategy
    , defaultStrategy
    )
    where

import Control.Monad (when,unless)
import qualified Data.ByteString as S
import Data.ByteString.Builder (Builder)
import Data.ByteString.Builder.Extra (runBuilder, BufferWriter, Next(Done, More, Chunk))
import Data.ByteString.Internal (mallocByteString, ByteString(PS))
import Data.ByteString.Lazy.Internal (defaultChunkSize)
import Data.IORef (newIORef, writeIORef, readIORef)
import Data.Word (Word8)
import Foreign.ForeignPtr (ForeignPtr, withForeignPtr)
import Foreign.Ptr (plusPtr, minusPtr)

import Data.Streaming.ByteString.Builder.Buffer

-- | Provides a series of @ByteString@s until empty, at which point it provides
-- an empty @ByteString@.
--
-- Since 0.1.10.0
--
type BuilderPopper = IO S.ByteString

type BuilderRecv = Builder -> IO BuilderPopper

type BuilderFinish = IO (Maybe S.ByteString)

newBuilderRecv :: BufferAllocStrategy -> IO (BuilderRecv, BuilderFinish)
newBuilderRecv :: BufferAllocStrategy -> IO (BuilderRecv, BuilderFinish)
newBuilderRecv = BufferAllocStrategy -> IO (BuilderRecv, BuilderFinish)
newByteStringBuilderRecv
{-# INLINE newBuilderRecv #-}

newByteStringBuilderRecv :: BufferAllocStrategy -> IO (BuilderRecv, BuilderFinish)
newByteStringBuilderRecv :: BufferAllocStrategy -> IO (BuilderRecv, BuilderFinish)
newByteStringBuilderRecv (IO Buffer
ioBufInit, Int -> Buffer -> IO (IO Buffer)
nextBuf) = do
    refBuf <- IO Buffer -> IO (IORef (IO Buffer))
forall a. a -> IO (IORef a)
newIORef IO Buffer
ioBufInit
    return (push refBuf, finish refBuf)
  where
    finish :: IORef (IO Buffer) -> BuilderFinish
finish IORef (IO Buffer)
refBuf = do
        ioBuf <- IORef (IO Buffer) -> IO (IO Buffer)
forall a. IORef a -> IO a
readIORef IORef (IO Buffer)
refBuf
        buf <- ioBuf
        return $ unsafeFreezeNonEmptyBuffer buf

    push :: IORef (IO Buffer) -> BuilderRecv
push IORef (IO Buffer)
refBuf Builder
builder = do
        refWri <- Either BufferWriter (IO ByteString)
-> IO (IORef (Either BufferWriter (IO ByteString)))
forall a. a -> IO (IORef a)
newIORef (Either BufferWriter (IO ByteString)
 -> IO (IORef (Either BufferWriter (IO ByteString))))
-> Either BufferWriter (IO ByteString)
-> IO (IORef (Either BufferWriter (IO ByteString)))
forall a b. (a -> b) -> a -> b
$ BufferWriter -> Either BufferWriter (IO ByteString)
forall a b. a -> Either a b
Left (BufferWriter -> Either BufferWriter (IO ByteString))
-> BufferWriter -> Either BufferWriter (IO ByteString)
forall a b. (a -> b) -> a -> b
$ Builder -> BufferWriter
runBuilder Builder
builder
        return $ popper refBuf refWri

    popper :: IORef (IO Buffer)
-> IORef (Either BufferWriter (IO ByteString)) -> IO ByteString
popper IORef (IO Buffer)
refBuf IORef (Either BufferWriter (IO ByteString))
refWri = do
        ioBuf <- IORef (IO Buffer) -> IO (IO Buffer)
forall a. IORef a -> IO a
readIORef IORef (IO Buffer)
refBuf
        ebWri <- readIORef refWri
        case ebWri of
            Left BufferWriter
bWri -> do
                !buf@(Buffer _ _ op ope) <- IO Buffer
ioBuf
                (bytes, next) <- bWri op (ope `minusPtr` op)
                let op' = Ptr Word8
op Ptr Word8 -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
bytes
                case next of
                    Next
Done -> do
                        IORef (IO Buffer) -> IO Buffer -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (IO Buffer)
refBuf (IO Buffer -> IO ()) -> IO Buffer -> IO ()
forall a b. (a -> b) -> a -> b
$ Buffer -> IO Buffer
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Buffer -> IO Buffer) -> Buffer -> IO Buffer
forall a b. (a -> b) -> a -> b
$ Buffer -> Ptr Word8 -> Buffer
updateEndOfSlice Buffer
buf Ptr Word8
forall {b}. Ptr b
op'
                        ByteString -> IO ByteString
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ByteString
S.empty
                    More Int
minSize BufferWriter
bWri' -> do
                        let buf' :: Buffer
buf' = Buffer -> Ptr Word8 -> Buffer
updateEndOfSlice Buffer
buf Ptr Word8
forall {b}. Ptr b
op'
                            {-# INLINE cont #-}
                            cont :: Maybe ByteString -> IO ByteString
cont Maybe ByteString
mbs = do
                                -- sequencing the computation of the next buffer
                                -- construction here ensures that the reference to the
                                -- foreign pointer `fp` is lost as soon as possible.
                                ioBuf' <- Int -> Buffer -> IO (IO Buffer)
nextBuf Int
minSize Buffer
buf'
                                writeIORef refBuf ioBuf'
                                writeIORef refWri $ Left bWri'
                                case mbs of
                                    Just ByteString
bs | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ ByteString -> Bool
S.null ByteString
bs -> ByteString -> IO ByteString
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ByteString
bs
                                    Maybe ByteString
_ -> IORef (IO Buffer)
-> IORef (Either BufferWriter (IO ByteString)) -> IO ByteString
popper IORef (IO Buffer)
refBuf IORef (Either BufferWriter (IO ByteString))
refWri
                        Maybe ByteString -> IO ByteString
cont (Maybe ByteString -> IO ByteString)
-> Maybe ByteString -> IO ByteString
forall a b. (a -> b) -> a -> b
$ Buffer -> Maybe ByteString
unsafeFreezeNonEmptyBuffer Buffer
buf'
                    Chunk ByteString
bs BufferWriter
bWri' -> do
                        let buf' :: Buffer
buf' = Buffer -> Ptr Word8 -> Buffer
updateEndOfSlice Buffer
buf Ptr Word8
forall {b}. Ptr b
op'
                        let yieldBS :: IO ByteString
yieldBS = do
                                Int -> Buffer -> IO (IO Buffer)
nextBuf Int
1 Buffer
buf' IO (IO Buffer) -> (IO Buffer -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= IORef (IO Buffer) -> IO Buffer -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (IO Buffer)
refBuf
                                IORef (Either BufferWriter (IO ByteString))
-> Either BufferWriter (IO ByteString) -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (Either BufferWriter (IO ByteString))
refWri (Either BufferWriter (IO ByteString) -> IO ())
-> Either BufferWriter (IO ByteString) -> IO ()
forall a b. (a -> b) -> a -> b
$ BufferWriter -> Either BufferWriter (IO ByteString)
forall a b. a -> Either a b
Left BufferWriter
bWri'
                                if ByteString -> Bool
S.null ByteString
bs
                                    then IORef (IO Buffer)
-> IORef (Either BufferWriter (IO ByteString)) -> IO ByteString
popper IORef (IO Buffer)
refBuf IORef (Either BufferWriter (IO ByteString))
refWri
                                    else ByteString -> IO ByteString
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ByteString
bs
                        case Buffer -> Maybe ByteString
unsafeFreezeNonEmptyBuffer Buffer
buf' of
                            Maybe ByteString
Nothing -> IO ByteString
yieldBS
                            Just ByteString
bs' -> do
                                IORef (Either BufferWriter (IO ByteString))
-> Either BufferWriter (IO ByteString) -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (Either BufferWriter (IO ByteString))
refWri (Either BufferWriter (IO ByteString) -> IO ())
-> Either BufferWriter (IO ByteString) -> IO ()
forall a b. (a -> b) -> a -> b
$ IO ByteString -> Either BufferWriter (IO ByteString)
forall a b. b -> Either a b
Right IO ByteString
yieldBS
                                ByteString -> IO ByteString
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ByteString
bs'
            Right IO ByteString
action -> IO ByteString
action

-- | Use a pre-existing buffer to 'toByteStringIOWith'.
--
-- Since 0.1.9
--
toByteStringIOWithBuffer :: Int
                         -> (ByteString -> IO ())
                         -> Builder
                         -> ForeignPtr Word8
                         -> IO ()
toByteStringIOWithBuffer :: Int
-> (ByteString -> IO ()) -> Builder -> ForeignPtr Word8 -> IO ()
toByteStringIOWithBuffer Int
initBufSize ByteString -> IO ()
io Builder
b ForeignPtr Word8
initBuf = do
    Int -> ForeignPtr Word8 -> BufferWriter -> IO ()
go Int
initBufSize ForeignPtr Word8
initBuf (Builder -> BufferWriter
runBuilder Builder
b)
  where
    go :: Int -> ForeignPtr Word8 -> BufferWriter -> IO ()
go Int
bufSize ForeignPtr Word8
buf = BufferWriter -> IO ()
loop
      where
        loop :: BufferWriter -> IO ()
        loop :: BufferWriter -> IO ()
loop BufferWriter
wr = do
            (len, 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
buf (BufferWriter -> Int -> Ptr Word8 -> IO (Int, Next)
forall a b c. (a -> b -> c) -> b -> a -> c
flip BufferWriter
wr Int
bufSize)
            when (len > 0) (io $! PS buf 0 len)
            case next of
                Next
Done -> () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
                More Int
newBufSize BufferWriter
nextWr
                    | Int
newBufSize Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
bufSize -> do
                        newBuf <- Int -> IO (ForeignPtr Word8)
forall a. Int -> IO (ForeignPtr a)
mallocByteString Int
newBufSize
                        go newBufSize newBuf nextWr
                    | Bool
otherwise -> BufferWriter -> IO ()
loop BufferWriter
nextWr
                Chunk ByteString
s BufferWriter
nextWr -> do
                    Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (ByteString -> Bool
S.null ByteString
s) (ByteString -> IO ()
io ByteString
s)
                    BufferWriter -> IO ()
loop BufferWriter
nextWr

-- | @toByteStringIOWith bufSize io b@ runs the builder @b@ with a buffer of
-- at least the size @bufSize@ and executes the 'IO' action @io@ whenever the
-- buffer is full.
--
-- Compared to 'toLazyByteStringWith' this function requires less allocation,
-- as the output buffer is only allocated once at the start of the
-- serialization and whenever something bigger than the current buffer size has
-- to be copied into the buffer, which should happen very seldomly for the
-- default buffer size of 32kb. Hence, the pressure on the garbage collector is
-- reduced, which can be an advantage when building long sequences of bytes.
--
-- Since 0.1.9
--
toByteStringIOWith :: Int                    -- ^ Buffer size (upper bounds
                                             -- the number of bytes forced
                                             -- per call to the 'IO' action).
                   -> (ByteString -> IO ())  -- ^ 'IO' action to execute per
                                             -- full buffer, which is
                                             -- referenced by a strict
                                             -- 'S.ByteString'.
                   -> Builder                -- ^ 'Builder' to run.
                   -> IO ()
toByteStringIOWith :: Int -> (ByteString -> IO ()) -> Builder -> IO ()
toByteStringIOWith Int
bufSize ByteString -> IO ()
io Builder
b =
    Int
-> (ByteString -> IO ()) -> Builder -> ForeignPtr Word8 -> IO ()
toByteStringIOWithBuffer Int
bufSize ByteString -> IO ()
io Builder
b (ForeignPtr Word8 -> IO ()) -> IO (ForeignPtr Word8) -> IO ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Int -> IO (ForeignPtr Word8)
forall a. Int -> IO (ForeignPtr a)
mallocByteString Int
bufSize
{-# INLINE toByteStringIOWith #-}

-- | Run the builder with a 'defaultChunkSize'd buffer and execute the given
-- 'IO' action whenever the buffer is full or gets flushed.
--
-- @ 'toByteStringIO' = 'toByteStringIOWith' 'defaultChunkSize'@
--
-- Since 0.1.9
--
toByteStringIO :: (ByteString -> IO ())
               -> Builder
               -> IO ()
toByteStringIO :: (ByteString -> IO ()) -> Builder -> IO ()
toByteStringIO = Int -> (ByteString -> IO ()) -> Builder -> IO ()
toByteStringIOWith Int
defaultChunkSize
{-# INLINE toByteStringIO #-}