module ByteString.TreeBuilder
(
  Builder,
  -- * Declaration
  -- ** Primitives
  byteString,
  byte,
  -- ** Extras
  asciiIntegral,
  asciiChar,
  utf8Char,
  utf8Ord,
  utf8Text,
  utf8LazyText,
  intercalate,
  -- * Execution
  length,
  toByteString,
  toLazyByteString,
)
where

import ByteString.TreeBuilder.Prelude hiding (foldl, foldr, length, intercalate)
import qualified ByteString.TreeBuilder.Tree as A
import qualified ByteString.TreeBuilder.Poker as D
import qualified ByteString.TreeBuilder.Prelude as F
import qualified Data.ByteString as B
import qualified Data.ByteString.Internal as C
import qualified Data.ByteString.Lazy.Internal as E
import qualified Data.Text
import qualified Data.Text.Lazy


-- |
-- A binary-tree-based datastructure optimized for aggregation of bytestrings
-- using the /O(1)/ appending operation.
data Builder =
  Builder !Int !A.Tree

-- |
-- Implements `mappend` with /O(1)/ complexity.
instance Monoid Builder where
  {-# INLINE mempty #-}
  mempty :: Builder
mempty =
    Int -> Tree -> Builder
Builder Int
0 Tree
A.Empty
  {-# INLINABLE mappend #-}
  mappend :: Builder -> Builder -> Builder
mappend (Builder Int
length1 Tree
tree1) (Builder Int
length2 Tree
tree2) =
    Int -> Tree -> Builder
Builder (Int
length1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
length2) (Tree -> Tree -> Tree
A.Branch Tree
tree1 Tree
tree2)
  {-# INLINE mconcat #-}
  mconcat :: [Builder] -> Builder
mconcat =
    (Builder -> Builder -> Builder) -> Builder -> [Builder] -> Builder
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Builder -> Builder -> Builder
forall a. Monoid a => a -> a -> a
mappend Builder
forall a. Monoid a => a
mempty

instance Semigroup Builder where
  {-# INLINE sconcat #-}
  sconcat :: NonEmpty Builder -> Builder
sconcat =
    (Builder -> Builder -> Builder)
-> Builder -> NonEmpty Builder -> Builder
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Builder -> Builder -> Builder
forall a. Monoid a => a -> a -> a
mappend Builder
forall a. Monoid a => a
mempty

  {-# INLINABLE (<>) #-}
  <> :: Builder -> Builder -> Builder
(<>) = Builder -> Builder -> Builder
forall a. Monoid a => a -> a -> a
mappend

instance IsString Builder where
  {-# INLINE fromString #-}
  fromString :: String -> Builder
fromString String
string =
    Int -> Tree -> Builder
Builder (ByteString -> Int
B.length ByteString
bytes) (ByteString -> Tree
A.Leaf ByteString
bytes)
    where
      bytes :: ByteString
bytes =
        String -> ByteString
forall a. IsString a => String -> a
fromString String
string


-- * Declaration
-------------------------

-- ** Primitives
-------------------------

-- |
-- Lifts a bytestring into the builder.
{-# INLINE byteString #-}
byteString :: ByteString -> Builder
byteString :: ByteString -> Builder
byteString ByteString
bytes =
  Int -> Tree -> Builder
Builder (ByteString -> Int
B.length ByteString
bytes) (ByteString -> Tree
A.Leaf ByteString
bytes)

-- |
-- Lifts a single byte into the builder.
{-# INLINE byte #-}
byte :: Word8 -> Builder
byte :: Word8 -> Builder
byte Word8
byte =
  Int -> Tree -> Builder
Builder Int
1 (ByteString -> Tree
A.Leaf (Word8 -> ByteString
B.singleton Word8
byte))

-- ** Extras
-------------------------

{-# INLINABLE asciiIntegral #-}
asciiIntegral :: Integral a => a -> Builder
asciiIntegral :: a -> Builder
asciiIntegral =
  \case
    a
0 ->
      Word8 -> Builder
byte Word8
48
    a
x ->
      (Builder -> Builder)
-> (Builder -> Builder) -> Bool -> Builder -> Builder
forall a. a -> a -> Bool -> a
bool (Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
(<>) (Word8 -> Builder
byte Word8
45)) Builder -> Builder
forall k (cat :: k -> k -> *) (a :: k). Category cat => cat a a
id (a
x a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>= a
0) (Builder -> Builder) -> Builder -> Builder
forall a b. (a -> b) -> a -> b
$
      Builder -> a -> Builder
forall t. Integral t => Builder -> t -> Builder
loop Builder
forall a. Monoid a => a
mempty (a -> Builder) -> a -> Builder
forall a b. (a -> b) -> a -> b
$
      a -> a
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 t -> t -> (t, t)
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 Word8 -> Word8 -> Word8
forall a. Num a => a -> a -> a
+ t -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral t
rem) Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Builder
builder) t
quot

{-# INLINE asciiChar #-}
asciiChar :: Char -> Builder
asciiChar :: Char -> Builder
asciiChar =
  Word8 -> Builder
byte (Word8 -> Builder) -> (Char -> Word8) -> Char -> Builder
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Word8) -> (Char -> Int) -> Char -> Word8
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 utf8Char #-}
utf8Char :: Char -> Builder
utf8Char :: Char -> Builder
utf8Char =
  Int -> Builder
utf8Ord (Int -> Builder) -> (Char -> Int) -> Char -> Builder
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 utf8Ord #-}
utf8Ord :: Int -> Builder
utf8Ord :: Int -> Builder
utf8Ord Int
x =
  if Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0x7F
    then
      Word8 -> Builder
byte (Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
x)
    else 
      if Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0x07FF
        then
          Word8 -> Builder
byte (Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral ((Int
x Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
6) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
0xC0)) Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<>
          Word8 -> Builder
byte (Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral ((Int
x Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
0x3F) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
0x80))
        else
          if Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0xFFFF
            then
              Word8 -> Builder
byte (Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
x Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
12) Word8 -> Word8 -> Word8
forall a. Num a => a -> a -> a
+ Word8
0xE0) Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<>
              Word8 -> Builder
byte (Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral ((Int
x Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
6) Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
0x3F) Word8 -> Word8 -> Word8
forall a. Num a => a -> a -> a
+ Word8
0x80) Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<>
              Word8 -> Builder
byte (Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
x Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
0x3F) Word8 -> Word8 -> Word8
forall a. Num a => a -> a -> a
+ Word8
0x80)
            else
              Word8 -> Builder
byte (Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
x Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
18) Word8 -> Word8 -> Word8
forall a. Num a => a -> a -> a
+ Word8
0xF0) Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<>
              Word8 -> Builder
byte (Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral ((Int
x Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
12) Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
0x3F) Word8 -> Word8 -> Word8
forall a. Num a => a -> a -> a
+ Word8
0x80) Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<>
              Word8 -> Builder
byte (Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral ((Int
x Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
6) Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
0x3F) Word8 -> Word8 -> Word8
forall a. Num a => a -> a -> a
+ Word8
0x80) Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<>
              Word8 -> Builder
byte (Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
x Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
0x3F) Word8 -> Word8 -> Word8
forall a. Num a => a -> a -> a
+ Word8
0x80)

{-# INLINE utf8Text #-}
utf8Text :: Data.Text.Text -> Builder
utf8Text :: Text -> Builder
utf8Text =
  (Builder -> Char -> Builder) -> Builder -> Text -> Builder
forall a. (a -> Char -> a) -> a -> Text -> a
Data.Text.foldl' (\Builder
builder -> Builder -> Builder -> Builder
forall a. Monoid a => a -> a -> a
mappend Builder
builder (Builder -> Builder) -> (Char -> Builder) -> Char -> Builder
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Char -> Builder
utf8Char) Builder
forall a. Monoid a => a
mempty

{-# INLINE utf8LazyText #-}
utf8LazyText :: Data.Text.Lazy.Text -> Builder
utf8LazyText :: Text -> Builder
utf8LazyText =
  (Builder -> Char -> Builder) -> Builder -> Text -> Builder
forall a. (a -> Char -> a) -> a -> Text -> a
Data.Text.Lazy.foldl' (\Builder
builder -> Builder -> Builder -> Builder
forall a. Monoid a => a -> a -> a
mappend Builder
builder (Builder -> Builder) -> (Char -> Builder) -> Char -> Builder
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Char -> Builder
utf8Char) Builder
forall a. Monoid a => a
mempty

{-# INLINABLE intercalate #-}
intercalate :: (Foldable f, Monoid m) => m -> f m -> m
intercalate :: m -> f m -> m
intercalate m
incut =
  (m, m -> m) -> m
forall a b. (a, b) -> a
fst ((m, m -> m) -> m) -> (f m -> (m, m -> m)) -> f m -> m
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
.
  ((m, m -> m) -> m -> (m, m -> m))
-> (m, m -> m) -> f m -> (m, m -> m)
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\(m
acc, m -> m
incutFn) m
x -> (m -> m
incutFn (m -> m -> m
forall a. Monoid a => a -> a -> a
mappend m
x m
acc), m -> m -> m
forall a. Monoid a => a -> a -> a
mappend m
incut)) (m
forall a. Monoid a => a
mempty, m -> m
forall k (cat :: k -> k -> *) (a :: k). Category cat => cat a a
id)


-- * Execution
-------------------------

-- |
-- Performs a left-fold over the aggregated chunks.
{-# INLINE foldl #-}
foldl :: (a -> ByteString -> a) -> a -> Builder -> a
foldl :: (a -> ByteString -> a) -> a -> Builder -> a
foldl a -> ByteString -> a
step a
init (Builder Int
length Tree
tree) =
  (a -> ByteString -> a) -> a -> Tree -> a
forall a. (a -> ByteString -> a) -> a -> Tree -> a
A.foldl a -> ByteString -> a
step a
init Tree
tree

-- |
-- Performs a right-fold over the aggregated chunks.
{-# INLINE foldr #-}
foldr :: (ByteString -> a -> a) -> a -> Builder -> a
foldr :: (ByteString -> a -> a) -> a -> Builder -> a
foldr ByteString -> a -> a
step a
init (Builder Int
length Tree
tree) =
  (ByteString -> a -> a) -> a -> Tree -> a
forall a. (ByteString -> a -> a) -> a -> Tree -> a
A.foldr ByteString -> a -> a
step a
init Tree
tree

-- |
-- /O(1)/. Gets the total length.
{-# INLINE length #-}
length :: Builder -> Int
length :: Builder -> Int
length (Builder Int
length Tree
tree) =
  Int
length

-- |
-- /O(n)/. Converts the builder into a strict bytestring.
{-# INLINABLE toByteString #-}
toByteString :: Builder -> ByteString
toByteString :: Builder -> ByteString
toByteString (Builder Int
length Tree
tree) =
  Int -> (Ptr Word8 -> IO ()) -> ByteString
C.unsafeCreate Int
length ((Ptr Word8 -> IO ()) -> ByteString)
-> (Ptr Word8 -> IO ()) -> ByteString
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
ptr -> 
    IO (Ptr Word8) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Ptr Word8) -> IO ()) -> IO (Ptr Word8) -> IO ()
forall a b. (a -> b) -> a -> b
$ Tree -> Ptr Word8 -> IO (Ptr Word8)
D.pokeTree Tree
tree Ptr Word8
ptr

-- |
-- /O(n)/. Converts the builder into a lazy bytestring.
{-# INLINABLE toLazyByteString #-}
toLazyByteString :: Builder -> E.ByteString
toLazyByteString :: Builder -> ByteString
toLazyByteString =
  (ByteString -> ByteString -> ByteString)
-> ByteString -> Builder -> ByteString
forall a. (ByteString -> a -> a) -> a -> Builder -> a
foldr ByteString -> ByteString -> ByteString
E.Chunk ByteString
E.Empty