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 ConstraintKinds, CPP #-}
  module Thrift.Protocol
    ( ThriftSerializable, serializeGen, deserializeGen
    , ThriftStruct(..), ProtocolException(..)
    , Protocol(..), parseStop, parseEnum, parseEnumNoUnknown
    , FieldId, Length, Type, MsgType, SeqNum
    , FieldBegin(..), MsgBegin(..)
    , genByte, genI16, genI32, genI64, genBool
    , ThriftEnum(..)
    , unknownThriftEnumErrorMsg
    ) where
  
  import Control.Exception
  import Data.Aeson hiding (Bool, String)
  import Thrift.Binary.Parser
  import Data.ByteString (ByteString)
  import Data.ByteString.Builder
  import Data.ByteString.Builder.Prim
  import Data.HashMap.Strict as HashMap
  import Data.Int
  import Data.Proxy
  import Data.Text (Text)
  import Data.Text.Encoding
  import Data.Typeable
  import Data.Word
  import qualified Data.ByteString.Lazy as LBS
  
  #if __GLASGOW_HASKELL__ == 806
  import Prelude hiding (fail)
  import Control.Monad.Fail (fail)
  #endif
  
  
  -- Constraint Kind that every Thrift generated datatype must satisfy
  type ThriftSerializable a = (ToJSON a, ThriftStruct a)
  
  class ThriftStruct a where
    buildStruct :: Protocol p => Proxy p -> a -> Builder
    parseStruct :: Protocol p => Proxy p -> Parser a
  
  class ThriftEnum a where
    toThriftEnum        :: Int -> a
    fromThriftEnum      :: a -> Int
    allThriftEnumValues :: [a]
    toThriftEnumEither  :: Int -> Either String a
  
  instance ThriftEnum Bool where
    toThriftEnum n | n == 0    = False
                   | n == 1    = True
                   | otherwise = errorWithoutStackTrace
                     "Thrift.Protocol.ThriftEnum.Bool.toThriftEnum: bad argument"
    fromThriftEnum False = 0
    fromThriftEnum True  = 1
    allThriftEnumValues = [False, True]
    toThriftEnumEither n | n == 0    = Right False
                         | n == 1    = Right True
                         | otherwise = Left "Thrift.Protocol.ThriftEnum.Bool.toThriftEnumEither: bad argument"
  
  {-# INLINE serializeGen #-}
  serializeGen
    :: (Protocol p, ThriftStruct a)
    => Proxy p
    -> a
    -> ByteString
  serializeGen proxy = LBS.toStrict . toLazyByteString . buildStruct proxy
  
  {-# INLINE deserializeGen #-}
  deserializeGen
    :: (Protocol p, ThriftStruct a)
    => Proxy p
    -> ByteString
    -> Either String a
  deserializeGen proxy input = parse (parseStruct proxy) input
  
  newtype ProtocolException = ProtocolException String
    deriving (Show, Eq)
  instance Exception ProtocolException
  
  -- Binary Protocol type class --------------------------------------------------
  
  type FieldId = Int16
  type Length  = Word32
  type Type    = Word8
  type IdMap   = HashMap Text FieldId
  
  type MsgType = Word8
  type SeqNum  = Int32
  
  data FieldBegin
    = FieldBegin {-# UNPACK #-} !Type {-# UNPACK #-} !FieldId !Bool
    | FieldEnd
  
  data MsgBegin
    = MsgBegin !Text {-# UNPACK #-} !MsgType {-# UNPACK #-} !SeqNum
      deriving (Show, Eq)
  
  class Protocol p where
    -- Generators for Types
    getByteType   :: Proxy p -> Type
    getI16Type    :: Proxy p -> Type
    getI32Type    :: Proxy p -> Type
    getI64Type    :: Proxy p -> Type
    getFloatType  :: Proxy p -> Type
    getDoubleType :: Proxy p -> Type
    getBoolType   :: Proxy p -> Type
    getStringType :: Proxy p -> Type
    getListType   :: Proxy p -> Type
    getSetType    :: Proxy p -> Type
    getMapType    :: Proxy p -> Type
    getStructType :: Proxy p -> Type
  
    -- Generators for tokens
    genMsgBegin   :: Proxy p -> Text -> MsgType -> SeqNum -> Builder
  
    genStruct  :: Proxy p -> [Builder] -> Builder
  
    genField
      :: Proxy p -> Text -> Type -> FieldId -> FieldId -> Builder -> Builder
    genFieldPrim
      :: Proxy p -> Text -> Type -> FieldId -> FieldId
      -> BoundedPrim a -> a -> Builder
  
    -- Only Compact protocol needs a custom implementation for this
    genFieldBool
      :: Proxy p -> Text -> FieldId -> FieldId -> Bool -> Builder
    genFieldBool proxy name fid lastId =
      genFieldPrim proxy name (getBoolType proxy) fid lastId (genBoolPrim proxy)
  
    genList :: Proxy p -> Type -> (a -> Builder) -> [a] -> Builder
    genListPrim :: Proxy p -> Type -> BoundedPrim a -> [a] -> Builder
  
    genMap  :: Proxy p -> Type -> Type -> Bool
            -> (k -> Builder) -> (v -> Builder)
            -> [(k, v)] -> Builder
    genMapPrim :: Proxy p -> Type -> Type -> Bool
               -> BoundedPrim k -> BoundedPrim v
               -> [(k, v)] -> Builder
  
    genMsgEnd :: Proxy p -> Builder
    genMsgEnd _ = mempty
  
    -- Primitive Generators (for base types)
    genBytePrim   :: Proxy p -> BoundedPrim Int8
    genI16Prim    :: Proxy p -> BoundedPrim Int16
    genI32Prim    :: Proxy p -> BoundedPrim Int32
    genI64Prim    :: Proxy p -> BoundedPrim Int64
    genBoolPrim   :: Proxy p -> BoundedPrim Bool
  
  
    -- Generators for string types (implement 1)
    genText   :: Proxy p -> Text -> Builder
    genText p = genByteString p . encodeUtf8
    genByteString :: Proxy p -> ByteString -> Builder
    genByteString p = genText p . decodeUtf8
  
    -- Generator for binary
    genBytes :: Proxy p -> ByteString -> Builder
    genBytes = genByteString
  
    -- Generators for other base types
    genFloat  :: Proxy p -> Float  -> Builder
    genDouble :: Proxy p -> Double -> Builder
  
    -- Parsers for tokens
    parseMsgBegin   :: Proxy p -> Parser MsgBegin
    parseMsgEnd     :: Proxy p -> Parser ()
    parseMsgEnd _ = return ()
  
    parseFieldBegin :: Proxy p -> FieldId -> IdMap -> Parser FieldBegin
    parseList :: Proxy p -> Parser a -> Parser (Int,[a])
    parseMap  :: Proxy p -> Parser k -> Parser v -> Bool -> Parser [(k, v)]
  
    -- Parser for base types
    parseByte   :: Proxy p -> Parser Int8
    parseI16    :: Proxy p -> Parser Int16
    parseI32    :: Proxy p -> Parser Int32
    parseI64    :: Proxy p -> Parser Int64
    parseFloat  :: Proxy p -> Parser Float
    parseDouble :: Proxy p -> Parser Double
    parseBool   :: Proxy p -> Parser Bool
    parseBoolF  :: Proxy p -> Bool -> Parser Bool
  
    -- Parsers for strings
    parseText :: Proxy p -> Parser Text
    parseText p = decodeUtf8 <$> parseByteString p
    parseByteString :: Proxy p -> Parser ByteString
    parseByteString p = encodeUtf8 <$> parseText p
  
    -- Parser for binary
    parseBytes :: Proxy p -> Parser ByteString
    parseBytes = parseByteString
  
    parseSkip :: Proxy p -> Type -> Maybe Bool -> Parser ()
  
  genByte   :: Protocol p => Proxy p -> Int8   -> Builder
  genByte p = primBounded (genBytePrim p)
  
  genI16    :: Protocol p => Proxy p -> Int16  -> Builder
  genI16 p = primBounded (genI16Prim p)
  
  genI32    :: Protocol p => Proxy p -> Int32  -> Builder
  genI32 p = primBounded (genI32Prim p)
  
  genI64    :: Protocol p => Proxy p -> Int64  -> Builder
  genI64 p = primBounded (genI64Prim p)
  
  genBool   :: Protocol p => Proxy p -> Bool   -> Builder
  genBool p = primBounded (genBoolPrim p)
  
  parseStop :: Protocol p => Proxy p -> Parser ()
  parseStop token = do
    end <- parseFieldBegin token 0 HashMap.empty
    case end of
      FieldBegin{} -> fail "parseStop: expected end of struct"
      FieldEnd -> return ()
  
  parseEnum
    :: (Protocol p, ThriftEnum a)
    => Proxy p
    -> String
    -> Parser a
  parseEnum proxy _ = do
    n <- parseI32 proxy
    return $ toThriftEnum (fromIntegral n)
  
  parseEnumNoUnknown
    :: (Protocol p, ThriftEnum a)
    => Proxy p
    -> String
    -> Parser a
  parseEnumNoUnknown proxy name = do
    n <- parseI32 proxy
    case toThriftEnumEither (fromIntegral n) of
      Left _ -> fail $
        "parseEnum: not a valid identifier for thrift enum '" ++ name ++ "': "
          ++ show n
      Right enum -> return enum
  
  unknownThriftEnumErrorMsg :: (Typeable a, Show a) => a -> String
  unknownThriftEnumErrorMsg x = "Unknown value "
    ++ show x
    ++ " for enum "
    ++ show (tyConModule $ typeRepTyCon $ typeOf x)
    ++ "."
    ++ show (typeOf x)
  