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.

9 files changed, 162 insertions, 66 deletions

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 6cf72fdee..9e0b03dbc 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -166,7 +166,8 @@ jobs:
               -use-test-server \
               -category ${{ matrix.test-category }} \
               -api all-dx12 \
-              -expected-failure-list tests/expected-failure-github.txt
+              -expected-failure-list tests/expected-failure-github.txt \
+              -expected-failure-list tests/expected-failure-record-replay-tests.txt
           else
             "$bin_dir/slang-test" \
               -use-test-server \
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;
     }
 }
diff --git a/tests/diagnostics/discard-in-rt.slang b/tests/diagnostics/discard-in-rt.slang
new file mode 100644
index 000000000..93c3d1038
--- /dev/null
+++ b/tests/diagnostics/discard-in-rt.slang
@@ -0,0 +1,24 @@
+//TEST:SIMPLE(filecheck=CHECK): -target spirv
+
+// CHECK: 'closestHit' uses features that are not available in 'closesthit' stage.
+// CHECK: see using of 'discard'
+
+struct PrimaryRayPayload{}
+[require(spvRayTracingKHR)]
+[shader("closesthit")]
+void closestHit(
+    inout PrimaryRayPayload rayData : SV_RayPayload,
+    BuiltInTriangleIntersectionAttributes attribs : SV_IntersectionAttributes)
+{
+    discard;
+}
+
+// CHECK: 'closestHit1' uses features that are not available in 'closesthit' stage
+// CHECK: see using of 'discard'
+[shader("closesthit")]
+void closestHit1(
+    inout PrimaryRayPayload rayData : SV_RayPayload,
+    BuiltInTriangleIntersectionAttributes attribs : SV_IntersectionAttributes)
+{
+    discard;
+}
diff --git a/tests/language-feature/capability/capabilitySimplification1.slang b/tests/language-feature/capability/capabilitySimplification1.slang
index 1d781a45e..440ac1ced 100644
--- a/tests/language-feature/capability/capabilitySimplification1.slang
+++ b/tests/language-feature/capability/capabilitySimplification1.slang
@@ -5,7 +5,7 @@
 // CHECK_IGNORE_CAPS-NOT: error 36107
 
 // CHECK: error 36107
-// CHECK-SAME: entrypoint 'computeMain' does not support compilation target 'glsl' with stage 'compute'
+// CHECK-SAME: entrypoint 'computeMain' uses features that are not available in 'compute' stage for 'glsl' target.
 // CHECK: capabilitySimplification1.slang(21): note: see using of 'WaveMultiPrefixCountBits'
 // CHECK-NOT: see using of 'WaveMultiPrefixCountBits'
 // CHECK: {{.*}}.meta.slang({{.*}}): note: see definition of 'WaveMultiPrefixCountBits'
diff --git a/tests/language-feature/capability/capabilitySimplification3.slang b/tests/language-feature/capability/capabilitySimplification3.slang
index 808c19bf6..150f926a9 100644
--- a/tests/language-feature/capability/capabilitySimplification3.slang
+++ b/tests/language-feature/capability/capabilitySimplification3.slang
@@ -4,7 +4,7 @@
 
 // CHECK_IGNORE_CAPS-NOT: error 36107
 
-// CHECK: error 36107: entrypoint 'computeMain' does not support compilation target 'glsl' with stage 'compute'
+// CHECK: error 36107: entrypoint 'computeMain' uses features that are not available in 'compute' stage for 'glsl' target.
 // CHECK: capabilitySimplification3.slang(16): note: see using of 'WaveMultiPrefixCountBits'
 // CHECK-NOT: see using of 'WaveMultiPrefixCountBits'
 // CHECK: {{.*}}.meta.slang({{.*}}): note: see definition of 'WaveMultiPrefixCountBits'
diff --git a/tests/language-feature/capability/explicit-shader-stage-2.slang b/tests/language-feature/capability/explicit-shader-stage-2.slang
index b4c6d3cfe..a01cff7c2 100644
--- a/tests/language-feature/capability/explicit-shader-stage-2.slang
+++ b/tests/language-feature/capability/explicit-shader-stage-2.slang
@@ -1,7 +1,7 @@
 //TEST:SIMPLE(filecheck=CHECK): -target hlsl -entry main -allow-glsl -profile sm_5_0
 //TEST:SIMPLE(filecheck=CHECK_IGNORE_CAPS): -target hlsl -entry main -allow-glsl -profile sm_5_0 -ignore-capabilities
 
-//CHECK: error 36107: entrypoint 'main' does not support compilation target 'hlsl' with stage 'fragment'
+//CHECK: error 36107: entrypoint 'main' uses features that are not available in 'fragment' stage for 'hlsl' target.
 //CHECK_IGNORE_CAPS-NOT: error 36100
 [shader("fragment")]
 float4 main()