module Hpgsql.Builder where

-- \| This module replicates parts of the API of Data.ByteString.Builder but its own
-- builder is length-aware, which makes other parts of the code a little bit nicer.
-- In COPY benchmarks, this module was introduced in a commit (together with other
-- changes, like replacing `Maybe` with `BinaryField` in `ToPgField`) that barely
-- changed memory usage and runtime.
-- The benefits are exclusively for code readability, then.
-- \|

import Data.ByteString (ByteString)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Builder as Builder
import qualified Data.ByteString.Builder.Extra as Builder
import qualified Data.ByteString.Lazy as LBS
import Data.Int (Int16, Int32, Int64)
import Data.Word (Word8)

data BinaryField = SqlNull | NotNull !ByteString
  deriving stock (BinaryField -> BinaryField -> Bool
(BinaryField -> BinaryField -> Bool)
-> (BinaryField -> BinaryField -> Bool) -> Eq BinaryField
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: BinaryField -> BinaryField -> Bool
== :: BinaryField -> BinaryField -> Bool
$c/= :: BinaryField -> BinaryField -> Bool
/= :: BinaryField -> BinaryField -> Bool
Eq)

instance Show BinaryField where
  show :: BinaryField -> String
show BinaryField
SqlNull = String
"NULL"
  show (NotNull ByteString
bs) = ByteString -> String
forall a. Show a => a -> String
show ByteString
bs

data LengthAwareBuilder = LengthAwareBuilder !Int32 !Builder.Builder

type Builder = LengthAwareBuilder

instance Semigroup LengthAwareBuilder where
  LengthAwareBuilder Int32
l1 Builder
b1 <> :: LengthAwareBuilder -> LengthAwareBuilder -> LengthAwareBuilder
<> LengthAwareBuilder Int32
l2 Builder
b2 = Int32 -> Builder -> LengthAwareBuilder
LengthAwareBuilder (Int32
l1 Int32 -> Int32 -> Int32
forall a. Num a => a -> a -> a
+ Int32
l2) (Builder
b1 Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Builder
b2)

instance Monoid LengthAwareBuilder where
  mempty :: LengthAwareBuilder
mempty = Int32 -> Builder -> LengthAwareBuilder
LengthAwareBuilder Int32
0 Builder
forall a. Monoid a => a
mempty

builderLength :: LengthAwareBuilder -> Int32
builderLength :: LengthAwareBuilder -> Int32
builderLength (LengthAwareBuilder Int32
n Builder
_) = Int32
n

binaryField :: BinaryField -> LengthAwareBuilder
binaryField :: BinaryField -> LengthAwareBuilder
binaryField = \case
  BinaryField
SqlNull -> Int32 -> Builder -> LengthAwareBuilder
LengthAwareBuilder Int32
4 (Int32 -> Builder
Builder.int32BE (-Int32
1))
  NotNull ByteString
val -> let valLen :: Int32
valLen = Int -> Int32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (ByteString -> Int
BS.length ByteString
val) in Int32 -> Builder -> LengthAwareBuilder
LengthAwareBuilder (Int32
4 Int32 -> Int32 -> Int32
forall a. Num a => a -> a -> a
+ Int32
valLen) (Int32 -> Builder
Builder.int32BE Int32
valLen Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> ByteString -> Builder
Builder.byteString ByteString
val)

toLazyByteString :: LengthAwareBuilder -> LBS.ByteString
toLazyByteString :: LengthAwareBuilder -> LazyByteString
toLazyByteString (LengthAwareBuilder (Int32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral -> Int
totalLen) Builder
builder) = AllocationStrategy -> LazyByteString -> Builder -> LazyByteString
Builder.toLazyByteStringWith (Int -> Int -> AllocationStrategy
Builder.untrimmedStrategy Int
totalLen Int
0) LazyByteString
forall a. Monoid a => a
mempty Builder
builder

toStrictByteString :: LengthAwareBuilder -> ByteString
toStrictByteString :: LengthAwareBuilder -> ByteString
toStrictByteString = LazyByteString -> ByteString
LBS.toStrict (LazyByteString -> ByteString)
-> (LengthAwareBuilder -> LazyByteString)
-> LengthAwareBuilder
-> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. LengthAwareBuilder -> LazyByteString
toLazyByteString

char7 :: Char -> LengthAwareBuilder
char7 :: Char -> LengthAwareBuilder
char7 Char
c = Int32 -> Builder -> LengthAwareBuilder
LengthAwareBuilder Int32
1 (Char -> Builder
Builder.char7 Char
c)

string7 :: String -> LengthAwareBuilder
string7 :: String -> LengthAwareBuilder
string7 String
s = Int32 -> Builder -> LengthAwareBuilder
LengthAwareBuilder (Int -> Int32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Int32) -> Int -> Int32
forall a b. (a -> b) -> a -> b
$ String -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length String
s) (String -> Builder
Builder.string7 String
s)

word8 :: Word8 -> LengthAwareBuilder
word8 :: Word8 -> LengthAwareBuilder
word8 Word8
s = Int32 -> Builder -> LengthAwareBuilder
LengthAwareBuilder Int32
1 (Word8 -> Builder
Builder.word8 Word8
s)

int16BE :: Int16 -> LengthAwareBuilder
int16BE :: Int16 -> LengthAwareBuilder
int16BE Int16
n = Int32 -> Builder -> LengthAwareBuilder
LengthAwareBuilder Int32
2 (Int16 -> Builder
Builder.int16BE Int16
n)

int32BE :: Int32 -> LengthAwareBuilder
int32BE :: Int32 -> LengthAwareBuilder
int32BE Int32
n = Int32 -> Builder -> LengthAwareBuilder
LengthAwareBuilder Int32
4 (Int32 -> Builder
Builder.int32BE Int32
n)

int64BE :: Int64 -> LengthAwareBuilder
int64BE :: Int64 -> LengthAwareBuilder
int64BE Int64
n = Int32 -> Builder -> LengthAwareBuilder
LengthAwareBuilder Int32
8 (Int64 -> Builder
Builder.int64BE Int64
n)

byteString :: ByteString -> LengthAwareBuilder
byteString :: ByteString -> LengthAwareBuilder
byteString ByteString
bs = Int32 -> Builder -> LengthAwareBuilder
LengthAwareBuilder (Int -> Int32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Int32) -> Int -> Int32
forall a b. (a -> b) -> a -> b
$ ByteString -> Int
BS.length ByteString
bs) (ByteString -> Builder
Builder.byteString ByteString
bs)

lazyByteString :: LBS.ByteString -> LengthAwareBuilder
lazyByteString :: LazyByteString -> LengthAwareBuilder
lazyByteString LazyByteString
bs = Int32 -> Builder -> LengthAwareBuilder
LengthAwareBuilder (Int64 -> Int32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int64 -> Int32) -> Int64 -> Int32
forall a b. (a -> b) -> a -> b
$ LazyByteString -> Int64
LBS.length LazyByteString
bs) (LazyByteString -> Builder
Builder.lazyByteString LazyByteString
bs)