From 65240d074b4ddec55e56962ebf8de46207bcf5fa Mon Sep 17 00:00:00 2001
From: ArielG-NV <159081215+ArielG-NV@users.noreply.github.com>
Date: Wed, 28 Aug 2024 15:06:23 -0400
Subject: Allow capabilities to be used with `[shader("...")]` (#4928)

* Allow capabilities to be used with `[shader("...")]`

Fixes: #4917

Changes:
1. Allow using capabilities instead of `Stage`s with `EntryPointAttribute`.
2. When resolving capabilities for an entrypoint+profile (per entrypoint) in `resolveStageOfProfileWithEntryPoint` add our `EntryPointAttribute` and resolved capability
3. Added tests and some capabilities related clean-up

* fix a warning made by a mistake in syntax

* change fineStageByName to assume it is passed a stage without a '_'

* test with and without prefix '_'

* cleanup some profiles and reprisentation to work better with 'Stage' and 'Profile'

This use case is why we need to clean all profile-usage into `CapabilityName`s directly.

* change how we compare

* only change profiles

* let all capabilities be resolved by 'shader' profile for now

* fix warning checks I accidently broke

* meshshading_internal to _meshshading

---------

Co-authored-by: Yong He <yonghe@outlook.com>
---
 source/slang/slang-check-modifier.cpp | 42 +++++++++++++++++++++++++++++------
 1 file changed, 35 insertions(+), 7 deletions(-)

(limited to 'source/slang/slang-check-modifier.cpp')

diff --git a/source/slang/slang-check-modifier.cpp b/source/slang/slang-check-modifier.cpp
index 705d0bb3b..d7f879c51 100644
--- a/source/slang/slang-check-modifier.cpp
+++ b/source/slang/slang-check-modifier.cpp
@@ -548,19 +548,47 @@ namespace Slang
         {
             SLANG_ASSERT(attr->args.getCount() == 1);
 
-            String stageName;
-            if (!checkLiteralStringVal(attr->args[0], &stageName))
+            String capNameString;
+            if (!checkLiteralStringVal(attr->args[0], &capNameString))
             {
                 return false;
             }
 
-            auto stage = findStageByName(stageName);
-            if (stage == Stage::Unknown)
+            CapabilityName capName = findCapabilityName(capNameString.getUnownedSlice());
+            if (capName != CapabilityName::Invalid)
             {
-                getSink()->diagnose(attr->args[0], Diagnostics::unknownStageName, stageName);
-            }
+                if (isInternalCapabilityName(capName))
+                    maybeDiagnose(getSink(), this->getOptionSet(), DiagnosticCategory::Capability, attr, Diagnostics::usingInternalCapabilityName, attr, capName);
+                
+                // Ensure this capability only defines 1 stage per target, else diagnose an error.
+                // This is a fatal error, do not allow toggling this error off.
+                entryPointAttr->capabilitySet = CapabilitySet(capName);
+                HashSet<CapabilityAtom> stageToBeUsed;
+                for (auto& targetSet : entryPointAttr->capabilitySet.getCapabilityTargetSets())
+                {
+                    for(auto& stageSet : targetSet.second.shaderStageSets)
+                        stageToBeUsed.add(stageSet.first);
+                }
 
-            entryPointAttr->stage = stage;
+                // TODO: Once profiles are removed in favor for `CapabilitySet`s we will beable to use more complex relationships,
+                // Until then we have an artificial limit that any capabilites used inside '[shader(...)]' must only specify 1 stage type 
+                // uniformly across targets.
+                if (stageToBeUsed.getCount() > 1)
+                {
+                    List<CapabilityAtom> atomsToPrint;
+                    atomsToPrint.reserve(stageToBeUsed.getCount());
+                    for (auto i : stageToBeUsed)
+                        atomsToPrint.add(i);
+                    getSink()->diagnose(attr, Diagnostics::capabilityHasMultipleStages, capNameString, atomsToPrint);
+                }
+                return entryPointAttr;
+            }
+            else
+            {
+                // always diagnose this error since nothing can compile with an invalid capability
+                getSink()->diagnose(attr, Diagnostics::unknownCapability, capNameString);
+                return false;
+            }
         }
         else if ((as<DomainAttribute>(attr)) ||
             (as<MaxTessFactorAttribute>(attr)) ||
-- 
cgit v1.2.3