module Data.ByteString.SuperBuffer.Pure
( SuperBuffer, withBuffer, appendBuffer, appendBufferT, size )
where
import Control.Concurrent.MVar
import Control.Exception
import Data.Bits
import Data.IORef
import Data.Word
import Foreign.Marshal.Alloc
import Foreign.Marshal.Utils
import Foreign.Ptr
import qualified Data.ByteString as BS
import qualified Data.ByteString.Unsafe as BS
data SuperBuffer
= SuperBuffer
{ sb_buffer :: !(IORef (Ptr Word8))
, sb_currentSize :: !(IORef Int)
, sb_maxSize :: !(IORef Int)
, sb_lock :: !(MVar ())
}
withBuffer :: Int -> (SuperBuffer -> IO ()) -> IO BS.ByteString
withBuffer sz action =
do ptr <- mallocBytes sz
ptrRef <- newIORef ptr
go ptrRef `onException` freeOnException ptrRef
where
freeOnException ref =
do ptr <- readIORef ref
free ptr
go ptrRef =
do sizeRef <- newIORef 0
maxSizeRef <- newIORef sz
lock <- newEmptyMVar
let sb = SuperBuffer ptrRef sizeRef maxSizeRef lock
action sb
readBuffer sb
appendBuffer :: SuperBuffer -> BS.ByteString -> IO ()
appendBuffer sb bs
| BS.null bs = pure ()
| otherwise =
BS.unsafeUseAsCStringLen bs $ \(cstr, len) ->
do currentSize <- readIORef (sb_currentSize sb)
maxSize <- readIORef (sb_maxSize sb)
let nextSize = currentSize + len
writePtr <-
if nextSize > maxSize
then do let maxSize' = nextSize + unsafeShiftR nextSize 1
writeIORef (sb_maxSize sb) maxSize'
buff <- readIORef (sb_buffer sb)
buff' <- reallocBytes buff maxSize'
writeIORef (sb_buffer sb) buff'
pure buff'
else readIORef (sb_buffer sb)
let copyTarget = writePtr `plusPtr` currentSize
copyBytes copyTarget cstr len
writeIORef (sb_currentSize sb) (currentSize + len)
appendBufferT :: SuperBuffer -> BS.ByteString -> IO ()
appendBufferT sb bs =
bracket_ (putMVar (sb_lock sb) ()) (takeMVar (sb_lock sb)) $
appendBuffer sb bs
readBuffer :: SuperBuffer -> IO BS.ByteString
readBuffer sb =
do (buff, currentSize, maxSize) <-
(,,)
<$> readIORef (sb_buffer sb)
<*> readIORef (sb_currentSize sb)
<*> readIORef (sb_maxSize sb)
finalPtr <-
if currentSize < maxSize
then reallocBytes buff currentSize
else pure buff
BS.unsafePackCStringFinalizer finalPtr currentSize (free finalPtr)
size :: SuperBuffer -> IO Int
size sb = readIORef $ sb_currentSize sb