Add `IRThisTypeWitness` to stand in for witness lookups inside an interface definition. (#3316)

Co-authored-by: Yong He <yhe@nvidia.com>

5 files changed, 49 insertions, 21 deletions

diff --git a/source/slang/slang-ir-inst-defs.h b/source/slang/slang-ir-inst-defs.h
index 026b8b110..e32f0d084 100644
--- a/source/slang/slang-ir-inst-defs.h
+++ b/source/slang/slang-ir-inst-defs.h
@@ -261,6 +261,9 @@ INST(StructKey, key, 0, GLOBAL)
 INST(GlobalGenericParam, global_generic_param, 0, GLOBAL)
 INST(WitnessTable, witness_table, 0, 0)
 
+// A placeholder witness that ThisType implements the enclosing interface.
+// Used only in interface definitions.
+INST(ThisTypeWitness, thisTypeWitness, 1, 0)
 INST(GlobalHashedStringLiterals, global_hashed_string_literals, 0, 0)
 
 INST(Module, module, 0, PARENT)
diff --git a/source/slang/slang-ir-insts.h b/source/slang/slang-ir-insts.h
index be6228f2d..a2662d53f 100644
--- a/source/slang/slang-ir-insts.h
+++ b/source/slang/slang-ir-insts.h
@@ -3703,6 +3703,8 @@ public:
         IRInst*        requirementKey,
         IRInst*        satisfyingVal);
 
+    IRInst* createThisTypeWitness(IRType* interfaceType);
+
     IRInterfaceRequirementEntry* createInterfaceRequirementEntry(
         IRInst* requirementKey,
         IRInst* requirementVal);
diff --git a/source/slang/slang-ir.cpp b/source/slang/slang-ir.cpp
index 2f603ac17..c40a62941 100644
--- a/source/slang/slang-ir.cpp
+++ b/source/slang/slang-ir.cpp
@@ -4381,6 +4381,17 @@ namespace Slang
         return entry;
     }
 
+    IRInst* IRBuilder::createThisTypeWitness(IRType* interfaceType)
+    {
+        IRInst* witness = createInst<IRThisTypeWitness>(
+            this,
+            kIROp_ThisTypeWitness,
+            getWitnessTableType(interfaceType));
+        addGlobalValue(this, witness);
+        return witness;
+    }
+
+
     IRStructType* IRBuilder::createStructType()
     {
         IRStructType* structType = createInst<IRStructType>(
@@ -7691,6 +7702,7 @@ namespace Slang
         case kIROp_GlobalParam:
         case kIROp_StructKey:
         case kIROp_GlobalGenericParam:
+        case kIROp_ThisTypeWitness:
         case kIROp_WitnessTable:
         case kIROp_WitnessTableEntry:
         case kIROp_InterfaceRequirementEntry:
diff --git a/source/slang/slang-ir.h b/source/slang/slang-ir.h
index cb744d4da..8d7cad0b4 100644
--- a/source/slang/slang-ir.h
+++ b/source/slang/slang-ir.h
@@ -1802,6 +1802,16 @@ struct IRThisType : IRType
     }
 };
 
+struct IRThisTypeWitness : IRInst
+{
+    IR_LEAF_ISA(ThisTypeWitness)
+
+    IRInst* getConstraintType()
+    {
+        return getOperand(0);
+    }
+};
+
 struct IRInterfaceRequirementEntry : IRInst
 {
     IRInst* getRequirementKey() { return getOperand(0); }
diff --git a/source/slang/slang-lower-to-ir.cpp b/source/slang/slang-lower-to-ir.cpp
index 8851195b3..b2c0dbb18 100644
--- a/source/slang/slang-lower-to-ir.cpp
+++ b/source/slang/slang-lower-to-ir.cpp
@@ -1659,17 +1659,7 @@ struct ValLoweringVisitor : ValVisitor<ValLoweringVisitor, LoweredValInfo, Lower
         // produce transitive witnesses in shapes that will cuase us
         // problems here.
         //
-        if (!baseWitnessTable)
-        {
-            // If we don't have a valid baseWitnessTable,
-            // This can happen when we are looking up an associatedtype defined in the base interface from
-            // a derived interface that inherits the base interface.
-            // For now, we just emit a null witness.
-            // In the future, we may want to consider lower `ThisTypeConstraint` into IR as something like
-            // `IRThisTypeWitness`, and emit an explicit lookup through that witness instead.
-            SLANG_RELEASE_ASSERT(as<ThisType>(val->getSub()));
-            return LoweredValInfo();
-        }
+        SLANG_RELEASE_ASSERT(baseWitnessTable);
 
         if (auto declaredMidToSup = as<DeclaredSubtypeWitness>(val->getMidToSup()))
         {
@@ -1925,7 +1915,9 @@ struct ValLoweringVisitor : ValVisitor<ValLoweringVisitor, LoweredValInfo, Lower
         // add them as operands of the IR type.
         SubstitutionSet(type->getDeclRef()).forEachSubstitutionArg([&](Val* arg)
             {
-                operands.add(lowerVal(context, arg).val);
+                auto argVal = lowerVal(context, arg).val;
+                SLANG_ASSERT(argVal);
+                operands.add(argVal);
             });
         return getBuilder()->getType(
             op,
@@ -7810,9 +7802,8 @@ struct DeclLoweringVisitor : DeclVisitor<DeclLoweringVisitor, LoweredValInfo>
             context->astBuilder,
             createDefaultSpecializedDeclRef(subContext, nullptr, decl->getThisTypeDecl()));
         subContext->thisType = thisType;
-
-        // TODO: Need to add an appropriate stand-in witness here.
-        subContext->thisTypeWitness = nullptr;
+        // Create a stand-in witness that represents `ThisType` conforms to the interface.
+        subContext->thisTypeWitness = subBuilder->createThisTypeWitness((IRType*)finalVal);
 
         // Lower associated types first, so they can be referred to when lowering functions.
         for (auto assocTypeDecl : decl->getMembersOfType<AssocTypeDecl>())
@@ -9712,6 +9703,15 @@ static void _addFlattenedTupleArgs(
     }
 }
 
+bool isAbstractWitnessTable(IRInst* inst)
+{
+    if (as<IRThisTypeWitness>(inst))
+        return true;
+    if (auto lookup = as<IRLookupWitnessMethod>(inst))
+        return isAbstractWitnessTable(lookup->getWitnessTable());
+    return false;
+}
+
 LoweredValInfo emitDeclRef(
     IRGenContext*           context,
     Decl*            decl,
@@ -9853,25 +9853,26 @@ LoweredValInfo emitDeclRef(
             // witness table for the concrete type that conforms to `ISomething<Foo>`.
             //
             auto irWitnessTable = lowerSimpleVal(context, thisTypeSubst->getWitness());
-            if (!irWitnessTable)
+            if (isAbstractWitnessTable(irWitnessTable))
             {
-                // If `thisTypeSubst` doesn't lower into an IRWitnessTable,
+                // If `thisTypeSubst` doesn't lower into a concrete IRWitnessTable,
                 // this is a lookup of an interface requirement
                 // defined in some base interface from an interface type.
                 // For now we just lower that decl as if it is referenced
                 // from the same interface directly, e.g. a reference to
                 // IBase.AssocType from IDerived:IBase will be lowered as
                 // IRAssocType(IBase).
-                // We may want to consider extend our IR representation to
-                // have a `IRThisTypeWitness` object, so we can lower this case
-                // into an explicit lookup from `IRThisTypeWitness`,
-                // just like any other cases.
+                // We may want to consider unifying our IR representation to
+                // represent associated types with lookupWitness inst even inside
+                // interface definitions.
                 return emitDeclRef(
                     context,
                     createDefaultSpecializedDeclRef(context, nullptr, decl),
                     context->irBuilder->getTypeKind());
             }
 
+            SLANG_RELEASE_ASSERT(irWitnessTable);
+
             //
             // The key to use for looking up the interface member is
             // derived from the declaration.