{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE Rank2Types #-}

-- | Buffers for 'Builder's.  This is a partial copy of blaze-builder-0.3.3.4's
-- "Blaze.ByteString.Builder.Internal.Buffer" module, which was removed in
-- blaze-builder-0.4.
--
-- If you are using blaze-builder 0.3.*, this module just re-exports from
-- "Blaze.ByteString.Builder.Internal.Buffer".
--
-- Since 0.1.10.0
--
module Data.Streaming.ByteString.Builder.Buffer
    (
    -- * Buffers
      Buffer (..)

    -- ** Status information
    , freeSize
    , sliceSize
    , bufferSize

    -- ** Creation and modification
    , allocBuffer
    , reuseBuffer
    , nextSlice
    , updateEndOfSlice

    -- ** Conversion to bytestings
    , unsafeFreezeBuffer
    , unsafeFreezeNonEmptyBuffer

    -- * Buffer allocation strategies
    , BufferAllocStrategy
    , allNewBuffersStrategy
    , reuseBufferStrategy
    , defaultStrategy
    ) where

import Data.ByteString.Lazy.Internal (defaultChunkSize)

import qualified Data.ByteString as S
import qualified Data.ByteString.Internal as S
import Foreign (Word8, ForeignPtr, Ptr, plusPtr, minusPtr)
import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)

------------------------------------------------------------------------------
-- Buffers
------------------------------------------------------------------------------

-- | A buffer @Buffer fpbuf p0 op ope@ describes a buffer with the underlying
-- byte array @fpbuf..ope@, the currently written slice @p0..op@ and the free
-- space @op..ope@.
--
-- Since 0.1.10.0
--
data Buffer = Buffer {-# UNPACK #-} !(ForeignPtr Word8) -- underlying pinned array
                     {-# UNPACK #-} !(Ptr Word8)        -- beginning of slice
                     {-# UNPACK #-} !(Ptr Word8)        -- next free byte
                     {-# UNPACK #-} !(Ptr Word8)        -- first byte after buffer

-- | The size of the free space of the buffer.
--
-- Since 0.1.10.0
--
freeSize :: Buffer -> Int
freeSize :: Buffer -> Int
freeSize (Buffer ForeignPtr Word8
_ Ptr Word8
_ Ptr Word8
op Ptr Word8
ope) = Ptr Word8
ope Ptr Word8 -> Ptr Word8 -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr Word8
op

-- | The size of the written slice in the buffer.
--
-- Since 0.1.10.0
--
sliceSize :: Buffer -> Int
sliceSize :: Buffer -> Int
sliceSize (Buffer ForeignPtr Word8
_ Ptr Word8
p0 Ptr Word8
op Ptr Word8
_) = Ptr Word8
op Ptr Word8 -> Ptr Word8 -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr Word8
p0

-- | The size of the whole byte array underlying the buffer.
--
-- Since 0.1.10.0
--
bufferSize :: Buffer -> Int
bufferSize :: Buffer -> Int
bufferSize (Buffer ForeignPtr Word8
fpbuf Ptr Word8
_ Ptr Word8
_ Ptr Word8
ope) =
    Ptr Word8
ope Ptr Word8 -> Ptr Word8 -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` ForeignPtr Word8 -> Ptr Word8
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr Word8
fpbuf

-- | @allocBuffer size@ allocates a new buffer of size @size@.
--
-- Since 0.1.10.0
--
{-# INLINE allocBuffer #-}
allocBuffer :: Int -> IO Buffer
allocBuffer :: Int -> IO Buffer
allocBuffer Int
size = do
    ForeignPtr Word8
fpbuf <- Int -> IO (ForeignPtr Word8)
forall a. Int -> IO (ForeignPtr a)
S.mallocByteString Int
size
    let !pbuf :: Ptr Word8
pbuf = ForeignPtr Word8 -> Ptr Word8
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr Word8
fpbuf
    Buffer -> IO Buffer
forall (m :: * -> *) a. Monad m => a -> m a
return (Buffer -> IO Buffer) -> Buffer -> IO Buffer
forall a b. (a -> b) -> a -> b
$! ForeignPtr Word8 -> Ptr Word8 -> Ptr Word8 -> Ptr Word8 -> Buffer
Buffer ForeignPtr Word8
fpbuf Ptr Word8
pbuf Ptr Word8
pbuf (Ptr Word8
pbuf Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
size)

-- | Resets the beginning of the next slice and the next free byte such that
-- the whole buffer can be filled again.
--
-- Since 0.1.10.0
--
{-# INLINE reuseBuffer #-}
reuseBuffer :: Buffer -> Buffer
reuseBuffer :: Buffer -> Buffer
reuseBuffer (Buffer ForeignPtr Word8
fpbuf Ptr Word8
_ Ptr Word8
_ Ptr Word8
ope) = ForeignPtr Word8 -> Ptr Word8 -> Ptr Word8 -> Ptr Word8 -> Buffer
Buffer ForeignPtr Word8
fpbuf Ptr Word8
p0 Ptr Word8
p0 Ptr Word8
ope
  where
    p0 :: Ptr Word8
p0 = ForeignPtr Word8 -> Ptr Word8
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr Word8
fpbuf

-- | Convert the buffer to a bytestring. This operation is unsafe in the sense
-- that created bytestring shares the underlying byte array with the buffer.
-- Hence, depending on the later use of this buffer (e.g., if it gets reset and
-- filled again) referential transparency may be lost.
--
-- Since 0.1.10.0
--
{-# INLINE unsafeFreezeBuffer #-}
unsafeFreezeBuffer :: Buffer -> S.ByteString
unsafeFreezeBuffer :: Buffer -> ByteString
unsafeFreezeBuffer (Buffer ForeignPtr Word8
fpbuf Ptr Word8
p0 Ptr Word8
op Ptr Word8
_) =
    ForeignPtr Word8 -> Int -> Int -> ByteString
S.PS ForeignPtr Word8
fpbuf (Ptr Word8
p0 Ptr Word8 -> Ptr Word8 -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` ForeignPtr Word8 -> Ptr Word8
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr Word8
fpbuf) (Ptr Word8
op Ptr Word8 -> Ptr Word8 -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr Word8
p0)

-- | Convert a buffer to a non-empty bytestring. See 'unsafeFreezeBuffer' for
-- the explanation of why this operation may be unsafe.
--
-- Since 0.1.10.0
--
{-# INLINE unsafeFreezeNonEmptyBuffer #-}
unsafeFreezeNonEmptyBuffer :: Buffer -> Maybe S.ByteString
unsafeFreezeNonEmptyBuffer :: Buffer -> Maybe ByteString
unsafeFreezeNonEmptyBuffer Buffer
buf
  | Buffer -> Int
sliceSize Buffer
buf Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = Maybe ByteString
forall a. Maybe a
Nothing
  | Bool
otherwise          = ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just (ByteString -> Maybe ByteString) -> ByteString -> Maybe ByteString
forall a b. (a -> b) -> a -> b
$ Buffer -> ByteString
unsafeFreezeBuffer Buffer
buf

-- | Update the end of slice pointer.
--
-- Since 0.1.10.0
--
{-# INLINE updateEndOfSlice #-}
updateEndOfSlice :: Buffer    -- Old buffer
                 -> Ptr Word8 -- New end of slice
                 -> Buffer    -- Updated buffer
updateEndOfSlice :: Buffer -> Ptr Word8 -> Buffer
updateEndOfSlice (Buffer ForeignPtr Word8
fpbuf Ptr Word8
p0 Ptr Word8
_ Ptr Word8
ope) Ptr Word8
op' = ForeignPtr Word8 -> Ptr Word8 -> Ptr Word8 -> Ptr Word8 -> Buffer
Buffer ForeignPtr Word8
fpbuf Ptr Word8
p0 Ptr Word8
op' Ptr Word8
ope

-- | Move the beginning of the slice to the next free byte such that the
-- remaining free space of the buffer can be filled further. This operation
-- is safe and can be used to fill the remaining part of the buffer after a
-- direct insertion of a bytestring or a flush.
--
-- Since 0.1.10.0
--
{-# INLINE nextSlice #-}
nextSlice :: Int -> Buffer -> Maybe Buffer
nextSlice :: Int -> Buffer -> Maybe Buffer
nextSlice Int
minSize (Buffer ForeignPtr Word8
fpbuf Ptr Word8
_ Ptr Word8
op Ptr Word8
ope)
  | Ptr Word8
ope Ptr Word8 -> Ptr Word8 -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr Word8
op Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
minSize = Maybe Buffer
forall a. Maybe a
Nothing
  | Bool
otherwise                    = Buffer -> Maybe Buffer
forall a. a -> Maybe a
Just (ForeignPtr Word8 -> Ptr Word8 -> Ptr Word8 -> Ptr Word8 -> Buffer
Buffer ForeignPtr Word8
fpbuf Ptr Word8
op Ptr Word8
op Ptr Word8
ope)

------------------------------------------------------------------------------
-- Buffer allocation strategies
------------------------------------------------------------------------------

-- | A buffer allocation strategy @(buf0, nextBuf)@ specifies the initial
-- buffer to use and how to compute a new buffer @nextBuf minSize buf@ with at
-- least size @minSize@ from a filled buffer @buf@. The double nesting of the
-- @IO@ monad helps to ensure that the reference to the filled buffer @buf@ is
-- lost as soon as possible, but the new buffer doesn't have to be allocated
-- too early.
--
-- Since 0.1.10.0
--
type BufferAllocStrategy = (IO Buffer, Int -> Buffer -> IO (IO Buffer))

-- | The simplest buffer allocation strategy: whenever a buffer is requested,
-- allocate a new one that is big enough for the next build step to execute.
--
-- NOTE that this allocation strategy may spill quite some memory upon direct
-- insertion of a bytestring by the builder. Thats no problem for garbage
-- collection, but it may lead to unreasonably high memory consumption in
-- special circumstances.
--
-- Since 0.1.10.0
--
allNewBuffersStrategy :: Int                 -- Minimal buffer size.
                      -> BufferAllocStrategy
allNewBuffersStrategy :: Int -> BufferAllocStrategy
allNewBuffersStrategy Int
bufSize =
    ( Int -> IO Buffer
allocBuffer Int
bufSize
    , \Int
reqSize Buffer
_ -> IO Buffer -> IO (IO Buffer)
forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> IO Buffer
allocBuffer (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
reqSize Int
bufSize)) )

-- | An unsafe, but possibly more efficient buffer allocation strategy:
-- reuse the buffer, if it is big enough for the next build step to execute.
--
-- Since 0.1.10.0
--
reuseBufferStrategy :: IO Buffer
                    -> BufferAllocStrategy
reuseBufferStrategy :: IO Buffer -> BufferAllocStrategy
reuseBufferStrategy IO Buffer
buf0 =
    (IO Buffer
buf0, Int -> Buffer -> IO (IO Buffer)
forall (m :: * -> *). Monad m => Int -> Buffer -> m (IO Buffer)
tryReuseBuffer)
  where
    tryReuseBuffer :: Int -> Buffer -> m (IO Buffer)
tryReuseBuffer Int
reqSize Buffer
buf
      | Buffer -> Int
bufferSize Buffer
buf Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
reqSize = IO Buffer -> m (IO Buffer)
forall (m :: * -> *) a. Monad m => a -> m a
return (IO Buffer -> m (IO Buffer)) -> IO Buffer -> m (IO Buffer)
forall a b. (a -> b) -> a -> b
$ Buffer -> IO Buffer
forall (m :: * -> *) a. Monad m => a -> m a
return (Buffer -> Buffer
reuseBuffer Buffer
buf)
      | Bool
otherwise                 = IO Buffer -> m (IO Buffer)
forall (m :: * -> *) a. Monad m => a -> m a
return (IO Buffer -> m (IO Buffer)) -> IO Buffer -> m (IO Buffer)
forall a b. (a -> b) -> a -> b
$ Int -> IO Buffer
allocBuffer Int
reqSize

defaultStrategy :: BufferAllocStrategy
defaultStrategy :: BufferAllocStrategy
defaultStrategy = Int -> BufferAllocStrategy
allNewBuffersStrategy Int
defaultChunkSize