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 TypeApplications, ApplicativeDo #-}
  
  module Glean.Glass.Pretty.Python
    (
      prettyPythonSignature
    ) where
  
  import Data.Maybe ( mapMaybe, isNothing )
  import Data.Map.Strict ( Map )
  import qualified Data.Map.Strict as Map
  import Data.Text ( Text, takeWhileEnd )
  import Compat.Prettyprinter
  import Control.Monad.Trans.Maybe (MaybeT (..))
  import Control.Monad ( forM )
  import Util.List ( uniq )
  
  import Glean.Angle as Angle
  import Glean.Glass.Path ( fromGleanPath )
  import Glean.Glass.Base ( GleanPath(GleanPath) )
  import Glean.Glass.Types ( SymbolId(..), RepoName(..) )
  import Glean.Glass.Utils
  import Glean.Glass.SymbolId ( toSymbolId )
  import Glean.Util.ToAngle ( ToAngle(toAngle) )
  import qualified Glean
  import qualified Glean.Haxl.Repos as Glean
  
  import Glean.Schema.CodePython.Types as Python ( Entity(..) )
  import qualified Glean.Schema.Python.Types as Python
  import qualified Glean.Schema.Src.Types as Src
  import qualified Glean.Schema.Code.Types as Code
  import qualified Glean.Schema.Codemarkup.Types as Code
  
  -- Type of Python entity signatures to capture the the subset we generate
  -- We use this to separate the processing of the Glean type from the text
  -- we wish to generate
  
  -- Python symbol kinds
  data Definition
    = Function {
        funModifier :: !AsyncModifier,
        funName :: !Name,
        funParams :: [Parameter], -- regular parameters
        funPosOnlyParams :: [Parameter], -- posonly
        funKWOnlyParams :: [Parameter], -- kw only
        funStarArg :: Maybe Parameter, -- star args
        funStarKWArg :: Maybe Parameter, --  star star kwargs
        funRetType :: !AType
      }
    | Module !Name
    | Variable !Name !AType
    | Class {
        _clsName :: !Name,
        _baseNames :: [Name]
      }
  
  data AType
    = NoType
    | AType !PyType XRefs
  
  data Parameter = Parameter !Name (Maybe ExprText) !(Maybe PyType) XRefs
  
  -- modifiers
  data AsyncModifier = Async | NotAsync
   deriving (Eq, Show)
  
  newtype ExprText = ExprText Text
  
  -- names
  newtype Name = Name Text
  
  newtype PyType = PyType Text
  
  type XRef = (Python.Declaration, Src.ByteSpan, GleanPath)
  type XRefs = [XRef]
  
  data Ann
    = None
    | BareDecl !Python.Declaration !GleanPath
    | SymId !SymbolId
  
  prettyPythonSignature
    :: LayoutOptions
    -> RepoName
    -> SymbolId
    -> Python.Entity
    -> Glean.RepoHaxl u w (Maybe (SimpleDocStream (Maybe SymbolId)))
  prettyPythonSignature opts repo sym (Python.Entity_decl decl) = runMaybeT $ do
      pyDef <- maybeT $ fetchDataRecursive (angleDeclToDef (toAngle decl))
      def <- maybeT $ fmap (pprDefinition sym) <$> fromAngleDefinition pyDef
      maybeT $ Just <$> sequence (annotateDocs def)
    where
      annotateDocs doc = reAnnotateS (declToSymbolId repo) (layoutSmart opts doc)
  prettyPythonSignature _ _ _ Python.Entity_EMPTY = return Nothing
  
  declToSymbolId :: RepoName -> Ann -> Glean.RepoHaxl u w (Maybe SymbolId)
  declToSymbolId _repo None = return Nothing
  declToSymbolId _repo (SymId symId) = return (Just symId)
  declToSymbolId repo (BareDecl decl filepath) = Just <$>
      toSymbolId (fromGleanPath repo filepath) entity
    where
      entity = Code.Entity_python (Python.Entity_decl decl)
  
  fromAngleDefinition
    :: Python.Definition -> Glean.RepoHaxl u w (Maybe Definition)
  fromAngleDefinition def = case def of
    Python.Definition_func fn -> Just <$>
      (fromFunctionDefinition =<< Glean.keyOf fn)
    Python.Definition_module m -> Just <$>
      (fromModuleDefinition =<< Glean.keyOf m)
    Python.Definition_cls c -> Just <$>
      (fromClassDefinition =<< Glean.keyOf c)
    Python.Definition_variable v -> Just <$>
      (fromVariableDefinition =<< Glean.keyOf v)
    _ -> pure Nothing
  
  fromClassDefinition
    :: Python.ClassDefinition_key -> Glean.RepoHaxl u w Definition
  fromClassDefinition def = do
    Python.ClassDeclaration_key name _bases <- Glean.keyOf decl
    -- get class xrefs from bases
    baseNames <- case mBases of
      Nothing -> pure []
      Just decls ->
        mapM (fmap Python.classDeclaration_key_name <$> Glean.keyOf) decls
  
    nameStr <- trimModule <$> Glean.keyOf name
    baseStrs <- mapM (fmap trimModule <$> Glean.keyOf) baseNames
    return $ Class nameStr baseStrs
    where
      Python.ClassDefinition_key {
        classDefinition_key_bases = mBases,
        classDefinition_key_declaration = decl
      } = def
  
  fromVariableDefinition
    :: Python.VariableDefinition_key -> Glean.RepoHaxl u w Definition
  fromVariableDefinition def = do
    Python.VariableDeclaration_key name <- Glean.keyOf decl
    let mVarTy = case mTypeInfo of
          Nothing -> Nothing
          Just tyInfo ->
            let tyNames = map Python.xRefViaName_target
                  (Python.typeInfo_xrefs tyInfo)
            in Just (tyNames, tyInfo)
    declMap <- fetchDeclWithNames $ maybe [] fst mVarTy
    varType <- case mVarTy of
      Nothing -> pure NoType
      Just (_, retTyInfo) -> do
        retTyText <- fromTypeInfo retTyInfo
        pure $ AType retTyText (mkXRefs declMap retTyInfo)
    varName <- trimModule <$> Glean.keyOf name
    return $ Variable varName varType
    where
      Python.VariableDefinition_key {
        variableDefinition_key_typeInfo = mTypeInfo,
        variableDefinition_key_declaration = decl
      } = def
  
  fromModuleDefinition
    :: Python.ModuleDefinition_key -> Glean.RepoHaxl u w Definition
  fromModuleDefinition (Python.ModuleDefinition_key mod) = do
    Python.Module_key name <- Glean.keyOf mod
    modName <- Glean.keyOf name
    return $ Module (Name modName)
  
  fromFunctionDefinition
    :: Python.FunctionDefinition_key -> Glean.RepoHaxl u w Definition
  fromFunctionDefinition def = do
    Python.FunctionDeclaration_key name <- Glean.keyOf decl
  
    let returnTy = case mReturnTy of
          Nothing -> Nothing
          Just tyInfo ->
            let retTypeDeclNames = map Python.xRefViaName_target
                  (Python.typeInfo_xrefs tyInfo)
            in Just (retTypeDeclNames, tyInfo)
  
    -- resolve all names to their xref decls in one shot
    declMap <- fetchDeclWithNames $ concat
      [ map Python.xRefViaName_target (Python.typeInfo_xrefs tyInfo)
      | Just tyInfo <- map Python.parameter_typeInfo pyParams
      ] ++ maybe [] fst returnTy
        ++ maybe []
              (map Python.xRefViaName_target . Python.typeInfo_xrefs)
              (Python.parameter_typeInfo =<< mStarKWArg)
        ++ maybe []
              (map Python.xRefViaName_target . Python.typeInfo_xrefs)
              (Python.parameter_typeInfo =<< mStarArg)
        ++ maybe [] (\ps -> concat
            [ map Python.xRefViaName_target (Python.typeInfo_xrefs tyInfo)
            | Just tyInfo <- map Python.parameter_typeInfo ps
            ]) mKWOnlyParams
        ++ maybe [] (\ps -> concat
            [ map Python.xRefViaName_target (Python.typeInfo_xrefs tyInfo)
            | Just tyInfo <- map Python.parameter_typeInfo ps
            ]) mPosOnlyParams
  
    let regParamsAndXRefs =
          [ case parameter_typeInfo of
              Nothing -> (param, [])
              Just tyInfo -> (param, mkXRefs declMap tyInfo)
          | param@Python.Parameter{..} <- pyParams
          ]
    let kwOnlyParamsAndXRefs = case mKWOnlyParams of
          Nothing -> []
          Just ps ->
            [ case parameter_typeInfo of
                Nothing -> (param, [])
                Just tyInfo -> (param, mkXRefs declMap tyInfo)
            | param@Python.Parameter{..} <- ps
            ]
    let posOnlyParamsAndXRefs = case mPosOnlyParams of
          Nothing -> []
          Just ps ->
            [ case parameter_typeInfo of
                Nothing -> (param, [])
                Just tyInfo -> (param, mkXRefs declMap tyInfo)
            | param@Python.Parameter{..} <- ps
            ]
    let starKWArgsAndXRefs = case mStarKWArg of
          Nothing -> Nothing
          Just param@Python.Parameter{..} -> Just $ case parameter_typeInfo of
              Nothing -> (param, []) -- has no type info
              Just tyInfo -> (param, mkXRefs declMap tyInfo)
    let starArgAndXRefs = case mStarArg of
          Nothing -> Nothing
          Just param@Python.Parameter{..} -> Just $ case parameter_typeInfo of
              Nothing -> (param, []) -- has no type info
              Just tyInfo -> (param, mkXRefs declMap tyInfo)
  
    funRetType <- case returnTy of
      Nothing -> pure NoType
      Just (_, retTyInfo) -> do
        retTyText <- fromTypeInfo retTyInfo
        pure $ AType retTyText (mkXRefs declMap retTyInfo)
  
    funName <- trimModule <$> Glean.keyOf name
    funParams <- mapM fromParameter regParamsAndXRefs
    funPosOnlyParams <- mapM fromParameter posOnlyParamsAndXRefs
    funKWOnlyParams <- mapM fromParameter kwOnlyParamsAndXRefs
    funStarKWArg <- mapM fromParameter starKWArgsAndXRefs
    funStarArg <- mapM fromParameter starArgAndXRefs
  
    return $ Function { funModifier = if async then Async else NotAsync, .. }
    where
      Python.FunctionDefinition_key {
        functionDefinition_key_declaration = decl,
        functionDefinition_key_is_async = async,
        functionDefinition_key_params = pyParams,
        functionDefinition_key_kwonly_params = mKWOnlyParams,
        functionDefinition_key_posonly_params = mPosOnlyParams,
        functionDefinition_key_star_arg = mStarArg,
        functionDefinition_key_star_kwarg = mStarKWArg,
        functionDefinition_key_returnsInfo = mReturnTy
      } = def
  
  mkXRefs
    :: Map.Map Python.Name (Python.Declaration, GleanPath)
    -> Python.TypeInfo
    -> XRefs
  mkXRefs declMap Python.TypeInfo{..} =
    mapMaybe (\Python.XRefViaName{..} ->
      case Map.lookup xRefViaName_target declMap of
        Nothing -> Nothing
        Just (decl, path) -> Just (decl, xRefViaName_source, path)
    ) typeInfo_xrefs
  
  fromTypeInfo  :: Python.TypeInfo -> Glean.RepoHaxl u w PyType
  fromTypeInfo Python.TypeInfo{..} = PyType <$> Glean.keyOf
    typeInfo_displayType
  
  fromParameter :: (Python.Parameter, XRefs) -> Glean.RepoHaxl u w Parameter
  fromParameter (Python.Parameter{..}, xrefs) = do
    nameStr <- Glean.keyOf parameter_name
    tyInfo <- case parameter_typeInfo of
      Nothing -> return Nothing
      Just ty -> Just <$> fromTypeInfo ty
    return $ Parameter (Name nameStr) (ExprText <$> parameter_value) tyInfo xrefs
  
  pprDefinition :: SymbolId -> Definition -> Doc Ann
  pprDefinition self (Module name) =
    "module" <+> annotate (SymId self) (pprName name)
  pprDefinition self (Class name _) =
    "class" <+> annotate (SymId self) (pprName name)
  
  pprDefinition self (Variable name mTyInfo) =
    annotate (SymId self) (pprName name) <> case mTyInfo of
      NoType -> emptyDoc
      AType ty xrefs -> colon <+> pprTypeXRefs ty xrefs
  
  -- empty param case
  pprDefinition self (Function async name [] [] [] Nothing Nothing returnTy) =
    hcat [
     pprAsync async, "def" <+>
        annotate (SymId self) (pprName name), "()",
     pprReturnType returnTy
    ]
  -- full param list
  pprDefinition self (Function async name params posOnlyParams
         kwOnlyParams starArg starKWArg returnTy) =
    vcat [
      nest 4 (vsep (
        hcat [pprAsync async, "def" <> space,
          annotate (SymId self) (pprName name) <> lparen] :
        punctuate comma ( -- ordering is quite semantically sensitive
          concat [
            pprPosOnlyParams posOnlyParams,
            map pprParam params,
            maybe [] (pure . pprStarArg) starArg,
            pprKWOnlyParams (isNothing starArg) kwOnlyParams,
            maybe [] (pure . pprStarKWArg) starKWArg
          ]))),
      rparen <> pprReturnType returnTy
      ]
  
  pprAsync :: AsyncModifier -> Doc Ann
  pprAsync Async = "async" <> space
  pprAsync _ = emptyDoc
  
  pprReturnType :: AType -> Doc Ann
  pprReturnType NoType = emptyDoc
  pprReturnType (AType ty xrefs) = space <> "->" <+> pprTypeXRefs ty xrefs
  
  pprKWOnlyParams :: Bool -> [Parameter] -> [Doc Ann]
  pprKWOnlyParams _ [] = []
  pprKWOnlyParams True xs  = "*" : map pprParam xs
  pprKWOnlyParams False xs =       map pprParam xs
  
  pprPosOnlyParams :: [Parameter] -> [Doc Ann]
  pprPosOnlyParams [] = []
  pprPosOnlyParams xs = map pprParam xs ++ ["/"]
  
  pprParam :: Parameter -> Doc Ann
  pprParam (Parameter name mDefValue mty xrefs) = hcat
    [ pprName name
    , case mty of
      Nothing -> emptyDoc
      Just ty -> colon <+> pprTypeXRefs ty xrefs
    , case mDefValue of
      Nothing -> emptyDoc
      Just (ExprText val) -> space <> equals <+> pretty val
    ]
  
  pprStarArg :: Parameter -> Doc Ann
  pprStarArg param = "*" <> pprParam param
  
  pprStarKWArg :: Parameter -> Doc Ann
  pprStarKWArg param = "**" <> pprParam param
  
  pprName :: Name -> Doc Ann
  pprName (Name name) = pretty name
  
  pprTypeXRefs :: PyType -> XRefs -> Doc Ann
  pprTypeXRefs (PyType ty) xrefs =
      mconcat $ (\(frag, ann) -> annotate (toAnn ann) $ pretty frag) <$>
        splitString ty spans
    where
      toAnn Nothing = None
      toAnn (Just (decl, path)) = BareDecl decl path
  
      spans = map (\(decl, Src.ByteSpan{..}, filepath) ->
                 ((decl,filepath), fromIntegral (Glean.fromNat byteSpan_start)
                     , fromIntegral (Glean.fromNat byteSpan_length))) xrefs
  
  --
  -- | Convert a list of python.Names into their corresponding decl/location pairs
  --
  fetchDeclWithNames
    :: [Python.Name]
    -> Glean.RepoHaxl u w (Map Python.Name (Python.Declaration, GleanPath))
  fetchDeclWithNames [] = pure mempty
  fetchDeclWithNames names = do
    (result,_truncated) <- searchRecursiveWithLimit maxXRefs
      (angleDeclsByNames ids)
    Map.fromList <$> forM result (\(decl,srcFile) -> do
      Python.DeclarationWithName_key{..} <- Glean.keyOf decl
      filepath <- GleanPath <$> Glean.keyOf srcFile
      let decl = (declarationWithName_key_declaration, filepath)
          name = declarationWithName_key_name
      return (name, decl)
     )
    where
      maxXRefs = Just (length ids)
  
      ids :: [Glean.IdOf Python.Name]
      ids = uniq (map Glean.getId names)
  
  angleDeclToDef :: Angle Python.Declaration -> Angle Python.Definition
  angleDeclToDef decl = var $ \(def :: Angle Python.Definition) ->
    def `where_` [
      wild .= predicate @Python.DeclarationDefinition (
        rec $
          field @"declaration" decl $
          field @"definition" def
        end
      )
    ]
  
  -- Bulk convert each name to its definition entity and file location
  -- to build a symbol id later
  angleDeclsByNames
    :: [Glean.IdOf Python.Name] -> Angle (Python.DeclarationWithName, Src.File)
  angleDeclsByNames names = vars $ \(decl :: Angle Python.Declaration)
      (file :: Angle Src.File) (p :: Angle Python.DeclarationWithName) ->
    tuple (p, file) `where_` [
      p .= predicate @Python.DeclarationWithName (
        rec $
          field @"name" (asPredicate (elementsOf (factIdsArray names))) $
          field @"declaration" decl
        end),
      wild .= predicate @Code.EntityLocation (
        rec $
          field @"entity" (alt @"python" (alt @"decl" decl)) $
          field @"location" (rec $ field @"file" (asPredicate file) $ end)
        end)
      ]
  
  -- | we could use the sname here to lookup the associated decl fact
  -- as an xref in the type signature (c.f how Hack does this)
  trimModule :: Text -> Name
  trimModule qname = Name (takeWhileEnd (/= '.') qname)
  