Add sequential ID cache in Linkage for witness tables and RTTI objects. (#1590)

5 files changed, 127 insertions, 0 deletions

diff --git a/source/slang/slang-compiler.h b/source/slang/slang-compiler.h
index e120d5849..79ef9df27 100755
--- a/source/slang/slang-compiler.h
+++ b/source/slang/slang-compiler.h
@@ -1216,6 +1216,10 @@ namespace Slang
             slang::TypeReflection* type,
             slang::TypeReflection* interfaceType,
             ISlangBlob** outNameBlob) override;
+        SLANG_NO_THROW SlangResult SLANG_MCALL getTypeConformanceWitnessSequentialID(
+            slang::TypeReflection* type,
+            slang::TypeReflection* interfaceType,
+            uint32_t*              outId) override;
         SLANG_NO_THROW SlangResult SLANG_MCALL createCompileRequest(
             SlangCompileRequest**   outCompileRequest) override;
 
@@ -1286,6 +1290,13 @@ namespace Slang
         // Map from the logical name of a module to its definition
         Dictionary<Name*, RefPtr<LoadedModule>> mapNameToLoadedModules;
 
+        // Map from the mangled name of RTTI objects to sequential IDs
+        // used by `switch`-based dynamic dispatch.
+        Dictionary<String, uint32_t> mapMangledNameToRTTIObjectIndex;
+
+        // Counters for allocating sequential IDs to witness tables conforming to each interface type.
+        Dictionary<String, uint32_t> mapInterfaceMangledNameToSequentialIDCounters;
+
         // The resulting specialized IR module for each entry point request
         List<RefPtr<IRModule>> compiledModules;
 
diff --git a/source/slang/slang-ir-inst-defs.h b/source/slang/slang-ir-inst-defs.h
index 340f18c8c..b225cbf0b 100644
--- a/source/slang/slang-ir-inst-defs.h
+++ b/source/slang/slang-ir-inst-defs.h
@@ -560,6 +560,7 @@ INST(HighLevelDeclDecoration,               highLevelDecl,          1, 0)
     /* Decorations for RTTI objects */
         INST(RTTITypeSizeDecoration, RTTI_typeSize, 1, 0)
     INST(AnyValueSizeDecoration, AnyValueSize, 1, 0)
+    INST(SequentialIDDecoration, SequentialIDDecoration, 1, 0)
 
     INST(TypeConstraintDecoration, TypeConstraintDecoration, 1, 0)
     INST_RANGE(LinkageDecoration, ImportDecoration, ExportDecoration)
diff --git a/source/slang/slang-ir-insts.h b/source/slang/slang-ir-insts.h
index af0f737ea..aedcea232 100644
--- a/source/slang/slang-ir-insts.h
+++ b/source/slang/slang-ir-insts.h
@@ -454,6 +454,18 @@ struct IRBuiltinDecoration : IRDecoration
     IR_LEAF_ISA(BuiltinDecoration)
 };
 
+struct IRSequentialIDDecoration : IRDecoration
+{
+    enum
+    {
+        kOp = kIROp_SequentialIDDecoration
+    };
+    IR_LEAF_ISA(SequentialIDDecoration)
+
+    IRIntLit* getSequentialIDOperand() { return cast<IRIntLit>(getOperand(0)); }
+    IRIntegerValue getSequentialID() { return getSequentialIDOperand()->getValue(); }
+};
+
 // An instruction that specializes another IR value
 // (representing a generic) to a particular set of generic arguments 
 // (instructions representing types, witness tables, etc.)
@@ -2533,6 +2545,11 @@ struct IRBuilder
     {
         addDecoration(inst, kIROp_BuiltinDecoration);
     }
+
+    void addSequentialIDDecoration(IRInst* inst, IRIntegerValue id)
+    {
+        addDecoration(inst, kIROp_SequentialIDDecoration, getIntValue(getUIntType(), id));
+    }
 };
 
 void addHoistableInst(
diff --git a/source/slang/slang-ir-specialize-dispatch.cpp b/source/slang/slang-ir-specialize-dispatch.cpp
index 0c519427d..05ed867bd 100644
--- a/source/slang/slang-ir-specialize-dispatch.cpp
+++ b/source/slang/slang-ir-specialize-dispatch.cpp
@@ -114,10 +114,78 @@ void specializeDispatchFunction(SharedGenericsLoweringContext* sharedContext, IR
     returnInst->removeAndDeallocate();
 }
 
+// Ensures every witness table object has been assigned a sequential ID.
+// All witness tables will have a SequentialID decoration after this function is run.
+// The sequantial ID in the decoration will be the same as the one specified in the Linkage.
+// Otherwise, a new ID will be generated and assigned to the witness table object, and
+// the sequantial ID map in the Linkage will be updated to include the new ID, so they
+// can be looked up by the user via future Slang API calls.
+void ensureWitnessTableSequentialIDs(SharedGenericsLoweringContext* sharedContext)
+{
+    auto linkage = sharedContext->targetReq->getLinkage();
+    for (auto inst : sharedContext->module->getGlobalInsts())
+    {
+        if (inst->op == kIROp_WitnessTable)
+        {
+            UnownedStringSlice witnessTableMangledName;
+            if (auto instLinkage = inst->findDecoration<IRLinkageDecoration>())
+            {
+                witnessTableMangledName = instLinkage->getMangledName();
+            }
+            else
+            {
+                // If this witness table entry does not have a linkage,
+                // don't assign sequential ID for it.
+                continue;
+            }
+
+            // If the inst already has a SequentialIDDecoration, stop now.
+            if (inst->findDecoration<IRSequentialIDDecoration>())
+                continue;
+
+            // Get a sequential ID for the witness table using the map from the Linkage.
+            uint32_t seqID = 0;
+            if (!linkage->mapMangledNameToRTTIObjectIndex.TryGetValue(
+                witnessTableMangledName, seqID))
+            {
+                auto interfaceType =
+                    cast<IRWitnessTableType>(inst->getDataType())->getConformanceType();
+                auto interfaceLinkage = interfaceType->findDecoration<IRLinkageDecoration>();
+                SLANG_ASSERT(
+                    interfaceLinkage && "An interface type does not have a linkage,"
+                                        "but a witness table associated with it has one.");
+                auto interfaceName = interfaceLinkage->getMangledName();
+                auto idAllocator =
+                    linkage->mapInterfaceMangledNameToSequentialIDCounters.TryGetValue(
+                        interfaceName);
+                if (!idAllocator)
+                {
+                    linkage->mapInterfaceMangledNameToSequentialIDCounters[interfaceName] = 0;
+                    idAllocator =
+                        linkage->mapInterfaceMangledNameToSequentialIDCounters.TryGetValue(
+                            interfaceName);
+                }
+                seqID = *idAllocator;
+                ++(*idAllocator);
+                linkage->mapMangledNameToRTTIObjectIndex[witnessTableMangledName] = seqID;
+            }
+
+            // Add a decoration to the inst.
+            IRBuilder builder;
+            builder.sharedBuilder = &sharedContext->sharedBuilderStorage;
+            builder.setInsertBefore(inst);
+            builder.addSequentialIDDecoration(inst, seqID);
+        }
+    }
+}
+
 void specializeDispatchFunctions(SharedGenericsLoweringContext* sharedContext)
 {
     sharedContext->sharedBuilderStorage.deduplicateAndRebuildGlobalNumberingMap();
 
+    // First we ensure that all witness table objects has a sequential ID assigned.
+    ensureWitnessTableSequentialIDs(sharedContext);
+
     for (auto kv : sharedContext->mapInterfaceRequirementKeyToDispatchMethods)
     {
         auto dispatchFunc = kv.Value;
diff --git a/source/slang/slang.cpp b/source/slang/slang.cpp
index 376376f24..ad4e82d4f 100644
--- a/source/slang/slang.cpp
+++ b/source/slang/slang.cpp
@@ -658,6 +658,36 @@ SLANG_NO_THROW SlangResult SLANG_MCALL Linkage::getTypeConformanceWitnessMangled
     return SLANG_OK;
 }
 
+SLANG_NO_THROW SlangResult SLANG_MCALL Linkage::getTypeConformanceWitnessSequentialID(
+    slang::TypeReflection* type,
+    slang::TypeReflection* interfaceType,
+    uint32_t* outId)
+{
+    auto subType = asInternal(type);
+    auto supType = asInternal(interfaceType);
+    auto name = getMangledNameForConformanceWitness(subType->getASTBuilder(), subType, supType);
+    auto interfaceName = getMangledTypeName(supType->getASTBuilder(), supType);
+    uint32_t resultIndex = 0;
+    if (mapMangledNameToRTTIObjectIndex.TryGetValue(name, resultIndex))
+    {
+        if (outId)
+            *outId = resultIndex;
+        return SLANG_OK;
+    }
+    auto idAllocator = mapInterfaceMangledNameToSequentialIDCounters.TryGetValue(interfaceName);
+    if (!idAllocator)
+    {
+        mapInterfaceMangledNameToSequentialIDCounters[interfaceName] = 0;
+        idAllocator = mapInterfaceMangledNameToSequentialIDCounters.TryGetValue(interfaceName);
+    }
+    resultIndex = (*idAllocator);
+    ++(*idAllocator);
+    mapMangledNameToRTTIObjectIndex[name] = resultIndex;
+    if (outId)
+        *outId = resultIndex;
+    return SLANG_OK;
+}
+
 SLANG_NO_THROW SlangResult SLANG_MCALL Linkage::createCompileRequest(
     SlangCompileRequest**   outCompileRequest)
 {