{-# 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
    IORef (IO Buffer)
refBuf <- forall a. a -> IO (IORef a)
newIORef IO Buffer
ioBufInit
    forall (m :: * -> *) a. Monad m => a -> m a
return (IORef (IO Buffer) -> BuilderRecv
push IORef (IO Buffer)
refBuf, IORef (IO Buffer) -> BuilderFinish
finish IORef (IO Buffer)
refBuf)
  where
    finish :: IORef (IO Buffer) -> BuilderFinish
finish IORef (IO Buffer)
refBuf = do
        IO Buffer
ioBuf <- forall a. IORef a -> IO a
readIORef IORef (IO Buffer)
refBuf
        Buffer
buf <- IO Buffer
ioBuf
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ Buffer -> Maybe ByteString
unsafeFreezeNonEmptyBuffer Buffer
buf

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

    popper :: IORef (IO Buffer)
-> IORef (Either BufferWriter (IO ByteString)) -> IO ByteString
popper IORef (IO Buffer)
refBuf IORef (Either BufferWriter (IO ByteString))
refWri = do
        IO Buffer
ioBuf <- forall a. IORef a -> IO a
readIORef IORef (IO Buffer)
refBuf
        Either BufferWriter (IO ByteString)
ebWri <- forall a. IORef a -> IO a
readIORef IORef (Either BufferWriter (IO ByteString))
refWri
        case Either BufferWriter (IO ByteString)
ebWri of
            Left BufferWriter
bWri -> do
                !buf :: Buffer
buf@(Buffer ForeignPtr Word8
_ Ptr Word8
_ Ptr Word8
op Ptr Word8
ope) <- IO Buffer
ioBuf
                (Int
bytes, Next
next) <- BufferWriter
bWri Ptr Word8
op (Ptr Word8
ope forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr Word8
op)
                let op' :: Ptr b
op' = Ptr Word8
op forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
bytes
                case Next
next of
                    Next
Done -> do
                        forall a. IORef a -> a -> IO ()
writeIORef IORef (IO Buffer)
refBuf forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ Buffer -> Ptr Word8 -> Buffer
updateEndOfSlice Buffer
buf forall {b}. Ptr b
op'
                        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 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.
                                IO Buffer
ioBuf' <- Int -> Buffer -> IO (IO Buffer)
nextBuf Int
minSize Buffer
buf'
                                forall a. IORef a -> a -> IO ()
writeIORef IORef (IO Buffer)
refBuf IO Buffer
ioBuf'
                                forall a. IORef a -> a -> IO ()
writeIORef IORef (Either BufferWriter (IO ByteString))
refWri forall a b. (a -> b) -> a -> b
$ forall a b. a -> Either a b
Left BufferWriter
bWri'
                                case Maybe ByteString
mbs of
                                    Just ByteString
bs | Bool -> Bool
not forall a b. (a -> b) -> a -> b
$ ByteString -> Bool
S.null ByteString
bs -> 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 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 forall {b}. Ptr b
op'
                        let yieldBS :: IO ByteString
yieldBS = do
                                Int -> Buffer -> IO (IO Buffer)
nextBuf Int
1 Buffer
buf' forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall a. IORef a -> a -> IO ()
writeIORef IORef (IO Buffer)
refBuf
                                forall a. IORef a -> a -> IO ()
writeIORef IORef (Either BufferWriter (IO ByteString))
refWri forall a b. (a -> b) -> a -> b
$ 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 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
                                forall a. IORef a -> a -> IO ()
writeIORef IORef (Either BufferWriter (IO ByteString))
refWri forall a b. (a -> b) -> a -> b
$ forall a b. b -> Either a b
Right IO ByteString
yieldBS
                                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
            (Int
len, Next
next) <- forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr Word8
buf (forall a b c. (a -> b -> c) -> b -> a -> c
flip BufferWriter
wr Int
bufSize)
            forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
len forall a. Ord a => a -> a -> Bool
> Int
0) (ByteString -> IO ()
io forall a b. (a -> b) -> a -> b
$! ForeignPtr Word8 -> Int -> Int -> ByteString
PS ForeignPtr Word8
buf Int
0 Int
len)
            case Next
next of
                Next
Done -> forall (m :: * -> *) a. Monad m => a -> m a
return ()
                More Int
newBufSize BufferWriter
nextWr
                    | Int
newBufSize forall a. Ord a => a -> a -> Bool
> Int
bufSize -> do
                        ForeignPtr Word8
newBuf <- forall a. Int -> IO (ForeignPtr a)
mallocByteString Int
newBufSize
                        Int -> ForeignPtr Word8 -> BufferWriter -> IO ()
go Int
newBufSize ForeignPtr Word8
newBuf BufferWriter
nextWr
                    | Bool
otherwise -> BufferWriter -> IO ()
loop BufferWriter
nextWr
                Chunk ByteString
s BufferWriter
nextWr -> do
                    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 forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< 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 #-}