module Buffer
(
Buffer,
new,
push,
pushBytes,
pushStorable,
pull,
pullBytes,
pullStorable,
getBytes,
getSpace,
)
where
import Buffer.Prelude hiding (State, Buffer)
import Foreign.C
import qualified Data.ByteString.Internal as C
foreign import ccall unsafe "memmove"
memmove :: Ptr a -> Ptr a -> CSize -> IO (Ptr a)
foreign import ccall unsafe "memcpy"
memcpy :: Ptr a -> Ptr a -> CSize -> IO (Ptr a)
newtype Buffer =
Buffer (IORef State)
data State =
State !(ForeignPtr Word8) !Int !Int !Int
new :: Int -> IO Buffer
new capacity =
do
fptr <- mallocForeignPtrBytes capacity
stateIORef <- newIORef (State fptr 0 0 capacity)
return (Buffer stateIORef)
push :: Buffer -> Int -> (Ptr Word8 -> IO (Int, result)) -> IO result
push (Buffer stateIORef) space ptrIO =
do
State fptr start end capacity <- readIORef stateIORef
let
!remainingSpace = capacity end
!capacityDelta = space remainingSpace
!occupiedSpace = end start
in
if capacityDelta <= 0
then
do
(!actualSpace, !output) <- withForeignPtr fptr $ \ptr -> ptrIO (plusPtr ptr end)
writeIORef stateIORef $! State fptr start (end + actualSpace) capacity
return output
else
if capacityDelta > start
then
do
let !newCapacity = occupiedSpace + space
newFPtr <- mallocForeignPtrBytes newCapacity
(!actualSpace, !output) <- withForeignPtr newFPtr $ \newPtr -> do
withForeignPtr fptr $ \ptr -> do
memcpy newPtr (plusPtr ptr start) (fromIntegral occupiedSpace)
ptrIO (plusPtr newPtr occupiedSpace)
let !newOccupiedSpace = occupiedSpace + actualSpace
writeIORef stateIORef $! State newFPtr 0 newOccupiedSpace newCapacity
return output
else
do
(!actualSpace, !output) <- withForeignPtr fptr $ \ptr -> do
memmove ptr (plusPtr ptr start) (fromIntegral occupiedSpace)
ptrIO (plusPtr ptr occupiedSpace)
writeIORef stateIORef $! State fptr 0 (occupiedSpace + actualSpace) capacity
return output
pull :: Buffer -> Int -> (Ptr Word8 -> IO result) -> (Int -> IO result) -> IO result
pull (Buffer stateIORef) pulledAmount ptrIO refill =
do
State fptr start end capacity <- readIORef stateIORef
let !newStart = start + pulledAmount
if newStart > end
then refill $! newStart end
else do
!pulled <- withForeignPtr fptr $ \ptr -> ptrIO (plusPtr ptr start)
writeIORef stateIORef $! State fptr newStart end capacity
return pulled
pushBytes :: Buffer -> ByteString -> IO ()
pushBytes buffer (C.PS bytesFPtr offset length) =
push buffer length $ \ptr ->
withForeignPtr bytesFPtr $ \bytesPtr ->
C.memcpy ptr (plusPtr bytesPtr offset) length $> (length, ())
pullBytes :: Buffer -> Int -> (ByteString -> result) -> (Int -> IO result) -> IO result
pullBytes buffer amount bytesResult =
pull buffer amount (\ptr -> fmap bytesResult (C.create amount (\destPtr -> C.memcpy destPtr ptr amount)))
pushStorable :: (Storable storable) => Buffer -> storable -> IO ()
pushStorable buffer storable =
push buffer amount (\ptr -> poke (castPtr ptr) storable $> (amount, ()))
where
amount = sizeOf storable
pullStorable :: (Storable storable) => Buffer -> (storable -> result) -> (Int -> IO result) -> IO result
pullStorable buffer storableResult =
pull buffer amount (\ptr -> fmap storableResult (peek (castPtr ptr)))
where
amount =
sizeOf ((undefined :: (a -> b) -> a) storableResult)
getSpace :: Buffer -> IO Int
getSpace (Buffer stateIORef) =
do
State fptr start end capacity <- readIORef stateIORef
return $! end start
getBytes :: Buffer -> IO ByteString
getBytes (Buffer stateIORef) =
do
State fptr start end capacity <- readIORef stateIORef
let size = end start
withForeignPtr fptr $ \ptr -> C.create size $ \destPtr -> C.memcpy destPtr (plusPtr ptr start) size