{-# 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 (countLeadingZeros,complement,unsafeShiftR,(.&.),(.|.))
import Data.Bits (testBit)
import Data.Bytes.Builder (Builder)
import Data.Int (Int32)
import Data.Primitive (ByteArray,newByteArray)
import Data.Primitive (writeByteArray,byteArrayFromListN,sizeofByteArray)
import Data.Primitive (MutableByteArray(..),unsafeFreezeByteArray)
import Data.Primitive (readByteArray,copyMutableByteArray)
import Data.Text.Short (ShortText)
import Data.Word (Word8,Word32,Word64)
import Data.Word.Zigzag (toZigzag32,toZigzag64)
import GHC.Exts (RealWorld,Word#,State#)
import GHC.IO (IO(IO))
import GHC.Word (Word(..))
import Json (Value(..), Member(..))
import Numeric.Natural (Natural)
import System.IO.Unsafe (unsafeDupablePerformIO)

import qualified Arithmetic.Nat as Nat
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.ByteString.Short as SBS
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' 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 forall a b. (a -> b) -> a -> b
$ case Value
val of
  Object SmallArray Member
obj ->
    Word8 -> Builder
B.word8 Word8
0xFA
    forall a. Semigroup a => a -> a -> a
<>
    forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (\Member{ShortText
key :: Member -> ShortText
key :: ShortText
key,Value
value :: Member -> Value
value :: Value
value} -> ShortText -> Builder
encodeKey ShortText
key forall a. Semigroup a => a -> a -> a
<> Value -> Builder
encodeNoHeader Value
value) SmallArray Member
obj
    forall a. Semigroup a => a -> a -> a
<>
    Word8 -> Builder
B.word8 Word8
0xFB
  Array SmallArray Value
arr -> Word8 -> Builder
B.word8 Word8
0xF8 forall a. Semigroup a => a -> a -> a
<> forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Value -> Builder
encodeNoHeader SmallArray Value
arr 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 forall a. Ord a => a -> a -> Bool
<= Int32
i32 Bool -> Bool -> Bool
&& Int32
i32 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 forall a. Num a => a -> a -> a
+ Word8
w5)
    | Just Int32
i32 <- Scientific -> Maybe Int32
Sci.toInt32 Scientific
x
      -> forall (n :: Natural). Nat n -> Builder n -> Builder
B.fromBounded forall (n :: Natural). KnownNat n => Nat n
Nat.constant (Word8 -> Builder 1
Bounded.word8 Word8
0x24 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
      -> forall (n :: Natural). Nat n -> Builder n -> Builder
B.fromBounded forall (n :: Natural). KnownNat n => Nat n
Nat.constant (Word8 -> Builder 1
Bounded.word8 Word8
0x25 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 -> 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 (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
c) (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
    forall a. Semigroup a => a -> a -> a
<> Natural -> Builder
vlqSmile ( forall a b. (Integral a, Num b) => a -> b
fromIntegral @Word32 @Natural
                forall a b. (a -> b) -> a -> b
$ Int32 -> Word32
toZigzag32 Int32
scale)
    forall a. Semigroup a => a -> a -> a
<> Natural -> Builder
vlqSmile (forall a b. (Integral a, Num b) => a -> b
fromIntegral @Int @Natural forall a b. (a -> b) -> a -> b
$ ByteArray -> Int
sizeofByteArray ByteArray
raw) -- size of byte digits
    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
  forall a. Semigroup a => a -> a -> a
<> Natural -> Builder
vlqSmile (forall a b. (Integral a, Num b) => a -> b
fromIntegral @Int @Natural forall a b. (a -> b) -> a -> b
$ ByteArray -> Int
sizeofByteArray ByteArray
raw) -- size of byte digits
  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 forall a. Eq a => a -> a -> Bool
== Integer
0
  then forall a. Prim a => Int -> [a] -> ByteArray
byteArrayFromListN Int
1 [Word8
0::Word8]
  else case Integer
c of
    Integer.IP ByteArray#
bn -> forall a. IO a -> a
unsafeDupablePerformIO 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 <- 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 (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 <- forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> m a
readByteArray MutableByteArray RealWorld
mut Int
0
      if 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 <- forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableByteArray (PrimState m))
newByteArray (Int
nDigits256 forall a. Num a => a -> a -> a
+ Int
1)
          forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> a -> m ()
writeByteArray MutableByteArray RealWorld
dst Int
0 (Word8
0x00 :: Word8)
          forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m)
-> Int -> MutableByteArray (PrimState m) -> Int -> Int -> m ()
copyMutableByteArray MutableByteArray RealWorld
dst Int
1 MutableByteArray RealWorld
mut Int
0 Int
nDigits256
          forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> m ByteArray
unsafeFreezeByteArray MutableByteArray RealWorld
dst
        else forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> m ByteArray
unsafeFreezeByteArray MutableByteArray RealWorld
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 = 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 = forall a. IO a -> a
unsafeDupablePerformIO 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 <- 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 (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 forall a. Ord a => a -> a -> Bool
>= Int
0
        then do
          Word8
w :: Word8 <- forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> m a
readByteArray MutableByteArray RealWorld
mut Int
ix
          forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> a -> m ()
writeByteArray MutableByteArray RealWorld
mut Int
ix (forall a. Bits a => a -> a
complement Word8
w)
          Int -> IO ()
goComplement (Int
ix forall a. Num a => a -> a -> a
- Int
1)
        else forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  Int -> IO ()
goComplement (Int
nDigits256 forall a. Num a => a -> a -> a
- Int
1)
  -- Second, add one
  let goAddOne :: Int -> IO ()
goAddOne !Int
ix = if Int
ix forall a. Ord a => a -> a -> Bool
>= Int
0
        then do
          Word8
w :: Word8 <- forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> m a
readByteArray MutableByteArray RealWorld
mut Int
ix
          case Word8
w of
            Word8
0xFF -> do
              forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> a -> m ()
writeByteArray MutableByteArray RealWorld
mut Int
ix (Word8
0 :: Word8)
              Int -> IO ()
goAddOne (Int
ix forall a. Num a => a -> a -> a
- Int
1)
            Word8
_ -> forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> a -> m ()
writeByteArray MutableByteArray RealWorld
mut Int
ix (Word8
w forall a. Num a => a -> a -> a
+ Word8
1)
        else forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  Int -> IO ()
goAddOne (Int
nDigits256 forall a. Num a => a -> a -> a
- Int
1)
  Word8
leader :: Word8 <- forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> m a
readByteArray MutableByteArray RealWorld
mut Int
0
  if forall a. Bits a => a -> Int -> Bool
testBit Word8
leader Int
7
    then forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> m ByteArray
unsafeFreezeByteArray MutableByteArray RealWorld
mut
    else do
      MutableByteArray RealWorld
dst <- forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableByteArray (PrimState m))
newByteArray (Int
nDigits256 forall a. Num a => a -> a -> a
+ Int
1)
      forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m)
-> Int -> MutableByteArray (PrimState m) -> Int -> Int -> m ()
copyMutableByteArray MutableByteArray RealWorld
dst Int
1 MutableByteArray RealWorld
mut Int
0 Int
nDigits256
      forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> a -> m ()
writeByteArray MutableByteArray RealWorld
dst Int
0 (Word8
0xFF :: Word8)
      forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> m ByteArray
unsafeFreezeByteArray MutableByteArray RealWorld
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 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 = forall a. Integral a => a -> a -> a
quot (Int
72 forall a. Num a => a -> a -> a
- forall b. FiniteBits b => b -> Int
countLeadingZeros Word
w) Int
8
  MutableByteArray s
dst <- 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 forall a. Ord a => a -> a -> Bool
>= Int
0 
        then do
          forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> a -> m ()
writeByteArray MutableByteArray 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 forall a. Num a => a -> a -> a
- Int
1) (forall a. Bits a => a -> Int -> a
unsafeShiftR Word
acc Int
8)
        else forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> m ByteArray
unsafeFreezeByteArray MutableByteArray s
dst
  Int -> Word -> ST s ByteArray
go (Int
total 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 a b. (a -> b) -> a -> b
$ do
  let !total :: Int
total = forall a. Integral a => a -> a -> a
quot (Int
72 forall a. Num a => a -> a -> a
- forall b. FiniteBits b => b -> Int
countLeadingZeros (forall a. Bits a => a -> a
complement Word
w)) Int
8
  MutableByteArray s
dst <- 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 forall a. Ord a => a -> a -> Bool
>= Int
0 
        then do
          forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> a -> m ()
writeByteArray MutableByteArray 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 forall a. Num a => a -> a -> a
- Int
1) (forall a. Bits a => a -> Int -> a
unsafeShiftR Word
acc Int
8)
        else forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> m ByteArray
unsafeFreezeByteArray MutableByteArray s
dst
  Int -> Word -> ST s ByteArray
go (Int
total 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 forall a. Eq a => a -> a -> Bool
== Int
0 = Word8 -> Builder
B.word8 Word8
0x20
  | Int
n forall a. Ord a => a -> a -> Bool
<= Int
64 = Word8 -> Bytes -> Builder
B.copyCons (Word8
0x40 forall a. Num a => a -> a -> a
+ forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
n 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 forall a. Semigroup a => a -> a -> a
<> ShortText -> Builder
B.shortTextUtf8 ShortText
str 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 forall a. Ord a => a -> a -> Bool
<= Int
64 -> Word8 -> Bytes -> Builder
B.copyCons (Word8
0x40 forall a. Num a => a -> a -> a
+ forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
n 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 forall a. Semigroup a => a -> a -> a
<> ShortText -> Builder
B.shortTextUtf8 ShortText
str forall a. Semigroup a => a -> a -> a
<> Word8 -> Builder
B.word8 Word8
0xFC
    Bool
Prelude.False
      | Int
n forall a. Ord a => a -> a -> Bool
<= Int
65 -> Word8 -> Bytes -> Builder
B.copyCons (Word8
0x80 forall a. Num a => a -> a -> a
+ forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
n 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 forall a. Semigroup a => a -> a -> a
<> ShortText -> Builder
B.shortTextUtf8 ShortText
str 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 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 forall a. Num a => a -> a -> a
- Int
1)
    -> Word8 -> Bytes -> Builder
B.copyCons (Word8
0x80 forall a. Num a => a -> a -> a
+ Word8
w8) (ShortByteString -> Bytes
Bytes.fromShortByteString (ShortText -> ShortByteString
TS.toShortByteString ShortText
str))
  Int
n | Int
n 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 forall a. Num a => a -> a -> a
- Int
2)
    -> Word8 -> Bytes -> Builder
B.copyCons (Word8
0xC0 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 forall a. Semigroup a => a -> a -> a
<> ShortText -> Builder
B.shortTextUtf8 ShortText
str 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 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 forall a. Num a => a -> a -> a
- Int
1)
    -> Word8 -> Builder
B.word8 (Word8
0x80 forall a. Num a => a -> a -> a
+ Word8
w8) forall a. Semigroup a => a -> a -> a
<> ShortText -> Builder
B.shortTextUtf8 ShortText
str
    | Bool
otherwise -> Word8 -> Builder
B.word8 Word8
0x34 forall a. Semigroup a => a -> a -> a
<> ShortText -> Builder
B.shortTextUtf8 ShortText
str 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 forall a. Semigroup a => a -> a -> a
<> Word8 -> Builder
B.word8 (Word8
lastBits forall a. Bits a => a -> a -> a
.|. Word8
0x80)
  where
  loop :: Natural -> Builder
loop Natural
n
    | Natural
n forall a. Eq a => a -> a -> Bool
== Natural
0 = forall a. Monoid a => a
mempty
    | (Natural
rest, Word8
bits) <- Natural -> (Natural, Word8)
take7bits Natural
n
      = Natural -> Builder
loop Natural
rest 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 forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
7, forall a b. (Integral a, Num b) => a -> b
fromIntegral @Natural @Word8 Natural
n forall a. Bits a => a -> a -> a
.&. Word8
0x7F)
  take6bits :: Natural -> (Natural, Word8)
  take6bits :: Natural -> (Natural, Word8)
take6bits Natural
n = (Natural
n forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
6, forall a b. (Integral a, Num b) => a -> b
fromIntegral @Natural @Word8 Natural
n 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 (n :: Natural).
(forall s. MutableByteArray s -> Int -> ST s Int) -> Builder n
Unsafe.construct 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 forall a. Bits a => a -> a -> a
.&. Word8
0x3F
  forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> a -> m ()
writeByteArray MutableByteArray s
buf Int
ix0 (Word8
0x80 forall a. Bits a => a -> a -> a
.|. Word8
w0)
  let !acc0 :: Word64
acc0 = Word64
n0 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 -> 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 forall a. Bits a => a -> a -> a
.&. Word8
0x7F
          forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> a -> m ()
writeByteArray MutableByteArray s
buf Int
ix Word8
w
          let !acc' :: Word64
acc' = Word64
acc forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
7
          Word64 -> Int -> ST s Int
loop Word64
acc' (Int
ix forall a. Num a => a -> a -> a
+ Int
1)
  Int
ix1 <- Word64 -> Int -> ST s Int
loop Word64
acc0 (Int
ix0 forall a. Num a => a -> a -> a
+ Int
1)
  forall s. MutableByteArray s -> Int -> Int -> ST s ()
reverseBytes MutableByteArray s
buf Int
ix0 (Int
ix1 forall a. Num a => a -> a -> a
- Int
1)
  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 forall a. Ord a => a -> a -> Bool
< Int
ixB
    then do
      Word8
a :: Word8 <- forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> m a
readByteArray MutableByteArray s
arr Int
ixA
      Word8
b :: Word8 <- forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> m a
readByteArray MutableByteArray s
arr Int
ixB
      forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> a -> m ()
writeByteArray MutableByteArray s
arr Int
ixA Word8
b
      forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> a -> m ()
writeByteArray MutableByteArray s
arr Int
ixB Word8
a
      Int -> Int -> ST s ()
go (Int
ixA forall a. Num a => a -> a -> a
+ Int
1) (Int
ixB forall a. Num a => a -> a -> a
- Int
1)
    else forall (f :: * -> *) a. Applicative f => a -> f a
pure ()