module ByteString.StrictBuilder
  ( Builder,
    builderBytes,
    builderChunksBuilder,
    builderLength,
    builderPtrFiller,
    bytes,
    lazyBytes,
    asciiIntegral,
    asciiChar,
    utf8Char,
    storable,
    word8,
    word16BE,
    word32BE,
    word64BE,
    int8,
    int16BE,
    int32BE,
    int64BE,
  )
where

import qualified ByteString.StrictBuilder.Population as A
import ByteString.StrictBuilder.Prelude
import qualified ByteString.StrictBuilder.UTF8 as E
import qualified Data.ByteString as B
import qualified Data.ByteString.Builder.Internal as G
import qualified Data.ByteString.Internal as C
import qualified Data.ByteString.Lazy as F

-- * Rewrite rules

{-# RULES "builderBytes/bytes" forall a. builderBytes (bytes a) = a #-}

-- * Builder

data Builder
  = Builder !Int !A.Population

instance Semigroup Builder where
  <> :: Builder -> Builder -> Builder
(<>) (Builder Int
leftSize Population
leftPopulation) (Builder Int
rightSize Population
rightPopulation) =
    Int -> Population -> Builder
Builder (Int
leftSize forall a. Num a => a -> a -> a
+ Int
rightSize) (Population
leftPopulation forall a. Semigroup a => a -> a -> a
<> Population
rightPopulation)
  {-# INLINE sconcat #-}
  sconcat :: NonEmpty Builder -> Builder
sconcat NonEmpty Builder
builders =
    Int -> Population -> Builder
Builder Int
size Population
population
    where
      size :: Int
size =
        forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Int
acc (Builder Int
x Population
_) -> Int
acc forall a. Num a => a -> a -> a
+ Int
x) Int
0 NonEmpty Builder
builders
      population :: Population
population =
        forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (\(Builder Int
_ Population
x) -> Population
x) NonEmpty Builder
builders

instance Monoid Builder where
  {-# INLINE mempty #-}
  mempty :: Builder
mempty =
    Int -> Population -> Builder
Builder Int
0 forall a. Monoid a => a
mempty
  {-# INLINE mconcat #-}
  mconcat :: [Builder] -> Builder
mconcat [Builder]
builders =
    Int -> Population -> Builder
Builder Int
size Population
population
    where
      size :: Int
size =
        forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Int
acc (Builder Int
x Population
_) -> Int
acc forall a. Num a => a -> a -> a
+ Int
x) Int
0 [Builder]
builders
      population :: Population
population =
        forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (\(Builder Int
_ Population
x) -> Population
x) [Builder]
builders

instance IsString Builder where
  fromString :: String -> Builder
fromString =
    ByteString -> Builder
bytes forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall a. IsString a => String -> a
fromString

instance Show Builder where
  show :: Builder -> String
show =
    forall a. Show a => a -> String
show forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Builder -> ByteString
builderBytes

-------------------------

-- |
-- Efficiently constructs a strict bytestring.
{-# NOINLINE builderBytes #-}
builderBytes :: Builder -> ByteString
builderBytes :: Builder -> ByteString
builderBytes (Builder Int
size Population
population) =
  Int -> (Ptr Word8 -> IO ()) -> ByteString
C.unsafeCreate Int
size forall a b. (a -> b) -> a -> b
$ \Ptr Word8
ptr -> Population -> Ptr Word8 -> IO (Ptr Word8)
A.populationPtrUpdate Population
population Ptr Word8
ptr forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> ()

-- |
-- Converts into the standard lazy bytestring builder.
-- Does so efficiently using the internal APIs of \"bytestring\",
-- without producing any intermediate representation.
{-# INLINE builderChunksBuilder #-}
builderChunksBuilder :: Builder -> G.Builder
builderChunksBuilder :: Builder -> Builder
builderChunksBuilder (Builder Int
size Population
population) =
  Int -> Builder
G.ensureFree Int
size forall a. Semigroup a => a -> a -> a
<> Population -> Builder
A.populationChunksBuilder Population
population

-- |
-- /O(1)/. Gets the size of the bytestring that is to be produced.
{-# INLINE builderLength #-}
builderLength :: Builder -> Int
builderLength :: Builder -> Int
builderLength (Builder Int
size Population
population) =
  Int
size

-- |
-- Use the builder to populate a buffer.
-- It is your responsibility to ensure that the bounds are not exceeded.
{-# INLINE builderPtrFiller #-}
builderPtrFiller ::
  Builder ->
  -- | A continuation on the amount of bytes to be written and the action populating the pointer.
  (Int -> (Ptr Word8 -> IO ()) -> result) ->
  result
builderPtrFiller :: forall result.
Builder -> (Int -> (Ptr Word8 -> IO ()) -> result) -> result
builderPtrFiller (Builder Int
size (A.Population Ptr Word8 -> IO (Ptr Word8)
ptrUpdate)) Int -> (Ptr Word8 -> IO ()) -> result
cont =
  Int -> (Ptr Word8 -> IO ()) -> result
cont Int
size (forall (f :: * -> *) a. Functor f => f a -> f ()
void forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Ptr Word8 -> IO (Ptr Word8)
ptrUpdate)

-- * Primitives

-------------------------

{-# NOINLINE bytes #-}
bytes :: ByteString -> Builder
bytes :: ByteString -> Builder
bytes ByteString
bytes =
  Int -> Population -> Builder
Builder (ByteString -> Int
B.length ByteString
bytes) (ByteString -> Population
A.bytes ByteString
bytes)

{-# INLINE lazyBytes #-}
lazyBytes :: F.ByteString -> Builder
lazyBytes :: ByteString -> Builder
lazyBytes =
  forall a. (a -> ByteString -> a) -> a -> ByteString -> a
F.foldlChunks (\Builder
builder -> forall a. Monoid a => a -> a -> a
mappend Builder
builder forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. ByteString -> Builder
bytes) forall a. Monoid a => a
mempty

{-# INLINE byte #-}
byte :: Word8 -> Builder
byte :: Word8 -> Builder
byte =
  Word8 -> Builder
word8

-- * Extras

-------------------------

{-# INLINEABLE asciiIntegral #-}
asciiIntegral :: (Integral a) => a -> Builder
asciiIntegral :: forall a. Integral a => a -> Builder
asciiIntegral =
  \case
    a
0 ->
      Word8 -> Builder
byte Word8
48
    a
x ->
      forall a. a -> a -> Bool -> a
bool (forall a. Semigroup a => a -> a -> a
(<>) (Word8 -> Builder
byte Word8
45)) forall {k} (cat :: k -> k -> *) (a :: k). Category cat => cat a a
id (a
x forall a. Ord a => a -> a -> Bool
>= a
0)
        forall a b. (a -> b) -> a -> b
$ forall {t}. Integral t => Builder -> t -> Builder
loop forall a. Monoid a => a
mempty
        forall a b. (a -> b) -> a -> b
$ forall a. Num a => a -> a
abs a
x
  where
    loop :: Builder -> t -> Builder
loop Builder
builder t
remainder =
      case t
remainder of
        t
0 ->
          Builder
builder
        t
_ ->
          case forall a. Integral a => a -> a -> (a, a)
quotRem t
remainder t
10 of
            (t
quot, t
rem) ->
              Builder -> t -> Builder
loop (Word8 -> Builder
byte (Word8
48 forall a. Num a => a -> a -> a
+ forall a b. (Integral a, Num b) => a -> b
fromIntegral t
rem) forall a. Semigroup a => a -> a -> a
<> Builder
builder) t
quot

{-# INLINE asciiChar #-}
asciiChar :: Char -> Builder
asciiChar :: Char -> Builder
asciiChar =
  Word8 -> Builder
byte forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall a b. (Integral a, Num b) => a -> b
fromIntegral forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Char -> Int
ord

{-# INLINE CONLIKE storable #-}
storable :: (Storable a) => a -> Builder
storable :: forall a. Storable a => a -> Builder
storable a
value =
  Int -> Population -> Builder
Builder Int
size (forall a. Storable a => Int -> a -> Population
A.storable Int
size a
value)
  where
    size :: Int
size =
      forall a. Storable a => a -> Int
sizeOf a
value

{-# INLINE word8 #-}
word8 :: Word8 -> Builder
word8 :: Word8 -> Builder
word8 =
  Int -> Population -> Builder
Builder Int
1 forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Word8 -> Population
A.word8

{-# INLINE word16BE #-}
word16BE :: Word16 -> Builder
word16BE :: Word16 -> Builder
word16BE =
  Int -> Population -> Builder
Builder Int
2 forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Word16 -> Population
A.word16BE

{-# INLINE word32BE #-}
word32BE :: Word32 -> Builder
word32BE :: Word32 -> Builder
word32BE =
  Int -> Population -> Builder
Builder Int
4 forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Word32 -> Population
A.word32BE

{-# INLINE word64BE #-}
word64BE :: Word64 -> Builder
word64BE :: Word64 -> Builder
word64BE =
  Int -> Population -> Builder
Builder Int
8 forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Word64 -> Population
A.word64BE

{-# INLINE int8 #-}
int8 :: Int8 -> Builder
int8 :: Int8 -> Builder
int8 =
  Int -> Population -> Builder
Builder Int
1 forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Int8 -> Population
A.int8

{-# INLINE int16BE #-}
int16BE :: Int16 -> Builder
int16BE :: Int16 -> Builder
int16BE =
  Int -> Population -> Builder
Builder Int
2 forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Int16 -> Population
A.int16BE

{-# INLINE int32BE #-}
int32BE :: Int32 -> Builder
int32BE :: Int32 -> Builder
int32BE =
  Int -> Population -> Builder
Builder Int
4 forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Int32 -> Population
A.int32BE

{-# INLINE int64BE #-}
int64BE :: Int64 -> Builder
int64BE :: Int64 -> Builder
int64BE =
  Int -> Population -> Builder
Builder Int
8 forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Int64 -> Population
A.int64BE

{-# INLINE utf8Char #-}
utf8Char :: Char -> Builder
utf8Char :: Char -> Builder
utf8Char Char
x =
  Char -> UTF8Char
E.char Char
x Word8 -> Builder
bytes_1 Word8 -> Word8 -> Builder
bytes_2 Word8 -> Word8 -> Word8 -> Builder
bytes_3 Word8 -> Word8 -> Word8 -> Word8 -> Builder
bytes_4

{-# INLINE bytes_1 #-}
bytes_1 :: Word8 -> Builder
bytes_1 :: Word8 -> Builder
bytes_1 Word8
b1 =
  Int -> Population -> Builder
Builder Int
1 (Word8 -> Population
A.bytes_1 Word8
b1)

{-# INLINE bytes_2 #-}
bytes_2 :: Word8 -> Word8 -> Builder
bytes_2 :: Word8 -> Word8 -> Builder
bytes_2 Word8
b1 Word8
b2 =
  Int -> Population -> Builder
Builder Int
2 (Word8 -> Word8 -> Population
A.bytes_2 Word8
b1 Word8
b2)

{-# INLINE bytes_3 #-}
bytes_3 :: Word8 -> Word8 -> Word8 -> Builder
bytes_3 :: Word8 -> Word8 -> Word8 -> Builder
bytes_3 Word8
b1 Word8
b2 Word8
b3 =
  Int -> Population -> Builder
Builder Int
3 (Word8 -> Word8 -> Word8 -> Population
A.bytes_3 Word8
b1 Word8
b2 Word8
b3)

{-# INLINE bytes_4 #-}
bytes_4 :: Word8 -> Word8 -> Word8 -> Word8 -> Builder
bytes_4 :: Word8 -> Word8 -> Word8 -> Word8 -> Builder
bytes_4 Word8
b1 Word8
b2 Word8
b3 Word8
b4 =
  Int -> Population -> Builder
Builder Int
4 (Word8 -> Word8 -> Word8 -> Word8 -> Population
A.bytes_4 Word8
b1 Word8
b2 Word8
b3 Word8
b4)