@@ -5,9 +5,9 @@ import qualified Data.Streaming.ByteString.Builder.Buffer as B (Buffer (..))
55import Data.Word (Word8)
66import Data.IORef (IORef, newIORef, readIORef, writeIORef)
77import Data.ByteString (ByteString, empty, length)
8−import Data.ByteString.Internal (ByteString(PS), create)
8+import Data.ByteString.Internal (ByteString(PS), mallocByteString, c_free_finalizer)
99import Data.ByteString.Unsafe (unsafeTake, unsafeDrop)
10−import Foreign.ForeignPtr (newForeignPtr_, withForeignPtr)
10+import Foreign.ForeignPtr (newForeignPtr, newForeignPtr_, withForeignPtr)
1111import Foreign.Marshal.Alloc (mallocBytes, free)
1212import Foreign.Ptr (Ptr, castPtr, plusPtr)
1313
@@ -21,10 +21,18 @@ type BufferPool = IORef ByteString
2121newBufferPool :: IO BufferPool
2222newBufferPool = newIORef Data.ByteString.empty
2323
24−allocBuffer :: Int -> IO ByteString
25−allocBuffer size = create size initializer
26− where initializer = \_ -> return ()
27−{-# INLINE allocBuffer #-}
24+mallocBuffer :: Int -> IO ByteString
25+mallocBuffer size = do
26+ ptr <- mallocBytes size
27+ fptr <- newForeignPtr c_free_finalizer ptr
28+ return $! PS fptr 0 size
29+{-# INLINE mallocBuffer #-}
30+
31+createBuffer :: Int -> IO ByteString
32+createBuffer size = do
33+ fptr <- mallocByteString size
34+ return $! PS fptr 0 size
35+{-# INLINE createBuffer #-}
2836
2937usefulBuffer :: ByteString -> Bool
3038usefulBuffer buffer = Data.ByteString.length buffer >= minBufferSize
@@ -33,7 +41,7 @@ usefulBuffer buffer = Data.ByteString.length buffer >= minBufferSize
3341getBuffer :: BufferPool -> IO ByteString
3442getBuffer pool = do
3543 buffer <- readIORef pool
36− if usefulBuffer buffer then return buffer else allocBuffer largeBufferSize
44+ if usefulBuffer buffer then return buffer else createBuffer largeBufferSize
3745{-# INLINE getBuffer #-}
3846
3947putBuffer :: BufferPool -> ByteString -> IO ()