{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE RankNTypes #-}
module Data.Streaming.ByteString.Builder
( BuilderRecv
, BuilderPopper
, BuilderFinish
, newBuilderRecv
, newByteStringBuilderRecv
, toByteStringIO
, toByteStringIOWith
, toByteStringIOWithBuffer
, Buffer
, freeSize
, sliceSize
, bufferSize
, allocBuffer
, reuseBuffer
, nextSlice
, unsafeFreezeBuffer
, unsafeFreezeNonEmptyBuffer
, 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
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
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
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 :: Int
-> (ByteString -> IO ())
-> Builder
-> 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 #-}
toByteStringIO :: (ByteString -> IO ())
-> Builder
-> IO ()
toByteStringIO :: (ByteString -> IO ()) -> Builder -> IO ()
toByteStringIO = Int -> (ByteString -> IO ()) -> Builder -> IO ()
toByteStringIOWith Int
defaultChunkSize
{-# INLINE toByteStringIO #-}