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

-- | The buffer data structure.
data SuperBuffer
    = SuperBuffer
    { SuperBuffer -> IORef (Ptr Word8)
sb_buffer :: {-# UNPACK #-}!(IORef (Ptr Word8))
    , SuperBuffer -> IORef Int
sb_currentSize :: {-# UNPACK #-}!(IORef Int)
    , SuperBuffer -> IORef Int
sb_maxSize :: {-# UNPACK #-}!(IORef Int)
    , SuperBuffer -> MVar ()
sb_lock :: {-# UNPACK #-}!(MVar ())
    }

-- | Allocate a new buffer with a given initial size. The perfect starting point
-- depends on the expected total size and the average size for a single chunk
-- written with 'appendBuffer'. You can always start with 1024 and optimize from
-- there with benchmarks. Please note that the SuperBuffer will no longer be
-- valid after this function terminates, so do NOT pass it to some other
-- thread without waiting for it to finish in the action.
withBuffer :: Int -> (SuperBuffer -> IO ()) -> IO BS.ByteString
withBuffer :: Int -> (SuperBuffer -> IO ()) -> IO ByteString
withBuffer Int
sz SuperBuffer -> IO ()
action =
    do Ptr Word8
ptr <- Int -> IO (Ptr Word8)
forall a. Int -> IO (Ptr a)
mallocBytes Int
sz
       IORef (Ptr Word8)
ptrRef <- Ptr Word8 -> IO (IORef (Ptr Word8))
forall a. a -> IO (IORef a)
newIORef Ptr Word8
ptr
       IORef (Ptr Word8) -> IO ByteString
go IORef (Ptr Word8)
ptrRef IO ByteString -> IO () -> IO ByteString
forall a b. IO a -> IO b -> IO a
`onException` IORef (Ptr Word8) -> IO ()
forall a. IORef (Ptr a) -> IO ()
freeOnException IORef (Ptr Word8)
ptrRef
    where
        freeOnException :: IORef (Ptr a) -> IO ()
freeOnException IORef (Ptr a)
ref =
            do Ptr a
ptr <- IORef (Ptr a) -> IO (Ptr a)
forall a. IORef a -> IO a
readIORef IORef (Ptr a)
ref
               Ptr a -> IO ()
forall a. Ptr a -> IO ()
free Ptr a
ptr
        go :: IORef (Ptr Word8) -> IO ByteString
go IORef (Ptr Word8)
ptrRef =
            do IORef Int
sizeRef <- Int -> IO (IORef Int)
forall a. a -> IO (IORef a)
newIORef Int
0
               IORef Int
maxSizeRef <- Int -> IO (IORef Int)
forall a. a -> IO (IORef a)
newIORef Int
sz
               MVar ()
lock <- IO (MVar ())
forall a. IO (MVar a)
newEmptyMVar
               let sb :: SuperBuffer
sb = IORef (Ptr Word8)
-> IORef Int -> IORef Int -> MVar () -> SuperBuffer
SuperBuffer IORef (Ptr Word8)
ptrRef IORef Int
sizeRef IORef Int
maxSizeRef MVar ()
lock
               SuperBuffer -> IO ()
action SuperBuffer
sb
               SuperBuffer -> IO ByteString
readBuffer SuperBuffer
sb
{-# INLINE withBuffer #-}

-- | Write a bytestring to the buffer and grow the buffer if needed. Note that only
-- one thread at any given time may call this function. Use 'appendBufferT' when
-- accessing 'SuperBuffer' from multiple threads.
appendBuffer :: SuperBuffer -> BS.ByteString -> IO ()
appendBuffer :: SuperBuffer -> ByteString -> IO ()
appendBuffer SuperBuffer
sb ByteString
bs
    | ByteString -> Bool
BS.null ByteString
bs = () -> IO ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    | Bool
otherwise =
          ByteString -> (CStringLen -> IO ()) -> IO ()
forall a. ByteString -> (CStringLen -> IO a) -> IO a
BS.unsafeUseAsCStringLen ByteString
bs ((CStringLen -> IO ()) -> IO ()) -> (CStringLen -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \(Ptr CChar
cstr, Int
len) ->
          do Int
currentSize <- IORef Int -> IO Int
forall a. IORef a -> IO a
readIORef (SuperBuffer -> IORef Int
sb_currentSize SuperBuffer
sb)
             Int
maxSize <- IORef Int -> IO Int
forall a. IORef a -> IO a
readIORef (SuperBuffer -> IORef Int
sb_maxSize SuperBuffer
sb)
             let nextSize :: Int
nextSize = Int
currentSize Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len
             Ptr Word8
writePtr <-
                 if Int
nextSize Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
maxSize
                 then do let maxSize' :: Int
maxSize' = Int
nextSize Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int -> Int -> Int
forall a. Bits a => a -> Int -> a
unsafeShiftR Int
nextSize Int
1
                         IORef Int -> Int -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (SuperBuffer -> IORef Int
sb_maxSize SuperBuffer
sb) Int
maxSize'
                         Ptr Word8
buff <- IORef (Ptr Word8) -> IO (Ptr Word8)
forall a. IORef a -> IO a
readIORef (SuperBuffer -> IORef (Ptr Word8)
sb_buffer SuperBuffer
sb)
                         Ptr Word8
buff' <- Ptr Word8 -> Int -> IO (Ptr Word8)
forall a. Ptr a -> Int -> IO (Ptr a)
reallocBytes Ptr Word8
buff Int
maxSize'
                         IORef (Ptr Word8) -> Ptr Word8 -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (SuperBuffer -> IORef (Ptr Word8)
sb_buffer SuperBuffer
sb) Ptr Word8
buff'
                         Ptr Word8 -> IO (Ptr Word8)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Ptr Word8
buff'
                 else IORef (Ptr Word8) -> IO (Ptr Word8)
forall a. IORef a -> IO a
readIORef (SuperBuffer -> IORef (Ptr Word8)
sb_buffer SuperBuffer
sb)
             let copyTarget :: Ptr b
copyTarget = Ptr Word8
writePtr Ptr Word8 -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
currentSize
             Ptr CChar -> Ptr CChar -> Int -> IO ()
forall a. Ptr a -> Ptr a -> Int -> IO ()
copyBytes Ptr CChar
forall b. Ptr b
copyTarget Ptr CChar
cstr Int
len
             IORef Int -> Int -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (SuperBuffer -> IORef Int
sb_currentSize SuperBuffer
sb) (Int
currentSize Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len)
{-# INLINE appendBuffer #-}

-- | Write a bytestring to the buffer and grow the buffer if needed. This function
-- can be used accross different threads, but is slower than 'appendBuffer'.
appendBufferT :: SuperBuffer -> BS.ByteString -> IO ()
appendBufferT :: SuperBuffer -> ByteString -> IO ()
appendBufferT SuperBuffer
sb ByteString
bs =
    IO () -> IO () -> IO () -> IO ()
forall a b c. IO a -> IO b -> IO c -> IO c
bracket_ (MVar () -> () -> IO ()
forall a. MVar a -> a -> IO ()
putMVar (SuperBuffer -> MVar ()
sb_lock SuperBuffer
sb) ()) (MVar () -> IO ()
forall a. MVar a -> IO a
takeMVar (SuperBuffer -> MVar ()
sb_lock SuperBuffer
sb)) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
    SuperBuffer -> ByteString -> IO ()
appendBuffer SuperBuffer
sb ByteString
bs
{-# INLINE appendBufferT #-}

-- | Read the final buffer contents. This must only be called once
readBuffer :: SuperBuffer -> IO BS.ByteString
readBuffer :: SuperBuffer -> IO ByteString
readBuffer SuperBuffer
sb =
    do (Ptr Word8
buff, Int
currentSize, Int
maxSize) <-
           (,,)
           (Ptr Word8 -> Int -> Int -> (Ptr Word8, Int, Int))
-> IO (Ptr Word8) -> IO (Int -> Int -> (Ptr Word8, Int, Int))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IORef (Ptr Word8) -> IO (Ptr Word8)
forall a. IORef a -> IO a
readIORef (SuperBuffer -> IORef (Ptr Word8)
sb_buffer SuperBuffer
sb)
           IO (Int -> Int -> (Ptr Word8, Int, Int))
-> IO Int -> IO (Int -> (Ptr Word8, Int, Int))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> IORef Int -> IO Int
forall a. IORef a -> IO a
readIORef (SuperBuffer -> IORef Int
sb_currentSize SuperBuffer
sb)
           IO (Int -> (Ptr Word8, Int, Int))
-> IO Int -> IO (Ptr Word8, Int, Int)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> IORef Int -> IO Int
forall a. IORef a -> IO a
readIORef (SuperBuffer -> IORef Int
sb_maxSize SuperBuffer
sb)
       Ptr Word8
finalPtr <-
           if Int
currentSize Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
maxSize
           then Ptr Word8 -> Int -> IO (Ptr Word8)
forall a. Ptr a -> Int -> IO (Ptr a)
reallocBytes Ptr Word8
buff Int
currentSize
           else Ptr Word8 -> IO (Ptr Word8)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Ptr Word8
buff
       Ptr Word8 -> Int -> IO () -> IO ByteString
BS.unsafePackCStringFinalizer Ptr Word8
finalPtr Int
currentSize (Ptr Word8 -> IO ()
forall a. Ptr a -> IO ()
free Ptr Word8
finalPtr)
{-# INLINE readBuffer #-}

-- | Get current (filled) size of the buffer
size :: SuperBuffer -> IO Int
size :: SuperBuffer -> IO Int
size SuperBuffer
sb = IORef Int -> IO Int
forall a. IORef a -> IO a
readIORef (IORef Int -> IO Int) -> IORef Int -> IO Int
forall a b. (a -> b) -> a -> b
$ SuperBuffer -> IORef Int
sb_currentSize SuperBuffer
sb
{-# INLINE size #-}