[HLSL] Run availability diagnostic on exported functions (#97352)

Implements availability diagnostic on `export` functions. 

For shader libraries the HLSL availability diagnostic should run on all
entry points and export functions. Now that the `export` keyword is
implemented (llvm/llvm-project#96823), we can detect which functions are
exported and run the diagnostic on them.

Exported functions can be nested in namespaces and in export
declarations so we need to scan not just the current translation unit
but also namespace and export declarations contexts.

Fixes #92073

1 files changed, 48 insertions, 21 deletions

diff --git a/clang/lib/Sema/SemaHLSL.cpp b/clang/lib/Sema/SemaHLSL.cpp
index eebe17a..babb984 100644
--- a/clang/lib/Sema/SemaHLSL.cpp
+++ b/clang/lib/Sema/SemaHLSL.cpp
@@ -671,30 +671,58 @@ void DiagnoseHLSLAvailability::HandleFunctionOrMethodRef(FunctionDecl *FD,
 
 void DiagnoseHLSLAvailability::RunOnTranslationUnit(
     const TranslationUnitDecl *TU) {
+
   // Iterate over all shader entry functions and library exports, and for those
   // that have a body (definiton), run diag scan on each, setting appropriate
   // shader environment context based on whether it is a shader entry function
-  // or an exported function.
-  for (auto &D : TU->decls()) {
-    const FunctionDecl *FD = llvm::dyn_cast<FunctionDecl>(D);
-    if (!FD || !FD->isThisDeclarationADefinition())
-      continue;
+  // or an exported function. Exported functions can be in namespaces and in
+  // export declarations so we need to scan those declaration contexts as well.
+  llvm::SmallVector<const DeclContext *, 8> DeclContextsToScan;
+  DeclContextsToScan.push_back(TU);
+
+  while (!DeclContextsToScan.empty()) {
+    const DeclContext *DC = DeclContextsToScan.pop_back_val();
+    for (auto &D : DC->decls()) {
+      // do not scan implicit declaration generated by the implementation
+      if (D->isImplicit())
+        continue;
+
+      // for namespace or export declaration add the context to the list to be
+      // scanned later
+      if (llvm::dyn_cast<NamespaceDecl>(D) || llvm::dyn_cast<ExportDecl>(D)) {
+        DeclContextsToScan.push_back(llvm::dyn_cast<DeclContext>(D));
+        continue;
+      }
 
-    // shader entry point
-    auto ShaderAttr = FD->getAttr<HLSLShaderAttr>();
-    if (ShaderAttr) {
-      SetShaderStageContext(ShaderAttr->getType());
-      RunOnFunction(FD);
-      continue;
-    }
-    // exported library function with definition
-    // FIXME: tracking issue #92073
-#if 0
-    if (FD->getFormalLinkage() == Linkage::External) {
-      SetUnknownShaderStageContext();
-      RunOnFunction(FD);
+      // skip over other decls or function decls without body
+      const FunctionDecl *FD = llvm::dyn_cast<FunctionDecl>(D);
+      if (!FD || !FD->isThisDeclarationADefinition())
+        continue;
+
+      // shader entry point
+      if (HLSLShaderAttr *ShaderAttr = FD->getAttr<HLSLShaderAttr>()) {
+        SetShaderStageContext(ShaderAttr->getType());
+        RunOnFunction(FD);
+        continue;
+      }
+      // exported library function
+      // FIXME: replace this loop with external linkage check once issue #92071
+      // is resolved
+      bool isExport = FD->isInExportDeclContext();
+      if (!isExport) {
+        for (const auto *Redecl : FD->redecls()) {
+          if (Redecl->isInExportDeclContext()) {
+            isExport = true;
+            break;
+          }
+        }
+      }
+      if (isExport) {
+        SetUnknownShaderStageContext();
+        RunOnFunction(FD);
+        continue;
+      }
     }
-#endif
   }
 }
 
@@ -707,8 +735,7 @@ void DiagnoseHLSLAvailability::RunOnFunction(const FunctionDecl *FD) {
     // For any CallExpr found during the traversal add it's callee to the top of
     // the stack to be processed next. Functions already processed are stored in
     // ScannedDecls.
-    const FunctionDecl *FD = DeclsToScan.back();
-    DeclsToScan.pop_back();
+    const FunctionDecl *FD = DeclsToScan.pop_back_val();
 
     // Decl was already scanned
     const unsigned ScannedStages = GetScannedStages(FD);