content: {-
    Copyright (c) Meta Platforms, Inc. and affiliates.
    All rights reserved.
  
    This source code is licensed under the BSD-style license found in the
    LICENSE file in the root directory of this source tree.
  -}
  
  {-# LANGUAGE CPP #-}
  {-# LANGUAGE MagicHash #-}
  module Thrift.Protocol.Compact
    ( serializeCompact
    , deserializeCompact
    , Compact
    ) where
  
  import Control.Monad
  import Thrift.Binary.Parser as P
  import Data.Bits
  import Data.ByteString (ByteString)
  import Data.ByteString.Builder as B
  import Data.ByteString.Builder.Prim as Prim
  import Data.Int
  import Data.List
  import Data.Proxy
  import Data.Text.Encoding
  import Data.Word
  import GHC.Base hiding (build, Type)
  import GHC.Word
  import qualified Data.ByteString as BS
  import qualified Data.HashMap.Strict as HashMap
  
  import Thrift.Protocol
  import Thrift.Protocol.Binary.Internal
  
  data Compact
  
  serializeCompact :: ThriftStruct a => a -> ByteString
  serializeCompact = serializeGen (Proxy :: Proxy Compact)
  
  deserializeCompact :: ThriftStruct a => ByteString -> Either String a
  deserializeCompact = deserializeGen (Proxy :: Proxy Compact)
  
  -- Macros for ThriftTypes ------------------------------------------------------
  
  #define TSTOP   0x00
  #define TTRUE   0x01
  #define TFALSE  0x02
  #define TBYTE   0x03
  #define TI16    0x04
  #define TI32    0x05
  #define TI64    0x06
  #define TDOUBLE 0x07
  #define TSTRING 0x08
  #define TLIST   0x09
  #define TSET    0x0A
  #define TMAP    0x0B
  #define TSTRUCT 0x0C
  #define TFLOAT  0x0D
  
  protocolVersion :: Word8
  protocolVersion = 0x02
  
  protocolID :: Word8
  protocolID = 0x82
  
  typeMask :: Word8
  typeMask = 0xE0
  
  versionMask :: Word8
  versionMask = 0x1F
  
  shiftAmt :: Int
  shiftAmt = 5
  
  -- BinaryWriter instance -------------------------------------------------------
  
  instance Protocol Compact where
    -- Generators for Types
    getByteType   _ = TBYTE
    getI16Type    _ = TI16
    getI32Type    _ = TI32
    getI64Type    _ = TI64
    getFloatType  _ = TFLOAT
    getDoubleType _ = TDOUBLE
    getBoolType   _ = TTRUE
    getStringType _ = TSTRING
    getListType   _ = TLIST
    getSetType    _ = TSET
    getMapType    _ = TMAP
    getStructType _ = TSTRUCT
  
    -- Generators for tokens
    genMsgBegin proxy name msgType seqNum =
      primBounded
        (liftFixedToBounded Prim.word8 >*<
         liftFixedToBounded Prim.word8 >*<
         compactI32)
        (protocolID,
         (protocolVersion .|. ((msgType `shiftL` shiftAmt) .&. typeMask),
          seqNum)) <>
      genText proxy name
  
    genStruct _ fields = mconcat fields <> B.word8 TSTOP
  
    -- NOTE: This will not work correctly for generating boolean fields. Boolean
    -- fields MUST be generated using genFieldPrim
    genField _ _ ty fid lastId val =
      primBounded (compactFieldBegin ty) (fid, lastId) <> val
  
    genFieldPrim _ _ ty fid lastId prim val =
      primBounded (compactFieldBegin ty >*< prim) ((fid, lastId), val)
  
    genFieldBool _ _ fid lastId val =
      primBounded (compactFieldBegin ty) (fid, lastId)
      where
        ty = if val then TTRUE else TFALSE
  
    genList _ ty build elems =
      genListBegin ty elems <>
      mconcat (map build elems)
  
    genListPrim _ ty bounded elems =
      genListBegin ty elems <>
      primMapListBounded bounded elems
  
    genMap _ kt vt _ kbuild vbuild elems =
      genMapBegin kt vt elems <>
      mconcat (map (\(k, v) -> kbuild k <> vbuild v) elems)
  
    genMapPrim _ kt vt _ kbuild vbuild elems =
      genMapBegin kt vt elems <>
      primMapListBounded (kbuild >*< vbuild) elems
  
    -- Generators for base types
    genBytePrim _ = liftFixedToBounded Prim.int8
    genI16Prim  _ = compactI16
    genI32Prim  _ = compactI32
    genI64Prim  _ = compactI64
    genFloat  _ = B.floatBE
    genDouble _ = B.doubleBE
    genBoolPrim _ = liftFixedToBounded $ getVal >$< Prim.word8
      where getVal b = if b then TTRUE else TFALSE
  
    genByteString _ s =
      primBounded buildVarint (W# (int2Word# len)) <> byteString s
      where
        !(I# len) = BS.length s
  
    -- Parsers for tokens
    parseMsgBegin proxy = do
      pid <- P.anyWord8
      when (pid /= protocolID) $ fail "parseMsgBegin: invalid protocol id"
      versionAndType <- P.anyWord8
      let
        version = versionAndType .&. versionMask
        msgType = (versionAndType .&. typeMask) `shiftR` shiftAmt
      when (version /= protocolVersion) $
        fail "parseMsgBegin: invalid protocol version"
      seqNum <- parseCompactI32
      name   <- parseText proxy
      return $ MsgBegin name msgType seqNum
  
    parseFieldBegin _ lastId _ = do
      byte <- anyWord8
      let rawType = byte .&. 0x0F
      if rawType == TSTOP
      then pure FieldEnd
      else do
        -- Boolean Value is encoded in the type
        let (!ty, !bool) =
              case rawType of
                0x01 -> (TTRUE, True)
                0x02 -> (TTRUE, False)
                mask -> (mask, False)
        -- Field Id depends on previous field Id
        fid <- case byte `shiftR` 4 of
                 0x00 -> parseCompactI16
                 offs -> pure $ fromIntegral offs + lastId
        return $ FieldBegin ty fid bool
  
    parseList _ p = do
      byte <- anyWord8
      let _ty = byte .&. 0x0F
      len <- case byte `shiftR` 4 of
               0x0F -> parseVarint
               size -> pure $ fromIntegral size
      ps <- replicateM (fromIntegral len) p
      return (fromIntegral len, ps)
  
    parseMap _ pk pv _ = do
      len <- parseVarint
      if len == 0
      then pure []
      else do
        byte <- anyWord8
        let _ktype = byte `shiftR` 4
            _vtype = byte .&. 0x0F
        replicateM (fromIntegral len) $ (,) <$> pk <*> pv
  
    -- Parser for base types
    parseByte   _ = fromIntegral <$> anyWord8
    parseI16    _ = parseCompactI16
    parseI32    _ = parseCompactI32
    parseI64    _ = parseCompactI64
    parseFloat  _ = binaryFloat
    parseDouble _ = binaryDouble
    parseBool   _ = do
      byte <- P.anyWord8
      case byte of
        TTRUE -> pure True
        TFALSE -> pure False
        n -> fail $ "invalid boolean value: " ++ show n
  
    parseBoolF  _ = pure
    parseByteString _ = getBuffer parseVarint BS.copy
    parseText _ = getBuffer parseVarint decodeUtf8
  
    parseSkip _ TTRUE Nothing = P.skipN 1
    parseSkip _ TFALSE Nothing = P.skipN 1
    parseSkip _ TTRUE Just{}  = pure ()
    parseSkip _ TFALSE Just{}  = pure ()
    parseSkip _ TBYTE _ = P.skipN 1
    parseSkip _ TI16 _ = void parseCompactI16
    parseSkip _ TI32 _ = void parseCompactI32
    parseSkip _ TI64 _ = void parseCompactI64
    parseSkip _ TDOUBLE _ = P.skipN 8
    parseSkip _ TSTRING _ =
      void $ parseVarint >>= P.skipN . fromIntegral
    parseSkip proxy TLIST   _ = do
      byte <- anyWord8
      let ty = byte .&. 0x0F
      len <- case byte `shiftR` 4 of
               0x0F -> parseVarint
               size -> pure $ fromIntegral size
      void $ replicateM (fromIntegral len) (parseSkip proxy ty Nothing)
  
    parseSkip proxy TSET _ = parseSkip proxy TLIST Nothing
    parseSkip proxy TMAP _ = do
      len <- parseVarint
      if len == 0
      then pure ()
      else do
        byte <- anyWord8
        let ktype = byte `shiftR` 4
            vtype = byte .&. 0x0F
        void $ replicateM (fromIntegral len) $
          (,) <$> parseSkip proxy ktype Nothing <*> parseSkip proxy vtype Nothing
  
    parseSkip proxy TSTRUCT _ = do
      -- The last id deosn't matter since we do not need to correctly parse field
      -- ids. We don't recognize this field, therefore it will be discarded anyway
      fieldBegin <- parseFieldBegin proxy 0 HashMap.empty
      case fieldBegin of
        FieldBegin ty _ bool ->
          parseSkip proxy ty (Just bool) *> parseSkip proxy TSTRUCT Nothing
        FieldEnd -> pure ()
    parseSkip _ n _ = fail $ "unrecognized type code: " ++ show n
  
  -- Helpers ---------------------------------------------------------------------
  
  {-# INLINE compactFieldBegin #-}
  compactFieldBegin :: Type -> BoundedPrim (FieldId, FieldId)
  compactFieldBegin ty = condB
    (\(fid, lastId) -> fid > lastId && fid - lastId < 16)
    ((\(fid, lastId) -> fromIntegral (fid - lastId) `shiftL` 4 .|. ty) >$<
     liftFixedToBounded Prim.word8)
    ((\(fid, _) -> (ty, fid)) >$<
     liftFixedToBounded Prim.word8 >*< compactI16)
  
  compactI16 :: BoundedPrim Int16
  compactI16 = i16ToZigZag >$< buildVarint2
  
  compactI32 :: BoundedPrim Int32
  compactI32 = i32ToZigZag >$< buildVarint4
  
  compactI64 :: BoundedPrim Int64
  compactI64 = i64ToZigZag >$< buildVarint
  
  parseCompactI16 :: Parser Int16
  parseCompactI16 = zigZagToInt <$> parseVarint
  
  parseCompactI32 :: Parser Int32
  parseCompactI32 = zigZagToInt <$> parseVarint
  
  parseCompactI64 :: Parser Int64
  parseCompactI64 = zigZagToInt <$> parseVarint
  
  genListBegin :: Type -> [a] -> Builder
  genListBegin ty elems = (`primBounded` len) $
    condB (< 15)
      ((\l -> (fromIntegral l `shiftL` 4) .|. ty) >$<
       liftFixedToBounded Prim.word8)
      ((\l -> (0xF0 .|. ty, fromIntegral l)) >$<
       (liftFixedToBounded Prim.word8 >*< buildVarint))
    where len = length elems
  
  genMapBegin :: Type -> Type -> [a] -> Builder
  genMapBegin kt vt elems
    | null elems = B.word8 0x00
    | otherwise = (`primBounded` len) $
        (\l -> (l, kt `shiftL` 4 .|. vt)) >$<
        (buildVarint >*< liftFixedToBounded Prim.word8)
      where len = genericLength elems
  
  -- Variable Length Encoded Integers --------------------------------------------
  
  -- Signed numbers must be converted to "Zig Zag" format before they can be
  -- serialized in the Varint format
  
  {-# INLINE iToZigZag #-}
  iToZigZag :: (Bits i, Integral i) => Int -> i -> Word
  iToZigZag s n = fromIntegral $ (n `shiftL` 1) `xor` (n `shiftR` s)
  
  {-# INLINE i16ToZigZag #-}
  i16ToZigZag :: Int16 -> Word
  i16ToZigZag n = 0xFFFF .&. iToZigZag 15 n
  
  {-# INLINE i32ToZigZag #-}
  i32ToZigZag :: Int32 -> Word
  i32ToZigZag n = 0xFFFFFFFF .&. iToZigZag 31 n
  
  {-# INLINE i64ToZigZag #-}
  i64ToZigZag :: Int64 -> Word
  i64ToZigZag = iToZigZag 63
  
  -- This is janky, but BoundedPrim requires the size of the value it is building
  -- to be bounded. Varints *are* bounded, but the output size depends on their
  -- value. A 64-bit value can take up to 80 bits to encode, because each 7-bits
  -- in the input map to 8 bits in the output. We therefore need 10 builders, each
  -- one decreasing in potential maximum size.
  buildVarint :: BoundedPrim Word
  buildVarint =  buildVarintBase buildVarint8
  
  buildVarint8 :: BoundedPrim Word
  buildVarint8 = buildVarintBase buildVarint7
  
  buildVarint7 :: BoundedPrim Word
  buildVarint7 = buildVarintBase buildVarint6
  
  buildVarint6 :: BoundedPrim Word
  buildVarint6 = buildVarintBase buildVarint5
  
  buildVarint5 :: BoundedPrim Word
  buildVarint5 = buildVarintBase buildVarint4
  
  buildVarint4 :: BoundedPrim Word
  buildVarint4 = buildVarintBase buildVarint3
  
  buildVarint3 :: BoundedPrim Word
  buildVarint3 = buildVarintBase buildVarint2
  
  buildVarint2 :: BoundedPrim Word
  buildVarint2 = buildVarintBase buildVarint1
  
  buildVarint1 :: BoundedPrim Word
  buildVarint1 = buildVarintBase buildVarint0
  
  buildVarint0 :: BoundedPrim Word
  buildVarint0 = fromIntegral >$< liftFixedToBounded Prim.word8
  
  {-# INLINE buildVarintBase #-}
  buildVarintBase :: BoundedPrim Word -> BoundedPrim Word
  buildVarintBase base = condB (\n -> n .&. complement 0x7F == 0)
    buildVarint0
    ((\(W# n) ->
       (mkW8# (0x80## `or#` (n `and#` 0x7f##)), W# (n `uncheckedShiftRL#` 7#))) >$<
     (liftFixedToBounded Prim.word8 >*< base))
    where
  #if __GLASGOW_HASKELL__ >= 902
    mkW8# x = W8# (wordToWord8# x)
  #else
    mkW8# = W8#
  #endif
  
  {-# INLINE zigZagToInt #-}
  zigZagToInt :: (Bits i, Integral i) => Word -> i
  zigZagToInt w = fromIntegral (w `shiftR` 1) `xor` negate (fromIntegral w .&. 1)
  
  {-# INLINE parseVarint #-}
  parseVarint :: Parser Word
  parseVarint = go 0 0
    where
      go !val !n = do
        w <- anyWord8
        let newVal = val .|. (fromIntegral w .&. 0x7F) `shiftL` n
        if not (testBit w 7)
          then return newVal
          else go newVal (n + 7)
  