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 Thrift.Compiler.Plugins.Haskell
    ( Haskell, HSType, HS
    , SpecialType(..)
    , HsVectorKind(..), hsVectorImport, hsVectorQual
    , HsInterface(..), RenameMap
    , LangOpts(..), defaultHsOpts
    , toCamel
    ) where
  
  import Data.ByteString (ByteString)
  import qualified Data.Foldable as Foldable
  import qualified Data.Map as Map
  import Data.Maybe
  import Data.Some
  import Data.Text (Text)
  import qualified Data.Text as Text
  import qualified Data.Text.Encoding as Text
  import Data.Type.Equality
  
  import Language.Haskell.Exts.SrcLoc
  import Language.Haskell.Names hiding (None, resolve)
  import qualified Language.Haskell.Exts.Syntax as E -- TODO: t16933748 refactor
  
  import Thrift.Compiler.Options
  import Thrift.Compiler.Parser
  import Thrift.Compiler.Plugin
  import Thrift.Compiler.Typechecker
  import Thrift.Compiler.Typechecker.Monad
  import Thrift.Compiler.Types as Thrift hiding (noLoc)
  
  -- Haskell Types ---------------------------------------------------------------
  
  data Haskell
  
  type HSType = Type Haskell
  
  type HS t = t 'Resolved Haskell Thrift.Loc
  
  data HsVectorKind = HsVectorBoxed | HsVectorStorable
    deriving (Eq, Ord, Prelude.Enum, Bounded)
  
  hsVectorImport :: HsVectorKind -> Text
  hsVectorImport HsVectorBoxed = "Data.Vector"
  hsVectorImport HsVectorStorable = "Data.Vector.Storable"
  
  hsVectorQual :: HsVectorKind -> Text
  hsVectorQual HsVectorBoxed = "Vector"
  hsVectorQual HsVectorStorable = "VectorStorable"
  
  data instance SpecialType Haskell t where
    HsInt    :: SpecialType Haskell Int
    HsString :: SpecialType Haskell String
    HsByteString :: SpecialType Haskell ByteString
    HsVector :: HsVectorKind -> HSType t -> SpecialType Haskell (List Haskell t)
  
  data HsInterface = HsInterface Environment RenameMap
  
  instance Semigroup HsInterface where
    HsInterface e1 r1 <> HsInterface e2 r2 =
      HsInterface (Map.unionWith (++) e1 e2) (Map.union r1 r2)
  
  instance Monoid HsInterface where
    mempty = HsInterface Map.empty Map.empty
  
  
  -- | Map from Haskell name qualified Thrift name
  type RenameMap = Map.Map Symbol Text
  
  -- Haskell Options -------------------------------------------------------------
  
  data instance LangOpts Haskell = HsOpts
    { hsoptsEnableHaddock :: Bool
    , hsoptsUseInt :: Bool
    , hsoptsUseHashMap :: Bool
    , hsoptsUseHashSet :: Bool
    , hsoptsDupNames :: Bool
    , hsoptsExtensions :: [Text]
    , hsoptsGenPrefix :: FilePath
    , hsoptsExtraHasFields :: Bool
    }
  
  defaultHsOpts :: LangOpts Haskell
  defaultHsOpts = HsOpts
    { hsoptsEnableHaddock = False
    , hsoptsUseInt = False
    , hsoptsUseHashMap = False
    , hsoptsUseHashSet = False
    , hsoptsDupNames = False
    , hsoptsExtensions = []
    , hsoptsGenPrefix = "gen-hs2"
    , hsoptsExtraHasFields = False
    }
  
  -- Type Class Instance ---------------------------------------------------------
  
  instance Typecheckable Haskell where
    type Interface Haskell = HsInterface
  
    -- Annotation Processing
  
    resolveTypeAnnotations ty anns = do
      Env{ options = Options{..} } <- ask
      case optsLangSpecific of
        HsOpts{..} ->
          ifFlag hsoptsUseInt i64ToInt .
          ifFlag hsoptsUseHashMap map2HashMap .
          ifFlag hsoptsUseHashSet set2HashSet <$>
          resolve ty (getTypeAnns "hs" anns)
      where
        resolve
          :: HSType t
          -> [(Text, Annotation Thrift.Loc)]
          -> TC Haskell (Some HSType)
        resolve I64 [("Int",_)] = special HsInt
        resolve TText [("String",_)] = special HsString
        resolve (TMap k v) [("HashMap",_)] = pure $ Some $ THashMap k v
        resolve (TSet u) [("HashSet",_)] = pure $ Some $ THashSet u
        resolve TText [("ByteString",_)] = special HsByteString
        resolve (TList u) [(vec,_)]
          | Just kind <-
              lookup vec [(hsVectorQual x,x) | x <- [minBound .. maxBound]] =
                special $ HsVector kind u
        resolve u [] = pure $ Some u
        resolve u ((_,a):_) =
          typeError (annLoc a) $ AnnotationMismatch (AnnType u) a
  
        special = pure . Some . TSpecial
  
        ifFlag
          :: Bool
          -> (forall t. HSType t -> Some HSType)
          -> Some HSType
          -> Some HSType
        ifFlag flag fun (Some u)
          | flag      = fun u
          | otherwise = Some u
        i64ToInt :: HSType t -> Some HSType
        i64ToInt I64 = Some $ TSpecial HsInt
        i64ToInt u  = Some u
        map2HashMap (TMap k v) = Some $ THashMap k v
        map2HashMap u = Some u
        set2HashSet (TSet u) = Some $ THashSet u
        set2HashSet u = Some u
  
    -- Typechecking
  
    qualifySpecialType _ HsInt = HsInt
    qualifySpecialType _ HsString = HsString
    qualifySpecialType _ HsByteString = HsByteString
    qualifySpecialType m (HsVector kind ty) = HsVector kind $ qualifyType m ty
  
    typecheckSpecialConst HsInt (UntypedConst _ (IntConst i _)) =
      pure $ Literal $ fromIntegral i
    typecheckSpecialConst HsString (UntypedConst _ (StringConst s _)) =
      pure $ Literal $ Text.unpack s
    typecheckSpecialConst HsByteString (UntypedConst _ (StringConst s _)) =
      pure $ Literal $ Text.encodeUtf8 s
    typecheckSpecialConst (HsVector _ u) (UntypedConst _ ListConst{..}) =
      Literal . List <$> traverse (typecheckConst u . leElem) lvElems
    typecheckSpecialConst ty val@(UntypedConst Located{..} _) =
      typeError lLocation $ LiteralMismatch (TSpecial ty) val
  
    eqSpecial HsInt HsInt = Just Refl
    eqSpecial HsString HsString = Just Refl
    eqSpecial HsByteString HsByteString = Just Refl
    eqSpecial (HsVector a u) (HsVector b v)
      | a == b = apply Refl <$> eqOrAlias u v
    eqSpecial _ _ = Nothing
  
    -- Interfaces
  
    getInterface opts tf@ThriftFile{..} = mconcat $
      map (getDeclIface opts thriftName mname) thriftDecls
      where
        mname = E.ModuleName () $ Text.unpack $ renameModule opts tf <> ".Types"
  
    getExtraSymbols opts iface tf@ThriftFile{..} =
      maybe [] (getHsIncludeDeps opts iface tf) thriftSplice
  
    -- Renamers
  
    renameModule _ ThriftFile{..} = case getNamespace "hs" thriftHeaders of
      Just ns -> ns <> "." <> toCamel thriftName
      Nothing -> toCamel thriftName
  
    renameStruct _ Struct{..} = toConstructorName structName
  
    renameField Options{..} ann sname Field{..} =
      case optsLangSpecific of
        HsOpts{..} ->
          let basePrefix
                | hsoptsDupNames = ""
                | otherwise = sname <> "_"
          in  lowercase $ fromMaybe basePrefix (getPrefix ann) <> fieldName
  
    renameConst _ = lowercase
  
    renameService _ Service{..} = toConstructorName serviceName
  
    renameFunction _ Function{..} = lowercase $ prefix <> funName
      where
        prefix = fromMaybe "" $ getPrefix $ getAnns funAnns
  
    renameTypedef _ Typedef{..} = toConstructorName tdName
  
    renameEnum _ Enum{..} = toConstructorName enumName
  
    renameEnumAlt opts e@Enum{..} name =
      fixCase $ if
        | Just prefix <- getPrefix (getAnns enumAnns) -> prefix <> name
        | otherwise -> enumName <> "_" <> name
      where
        fixCase = case enumFlavourTag opts e of
          PseudoEnum{} -> lowercase
          _ -> uppercase
  
    renameUnion _ Union{..} = toConstructorName unionName
  
    renameUnionAlt _ Union{..} UnionAlt{..} =
      toConstructorName $ fromMaybe (unionName <> "_") (getPrefix $ getAnns unionAnns) <>
      altName
  
    getUnionEmptyName _ Union{..} =
      toConstructorName $ fromMaybe (unionName <> "_") (getPrefix $ getAnns unionAnns) <>
      "EMPTY"
  
    fieldsAreUnique Options{ optsLangSpecific = HsOpts{..} } = not hsoptsDupNames
    unionAltsAreUnique _ = True
    enumAltsAreUnique Options{} = True
  
    enumFlavourTag _ Enum{..}
      | hasSimpleAnn "hs.pseudoenum" = PseudoEnum False
      | hasValueAnn "hs.pseudoenum" "thriftenum" = PseudoEnum True
      | hasSimpleAnn "hs.nounknown" = SumTypeEnum True
      | otherwise = SumTypeEnum False
      where
        hasSimpleAnn t = or
          [ saTag == t
          | SimpleAnn{..} <- getAnns enumAnns
          ]
        hasValueAnn t v = or
          [ vaTag == t && av == v
          | ValueAnn{ vaVal=TextAnn av _, ..} <- getAnns enumAnns
          ]
  
    -- Back-Translators
  
    backTranslateType HsInt = (Some I64, "Int")
    backTranslateType HsString = (Some TText, "String")
    backTranslateType HsByteString = (Some TText, "ByteString")
    backTranslateType (HsVector kind u) = (Some (TList u), hsVectorQual kind)
  
    backTranslateLiteral HsInt i = ThisLit I64 (fromIntegral i)
    backTranslateLiteral HsString s = ThisLit TText (Text.pack s)
    backTranslateLiteral HsByteString s = ThisLit TText (Text.decodeUtf8 s)
    backTranslateLiteral (HsVector _ u) l = ThisLit (TList u) l
  
  -- Compute Decl Interfaces -----------------------------------------------------
  
  getDeclIface
    :: Options Haskell -> Text -> E.ModuleName () -> Parsed Decl -> HsInterface
  getDeclIface opts name mname decl = ifaceFromSymbols mname $ case decl of
    -- Structs
    D_Struct s@Struct{..} ->
      mkStruct (packT structName) (packHs $ renameStruct opts s) ++
      concatMap
      (\field ->
        mkSelector (packT structName)
        (packHs $ renameField opts (getAnns structAnns) structName field)
        (packHs $ renameStruct opts s))
      structMembers
    -- Unions
    D_Union u@Union{..} ->
      mkData (packT unionName) (packHs $ renameUnion opts u) ++
      concatMap
      (\alt ->
        mkConstructor (packT unionName)
        (packHs $ renameUnionAlt opts u alt)
        (packHs $ renameUnion opts u))
      unionAlts
    -- Enums
    D_Enum e@Enum{..} ->
      case enumFlavourTag opts e of
        PseudoEnum{} ->
          mkNewtype (packT enumName) (packHs $ renameEnum opts e)
            (packHs $ ("un" <>) $ renameEnum opts e) ++
          concatMap
          (\EnumValue{..} ->
            mkValue (packT enumName) (packHs $ renameEnumAlt opts e evName))
          enumConstants
        SumTypeEnum{} ->
          mkData (packT enumName) (packHs $ renameEnum opts e) ++
          concatMap
          (\EnumValue{..} ->
            mkConstructor (packT enumName)
            (packHs $ renameEnumAlt opts e evName)
            (packHs $ renameEnum opts e))
          enumConstants
    -- Typedefs
    D_Typedef t@Typedef{..}
      | isNewtype (getAnns tdAnns) ->
          mkNewtype (packT tdName) (packHs $ renameTypedef opts t)
          (packHs $ ("un" <>) $ renameTypedef opts t)
      | otherwise ->
          mkType (packT tdName) (packHs $ renameTypedef opts t)
    -- Constants
    D_Const Const{..} ->
      mkValue (packT constName) (packHs $ renameConst opts constName)
    -- Services are not supported yet
    D_Service{} -> []
    -- Interaction are not supported yet
    D_Interaction{} -> []
  
    where
      mkValue tname hsname =
        [ (Value mname hsname, tname) ]
      mkStruct tname hsname =
        [ (Data mname hsname, tname)
        , (Constructor mname hsname hsname, tname)
        ]
      mkSelector tname hsname tyname =
        [ (Selector mname hsname tyname [tyname], tname)
        ]
      mkData tname hsname =
        [ (Data mname hsname, tname) ]
      mkConstructor tname hsname tyname =
        [ (Constructor mname hsname tyname, tname) ]
      mkType tname hsname =
        [ (Type mname hsname, tname) ]
      mkNewtype tname hsname selname =
        [ (NewType mname hsname, tname)
        , (Constructor mname hsname hsname, tname)
        , (Selector mname selname hsname [hsname], tname)
        ]
  
      packHs = E.Ident () . Text.unpack
      packT t = name <> "." <> t
  
  ifaceFromSymbols :: E.ModuleName () -> [(Symbol, Text)] -> HsInterface
  ifaceFromSymbols mname ss = HsInterface
    (Map.singleton mname symbols)
    rmap
    where
      rmap = Map.fromList ss
      symbols = map fst ss
  
  -- HS Include Dependencies -----------------------------------------------------
  
  getHsIncludeDeps
    :: Options Haskell
    -> HsInterface
    -> ThriftFile a l
    -> E.Module SrcSpanInfo
    -> [Text]
  getHsIncludeDeps opts (HsInterface env rmap) tf (E.Module loc mhead ps is ds)
    | E.Module _ _ _ _ decls <- annotate env m' =
      [ thriftSym
      | decl <- decls
      , (Scoped (GlobalSymbol hsSymbol _) _) <- Foldable.toList decl
      , Just thriftSym <- [Map.lookup hsSymbol rmap]
      ]
    | otherwise = error "getHsIncludeDeps"
    where
      -- Add types module to imports so that haskell-names knows where the symbols
      -- come from
      m' = E.Module loc mhead ps (thriftImport : is) ds
      thriftImport = E.ImportDecl
        { importAnn = emptyLoc
        , importModule =
          E.ModuleName emptyLoc $ Text.unpack (renameModule opts tf) ++ ".Types"
        , importQualified = False
        , importSrc = False
        , importSafe = False
        , importPkg = Nothing
        , importAs = Nothing
        , importSpecs = Nothing
        }
      emptyLoc = toSrcInfo noLoc [] noLoc
  getHsIncludeDeps _ _ _ _ = []
  