{-# 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 <- IO Buffer -> IO (IORef (IO Buffer))
forall a. a -> IO (IORef a)
newIORef IO Buffer
ioBufInit
    (BuilderRecv, BuilderFinish) -> IO (BuilderRecv, BuilderFinish)
forall a. a -> IO a
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 <- IORef (IO Buffer) -> IO (IO Buffer)
forall a. IORef a -> IO a
readIORef IORef (IO Buffer)
refBuf
        Buffer
buf <- IO Buffer
ioBuf
        Maybe ByteString -> BuilderFinish
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe ByteString -> BuilderFinish)
-> Maybe ByteString -> BuilderFinish
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 <- 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
        IO ByteString -> IO (IO ByteString)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (IO ByteString -> IO (IO ByteString))
-> IO ByteString -> IO (IO ByteString)
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 <- IORef (IO Buffer) -> IO (IO Buffer)
forall a. IORef a -> IO a
readIORef IORef (IO Buffer)
refBuf
        Either BufferWriter (IO ByteString)
ebWri <- IORef (Either BufferWriter (IO ByteString))
-> IO (Either BufferWriter (IO ByteString))
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 Ptr Word8 -> Ptr Word8 -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr Word8
op)
                let op' :: Ptr b
op' = Ptr Word8
op Ptr Word8 -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
bytes
                case Next
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.
                                IO Buffer
ioBuf' <- Int -> Buffer -> IO (IO Buffer)
nextBuf Int
minSize Buffer
buf'
                                IORef (IO Buffer) -> IO Buffer -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (IO Buffer)
refBuf IO Buffer
ioBuf'
                                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'
                                case Maybe ByteString
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
            (Int
len, 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
buf (BufferWriter -> Int -> Ptr Word8 -> IO (Int, Next)
forall a b c. (a -> b -> c) -> b -> a -> c
flip BufferWriter
wr Int
bufSize)
            Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) (ByteString -> IO ()
io (ByteString -> IO ()) -> ByteString -> 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 -> () -> 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
                        ForeignPtr Word8
newBuf <- Int -> IO (ForeignPtr Word8)
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
                    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 #-}