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 TypeOperators #-}
  module Thrift.Compiler.Plugin
    ( Typecheckable(..), SomeLit(..)
    , qualify, qualifyType
    , getPrefix
    , lowercase, uppercase, toCamel
    , fixLeadingUnderscore
    , toConstructorName
    , getNamespace
    , isNewtype
    , filterHsAnns, getTypeAnns
    ) where
  
  import Data.Bifunctor
  import Data.List
  import qualified Data.Map.Strict as Map
  import Data.Maybe
  import Data.Some
  import Data.Text (Text)
  import qualified Data.Text as Text
  import Data.Type.Equality
  import Prelude hiding (Enum)
  
  import Thrift.Compiler.Options
  import Thrift.Compiler.Parser
  import Thrift.Compiler.Typechecker.Monad
  import Thrift.Compiler.Types
  
  data SomeLit l = forall t. ThisLit (Type l t) t
  
  -- | Typeclass to define Thrift typecheckable Languages. There are several
  -- sections of functions to define here, although many have default
  -- implementations
  class Monoid (Interface l) => Typecheckable l  where
    -- | Interface that gets generated by the Thrift compiler. This includes all
    -- the symbols exported by the generated code in the target language. It is
    -- used in order to add additional required symbols that get referenced in
    -- splice (hs_include) files. If your language does not support splice files,
    -- then set Interface l = ()
    type Interface l
  
    -- * Annotation Processing
  
    -- | Given a resolved Thrift type and a list of annotations, produce some new
    -- transformed type. This is how TSpecial types get generated
    resolveTypeAnnotations
      :: Type l t
      -> [Annotation Loc]
      -> TC l (Some (Type l))
    resolveTypeAnnotations ty _ = pure $ Some ty
  
    -- | Recursively qualify all of the named types in a SpecialType so that they
    -- can be properly identified in imports
    qualifySpecialType :: (Text, Text) -> SpecialType l t -> SpecialType l t
  
    -- * Typechecking
  
    -- | Typecheck literals of special, language specific types
    -- Note: this is *only* for literals, named constants are typechecked by the
    -- main typecheckConst function (even for special types)
    typecheckSpecialConst
      :: SpecialType l t
      -> UntypedConst Loc
      -> TC l (TypedConst l t)
  
    -- | Type equality for special types
    eqSpecial :: SpecialType l u -> SpecialType l v -> Maybe (u :~: v)
  
    -- | Generate an interface. See the Interface type family for more info on how
    -- this is used.
    getInterface :: Options l -> ThriftFile a Loc -> Interface l
    getInterface _ _ = mempty
  
    -- | Get extra Thrift symbols that are referenced in splice files
    getExtraSymbols ::
      Options l -> Interface l -> ThriftFile SpliceFile b -> [Text]
    getExtraSymbols _ _ _ = []
  
    -- * Renamers
  
    -- | Rename the thrift module, probably taking namespace into account
    renameModule :: Options l -> ThriftFile a b -> Text
    renameModule _ ThriftFile{..} = thriftName
  
    renameStruct :: Options l -> Struct s u a -> Text
    renameStruct _ Struct{..} = structName
  
    renameField
      :: Options l
      -> [Annotation a]
      -> Text
      -> Field u s m a
      -> Text
    renameField _ _ _ Field{..} = fieldName
  
    renameConst :: Options l -> Text -> Text
    renameConst _ = id
  
    renameService :: Options l -> Service s u a -> Text
    renameService _ Service{..} = serviceName
  
    renameFunction :: Options l -> Function s u a -> Text
    renameFunction _ Function{..} = funName
  
    renameTypedef :: Options l -> Typedef s u a -> Text
    renameTypedef _ Typedef{..} = tdName
  
    renameEnum :: Options l -> Enum s u a -> Text
    renameEnum _ Enum{..} = enumName
  
    renameEnumAlt :: Options l -> Enum s u a -> Text -> Text
    renameEnumAlt _ _ name = name
  
    renameUnion :: Options l -> Union s u a -> Text
    renameUnion _ Union{..} = unionName
  
    renameUnionAlt :: Options l -> Union s u a -> UnionAlt s u a -> Text
    renameUnionAlt _ _ UnionAlt{..} = altName
  
    getUnionEmptyName :: Options l -> Union s u a -> Text
    getUnionEmptyName _ Union{..} = unionName <> "_EMPTY"
  
    -- * Uniqueness options specify whether various renamed symbols need to be
    -- globally unique to the entire Thrift (if False, they are still unique
    -- within their structure)
  
    -- | Are field names globally unique? If yes, then we will report renamed
    -- field name collisions across structs
    fieldsAreUnique :: Options l -> Bool
    fieldsAreUnique _ = False
  
    -- | Are union alternatives unique?
    unionAltsAreUnique :: Options l -> Bool
    unionAltsAreUnique _ = False
  
    -- | Are enum values unique?
    enumAltsAreUnique :: Options l -> Bool
    enumAltsAreUnique _ = False
  
    -- | Get enum flavour from annotation tags
    enumFlavourTag :: Options l -> Enum s u a -> EnumFlavour
    enumFlavourTag _ _ = SumTypeEnum False
  
    -- * Back translators
    -- Translate Stuff Back to regular thrift for pretty printing and JSON output
  
    -- | Translate a special type back to its underlying Thrift type, and its
    -- special name in the target language.
    -- Note that this should be a shallow conversion. Do not backtranslate
    -- recursively
    backTranslateType :: SpecialType l t -> (Some (Type l), Text)
  
    -- | Translate a literal to its underlying thrift type and representation.
    -- Like @backTranslateType@, the conversion should be shallow
    backTranslateLiteral :: SpecialType l t -> t -> SomeLit l
  
  -- Env Qualifier ---------------------------------------------------------------
  
  -- | Qualify all the named types in an Env
  qualify :: Typecheckable l => (Text, Text) -> Env l -> Env l
  qualify m env = env
    { typeMap  = qualCtx qualST (typeMap env)
    , schemaMap = qualSchemas (schemaMap env)
    , unionMap = qualSchemas (unionMap env)
    , constMap = qualCtx qualC (constMap env)
    , enumMap  = qualEnums (enumMap env)
    , serviceMap = qualServices (serviceMap env)
    }
    where
      qualCtx
        :: (a -> a)
        -> Context a
        -> Context a
      qualCtx f ctx@Context{..} = ctx { cMap = Map.map f cMap }
      qualST :: Typecheckable l => Some (Type l) -> Some (Type l)
      qualST (Some ty) = Some $ qualifyType m ty
      qualC :: Typecheckable l => (Some (Type l), Name, Loc)
                               -> (Some (Type l), Name, Loc)
      qualC (st, n, loc) = (qualST st, qualName m n, loc)
      qualSchemas
        :: Typecheckable l
        => Map.Map Text (Some (SCHEMA l t)) -> Map.Map Text (Some (SCHEMA l t))
      qualSchemas = Map.map $ \(Some schema) -> Some (qualifySchema m schema)
      qualEnums = Map.map $ \(vs, ns) ->
        (Map.map (first (qualName m)) vs, Map.map (first (qualName m)) ns)
      qualServices = Map.map $ \(a, b, c) -> (qualName m a, b, c)
  
  qualifyType :: Typecheckable l => (Text, Text) -> Type l t -> Type l t
  qualifyType _ I8 = I8
  qualifyType _ I16 = I16
  qualifyType _ I32 = I32
  qualifyType _ I64 = I64
  qualifyType _ TFloat = TFloat
  qualifyType _ TDouble = TDouble
  qualifyType _ TBool = TBool
  qualifyType _ TBytes = TBytes
  qualifyType _ TText = TText
  qualifyType m (TSet u)       = TSet $ qualifyType m u
  qualifyType m (THashSet u)   = THashSet $ qualifyType m u
  qualifyType m (TList u)      = TList $ qualifyType m u
  qualifyType m (TMap k v)     = TMap (qualifyType m k) (qualifyType m v)
  qualifyType m (THashMap k v) = THashMap (qualifyType m k) (qualifyType m v)
  qualifyType m (TTypedef name t loc) =
    TTypedef (qualName m name) (qualifyType m t) loc
  qualifyType m (TStruct name loc)    = TStruct (qualName m name) loc
  qualifyType m (TException name loc) = TException (qualName m name) loc
  qualifyType m (TUnion name loc)     = TUnion (qualName m name) loc
  qualifyType m (TEnum name loc nounknown) =
    TEnum (qualName m name) loc nounknown
  qualifyType m (TNewtype name t loc) =
    TNewtype (qualName m name) (qualifyType m t) loc
  qualifyType m (TSpecial ty) = TSpecial $ qualifySpecialType m ty
  
  qualifySchema :: Typecheckable l => (Text, Text) -> SCHEMA l t s -> SCHEMA l t s
  qualifySchema _ SEmpty = SEmpty
  qualifySchema m (SField p n t s) =
    SField p n (qualifyType m t) (qualifySchema m s)
  qualifySchema m (SReqField p n t s) =
    SReqField p n (qualifyType m t) (qualifySchema m s)
  qualifySchema m (SOptField p n t s) =
    SOptField p n (qualifyType m t) (qualifySchema m s)
  
  qualName :: (Text, Text) -> Name -> Name
  qualName (tm, rm) Name{..} = Name
    { sourceName = qualName_ tm sourceName
    , resolvedName = qualName_ rm resolvedName
    }
  
  qualName_ :: Text -> Name_ s -> Name_ s
  qualName_ m (UName n) = QName m n
  qualName_ _ n@QName{} = n
  
  -- Renamer Helpers -------------------------------------------------------------
  
  getPrefix :: [Annotation a] -> Maybe Text
  getPrefix anns = listToMaybe
    [ pre
    | ValueAnn{vaVal=TextAnn pre _,..} <- anns
    , vaTag == "hs.prefix"
    ]
  
  lowercase :: Text -> Text
  lowercase = uncurry (<>) . first Text.toLower . Text.splitAt 1
  
  uppercase :: Text -> Text
  uppercase = uncurry (<>) . first Text.toUpper . Text.splitAt 1
  
  toCamel :: Text -> Text
  toCamel = Text.concat . map capitalize . Text.splitOn "_"
    where capitalize = uncurry (<>) . first Text.toUpper . Text.splitAt 1
  
  -- Prepend "TU" for "ThriftUnderscore"
  fixLeadingUnderscore :: Text -> Text
  fixLeadingUnderscore = uncurry (<>) . first fixUnderscore . Text.splitAt 1
    where
      fixUnderscore s = if s == "_" then "TU__" else s
  
  toConstructorName :: Text -> Text
  toConstructorName = fixLeadingUnderscore . uppercase
  
  -- | Select the last namespace header
  getNamespace :: Text -> [Header s l a] -> Maybe Text
  getNamespace l = foldl' getNS Nothing
    where
      getNS ns HInclude{} = ns
      getNS ns HPackage{} = ns
      getNS ns HNamespace{..}
        | nmLang == l = Just nmName
        | otherwise   = ns
  
  isNewtype :: [Annotation a] -> Bool
  isNewtype = any isNewtypeAnn
    where
      isNewtypeAnn SimpleAnn{..} = saTag == "hs.newtype"
      isNewtypeAnn _ = False
  
  filterHsAnns :: [Annotation a] -> [Annotation a]
  filterHsAnns = filter $ Text.isPrefixOf "hs." . \case
    SimpleAnn{..} -> saTag
    ValueAnn{..}  -> vaTag
  
  getTypeAnns :: Text -> [Annotation a] -> [(Text, Annotation a)]
  getTypeAnns lang anns =
    [ (ty, a) | a@ValueAnn{vaVal=TextAnn ty _,..} <- anns, vaTag == typeTag ]
    where
      typeTag = lang <> ".type"
  