Make capability diagnostic message more friendly. (#6474)

* Make capability diagnostic message more friendly.

* Fix.

* Fix.

* Fix.

* Fix test.

* Update expected fail setting for aarch64/linux

* Fix.

4 files changed, 133 insertions, 62 deletions

diff --git a/source/slang/slang-ast-base.h b/source/slang/slang-ast-base.h
index 9cc36e530..8f85334d6 100644
--- a/source/slang/slang-ast-base.h
+++ b/source/slang/slang-ast-base.h
@@ -741,9 +741,9 @@ class ModifiableSyntaxNode : public SyntaxNode
     }
 };
 
-struct DeclReferenceWithLoc
+struct ProvenenceNodeWithLoc
 {
-    Decl* referencedDecl;
+    NodeBase* referencedNode;
     SourceLoc referenceLoc;
 };
 
@@ -789,7 +789,7 @@ public:
     bool isChildOf(Decl* other) const;
 
     // Track the decl reference that caused the requirement of a capability atom.
-    SLANG_UNREFLECTED List<DeclReferenceWithLoc> capabilityRequirementProvenance;
+    SLANG_UNREFLECTED List<ProvenenceNodeWithLoc> capabilityRequirementProvenance;
 
     SLANG_UNREFLECTED bool hiddenFromLookup = false;
 
diff --git a/source/slang/slang-check-decl.cpp b/source/slang/slang-check-decl.cpp
index f12d11bdf..486ac6e9c 100644
--- a/source/slang/slang-check-decl.cpp
+++ b/source/slang/slang-check-decl.cpp
@@ -12744,11 +12744,11 @@ static void _propagateRequirement(
     // if stmt inside parent, set the provenance tracker to the calling function
     if (!decl)
         decl = visitor->getParentFuncOfVisitor();
-    if (referencedDecl && decl)
+    if (referencedNode && decl)
     {
         // Here we store a childDecl that added/removed capabilities from a parentDecl
         decl->capabilityRequirementProvenance.add(
-            DeclReferenceWithLoc{referencedDecl, referenceLoc});
+            ProvenenceNodeWithLoc{referencedNode, referenceLoc});
     }
 };
 
@@ -12977,15 +12977,19 @@ CapabilitySet SemanticsDeclCapabilityVisitor::getDeclaredCapabilitySet(Decl* dec
     // The requirement for `foo` should be glsl+glsl_ext_1 | spirv.
     //
     CapabilitySet declaredCaps;
+    CapabilityAtom stageToJoin = CapabilityAtom::Invalid;
     for (Decl* parent = decl; parent; parent = getParentDecl(parent))
     {
         CapabilitySet localDeclaredCaps;
         bool shouldBreak = false;
         if (!as<AggTypeDeclBase>(parent) || parent->inferredCapabilityRequirements.isEmpty())
         {
-            for (auto decoration : parent->getModifiersOfType<RequireCapabilityAttribute>())
+            for (auto mod : parent->modifiers)
             {
-                localDeclaredCaps.unionWith(decoration->capabilitySet);
+                if (auto decoration = as<RequireCapabilityAttribute>(mod))
+                    localDeclaredCaps.unionWith(decoration->capabilitySet);
+                else if (auto entrypoint = as<EntryPointAttribute>(mod))
+                    stageToJoin = entrypoint->capabilitySet.getTargetStage();
             }
         }
         else
@@ -13001,6 +13005,8 @@ CapabilitySet SemanticsDeclCapabilityVisitor::getDeclaredCapabilitySet(Decl* dec
         if (shouldBreak)
             break;
     }
+    if (!declaredCaps.isEmpty() && stageToJoin != CapabilityAtom::Invalid)
+        declaredCaps.join(CapabilitySet((CapabilityName)stageToJoin));
     return declaredCaps;
 }
 
@@ -13298,63 +13304,83 @@ void diagnoseMissingCapabilityProvenance(
     Decl* decl,
     CapabilitySet& setToFind)
 {
-    HashSet<Decl*> checkedDecls;
-    DeclReferenceWithLoc declWithRef;
-    declWithRef.referencedDecl = decl;
-    declWithRef.referenceLoc = (decl) ? decl->loc : SourceLoc();
+    HashSet<NodeBase*> checkedDecls;
+    ProvenenceNodeWithLoc provNode;
+    provNode.referencedNode = decl;
+    provNode.referenceLoc = (decl) ? decl->loc : SourceLoc();
     bool bottomOfProvenanceStack = false;
     // Find the bottom of the atom provenance stack which fails to contain `setToFind`
-    while (!bottomOfProvenanceStack && declWithRef.referencedDecl)
+    while (!bottomOfProvenanceStack && provNode.referencedNode)
     {
         bottomOfProvenanceStack = true;
-        for (auto& i : declWithRef.referencedDecl->capabilityRequirementProvenance)
+        if (auto referencedDecl = as<Decl>(provNode.referencedNode))
         {
-            if (checkedDecls.contains(i.referencedDecl))
-                continue;
-            checkedDecls.add(i.referencedDecl);
-
-            if (!i.referencedDecl->inferredCapabilityRequirements.implies(setToFind))
+            for (auto& i : referencedDecl->capabilityRequirementProvenance)
             {
-                // We found a source of the incompatible capability, follow this
-                // element inside the provenance stack until we are at the bottom
-                declWithRef = i;
-                bottomOfProvenanceStack = false;
-                break;
+                if (checkedDecls.contains(i.referencedNode))
+                    continue;
+                checkedDecls.add(i.referencedNode);
+                auto innerReferencedDecl = as<Decl>(i.referencedNode);
+                if (!innerReferencedDecl ||
+                    !innerReferencedDecl->inferredCapabilityRequirements.implies(setToFind))
+                {
+                    // We found a source of the incompatible capability, follow this
+                    // element inside the provenance stack until we are at the bottom
+                    provNode = i;
+                    bottomOfProvenanceStack = false;
+                    break;
+                }
             }
         }
+        else
+        {
+            bottomOfProvenanceStack = true;
+        }
     }
 
-    if (!declWithRef.referencedDecl)
+    if (!provNode.referencedNode)
         return;
 
-    // Diagnose the use-site
-    maybeDiagnose(
-        sink,
-        optionSet,
-        DiagnosticCategory::Capability,
-        declWithRef.referenceLoc,
-        Diagnostics::seeUsingOf,
-        declWithRef.referencedDecl);
-    // Diagnose the definition as the problem
-    maybeDiagnose(
-        sink,
-        optionSet,
-        DiagnosticCategory::Capability,
-        declWithRef.referencedDecl->loc,
-        Diagnostics::seeDefinitionOf,
-        declWithRef.referencedDecl);
-
-    // If we find a 'require' modifier, this is contributing to the overall capability
-    // incompatibility. We should hint to the user that this declaration is problematic.
-    if (auto requireCapabilityAttribute =
-            declWithRef.referencedDecl->findModifier<RequireCapabilityAttribute>())
+    if (auto referencedDecl = as<Decl>(provNode.referencedNode))
+    {
+        // Diagnose the use-site
         maybeDiagnose(
             sink,
             optionSet,
             DiagnosticCategory::Capability,
-            requireCapabilityAttribute->loc,
-            Diagnostics::seeDeclarationOf,
-            requireCapabilityAttribute);
+            provNode.referenceLoc,
+            Diagnostics::seeUsingOf,
+            referencedDecl);
+        // Diagnose the definition as the problem
+        maybeDiagnose(
+            sink,
+            optionSet,
+            DiagnosticCategory::Capability,
+            referencedDecl->loc,
+            Diagnostics::seeDefinitionOf,
+            referencedDecl);
+        // If we find a 'require' modifier, this is contributing to the overall capability
+        // incompatibility. We should hint to the user that this declaration is problematic.
+        if (auto requireCapabilityAttribute =
+                referencedDecl->findModifier<RequireCapabilityAttribute>())
+            maybeDiagnose(
+                sink,
+                optionSet,
+                DiagnosticCategory::Capability,
+                requireCapabilityAttribute->loc,
+                Diagnostics::seeDeclarationOf,
+                requireCapabilityAttribute);
+    }
+    else
+    {
+        maybeDiagnose(
+            sink,
+            optionSet,
+            DiagnosticCategory::Capability,
+            provNode.referenceLoc,
+            Diagnostics::seeUsingOf,
+            provNode.referencedNode->astNodeType);
+    }
 }
 
 void diagnoseCapabilityProvenance(
@@ -13372,7 +13398,20 @@ void diagnoseCapabilityProvenance(
         printedDecls.add(declToPrint);
         for (auto& provenance : declToPrint->capabilityRequirementProvenance)
         {
-            if (!provenance.referencedDecl->inferredCapabilityRequirements.implies(atomToFind))
+            auto referencedDecl = as<Decl>(provenance.referencedNode);
+            if (!referencedDecl)
+            {
+                maybeDiagnose(
+                    sink,
+                    optionSet,
+                    DiagnosticCategory::Capability,
+                    provenance.referenceLoc,
+                    Diagnostics::seeUsingOf,
+                    provenance.referencedNode->astNodeType);
+                break;
+            }
+
+            if (!referencedDecl->inferredCapabilityRequirements.implies(atomToFind))
                 continue;
             maybeDiagnose(
                 sink,
@@ -13380,8 +13419,8 @@ void diagnoseCapabilityProvenance(
                 DiagnosticCategory::Capability,
                 provenance.referenceLoc,
                 Diagnostics::seeUsingOf,
-                provenance.referencedDecl);
-            declToPrint = provenance.referencedDecl;
+                referencedDecl);
+            declToPrint = referencedDecl;
             if (printedDecls.contains(declToPrint))
                 break;
             if (declToPrint->findModifier<RequireCapabilityAttribute>())
@@ -13487,14 +13526,30 @@ void SemanticsDeclCapabilityVisitor::diagnoseUndeclaredCapability(
     for (auto i : simplifiedFailedAtomsSet)
     {
         CapabilityAtom formattedAtom = asAtom(i);
-        maybeDiagnose(
-            getSink(),
-            this->getOptionSet(),
-            DiagnosticCategory::Capability,
-            decl->loc,
-            diagnosticInfo,
-            decl,
-            formattedAtom);
+        CapabilityName canonicalName;
+        if (isStageAtom((CapabilityName)formattedAtom, canonicalName))
+        {
+            // Provide a more friendly message if atom is a stage.
+            maybeDiagnose(
+                getSink(),
+                this->getOptionSet(),
+                DiagnosticCategory::Capability,
+                decl->loc,
+                Diagnostics::declHasDependenciesNotCompatibleOnStage,
+                decl,
+                formattedAtom);
+        }
+        else
+        {
+            maybeDiagnose(
+                getSink(),
+                this->getOptionSet(),
+                DiagnosticCategory::Capability,
+                decl->loc,
+                diagnosticInfo,
+                decl,
+                formattedAtom);
+        }
         // Print provenances.
         diagnoseCapabilityProvenance(
             this->getOptionSet(),
diff --git a/source/slang/slang-diagnostic-defs.h b/source/slang/slang-diagnostic-defs.h
index 789eac88f..761dc4768 100644
--- a/source/slang/slang-diagnostic-defs.h
+++ b/source/slang/slang-diagnostic-defs.h
@@ -987,7 +987,7 @@ DIAGNOSTIC(
     36107,
     Error,
     entryPointUsesUnavailableCapability,
-    "entrypoint '$0' does not support compilation target '$1' with stage '$2'")
+    "entrypoint '$0' uses features that are not available in '$2' stage for '$1' target.")
 DIAGNOSTIC(
     36108,
     Error,
@@ -1025,7 +1025,11 @@ DIAGNOSTIC(
     Error,
     capabilityHasMultipleStages,
     "Capability '$0' is targeting stages '$1', only allowed to use 1 unique stage here.")
-
+DIAGNOSTIC(
+    36117,
+    Error,
+    declHasDependenciesNotCompatibleOnStage,
+    "'$0' uses features that are not available in '$1' stage.")
 
 // Attributes
 DIAGNOSTIC(31000, Warning, unknownAttributeName, "unknown attribute '$0'")
diff --git a/source/slang/slang-syntax.cpp b/source/slang/slang-syntax.cpp
index e479f1b9d..efb5814e6 100644
--- a/source/slang/slang-syntax.cpp
+++ b/source/slang/slang-syntax.cpp
@@ -223,8 +223,20 @@ void printDiagnosticArg(StringBuilder& sb, ASTNodeType nodeType)
     case ASTNodeType::RequireCapabilityDecl:
         sb << "__require_capability";
         break;
+    case ASTNodeType::DiscardStmt:
+        sb << "discard";
+        break;
     default:
-        sb << "decl";
+        if (ASTClassInfo::getInfo(nodeType)->isDerivedFrom((uint32_t)ASTNodeType::Expr))
+            sb << "expression";
+        else if (ASTClassInfo::getInfo(nodeType)->isDerivedFrom((uint32_t)ASTNodeType::Stmt))
+            sb << "statement";
+        else if (ASTClassInfo::getInfo(nodeType)->isDerivedFrom((uint32_t)ASTNodeType::Decl))
+            sb << "decl";
+        else if (ASTClassInfo::getInfo(nodeType)->isDerivedFrom((uint32_t)ASTNodeType::Val))
+            sb << "val";
+        else
+            sb << "node";
         break;
     }
 }