{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE UnboxedTuples #-}

module Json.Smile
  ( -- * Encode JSON Document
    encode

    -- * Encode JSON Atoms

    -- ** Integer
  , encodeBigInteger

    -- ** String
  , encodeString
  , encodeAsciiString

    -- ** Key
  , encodeKey
  , encodeAsciiKey
  ) where

import Prelude hiding (Bool (..))

import Control.Monad.ST (ST)
import Control.Monad.ST.Run (runByteArrayST)
import Data.Bits
  ( complement
  , countLeadingZeros
  , testBit
  , unsafeShiftR
  , (.&.)
  , (.|.)
  )
import Data.Bytes.Builder (Builder)
import Data.Int (Int32)
import Data.Primitive
  ( ByteArray
  , MutableByteArray (..)
  , byteArrayFromListN
  , copyMutableByteArray
  , newByteArray
  , readByteArray
  , sizeofByteArray
  , unsafeFreezeByteArray
  , writeByteArray
  )
import Data.Text.Short (ShortText)
import Data.Word (Word32, Word64, Word8)
import Data.Word.Zigzag (toZigzag32, toZigzag64)
import GHC.Exts (RealWorld, State#, Word#)
import GHC.IO (IO (IO))
import GHC.Word (Word (..))
import Json (Member (..), Value (..))
import Numeric.Natural (Natural)
import System.IO.Unsafe (unsafeDupablePerformIO)

import qualified Arithmetic.Nat as Nat
import qualified Data.ByteString.Short as SBS
import qualified Data.Bytes as Bytes
import qualified Data.Bytes.Builder as B
import qualified Data.Bytes.Builder.Bounded as Bounded
import qualified Data.Bytes.Builder.Bounded.Unsafe as Unsafe
import qualified Data.Number.Scientific as Sci
import qualified Data.Text.Short as TS
import qualified GHC.Exts as Exts
import qualified GHC.Num.BigNat as BN
import qualified GHC.Num.Integer as Integer
import qualified Prelude

{- | Encode a Json 'Value' to the Smile binary format.
This encoder does not produce backreferences.
-}
encode :: Value -> Builder
{-# INLINE encode #-}
encode :: Value -> Builder
encode Value
v0 = Char -> Char -> Char -> Char -> Builder
B.ascii4 Char
':' Char
')' Char
'\n' Char
'\x00' Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Value -> Builder
encodeNoHeader Value
v0

-- The "rebuild" trick was adapted from the fast-builder library. It
-- results in a 2x performance gain on the twitter benchmark.
-- This function is marked noinline to ensure that its performance is
-- stable.
encodeNoHeader :: Value -> Builder
{-# NOINLINE encodeNoHeader #-}
encodeNoHeader :: Value -> Builder
encodeNoHeader Value
val = Builder -> Builder
B.rebuild (Builder -> Builder) -> Builder -> Builder
forall a b. (a -> b) -> a -> b
$ case Value
val of
  Object SmallArray Member
obj ->
    Word8 -> Builder
B.word8 Word8
0xFA
      Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> (Member -> Builder) -> SmallArray Member -> Builder
forall m a. Monoid m => (a -> m) -> SmallArray a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (\Member {ShortText
key :: ShortText
key :: Member -> ShortText
key, Value
value :: Value
value :: Member -> Value
value} -> ShortText -> Builder
encodeKey ShortText
key Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Value -> Builder
encodeNoHeader Value
value) SmallArray Member
obj
      Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Word8 -> Builder
B.word8 Word8
0xFB
  Array SmallArray Value
arr -> Word8 -> Builder
B.word8 Word8
0xF8 Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> (Value -> Builder) -> SmallArray Value -> Builder
forall m a. Monoid m => (a -> m) -> SmallArray a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Value -> Builder
encodeNoHeader SmallArray Value
arr Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Word8 -> Builder
B.word8 Word8
0xF9
  String ShortText
str -> ShortText -> Builder
encodeString ShortText
str
  Number Scientific
x
    | Just Int32
i32 <- Scientific -> Maybe Int32
Sci.toInt32 Scientific
x
    , -Int32
16 Int32 -> Int32 -> Bool
forall a. Ord a => a -> a -> Bool
<= Int32
i32 Bool -> Bool -> Bool
&& Int32
i32 Int32 -> Int32 -> Bool
forall a. Ord a => a -> a -> Bool
<= Int32
15
    , Word8
w5 <- forall a b. (Integral a, Num b) => a -> b
fromIntegral @Word32 @Word8 (Int32 -> Word32
toZigzag32 Int32
i32) ->
        Word8 -> Builder
B.word8 (Word8
0xC0 Word8 -> Word8 -> Word8
forall a. Num a => a -> a -> a
+ Word8
w5)
    | Just Int32
i32 <- Scientific -> Maybe Int32
Sci.toInt32 Scientific
x ->
        Nat 11 -> Builder 11 -> Builder
forall (n :: Natural). Nat n -> Builder n -> Builder
B.fromBounded Nat 11
forall (n :: Natural). KnownNat n => Nat n
Nat.constant (Word8 -> Builder 1
Bounded.word8 Word8
0x24 Builder 1 -> Builder 10 -> Builder (1 + 10)
forall (m :: Natural) (n :: Natural).
Builder m -> Builder n -> Builder (m + n)
`Bounded.append` Word64 -> Builder 10
vlqSmile64 (forall a b. (Integral a, Num b) => a -> b
fromIntegral @Word32 @Word64 (Int32 -> Word32
toZigzag32 Int32
i32)))
    | Just Int64
i64 <- Scientific -> Maybe Int64
Sci.toInt64 Scientific
x ->
        Nat 11 -> Builder 11 -> Builder
forall (n :: Natural). Nat n -> Builder n -> Builder
B.fromBounded Nat 11
forall (n :: Natural). KnownNat n => Nat n
Nat.constant (Word8 -> Builder 1
Bounded.word8 Word8
0x25 Builder 1 -> Builder 10 -> Builder (1 + 10)
forall (m :: Natural) (n :: Natural).
Builder m -> Builder n -> Builder (m + n)
`Bounded.append` Word64 -> Builder 10
vlqSmile64 (Int64 -> Word64
toZigzag64 Int64
i64))
    | Bool
otherwise -> (Int -> Int -> Builder)
-> (Integer -> Integer -> Builder) -> Scientific -> Builder
forall a.
(Int -> Int -> a) -> (Integer -> Integer -> a) -> Scientific -> a
Sci.withExposed Int -> Int -> Builder
encodeSmallDecimal Integer -> Integer -> Builder
encodeBigDecimal Scientific
x
  Value
Null -> Word8 -> Builder
B.word8 Word8
0x21
  Value
False -> Word8 -> Builder
B.word8 Word8
0x22
  Value
True -> Word8 -> Builder
B.word8 Word8
0x23

encodeSmallDecimal :: Int -> Int -> Builder
encodeSmallDecimal :: Int -> Int -> Builder
encodeSmallDecimal !Int
c !Int
e = Integer -> Integer -> Builder
encodeBigDecimal (Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
c) (Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
e)

encodeBigDecimal :: Integer -> Integer -> Builder
encodeBigDecimal :: Integer -> Integer -> Builder
encodeBigDecimal Integer
c Integer
e = case Integer
e of
  Integer
0 -> Integer -> Builder
encodeBigInteger Integer
c
  Integer
_ ->
    Word8 -> Builder
B.word8 Word8
0x2A -- bigdecimal token tag
      Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Natural -> Builder
vlqSmile
        ( forall a b. (Integral a, Num b) => a -> b
fromIntegral @Word32 @Natural (Word32 -> Natural) -> Word32 -> Natural
forall a b. (a -> b) -> a -> b
$
            Int32 -> Word32
toZigzag32 Int32
scale
        )
      Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Natural -> Builder
vlqSmile (forall a b. (Integral a, Num b) => a -> b
fromIntegral @Int @Natural (Int -> Natural) -> Int -> Natural
forall a b. (a -> b) -> a -> b
$ ByteArray -> Int
sizeofByteArray ByteArray
raw) -- size of byte digits
      Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Bytes -> Builder
B.sevenEightSmile (ByteArray -> Bytes
Bytes.fromByteArray ByteArray
raw) -- 7/8 encoding of byte digits
   where
    scale :: Int32
    -- WARNING smile can't handle exponents outside int32_t, so this truncates
    -- WARNING "scale" is what Java BigDecimal thinks, which is
    -- negative of all mathematics since exponential notation was invented 💩
    scale :: Int32
scale = forall a b. (Integral a, Num b) => a -> b
fromIntegral @Integer @Int32 (-Integer
e)
    raw :: ByteArray
raw = Integer -> ByteArray
integerToBase256ByteArray Integer
c

-- | Encode a number using as SMILE @BigInteger@ token type (prefix @0x26@).
encodeBigInteger :: Integer -> Builder
encodeBigInteger :: Integer -> Builder
encodeBigInteger Integer
n =
  Word8 -> Builder
B.word8 Word8
0x26
    Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Natural -> Builder
vlqSmile (forall a b. (Integral a, Num b) => a -> b
fromIntegral @Int @Natural (Int -> Natural) -> Int -> Natural
forall a b. (a -> b) -> a -> b
$ ByteArray -> Int
sizeofByteArray ByteArray
raw) -- size of byte digits
    Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Bytes -> Builder
B.sevenEightSmile (ByteArray -> Bytes
Bytes.fromByteArray ByteArray
raw) -- 7/8 encoding of byte digits
 where
  !raw :: ByteArray
raw = Integer -> ByteArray
integerToBase256ByteArray Integer
n

integerToBase256ByteArray :: Integer -> ByteArray
integerToBase256ByteArray :: Integer -> ByteArray
integerToBase256ByteArray Integer
c =
  if Integer
c Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
0
    then Int -> [Word8] -> ByteArray
forall a. Prim a => Int -> [a] -> ByteArray
byteArrayFromListN Int
1 [Word8
0 :: Word8]
    else case Integer
c of
      Integer.IP ByteArray#
bn -> IO ByteArray -> ByteArray
forall a. IO a -> a
unsafeDupablePerformIO (IO ByteArray -> ByteArray) -> IO ByteArray -> ByteArray
forall a b. (a -> b) -> a -> b
$ do
        let nDigits256 :: Int
nDigits256 = forall a b. (Integral a, Num b) => a -> b
fromIntegral @Word @Int (Word# -> Word
W# (Word# -> ByteArray# -> Word#
BN.bigNatSizeInBase# Word#
256## ByteArray#
bn))
        MutableByteArray RealWorld
mut <- Int -> IO (MutableByteArray (PrimState IO))
forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableByteArray (PrimState m))
newByteArray Int
nDigits256
        let !(MutableByteArray MutableByteArray# RealWorld
mut#) = MutableByteArray RealWorld
mut
        !Word
_ <- (State# RealWorld -> (# State# RealWorld, Word# #)) -> IO Word
liftWordIO (ByteArray#
-> MutableByteArray# RealWorld
-> Word#
-> Bool#
-> State# RealWorld
-> (# State# RealWorld, Word# #)
forall s.
ByteArray#
-> MutableByteArray# s
-> Word#
-> Bool#
-> State# s
-> (# State# s, Word# #)
BN.bigNatToMutableByteArray# ByteArray#
bn MutableByteArray# RealWorld
mut# Word#
0## Bool#
1#)
        -- This is safe because Jp cannot have zero inside it.
        Word8
w0 :: Word8 <- MutableByteArray (PrimState IO) -> Int -> IO Word8
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> m a
readByteArray MutableByteArray RealWorld
MutableByteArray (PrimState IO)
mut Int
0
        if Word8 -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Word8
w0 Int
7
          then do
            -- If the upper bit is 1, then we must introduce a leading
            -- zero byte.
            MutableByteArray RealWorld
dst <- Int -> IO (MutableByteArray (PrimState IO))
forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableByteArray (PrimState m))
newByteArray (Int
nDigits256 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
            MutableByteArray (PrimState IO) -> Int -> Word8 -> IO ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> a -> m ()
writeByteArray MutableByteArray RealWorld
MutableByteArray (PrimState IO)
dst Int
0 (Word8
0x00 :: Word8)
            MutableByteArray (PrimState IO)
-> Int -> MutableByteArray (PrimState IO) -> Int -> Int -> IO ()
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m)
-> Int -> MutableByteArray (PrimState m) -> Int -> Int -> m ()
copyMutableByteArray MutableByteArray RealWorld
MutableByteArray (PrimState IO)
dst Int
1 MutableByteArray RealWorld
MutableByteArray (PrimState IO)
mut Int
0 Int
nDigits256
            MutableByteArray (PrimState IO) -> IO ByteArray
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> m ByteArray
unsafeFreezeByteArray MutableByteArray RealWorld
MutableByteArray (PrimState IO)
dst
          else MutableByteArray (PrimState IO) -> IO ByteArray
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> m ByteArray
unsafeFreezeByteArray MutableByteArray RealWorld
MutableByteArray (PrimState IO)
mut
      Integer.IN ByteArray#
bn -> ByteArray# -> ByteArray
twosComplementBigNat ByteArray#
bn
      Integer.IS Bool#
i -> case Bool#
i Bool# -> Bool# -> Bool#
Exts.># Bool#
0# of
        Bool#
1# -> Word -> ByteArray
encodePosWordBase256 (Word# -> Word
W# (Bool# -> Word#
Exts.int2Word# Bool#
i))
        Bool#
_ -> Word -> ByteArray
encodeNegWordBase256 (Word# -> Word
W# (Bool# -> Word#
Exts.int2Word# Bool#
i))

liftWordIO :: (State# RealWorld -> (# State# RealWorld, Word# #)) -> IO Word
{-# INLINE liftWordIO #-}
liftWordIO :: (State# RealWorld -> (# State# RealWorld, Word# #)) -> IO Word
liftWordIO State# RealWorld -> (# State# RealWorld, Word# #)
f =
  (State# RealWorld -> (# State# RealWorld, Word #)) -> IO Word
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO
    ( \State# RealWorld
s -> case State# RealWorld -> (# State# RealWorld, Word# #)
f State# RealWorld
s of
        (# State# RealWorld
s', Word#
w #) -> (# State# RealWorld
s', Word# -> Word
W# Word#
w #)
    )

twosComplementBigNat :: BN.BigNat# -> ByteArray
twosComplementBigNat :: ByteArray# -> ByteArray
twosComplementBigNat ByteArray#
bn = IO ByteArray -> ByteArray
forall a. IO a -> a
unsafeDupablePerformIO (IO ByteArray -> ByteArray) -> IO ByteArray -> ByteArray
forall a b. (a -> b) -> a -> b
$ do
  let nDigits256 :: Int
nDigits256 = forall a b. (Integral a, Num b) => a -> b
fromIntegral @Word @Int (Word# -> Word
W# (Word# -> ByteArray# -> Word#
BN.bigNatSizeInBase# Word#
256## ByteArray#
bn))
  MutableByteArray RealWorld
mut <- Int -> IO (MutableByteArray (PrimState IO))
forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableByteArray (PrimState m))
newByteArray Int
nDigits256
  let !(MutableByteArray MutableByteArray# RealWorld
mut#) = MutableByteArray RealWorld
mut
  !Word
_ <- (State# RealWorld -> (# State# RealWorld, Word# #)) -> IO Word
liftWordIO (ByteArray#
-> MutableByteArray# RealWorld
-> Word#
-> Bool#
-> State# RealWorld
-> (# State# RealWorld, Word# #)
forall s.
ByteArray#
-> MutableByteArray# s
-> Word#
-> Bool#
-> State# s
-> (# State# s, Word# #)
BN.bigNatToMutableByteArray# ByteArray#
bn MutableByteArray# RealWorld
mut# Word#
0## Bool#
1#)
  -- First, complement
  let goComplement :: Int -> IO ()
goComplement !Int
ix =
        if Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0
          then do
            Word8
w :: Word8 <- MutableByteArray (PrimState IO) -> Int -> IO Word8
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> m a
readByteArray MutableByteArray RealWorld
MutableByteArray (PrimState IO)
mut Int
ix
            MutableByteArray (PrimState IO) -> Int -> Word8 -> IO ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> a -> m ()
writeByteArray MutableByteArray RealWorld
MutableByteArray (PrimState IO)
mut Int
ix (Word8 -> Word8
forall a. Bits a => a -> a
complement Word8
w)
            Int -> IO ()
goComplement (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
          else () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  Int -> IO ()
goComplement (Int
nDigits256 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
  -- Second, add one
  let goAddOne :: Int -> IO ()
goAddOne !Int
ix =
        if Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0
          then do
            Word8
w :: Word8 <- MutableByteArray (PrimState IO) -> Int -> IO Word8
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> m a
readByteArray MutableByteArray RealWorld
MutableByteArray (PrimState IO)
mut Int
ix
            case Word8
w of
              Word8
0xFF -> do
                MutableByteArray (PrimState IO) -> Int -> Word8 -> IO ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> a -> m ()
writeByteArray MutableByteArray RealWorld
MutableByteArray (PrimState IO)
mut Int
ix (Word8
0 :: Word8)
                Int -> IO ()
goAddOne (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
              Word8
_ -> MutableByteArray (PrimState IO) -> Int -> Word8 -> IO ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> a -> m ()
writeByteArray MutableByteArray RealWorld
MutableByteArray (PrimState IO)
mut Int
ix (Word8
w Word8 -> Word8 -> Word8
forall a. Num a => a -> a -> a
+ Word8
1)
          else () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  Int -> IO ()
goAddOne (Int
nDigits256 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
  Word8
leader :: Word8 <- MutableByteArray (PrimState IO) -> Int -> IO Word8
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> m a
readByteArray MutableByteArray RealWorld
MutableByteArray (PrimState IO)
mut Int
0
  if Word8 -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Word8
leader Int
7
    then MutableByteArray (PrimState IO) -> IO ByteArray
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> m ByteArray
unsafeFreezeByteArray MutableByteArray RealWorld
MutableByteArray (PrimState IO)
mut
    else do
      MutableByteArray RealWorld
dst <- Int -> IO (MutableByteArray (PrimState IO))
forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableByteArray (PrimState m))
newByteArray (Int
nDigits256 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
      MutableByteArray (PrimState IO)
-> Int -> MutableByteArray (PrimState IO) -> Int -> Int -> IO ()
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m)
-> Int -> MutableByteArray (PrimState m) -> Int -> Int -> m ()
copyMutableByteArray MutableByteArray RealWorld
MutableByteArray (PrimState IO)
dst Int
1 MutableByteArray RealWorld
MutableByteArray (PrimState IO)
mut Int
0 Int
nDigits256
      MutableByteArray (PrimState IO) -> Int -> Word8 -> IO ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> a -> m ()
writeByteArray MutableByteArray RealWorld
MutableByteArray (PrimState IO)
dst Int
0 (Word8
0xFF :: Word8)
      MutableByteArray (PrimState IO) -> IO ByteArray
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> m ByteArray
unsafeFreezeByteArray MutableByteArray RealWorld
MutableByteArray (PrimState IO)
dst

-- Only works on 64-bit architectures.
-- Precondition: must be a bit pattern of a positive integer
encodePosWordBase256 :: Word -> ByteArray
encodePosWordBase256 :: Word -> ByteArray
encodePosWordBase256 !Word
w = (forall s. ST s ByteArray) -> ByteArray
runByteArrayST ((forall s. ST s ByteArray) -> ByteArray)
-> (forall s. ST s ByteArray) -> ByteArray
forall a b. (a -> b) -> a -> b
$ do
  -- If a positive number broken into its constituent bytes has
  -- leading byte with MSB 1, then we want to have a zero byte
  -- in front. That is, we want to produce:
  --
  -- 0000_0000 1101_1010 01010101
  --
  -- instead of
  --
  -- 1101_1010 01010101
  --
  -- Because the latter would be misinterpreted as a negative number.
  let !total :: Int
total = Int -> Int -> Int
forall a. Integral a => a -> a -> a
quot (Int
72 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Word -> Int
forall b. FiniteBits b => b -> Int
countLeadingZeros Word
w) Int
8
  MutableByteArray s
dst <- Int -> ST s (MutableByteArray (PrimState (ST s)))
forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableByteArray (PrimState m))
newByteArray Int
total
  let go :: Int -> Word -> ST s ByteArray
go !Int
ix !Word
acc =
        if Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0
          then do
            MutableByteArray (PrimState (ST s)) -> Int -> Word8 -> ST s ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> a -> m ()
writeByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
dst Int
ix (forall a b. (Integral a, Num b) => a -> b
fromIntegral @Word @Word8 Word
acc)
            Int -> Word -> ST s ByteArray
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (Word -> Int -> Word
forall a. Bits a => a -> Int -> a
unsafeShiftR Word
acc Int
8)
          else MutableByteArray (PrimState (ST s)) -> ST s ByteArray
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> m ByteArray
unsafeFreezeByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
dst
  Int -> Word -> ST s ByteArray
go (Int
total Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Word
w

-- Same deal as encodePosWordBase256.
-- Andrew Martin: I am not sure why we need to subtract from 72
-- instead of 71. But if we change it, a bunch of tests fail.
encodeNegWordBase256 :: Word -> ByteArray
encodeNegWordBase256 :: Word -> ByteArray
encodeNegWordBase256 !Word
w = (forall s. ST s ByteArray) -> ByteArray
runByteArrayST ((forall s. ST s ByteArray) -> ByteArray)
-> (forall s. ST s ByteArray) -> ByteArray
forall a b. (a -> b) -> a -> b
$ do
  let !total :: Int
total = Int -> Int -> Int
forall a. Integral a => a -> a -> a
quot (Int
72 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Word -> Int
forall b. FiniteBits b => b -> Int
countLeadingZeros (Word -> Word
forall a. Bits a => a -> a
complement Word
w)) Int
8
  MutableByteArray s
dst <- Int -> ST s (MutableByteArray (PrimState (ST s)))
forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableByteArray (PrimState m))
newByteArray Int
total
  let go :: Int -> Word -> ST s ByteArray
go !Int
ix !Word
acc =
        if Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0
          then do
            MutableByteArray (PrimState (ST s)) -> Int -> Word8 -> ST s ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> a -> m ()
writeByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
dst Int
ix (forall a b. (Integral a, Num b) => a -> b
fromIntegral @Word @Word8 Word
acc)
            Int -> Word -> ST s ByteArray
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (Word -> Int -> Word
forall a. Bits a => a -> Int -> a
unsafeShiftR Word
acc Int
8)
          else MutableByteArray (PrimState (ST s)) -> ST s ByteArray
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> m ByteArray
unsafeFreezeByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
dst
  Int -> Word -> ST s ByteArray
go (Int
total Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Word
w

{- | Encode a string in which all characters are ASCII. This precondition
is not checked. Resulting output will be corrupt if this condition
is not satisfied.
-}
encodeAsciiString :: ShortText -> Builder
encodeAsciiString :: ShortText -> Builder
encodeAsciiString !ShortText
str
  | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = Word8 -> Builder
B.word8 Word8
0x20
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
64 = Word8 -> Bytes -> Builder
B.copyCons (Word8
0x40 Word8 -> Word8 -> Word8
forall a. Num a => a -> a -> a
+ Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)) (ShortByteString -> Bytes
Bytes.fromShortByteString (ShortText -> ShortByteString
TS.toShortByteString ShortText
str))
  | Bool
otherwise = Word8 -> Builder
B.word8 Word8
0xe0 Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> ShortText -> Builder
B.shortTextUtf8 ShortText
str Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Word8 -> Builder
B.word8 Word8
0xFC
 where
  n :: Int
n = ShortByteString -> Int
SBS.length (ShortText -> ShortByteString
TS.toShortByteString ShortText
str)

-- | Encode a string.
encodeString :: ShortText -> Builder
encodeString :: ShortText -> Builder
encodeString !ShortText
str = case ShortByteString -> Int
SBS.length (ShortText -> ShortByteString
TS.toShortByteString ShortText
str) of
  Int
0 -> Word8 -> Builder
B.word8 Word8
0x20
  Int
n -> case ShortText -> Bool
TS.isAscii ShortText
str of
    Bool
Prelude.True
      | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
64 -> Word8 -> Bytes -> Builder
B.copyCons (Word8
0x40 Word8 -> Word8 -> Word8
forall a. Num a => a -> a -> a
+ Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)) (ShortByteString -> Bytes
Bytes.fromShortByteString (ShortText -> ShortByteString
TS.toShortByteString ShortText
str))
      | Bool
otherwise -> Word8 -> Builder
B.word8 Word8
0xe0 Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> ShortText -> Builder
B.shortTextUtf8 ShortText
str Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Word8 -> Builder
B.word8 Word8
0xFC
    Bool
Prelude.False
      | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
65 -> Word8 -> Bytes -> Builder
B.copyCons (Word8
0x80 Word8 -> Word8 -> Word8
forall a. Num a => a -> a -> a
+ Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2)) (ShortByteString -> Bytes
Bytes.fromShortByteString (ShortText -> ShortByteString
TS.toShortByteString ShortText
str))
      | Bool
otherwise -> Word8 -> Builder
B.word8 Word8
0xE4 Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> ShortText -> Builder
B.shortTextUtf8 ShortText
str Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Word8 -> Builder
B.word8 Word8
0xFC

-- | Encode a key.
encodeKey :: ShortText -> Builder
encodeKey :: ShortText -> Builder
encodeKey !ShortText
str = case ShortByteString -> Int
SBS.length (ShortText -> ShortByteString
TS.toShortByteString ShortText
str) of
  Int
0 -> Word8 -> Builder
B.word8 Word8
0x20
  Int
n
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
64
        Bool -> Bool -> Bool
&& ShortText -> Bool
TS.isAscii ShortText
str
    , Word8
w8 <- forall a b. (Integral a, Num b) => a -> b
fromIntegral @Int @Word8 (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) ->
        Word8 -> Bytes -> Builder
B.copyCons (Word8
0x80 Word8 -> Word8 -> Word8
forall a. Num a => a -> a -> a
+ Word8
w8) (ShortByteString -> Bytes
Bytes.fromShortByteString (ShortText -> ShortByteString
TS.toShortByteString ShortText
str))
  Int
n
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
56
    , Word8
w8 <- forall a b. (Integral a, Num b) => a -> b
fromIntegral @Int @Word8 (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2) ->
        Word8 -> Bytes -> Builder
B.copyCons (Word8
0xC0 Word8 -> Word8 -> Word8
forall a. Num a => a -> a -> a
+ Word8
w8) (ShortByteString -> Bytes
Bytes.fromShortByteString (ShortText -> ShortByteString
TS.toShortByteString ShortText
str))
    | Bool
otherwise -> Word8 -> Builder
B.word8 Word8
0x34 Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> ShortText -> Builder
B.shortTextUtf8 ShortText
str Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Word8 -> Builder
B.word8 Word8
0xFC

{- | Encode a key in which all characters are ASCII. This precondition
is not checked. Resulting output will be corrupt if this condition
is not satisfied.
-}
encodeAsciiKey :: ShortText -> Builder
encodeAsciiKey :: ShortText -> Builder
encodeAsciiKey ShortText
str = case ShortByteString -> Int
SBS.length (ShortText -> ShortByteString
TS.toShortByteString ShortText
str) of
  Int
0 -> Word8 -> Builder
B.word8 Word8
0x20
  Int
n
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
64
    , Word8
w8 <- forall a b. (Integral a, Num b) => a -> b
fromIntegral @Int @Word8 (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) ->
        Word8 -> Builder
B.word8 (Word8
0x80 Word8 -> Word8 -> Word8
forall a. Num a => a -> a -> a
+ Word8
w8) Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> ShortText -> Builder
B.shortTextUtf8 ShortText
str
    | Bool
otherwise -> Word8 -> Builder
B.word8 Word8
0x34 Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> ShortText -> Builder
B.shortTextUtf8 ShortText
str Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Word8 -> Builder
B.word8 Word8
0xFC

vlqSmile :: Natural -> Builder
vlqSmile :: Natural -> Builder
vlqSmile Natural
n0 =
  let (Natural
rest, Word8
lastBits) = Natural -> (Natural, Word8)
take6bits Natural
n0
   in Natural -> Builder
loop Natural
rest Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Word8 -> Builder
B.word8 (Word8
lastBits Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Word8
0x80)
 where
  loop :: Natural -> Builder
loop Natural
n
    | Natural
n Natural -> Natural -> Bool
forall a. Eq a => a -> a -> Bool
== Natural
0 = Builder
forall a. Monoid a => a
mempty
    | (Natural
rest, Word8
bits) <- Natural -> (Natural, Word8)
take7bits Natural
n =
        Natural -> Builder
loop Natural
rest Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Word8 -> Builder
B.word8 Word8
bits
  take7bits :: Natural -> (Natural, Word8)
  take7bits :: Natural -> (Natural, Word8)
take7bits Natural
n = (Natural
n Natural -> Int -> Natural
forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
7, forall a b. (Integral a, Num b) => a -> b
fromIntegral @Natural @Word8 Natural
n Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
0x7F)
  take6bits :: Natural -> (Natural, Word8)
  take6bits :: Natural -> (Natural, Word8)
take6bits Natural
n = (Natural
n Natural -> Int -> Natural
forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
6, forall a b. (Integral a, Num b) => a -> b
fromIntegral @Natural @Word8 Natural
n Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
0x3F)

-- Precondition: input is not zero
vlqSmile64 :: Word64 -> Bounded.Builder 10
vlqSmile64 :: Word64 -> Builder 10
vlqSmile64 !Word64
n0 = (forall s. MutableByteArray s -> Int -> ST s Int) -> Builder 10
forall (n :: Natural).
(forall s. MutableByteArray s -> Int -> ST s Int) -> Builder n
Unsafe.construct ((forall s. MutableByteArray s -> Int -> ST s Int) -> Builder 10)
-> (forall s. MutableByteArray s -> Int -> ST s Int) -> Builder 10
forall a b. (a -> b) -> a -> b
$ \MutableByteArray s
buf Int
ix0 -> do
  let !w0 :: Word8
w0 = forall a b. (Integral a, Num b) => a -> b
fromIntegral @Word64 @Word8 Word64
n0 Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
0x3F
  MutableByteArray (PrimState (ST s)) -> Int -> Word8 -> ST s ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> a -> m ()
writeByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
buf Int
ix0 (Word8
0x80 Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Word8
w0)
  let !acc0 :: Word64
acc0 = Word64
n0 Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
6
  let loop :: Word64 -> Int -> ST s Int
loop !Word64
acc !Int
ix = case Word64
acc of
        Word64
0 -> Int -> ST s Int
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
ix
        Word64
_ -> do
          let !w :: Word8
w = forall a b. (Integral a, Num b) => a -> b
fromIntegral @Word64 @Word8 Word64
acc Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
0x7F
          MutableByteArray (PrimState (ST s)) -> Int -> Word8 -> ST s ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> a -> m ()
writeByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
buf Int
ix Word8
w
          let !acc' :: Word64
acc' = Word64
acc Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
7
          Word64 -> Int -> ST s Int
loop Word64
acc' (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
  Int
ix1 <- Word64 -> Int -> ST s Int
loop Word64
acc0 (Int
ix0 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
  MutableByteArray s -> Int -> Int -> ST s ()
forall s. MutableByteArray s -> Int -> Int -> ST s ()
reverseBytes MutableByteArray s
buf Int
ix0 (Int
ix1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
  Int -> ST s Int
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
ix1

-- Reverse the bytes in the designated slice. This takes
-- an inclusive start offset and an inclusive end offset.
--
-- Copied from bytebuild
reverseBytes :: MutableByteArray s -> Int -> Int -> ST s ()
{-# INLINE reverseBytes #-}
reverseBytes :: forall s. MutableByteArray s -> Int -> Int -> ST s ()
reverseBytes MutableByteArray s
arr Int
begin Int
end = Int -> Int -> ST s ()
go Int
begin Int
end
 where
  go :: Int -> Int -> ST s ()
go Int
ixA Int
ixB =
    if Int
ixA Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
ixB
      then do
        Word8
a :: Word8 <- MutableByteArray (PrimState (ST s)) -> Int -> ST s Word8
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> m a
readByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
arr Int
ixA
        Word8
b :: Word8 <- MutableByteArray (PrimState (ST s)) -> Int -> ST s Word8
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> m a
readByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
arr Int
ixB
        MutableByteArray (PrimState (ST s)) -> Int -> Word8 -> ST s ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> a -> m ()
writeByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
arr Int
ixA Word8
b
        MutableByteArray (PrimState (ST s)) -> Int -> Word8 -> ST s ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> a -> m ()
writeByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
arr Int
ixB Word8
a
        Int -> Int -> ST s ()
go (Int
ixA Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
ixB Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
      else () -> ST s ()
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()