{-# LANGUAGE EmptyDataDecls #-}
{-# LANGUAGE ForeignFunctionInterface #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Data.ByteString.SuperBuffer
    ( SuperBuffer, withBuffer, appendBuffer, appendBufferT, size
    )
where

import Control.Concurrent.MVar
import Control.Exception
import Data.Coerce
import Foreign
import Foreign.C
import qualified Data.ByteString as BS
import qualified Data.ByteString.Unsafe as BS

-- | The buffer. Internally only a pointer to a C struct. Don't worry,
-- this module attempts to make usage of the SuperBuffer as safe as possible  in
-- terms of memory leaks (even when exceptions occur).
newtype SuperBuffer
    = SuperBuffer (SuperBufferP, 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 :: Int64 -> (SuperBuffer -> IO ()) -> IO BS.ByteString
withBuffer :: Int64 -> (SuperBuffer -> IO ()) -> IO ByteString
withBuffer Int64
sz SuperBuffer -> IO ()
action =
    IO SuperBuffer
-> (SuperBuffer -> IO ())
-> (SuperBuffer -> IO ByteString)
-> IO ByteString
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracket (Int64 -> IO SuperBuffer
newBuffer Int64
sz) SuperBuffer -> IO ()
destroyBuffer ((SuperBuffer -> IO ByteString) -> IO ByteString)
-> (SuperBuffer -> IO ByteString) -> IO ByteString
forall a b. (a -> b) -> a -> b
$ \SuperBuffer
buf ->
    do Either SomeException ()
ok <- IO () -> IO (Either SomeException ())
forall e a. Exception e => IO a -> IO (Either e a)
try (SuperBuffer -> IO ()
action SuperBuffer
buf)
       case Either SomeException ()
ok of
         Left (SomeException
exception :: SomeException) ->
             do SuperBuffer -> IO ()
destroyBufferContents SuperBuffer
buf
                SomeException -> IO ByteString
forall e a. Exception e => e -> IO a
throwIO SomeException
exception
         Right () ->
             SuperBuffer -> IO ByteString
readBuffer SuperBuffer
buf -- if something goes to shit here, we could be in trouble...
{-# INLINE withBuffer #-}

newBuffer :: Int64 -> IO SuperBuffer
newBuffer :: Int64 -> IO SuperBuffer
newBuffer Int64
sz = (SuperBufferP, MVar ()) -> SuperBuffer
SuperBuffer ((SuperBufferP, MVar ()) -> SuperBuffer)
-> IO (SuperBufferP, MVar ()) -> IO SuperBuffer
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((,) (SuperBufferP -> MVar () -> (SuperBufferP, MVar ()))
-> IO SuperBufferP -> IO (MVar () -> (SuperBufferP, MVar ()))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> CSize -> IO SuperBufferP
new_sbuf (Int64 -> CSize
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int64
sz) IO (MVar () -> (SuperBufferP, MVar ()))
-> IO (MVar ()) -> IO (SuperBufferP, MVar ())
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> IO (MVar ())
forall a. IO (MVar a)
newEmptyMVar)
{-# INLINE newBuffer #-}


-- | 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 (SuperBufferP
ptr, MVar ()
_)) ByteString
bs =
    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) ->
    SuperBufferP -> Ptr CChar -> CSize -> IO ()
append_sbuf SuperBufferP
ptr Ptr CChar
cstr (Int -> CSize
forall a b. (Integral a, Num b) => a -> b
fromIntegral 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 buf :: SuperBuffer
buf@(SuperBuffer (SuperBufferP
_, MVar ()
lock)) 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 MVar ()
lock ()) (MVar () -> IO ()
forall a. MVar a -> IO a
takeMVar MVar ()
lock) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
    SuperBuffer -> ByteString -> IO ()
appendBuffer SuperBuffer
buf ByteString
bs
{-# INLINE appendBufferT #-}

destroyBuffer :: SuperBuffer -> IO ()
destroyBuffer :: SuperBuffer -> IO ()
destroyBuffer (SuperBuffer (SuperBufferP
ptr, MVar ()
_)) = SuperBufferP -> IO ()
destroy_sbuf SuperBufferP
ptr
{-# INLINE destroyBuffer #-}

destroyBufferContents :: SuperBuffer -> IO ()
destroyBufferContents :: SuperBuffer -> IO ()
destroyBufferContents (SuperBuffer (SuperBufferP
ptr, MVar ()
_)) = SuperBufferP -> IO ()
destroyContents_sbuf SuperBufferP
ptr
{-# INLINE destroyBufferContents #-}

-- | Read the final buffer contents. This must only
-- be called once
readBuffer :: SuperBuffer -> IO BS.ByteString
readBuffer :: SuperBuffer -> IO ByteString
readBuffer (SuperBuffer (SuperBufferP
ptr, MVar ()
_)) =
    do (Ptr CChar
cstr, CSize
sz) <- IO (Ptr CChar, CSize)
readLocal
       Ptr Word8 -> Int -> IO () -> IO ByteString
BS.unsafePackCStringFinalizer (Ptr CChar -> Ptr Word8
coerce Ptr CChar
cstr) (CSize -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral CSize
sz) (Ptr CChar -> IO ()
forall a. Ptr a -> IO ()
free Ptr CChar
cstr)
    where
      readLocal :: IO (Ptr CChar, CSize)
readLocal =
          (Ptr CSize -> IO (Ptr CChar, CSize)) -> IO (Ptr CChar, CSize)
forall a b. Storable a => (Ptr a -> IO b) -> IO b
alloca ((Ptr CSize -> IO (Ptr CChar, CSize)) -> IO (Ptr CChar, CSize))
-> (Ptr CSize -> IO (Ptr CChar, CSize)) -> IO (Ptr CChar, CSize)
forall a b. (a -> b) -> a -> b
$ \Ptr CSize
sizePtr ->
          do Ptr CChar
cstr <- SuperBufferP -> Ptr CSize -> IO (Ptr CChar)
read_sbuf SuperBufferP
ptr Ptr CSize
sizePtr
             CSize
sz <- Ptr CSize -> IO CSize
forall a. Storable a => Ptr a -> IO a
peek Ptr CSize
sizePtr
             (Ptr CChar, CSize) -> IO (Ptr CChar, CSize)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Ptr CChar
cstr, CSize
sz)
{-# INLINE readBuffer #-}

-- | Get current (filled) size of the buffer
size :: SuperBuffer -> IO Int
size :: SuperBuffer -> IO Int
size (SuperBuffer (SuperBufferP
ptr, MVar ()
_)) =
    CSize -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (CSize -> Int) -> IO CSize -> IO Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SuperBufferP -> IO CSize
size_sbuf SuperBufferP
ptr
{-# INLINE size #-}

data SBuf
type SuperBufferP = Ptr SBuf

foreign import ccall unsafe "new_sbuf" new_sbuf :: CSize -> IO SuperBufferP
foreign import ccall unsafe "append_sbuf" append_sbuf :: SuperBufferP -> CString -> CSize -> IO ()
foreign import ccall unsafe "read_sbuf" read_sbuf :: SuperBufferP -> Ptr CSize -> IO CString
foreign import ccall unsafe "destroy_sbuf" destroy_sbuf :: SuperBufferP -> IO ()
foreign import ccall unsafe "destroyContents_sbuf" destroyContents_sbuf :: SuperBufferP -> IO ()
foreign import ccall unsafe "size_sbuf" size_sbuf :: SuperBufferP -> IO CSize