1 files changed, 140 insertions, 61 deletions
diff --git a/source/slang/slang-serialize-ast.cpp b/source/slang/slang-serialize-ast.cpp
index be959b8a1..a11575849 100644
--- a/source/slang/slang-serialize-ast.cpp
+++ b/source/slang/slang-serialize-ast.cpp
@@ -1093,95 +1093,174 @@ void serialize(S const& serializer, DeclAssociationList& value)
     serialize(serializer, value.associations);
 }
 
+// Intermediate type for CapabilitySet serialization.
 //
-// The various types used to store capabilities on declarations
-// are all semantically equivalent to simpler types.
+// This essentially flattens the hierarchy found in CapabilitySet, making it
+// about twice as fast to deserialize. Deserializing capability sets was taking
+// up a significant part of AST deserialization time.
 //
+// This follows the same pattern as ASTModuleInfo and ContainerDeclDirectMemberDeclsInfo
+FIDDLE()
+struct CapabilitySetInfo
+{
+    FIDDLE(...)
 
-// A `CapabilityAtomSet` is an optimized representation of a
-// set of a `CapabilityAtom`s (which we can encode as just
-// a sequence).
-//
-SLANG_DECLARE_FOSSILIZED_AS(CapabilityAtomSet, List<CapabilityAtom>);
+    // Store the capability data as a list of entries
+    // Each entry contains: target, stage, and buffer data
+    struct Entry
+    {
+        CapabilityAtom target;
+        CapabilityAtom stage;
+        List<UInt64> bufferData;
+    };
 
-// A `CapabilityStateSet` can simply be encoded using its `atomSet` member.
-//
-SLANG_DECLARE_FOSSILIZED_AS_MEMBER(CapabilityStageSet, atomSet);
+    FIDDLE() List<Entry> entries;
+};
 
-// A `CapabilityStageSet` is really just a wrapper around a `CapabilityStageSets`
-// (which is itself just a dictionary of `CapabilityStateSet`s).
-//
-SLANG_DECLARE_FOSSILIZED_AS(CapabilityTargetSet, CapabilityStageSets);
+// Forward declare the intermediate type
+template<typename S>
+void serialize(S const& serializer, CapabilitySetInfo::Entry& value);
 
-// A `CapabilitySet` is really just a wrapper around a `CapabilityTargetSets`
-// (which is itself just a dictionary of `CapabilityTargetSet`s).
-//
-SLANG_DECLARE_FOSSILIZED_AS(CapabilitySet, CapabilityTargetSets);
+template<typename S>
+void serialize(S const& serializer, CapabilitySetInfo& value);
 
+// Declare fossilized representation for CapabilitySetInfo::Entry
+template<>
+struct FossilizedTypeTraits<CapabilitySetInfo::Entry>
+{
+    struct FossilizedType
+    {
+        Fossilized<CapabilityAtom> target;
+        Fossilized<CapabilityAtom> stage;
+        Fossilized<List<UInt64>> bufferData;
+    };
+};
+
+// Serialize a CapabilitySetInfo::Entry
 template<typename S>
-void serialize(S const& serializer, CapabilityAtomSet& value)
+void serialize(S const& serializer, CapabilitySetInfo::Entry& value)
 {
-    SLANG_SCOPED_SERIALIZER_ARRAY(serializer);
-    if (isWriting(serializer))
+    SLANG_SCOPED_SERIALIZER_STRUCT(serializer);
+    serialize(serializer, value.target);
+    serialize(serializer, value.stage);
+    serialize(serializer, value.bufferData);
+}
+
+// Declare fossilized representation for CapabilitySetInfo
+template<>
+struct FossilizedTypeTraits<CapabilitySetInfo>
+{
+    struct FossilizedType
     {
-        for (auto rawAtom : value)
-        {
-            auto atom = CapabilityAtom(rawAtom);
-            serialize(serializer, atom);
-        }
-    }
-    else
+        Fossilized<List<CapabilitySetInfo::Entry>> entries;
+    };
+};
+
+// Serialize a CapabilitySetInfo
+template<typename S>
+void serialize(S const& serializer, CapabilitySetInfo& value)
+{
+    SLANG_SCOPED_SERIALIZER_STRUCT(serializer);
+    serialize(serializer, value.entries);
+}
+
+// A `CapabilitySet` is serialized via the intermediate `CapabilitySetInfo` type
+SLANG_DECLARE_FOSSILIZED_AS(CapabilitySet, CapabilitySetInfo);
+
+// Helper function to collect CapabilitySetInfo from a CapabilitySet
+static CapabilitySetInfo _collectCapabilitySetInfo(const CapabilitySet& capSet)
+{
+    CapabilitySetInfo info;
+
+    // Collect all non-empty entries
+    for (const auto& targetPair : capSet.getCapabilityTargetSets())
     {
-        while (hasElements(serializer))
+        CapabilityAtom target = targetPair.first;
+        const auto& targetSet = targetPair.second;
+
+        for (const auto& stagePair : targetSet.shaderStageSets)
         {
-            CapabilityAtom atom = CapabilityAtom(0);
-            serialize(serializer, atom);
-            value.add(UInt(atom));
+            CapabilityAtom stage = stagePair.first;
+            const auto& stageSet = stagePair.second;
+
+            if (stageSet.atomSet && !stageSet.atomSet->isEmpty())
+            {
+                CapabilitySetInfo::Entry entry;
+                entry.target = target;
+                entry.stage = stage;
+
+                // Copy buffer data
+                const auto& buffer = stageSet.atomSet->getBuffer();
+                for (Index i = 0; i < buffer.getCount(); i++)
+                {
+                    entry.bufferData.add(UInt64(buffer[i]));
+                }
+
+                info.entries.add(entry);
+            }
         }
     }
-}
 
-template<typename S>
-void serialize(S const& serializer, CapabilityStageSet& value)
-{
-    serialize(serializer, value.atomSet);
+    return info;
 }
 
-template<typename S>
-void serialize(S const& serializer, CapabilityTargetSet& value)
+// Helper function to reconstruct a CapabilitySet from CapabilitySetInfo
+static void _reconstructCapabilitySet(CapabilitySet& capSet, const CapabilitySetInfo& info)
 {
-    serialize(serializer, value.shaderStageSets);
+    // Clear existing data
+    capSet.getCapabilityTargetSets().clear();
 
-    // The value for each entry in `shaderStageSets` have
-    // a `stage` field that is redundant with the key for
-    // that entry. Rather than serialize the key as part
-    // of the `CapabilityStageSet` type, we instead copy
-    // it over from the key to the value in the case where
-    // we are reading.
-    //
-    if (isReading(serializer))
+    // Reconstruct from entries
+    for (const auto& entry : info.entries)
     {
-        for (auto& p : value.shaderStageSets)
-            p.second.stage = p.first;
+        // Get references to target structures and populate directly
+        auto& targetSet = capSet.getCapabilityTargetSets()[entry.target];
+        targetSet.target = entry.target;
+
+        auto& stageSet = targetSet.shaderStageSets[entry.stage];
+        stageSet.stage = entry.stage;
+
+        // Create CapabilityAtomSet and reconstruct from buffer data
+        CapabilityAtomSet atomSet;
+        if (entry.bufferData.getCount() > 0)
+        {
+            atomSet.resizeBackingBufferDirectly(entry.bufferData.getCount());
+
+            for (Index i = 0; i < entry.bufferData.getCount(); i++)
+            {
+                if (entry.bufferData[i] != 0)
+                {
+                    atomSet.addRawElement(UIntSet::Element(entry.bufferData[i]), i);
+                }
+            }
+        }
+
+        stageSet.atomSet = atomSet;
     }
 }
 
+// Simplified CapabilitySet serialization using intermediate data structure
 template<typename S>
 void serialize(S const& serializer, CapabilitySet& value)
 {
-    serialize(serializer, value.getCapabilityTargetSets());
+    SLANG_PROFILE_SECTION(serialize_CapabilitySet);
 
-    // The value for each entry in `getCapabilityTargetSets()` have
-    // a `target` field that is redundant with the key for
-    // that entry. Rather than serialize the key as part
-    // of the `CapabilityTargetSet` type, we instead copy
-    // it over from the key to the value in the case where
-    // we are reading.
-    //
-    if (isReading(serializer))
+    if (isWriting(serializer))
+    {
+        // Collect intermediate representation
+        CapabilitySetInfo info = _collectCapabilitySetInfo(value);
+
+        // Serialize the intermediate representation
+        serialize(serializer, info);
+    }
+    else
     {
-        for (auto& p : value.getCapabilityTargetSets())
-            p.second.target = p.first;
+        // Deserialize the intermediate representation
+        CapabilitySetInfo info;
+        serialize(serializer, info);
+
+        // Reconstruct the CapabilitySet
+        _reconstructCapabilitySet(value, info);
     }
 }