-- |
-- Module      :  Pinch.Internal.Builder
-- Copyright   :  (c) Abhinav Gupta 2015
-- License     :  BSD3
--
-- Maintainer  :  Abhinav Gupta <mail@abhinavg.net>
-- Stability   :  experimental
--
-- This module implements a ByteString builder very similar to
-- 'Data.ByteString.Builder' except that it keeps track of its final serialized
-- length. This allows it to allocate the target ByteString in one @malloc@ and
-- simply write the bytes to it.
module Pinch.Internal.Builder
    ( Builder
    , runBuilder
    , unsafeRunBuilder

    , append
    , int8
    , word8
    , int16BE
    , int32BE
    , int64BE
    , int64LE
    , doubleBE
    , doubleLE
    , byteString

    , getSize
    ) where

import Data.ByteString              (ByteString)
import Data.ByteString.Builder.Prim ((>*<))
import Data.Int
import Data.Semigroup
import Data.Word                    (Word8)
import Foreign.ForeignPtr           (withForeignPtr)
import Foreign.Ptr                  (Ptr, plusPtr)

import qualified Data.ByteString.Builder.Prim          as BP
import qualified Data.ByteString.Builder.Prim.Internal as BPI
import qualified Data.ByteString.Internal              as BI

-- | A ByteString Builder that knows its final size.
data Builder = B {-# UNPACK #-} !Int (Ptr Word8 -> IO ())

-- | Build a ByteString from the given ByteString builder.
runBuilder :: Builder -> ByteString
runBuilder :: Builder -> ByteString
runBuilder (B Int
size Ptr Word8 -> IO ()
fill) = Int -> (Ptr Word8 -> IO ()) -> ByteString
BI.unsafeCreate Int
size Ptr Word8 -> IO ()
fill
{-# INLINE runBuilder #-}

-- | Fill a buffer with the bytes of a builder. The buffer is assumed
-- to be big enough to hold the contents of the entire builder.
unsafeRunBuilder :: Builder -> Ptr Word8 -> IO ()
unsafeRunBuilder :: Builder -> Ptr Word8 -> IO ()
unsafeRunBuilder (B Int
_size Ptr Word8 -> IO ()
fill) = Ptr Word8 -> IO ()
fill
{-# INLINE unsafeRunBuilder #-}

-- | Append two Builders into one.
append :: Builder -> Builder -> Builder
append :: Builder -> Builder -> Builder
append (B Int
ll Ptr Word8 -> IO ()
lf) (B Int
rl Ptr Word8 -> IO ()
rf) = Int -> (Ptr Word8 -> IO ()) -> Builder
B (Int
ll Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
rl) (\Ptr Word8
p -> Ptr Word8 -> IO ()
lf Ptr Word8
p IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Ptr Word8 -> IO ()
rf (Ptr Word8
p Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
ll))
{-# INLINE [1] append #-}
    -- Don't inline append until phase 1. This ensures that the
    -- append/primFixed* rules have a chance to fire.

instance Semigroup Builder where
    <> :: Builder -> Builder -> Builder
(<>) = Builder -> Builder -> Builder
append
    sconcat :: NonEmpty Builder -> Builder
sconcat = (Builder -> Builder -> Builder)
-> Builder -> NonEmpty Builder -> Builder
forall a b. (a -> b -> b) -> b -> NonEmpty a -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
(<>) Builder
forall a. Monoid a => a
mempty

instance Monoid Builder where
    {-# INLINE mempty #-}
    mempty :: Builder
mempty = Int -> (Ptr Word8 -> IO ()) -> Builder
B Int
0 (\Ptr Word8
_ -> () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ())

primFixed :: BP.FixedPrim a -> a -> Builder
primFixed :: forall a. FixedPrim a -> a -> Builder
primFixed FixedPrim a
prim a
a = Int -> (Ptr Word8 -> IO ()) -> Builder
B (FixedPrim a -> Int
forall a. FixedPrim a -> Int
BPI.size FixedPrim a
prim) (FixedPrim a -> a -> Ptr Word8 -> IO ()
forall a. FixedPrim a -> a -> Ptr Word8 -> IO ()
BPI.runF FixedPrim a
prim a
a)
{-# INLINE [1] primFixed #-}
    -- Don't inline append until phase 1. This ensures that the
    -- append/primFixed* rules have a chance to fire.

-- The following rules try to join together instances of primFixed that are
-- being appended together. These were adapted almost as-is from
-- ByteString.Builder.Prim's rules around this.

{-# RULES

"append/primFixed" forall p1 p2 v1 v2.
    append (primFixed p1 v1) (primFixed p2 v2) =
        primFixed (p1 >*< p2) (v1, v2)

"append/primFixed/rightAssociative" forall p1 p2 v1 v2 b.
    append (primFixed p1 v1) (append (primFixed p2 v2) b) =
        append (primFixed (p1 >*< p2) (v1, v2)) b

"append/primFixed/leftAssociative" forall p1 p2 v1 v2 b.
    append (append b (primFixed p1 v1)) (primFixed p2 v2) =
        append b (primFixed (p1 >*< p2) (v1, v2))

  #-}

-- | Serialize a single signed byte.
int8 :: Int8 -> Builder
int8 :: Int8 -> Builder
int8 = FixedPrim Int8 -> Int8 -> Builder
forall a. FixedPrim a -> a -> Builder
primFixed FixedPrim Int8
BP.int8
{-# INLINE int8 #-}

-- | Serialize a single unsigned byte.
word8 :: Word8 -> Builder
word8 :: Word8 -> Builder
word8 = FixedPrim Word8 -> Word8 -> Builder
forall a. FixedPrim a -> a -> Builder
primFixed FixedPrim Word8
BP.word8
{-# INLINE word8 #-}

-- | Serialize a signed 16-bit integer in big endian format.
int16BE :: Int16 -> Builder
int16BE :: Int16 -> Builder
int16BE = FixedPrim Int16 -> Int16 -> Builder
forall a. FixedPrim a -> a -> Builder
primFixed FixedPrim Int16
BP.int16BE
{-# INLINE int16BE #-}

-- | Serialize a signed 32-bit integer in big endian format.
int32BE :: Int32 -> Builder
int32BE :: Int32 -> Builder
int32BE = FixedPrim Int32 -> Int32 -> Builder
forall a. FixedPrim a -> a -> Builder
primFixed FixedPrim Int32
BP.int32BE
{-# INLINE int32BE #-}

-- | Serialize a signed 64-bit integer in big endian format.
int64BE :: Int64 -> Builder
int64BE :: Int64 -> Builder
int64BE = FixedPrim Int64 -> Int64 -> Builder
forall a. FixedPrim a -> a -> Builder
primFixed FixedPrim Int64
BP.int64BE
{-# INLINE int64BE #-}

-- | Serialize a signed 64-bit integer in little endian format.
int64LE :: Int64 -> Builder
int64LE :: Int64 -> Builder
int64LE = FixedPrim Int64 -> Int64 -> Builder
forall a. FixedPrim a -> a -> Builder
primFixed FixedPrim Int64
BP.int64LE
{-# INLINE int64LE #-}

-- | Serialize a signed 64-bit floating point number in big endian format.
doubleBE :: Double -> Builder
doubleBE :: Double -> Builder
doubleBE = FixedPrim Double -> Double -> Builder
forall a. FixedPrim a -> a -> Builder
primFixed FixedPrim Double
BP.doubleBE
{-# INLINE doubleBE #-}

-- | Serialize a signed 64-bit floating point number in little endian format.
doubleLE :: Double -> Builder
doubleLE :: Double -> Builder
doubleLE = FixedPrim Double -> Double -> Builder
forall a. FixedPrim a -> a -> Builder
primFixed FixedPrim Double
BP.doubleLE
{-# INLINE doubleLE #-}


-- | Inlcude the given ByteString as-is in the builder.
--
-- Note that because this operation is applied lazily, we will maintain a
-- reference to the ByteString until the builder is executed.
byteString :: ByteString -> Builder
byteString :: ByteString -> Builder
byteString (BI.PS ForeignPtr Word8
fp Int
off Int
len) =
    Int -> (Ptr Word8 -> IO ()) -> Builder
B Int
len ((Ptr Word8 -> IO ()) -> Builder)
-> (Ptr Word8 -> IO ()) -> Builder
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
dst ->
    ForeignPtr Word8 -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr Word8
fp ((Ptr Word8 -> IO ()) -> IO ()) -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
src ->
        Ptr Word8 -> Ptr Word8 -> Int -> IO ()
BI.memcpy Ptr Word8
dst (Ptr Word8
src Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
off) Int
len
{-# INLINE byteString #-}

-- | Returns the number of bytes in the builder.
getSize :: Builder -> Int
getSize :: Builder -> Int
getSize (B Int
sz Ptr Word8 -> IO ()
_) = Int
sz