{-# 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
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
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
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 :: 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 forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< 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 #-}