1 files changed, 64 insertions, 0 deletions
diff --git a/tools/slang-capability-generator/capability-generator-main.cpp b/tools/slang-capability-generator/capability-generator-main.cpp
index 79582a8c5..f62f7ce20 100644
--- a/tools/slang-capability-generator/capability-generator-main.cpp
+++ b/tools/slang-capability-generator/capability-generator-main.cpp
@@ -7,6 +7,7 @@
 #include "../../source/core/slang-secure-crt.h"
 #include "../../source/core/slang-string-util.h"
 #include "../../source/core/slang-file-system.h"
+#include "../../source/core/slang-uint-set.h"
 
 using namespace Slang;
 
@@ -389,8 +390,42 @@ void calcCanonicalRepresentations(const List<RefPtr<CapabilityDef>>& defs, const
         calcCanonicalRepresentation(def, mapEnumValueToDef);
 }
 
+const Index kUnusedEnumValue = -1;
+
+// Check if "def" uses a named ("name"/enumValueOfAbstract) abstract atom. If true, assign 
+// the enumValue of the found abstract atom (cache the unique ID) and increment counter.
+bool maybeProcessConcreteAtomForAbstractCapability(CapabilityDef* def, Index& enumValueOfAbstract, const String& name, Index& counter)
+{
+    if (def->getAbstractBase()
+        && (
+            (enumValueOfAbstract == def->getAbstractBase()->enumValue)
+            || 
+            (enumValueOfAbstract == kUnusedEnumValue && def->getAbstractBase()->name.equals(name))
+            )
+        )
+    {
+        counter++;
+        enumValueOfAbstract = def->getAbstractBase()->enumValue;
+        return true;
+    }
+    return false;
+}
+
+void outputUIntSetAsBufferValues(const String& nameOfBuffer, StringBuilder& resultBuilder, UIntSet& set)
+{
+    // store UIntSet::Element as uint8_t to stay sizeof(UIntSet::Element) independent.
+    // underlying type may change, bits stay the same.
+    resultBuilder << "const static CapabilityAtomSet " << nameOfBuffer << " = CapabilityAtomSet({\n";
+    for (auto i : set.getBuffer())
+    {
+        resultBuilder << "    UIntSet::Element(" << i << "),\n";
+    }
+    resultBuilder << "    0\n});\n";
+}
+
 SlangResult generateDefinitions(const List<RefPtr<CapabilityDef>>& defs, StringBuilder& sbHeader, StringBuilder& sbCpp)
 {
+   
     sbHeader << "enum class CapabilityAtom\n{\n";
     sbHeader << "    Invalid,\n";
     for (auto def : defs)
@@ -402,6 +437,7 @@ SlangResult generateDefinitions(const List<RefPtr<CapabilityDef>>& defs, StringB
     }
     sbHeader << "    Count\n";
     sbHeader << "};\n";
+  
     CapabilityDef* firstAbstractDef = nullptr;
     CapabilityDef* firstAliasDef = nullptr;
     sbHeader << "enum class CapabilityName\n{\n";
@@ -446,6 +482,31 @@ SlangResult generateDefinitions(const List<RefPtr<CapabilityDef>>& defs, StringB
     sbHeader << "    Count\n";
     sbHeader << "};\n";
 
+    Index enumValueOfTarget = kUnusedEnumValue;
+    Index targetCount = 0;
+    Index enumValueOfStage = kUnusedEnumValue;
+    Index stageCount = 0;
+
+    UIntSet anyTargetAtomSet{};
+    UIntSet anyStageAtomSet{};
+    StringBuilder anyTargetUIntSetHash;
+    StringBuilder anyStageUIntSetHash;
+
+    for (auto def : defs)
+    {
+        if (maybeProcessConcreteAtomForAbstractCapability(def.get(), enumValueOfTarget, "target", targetCount))
+            anyTargetAtomSet.add(def->enumValue);
+        else if (maybeProcessConcreteAtomForAbstractCapability(def.get(), enumValueOfStage, "stage", stageCount))
+            anyStageAtomSet.add(def->enumValue);
+    }
+    outputUIntSetAsBufferValues("kAnyTargetUIntSetBuffer", anyTargetUIntSetHash, anyTargetAtomSet);
+    outputUIntSetAsBufferValues("kAnyStageUIntSetBuffer", anyStageUIntSetHash, anyStageAtomSet);
+    
+    sbHeader << "\nenum {\n";
+    sbHeader << "    kCapabilityTargetCount = " << targetCount << ",\n";
+    sbHeader << "    kCapabilityStageCount = " << stageCount << ",\n";
+    sbHeader << "};\n\n";
+
     calcCanonicalRepresentations(defs, mapEnumValueToDef);
 
     List<String> capabiltiyNameArray;
@@ -500,6 +561,9 @@ SlangResult generateDefinitions(const List<RefPtr<CapabilityDef>>& defs, StringB
             conjunctions.add(serializeConjunction(c));
         def->serializedCanonicalRepresentation = serializeDisjunction(conjunctions);
     }
+    
+    sbCpp << anyTargetUIntSetHash;
+    sbCpp << anyStageUIntSetHash;
 
     sbCpp << "static CapabilityName kCapabilityArray[] = {\n";
     Index arrayIndex = 0;