{-# 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 = newByteStringBuilderRecv
{-# INLINE newBuilderRecv #-}
newByteStringBuilderRecv :: BufferAllocStrategy -> IO (BuilderRecv, BuilderFinish)
newByteStringBuilderRecv (ioBufInit, nextBuf) = do
refBuf <- newIORef ioBufInit
return (push refBuf, finish refBuf)
where
finish refBuf = do
ioBuf <- readIORef refBuf
buf <- ioBuf
return $ unsafeFreezeNonEmptyBuffer buf
push refBuf builder = do
refWri <- newIORef $ Left $ runBuilder builder
return $ popper refBuf refWri
popper refBuf refWri = do
ioBuf <- readIORef refBuf
ebWri <- readIORef refWri
case ebWri of
Left bWri -> do
!buf@(Buffer _ _ op ope) <- ioBuf
(bytes, next) <- bWri op (ope `minusPtr` op)
let op' = op `plusPtr` bytes
case next of
Done -> do
writeIORef refBuf $ return $ updateEndOfSlice buf op'
return S.empty
More minSize bWri' -> do
let buf' = updateEndOfSlice buf op'
{-# INLINE cont #-}
cont mbs = do
ioBuf' <- nextBuf minSize buf'
writeIORef refBuf ioBuf'
writeIORef refWri $ Left bWri'
case mbs of
Just bs | not $ S.null bs -> return bs
_ -> popper refBuf refWri
cont $ unsafeFreezeNonEmptyBuffer buf'
Chunk bs bWri' -> do
let buf' = updateEndOfSlice buf op'
let yieldBS = do
nextBuf 1 buf' >>= writeIORef refBuf
writeIORef refWri $ Left bWri'
if S.null bs
then popper refBuf refWri
else return bs
case unsafeFreezeNonEmptyBuffer buf' of
Nothing -> yieldBS
Just bs' -> do
writeIORef refWri $ Right yieldBS
return bs'
Right action -> action
toByteStringIOWithBuffer :: Int
-> (ByteString -> IO ())
-> Builder
-> ForeignPtr Word8
-> IO ()
toByteStringIOWithBuffer initBufSize io b initBuf = do
go initBufSize initBuf (runBuilder b)
where
go bufSize buf = loop
where
loop :: BufferWriter -> IO ()
loop wr = do
(len, next) <- withForeignPtr buf (flip wr bufSize)
when (len > 0) (io $! PS buf 0 len)
case next of
Done -> return ()
More newBufSize nextWr
| newBufSize > bufSize -> do
newBuf <- mallocByteString newBufSize
go newBufSize newBuf nextWr
| otherwise -> loop nextWr
Chunk s nextWr -> do
unless (S.null s) (io s)
loop nextWr
toByteStringIOWith :: Int
-> (ByteString -> IO ())
-> Builder
-> IO ()
toByteStringIOWith bufSize io b =
toByteStringIOWithBuffer bufSize io b =<< mallocByteString bufSize
{-# INLINE toByteStringIOWith #-}
toByteStringIO :: (ByteString -> IO ())
-> Builder
-> IO ()
toByteStringIO = toByteStringIOWith defaultChunkSize
{-# INLINE toByteStringIO #-}