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.
  -}
  
  module Glean.Query.JSON
    ( factToJSON
    , factToCompact
    ) where
  
  import Control.Exception
  import Control.Monad
  import Control.Monad.ST (ST, stToIO)
  import Control.Monad.ST.Unsafe (unsafeIOToST)
  import Data.Bits
  import Data.ByteString (ByteString)
  import qualified Data.ByteString.Builder as Builder
  import qualified Data.ByteString.Lazy as LBS
  import Data.IntMap (IntMap)
  import qualified Data.IntMap as IntMap
  import Data.Maybe
  import Data.Text (Text)
  import qualified Data.Text.Encoding as Text
  import Data.Word (Word8, Word64)
  import Foreign.Marshal.Utils (copyBytes)
  import Safe (atMay)
  
  import Thrift.Protocol.JSON.Base64
  import Util.Buffer (ascii, liftST)
  import qualified Util.Buffer as Buffer
  
  import Glean.RTS as RTS
  import qualified Glean.RTS.Foreign.JSON as RTS hiding (String, Array)
  import qualified Glean.RTS.Foreign.Thrift as Thrift
  import Glean.RTS.Types (Type, ExpandedType(..))
  import Glean.Database.Schema
  import Glean.Angle.Types as Schema hiding (Type)
  
  newtype EncodingError = EncodingError String
    deriving(Show)
  
  instance Exception EncodingError
  
  encodingError :: String -> ST s a
  encodingError = unsafeIOToST . throwIO . EncodingError
  
  -- | Convert a fact to JSON
  factToJSON
    :: Bool
    -> IntMap ByteString
    -> PredicateDetails
    -> Fid
    -> ByteString
    -> ByteString
    -> IO ByteString
  factToJSON no_base64 expanded pred fid key value =
    encodeFact expanded pred fid key value (jsonEncoder no_base64)
  
  -- | Convert a fact to Thrift Compact
  factToCompact
    :: IntMap ByteString
    -> PredicateDetails
    -> Fid
    -> ByteString
    -> ByteString
    -> IO ByteString
  factToCompact expanded pred fid key value =
    encodeFact expanded pred fid key value compactEncoder
  
  type Enc = forall s . Buffer.Fill s ()
  
  -- | Specifies how facts should be encoded
  data Encoder = Encoder
    { -- | Bytes
      encByte :: Word8 -> Enc
  
      -- | Numbers
    , encNat :: Word64 -> Enc
  
      -- | Booleans
    , encBool :: Bool -> Enc
  
      -- | Mangled RTS strings
    , encMangledString :: RTS.MangledStringRef -> Int -> Enc
  
      -- | Binary
    , encBinary :: RTS.ByteStringRef -> Enc
  
      -- | Start an array
    , encArrayBegin :: Type -> Int -> Enc
  
      -- | Separate two array elements
    , encArraySep :: Enc
  
      -- | Finish an array
    , encArrayEnd :: Enc
  
      -- | Start a set
    , encSetBegin :: Type -> Int -> Enc
  
      -- | Separate two set elements
    , encSetSep :: Enc
  
      -- | Finish a set
    , encSetEnd :: Enc
  
      -- | Start an object
    , encObjectBegin :: Enc
  
      -- | Separate two object fields
    , encObjectSep :: Enc
  
      -- | Object field name/number - Thrift folds boolean values into the field
      -- header so pass in the 'Decoder' and an action to decode and write the
      -- value into the stream and let the encoder decide what to do with it.
    , encObjectField
        :: forall s.
        Maybe Int  -- previous field id (if any)
        -> Int  -- current field id
        -> Text  -- field name
        -> Type  -- field type
        -> Decoder s  -- Thrift folds bool fields into the header so pass in
        -> Buffer.Fill s ()  -- a Decoder and an action to write the field value
  
        -> Buffer.Fill s ()
  
      -- | Finish an object
    , encObjectEnd :: Enc
    }
  
  -- | Encode a fact with the given 'Encoder'
  encodeFact
    :: IntMap ByteString
    -> PredicateDetails
    -> Fid
    -> ByteString
    -> ByteString
    -> Encoder
    -> IO ByteString
  {-# INLINE encodeFact #-}
  encodeFact
    expanded
    PredicateDetails{..}
    fid
    key
    value
    encoder@Encoder{..} = stToIO $
      RTS.withDecoder key $ \dkey ->
      RTS.withDecoder value $ \dvalue ->
      Buffer.fillByteString 8 $ do
        -- NOTE: We hardcode the field numbers for id, key and value here.
        encObjectBegin
        encObjectField Nothing 1 "id" NatTy dkey
          $ encNat $ fromIntegral $ fromFid fid
        encObjectSep
        encObjectField (Just 1) 2 "key" predicateKeyType dkey
          $ encode expanded encoder dkey predicateKeyType
        case predicateValueType of
          Schema.RecordTy [] -> return () -- omit "value" field if value is unit
          _other -> do
            encObjectSep
            encObjectField (Just 2) 3 "value"
                predicateValueType dvalue
              $ encode expanded encoder dvalue predicateValueType
        encObjectEnd
  
  -- | Takes a Glean schema type and a value of that type, and fills a 'Buffer'
  -- with a blob that can be decoded into the appropriate native Client Type
  -- corresponding to the generated Thrift schema. The encoding (currently only
  -- JSON) is specified by the 'Encoder'.
  --
  -- NOTE: This is INLINE because we want a specialised version for each encoding
  -- that we support.
  encode
    :: IntMap ByteString
    -> Encoder
    -> RTS.Decoder s
    -> Type
    -> Buffer.Fill s ()
  {-# INLINE encode #-}
  encode expanded Encoder{..} !d = enc
    where
      enc typ = case typ of
        ByteTy -> do
          x <- liftST $ RTS.dByte d
          encByte x
        NatTy{} -> do
          x <- liftST $ RTS.dNat d
          encNat x
        StringTy{} -> do
          (x,n) <- liftST $ RTS.dTrustedStringRef d
          encMangledString x $ fromIntegral n
        ArrayTy elty -> do
          size <- liftST $ RTS.dArray d
          case elty of
            ByteTy{} -> do
              ref <- liftST $ RTS.dByteStringRef d size
              encBinary ref
            _ -> do
              encArrayBegin elty $ fromIntegral size
              case size of
                0 -> return ()
                1 -> enc elty
                _ -> do
                  enc elty
                  let go 0 = return ()
                      go n = do
                        encArraySep
                        enc elty
                        go (n-1)
                  go (size-1)
              encArrayEnd
        RecordTy fieldTys -> do
          encObjectBegin
          -- We need to carry both the current and previous field numbers as
          -- Thrift encodes deltas between those.
          let field prev (i, FieldDef name ety) = case ety of
                MaybeTy ety -> do
                  sel <- liftST $ RTS.dSelector d
                  if sel == 0
                    then return prev
                    else enc_field prev i name ety d
                ty -> enc_field prev i name ty d
          foldM_ field Nothing $ zip [1..] fieldTys
          encObjectEnd
        SumTy fieldTys -> do
          sel <- liftST $ RTS.dSelector d
          encObjectBegin
          void $ case atMay fieldTys (fromIntegral sel) of
            Just (FieldDef name ty) -> do
              enc_field Nothing (fromIntegral sel + 1) name ty d
            Nothing ->
              enc_field Nothing (length fieldTys + 1) "UNKNOWN" (RecordTy []) d
          encObjectEnd
        SetTy elty -> do
          size <- liftST $ RTS.dSet d
          case elty of
            ByteTy{} -> do
              ref <- liftST $ RTS.dByteStringRef d size
              encBinary ref
            _ -> do
              encSetBegin elty $ fromIntegral size
              case size of
                0 -> return ()
                1 -> enc elty
                _ -> do
                  enc elty
                  let go 0 = return ()
                      go n = do
                        encSetSep
                        enc elty
                        go (n-1)
                  go (size-1)
              encSetEnd
        NamedTy _ (ExpandedType _ ty) -> enc ty
        EnumeratedTy _ -> do
          x <- liftST $ RTS.dSelector d
          encNat $ fromIntegral x
        MaybeTy{} -> liftST $ encodingError "unsupported Maybe"
        BooleanTy{} -> do
          x <- liftST $ RTS.dSelector d
          encBool (x /= 0)
        PredicateTy{} -> do
          fid <- liftST $ RTS.dFact d
          let id = fromIntegral (fromFid fid)
          case IntMap.lookup id expanded of
            Just bs ->  Buffer.byteString bs
            Nothing -> do
              encObjectBegin
              encObjectField Nothing 1 "id" NatTy d
                $ encNat $ fromIntegral id
              encObjectEnd
        TyVar{} -> error "JSON.encode: TyVar"
        HasTy{} -> error "JSON.encode: HasTy"
        HasKey{} -> error "JSON.encode: HasKey"
        ElementsOf{} -> error "JSON.encode: ElementsOf"
  
      {-# INLINE enc_field #-}
      enc_field prev i name ty d = do
        when (isJust prev) encObjectSep
        encObjectField prev i name ty d $ enc ty
        return $ Just i
  
  jsonEncoder :: Bool -> Encoder
  jsonEncoder no_base64 = Encoder
    { encByte = number . fromIntegral
    , encNat = number
    , encBool = \b -> if b
        then do
          ascii 't'
          ascii 'r'
          ascii 'u'
          ascii 'e'
        else do
          ascii 'f'
          ascii 'a'
          ascii 'l'
          ascii 's'
          ascii 'e'
    , encMangledString = \ref _ -> do
        n <- liftST $ RTS.getMangledStringEscapedSize ref
        do
          ascii '\"'
          Buffer.alloc n $ \p -> do
            RTS.escapeMangledString ref p n
            return n
          ascii '\"'
    , encBinary = \ !ref ->
        if no_base64
          then do
            n <- liftST $ RTS.getStringEscapedSize ref
            do
              ascii '\"'
              Buffer.alloc n $ \p -> do
                RTS.escapeString ref p n
                return n
              ascii '\"'
          else do
            -- TODO: can be optimised if necessary
            bytes <- liftST $ unsafeIOToST $ RTS.derefByteString ref
            ascii '\"'
            Buffer.byteString
              $ LBS.toStrict $ Builder.toLazyByteString $ encodeBase64 bytes
            ascii '\"'
    , encArrayBegin = \_ _ -> ascii '['
    , encArraySep = ascii ','
    , encArrayEnd = ascii ']'
    -- Sets are encoded as ordered arrays.
    , encSetBegin = \_ _ -> ascii '['
    , encSetSep = ascii ','
    , encSetEnd = ascii ']'
    , encObjectBegin = ascii '{'
    , encObjectSep = ascii ','
    , encObjectField = \_ _ !name _ _ enc -> do
        ascii '\"'
        Buffer.byteString (Text.encodeUtf8 name)
        ascii '\"'
        ascii ':'
        enc
    , encObjectEnd = ascii '}'
    }
    where
      number :: Word64 -> Enc
      number !n = Buffer.alloc 21 $ RTS.encodeNumber $ fromIntegral n
        -- max length of 64 bit number is sign + 20 digits
  
  type ThriftType = Word8
  
  -- | Compute the Thrift type for a 'Type'
  -- See https://github.com/apache/thrift/blob/master/doc/specs/thrift-compact-protocol.md#struct-encoding
  thriftType :: Type -> ThriftType
  thriftType ByteTy{} = 3
  thriftType NatTy{} = 6
  thriftType StringTy{} = 8
  thriftType (ArrayTy ByteTy{}) = 8
  thriftType ArrayTy{} = 9
  thriftType (SetTy ByteTy{}) = 8
  thriftType SetTy{} = 9
  thriftType RecordTy{} = 12
  thriftType SumTy{} = 12
  thriftType (NamedTy _ (ExpandedType _ ty)) = thriftType ty
  thriftType EnumeratedTy{} = 5 -- spec says enums are i32
  thriftType BooleanTy{} = 1
  thriftType PredicateTy{} = 12
  thriftType typ =
    throw $ EncodingError $ "thrifTType: invalid type " ++ show typ
  
  
  compactEncoder :: Encoder
  compactEncoder = Encoder
    { encByte = Buffer.byte
    , encNat = Thrift.encodeZigZag . fromIntegral
    , encBool = Buffer.byte . bool
    , encMangledString = \ref n -> do
        -- length
        Thrift.encodeVarint (fromIntegral n)
        -- bytes
        Buffer.alloc (fromIntegral n) $ \p -> do
          m <- unsafeIOToST (RTS.demangle ref p)
          assert (m == fromIntegral n) $ return (fromIntegral n)
    , encBinary = \(RTS.ByteStringRef s n) -> do
        -- length
        Thrift.encodeVarint (fromIntegral n)
        -- bytes
        Buffer.alloc (fromIntegral n) $ \p -> do
          -- s might be NULL if the range is empty so check for that
          when (n /= 0) $ unsafeIOToST $ copyBytes p s $ fromIntegral n
          return (fromIntegral n)
    , encArrayBegin = encArray
    , encArraySep = return ()
    , encArrayEnd = return ()
    , encSetBegin = encArray
    , encSetSep = return ()
    , encSetEnd = return ()
    , encObjectBegin = return ()
    , encObjectSep = return ()
    , encObjectField = \prev i _ typ d enc -> do
        -- field types - different from array types
        !ety <- case thriftType typ of
          1 -> do
            -- bool values folded into the header
            sel <- liftST $ dSelector d
            return $ bool $ sel /= 0
          ety -> return ety
        let delta = i - fromMaybe 0 prev
        if delta /= 0 && delta <= 15
          -- small field number deltas folded into the field header
          then Buffer.byte $ ((fromIntegral delta .&. 0x0F) `shiftL` 4) .|. ety
          else do
            -- field number delta 0 indicates that the real field number follows
            Buffer.byte ety
            Thrift.encodeZigZag $ fromIntegral i
        when (ety > 2) enc
      -- object terminator
    , encObjectEnd = Buffer.byte 0
    }
    where
      bool True = 1
      bool False = 2
      encArray typ n = do
        -- array types
        let !ety = thriftType typ
        if n < 15
          -- small lengths are folded into the header byte
          then Buffer.byte $ ((fromIntegral n .&. 0x0F) `shiftL` 4) .|. ety
          else do
            Buffer.byte $ 0xF0 .|. ety
            Thrift.encodeVarint $ fromIntegral n
  