Lookup refactor (#8467)

Close #8201.

This PR unify the lowering logic for LookupDeclRef of an interface
requirement. We will always lower this AST node to a
LookupWitness IR. The key of this IR is the special witnessTableType
`ThisTypeWitness`, this witness Table is simply a wrapper for an
interface type. Our current specialization pass doesn't handle this kind
of LookupWitness IR at all, so we will also add the specialization of
this_type IR as well.

4 files changed, 82 insertions, 30 deletions

diff --git a/source/slang/slang-ast-val.h b/source/slang/slang-ast-val.h
index a381410b6..7b9462f2f 100644
--- a/source/slang/slang-ast-val.h
+++ b/source/slang/slang-ast-val.h
@@ -56,8 +56,8 @@ public:
 };
 
 
-// Represent a lookup of SuperType::`m_decl` from `lookupSourceType` type that we know conforms to
-// SuperType.
+// Represent a lookup of witness of SuperType::`m_decl` from `lookupSource` type that we know
+// conforms to SuperType.
 FIDDLE()
 class LookupDeclRef : public DeclRefBase
 {
@@ -68,7 +68,7 @@ public:
     // The source type that we are looking up from.
     Type* getLookupSource() { return as<Type>(getOperand(1)); }
 
-    // Witness that `lookupSourceType`:SuperType.
+    // Witness that `lookupSource`:SuperType.
     SubtypeWitness* getWitness() { return as<SubtypeWitness>(getOperand(2)); }
 
     LookupDeclRef(Decl* declToLookup, Type* lookupSource, SubtypeWitness* witness)
diff --git a/source/slang/slang-ir-lower-witness-lookup.cpp b/source/slang/slang-ir-lower-witness-lookup.cpp
index 8aa42ab41..8d6979913 100644
--- a/source/slang/slang-ir-lower-witness-lookup.cpp
+++ b/source/slang/slang-ir-lower-witness-lookup.cpp
@@ -98,6 +98,7 @@ struct WitnessLookupLoweringContext
         case kIROp_StructType:
         case kIROp_ClassType:
         case kIROp_InterfaceType:
+        case kIROp_LookupWitnessMethod:
             return (IRType*)type;
         default:
             {
diff --git a/source/slang/slang-ir-specialize.cpp b/source/slang/slang-ir-specialize.cpp
index ed96b23c1..bc7bcab12 100644
--- a/source/slang/slang-ir-specialize.cpp
+++ b/source/slang/slang-ir-specialize.cpp
@@ -845,10 +845,31 @@ struct SpecializationContext
         // the result of a `specialize` instruction or other
         // operation that will yield such a table.
         //
+        // Since we unify the frontend such that all LookupDeclRef node
+        // on a interface requirement will always be lowered to lookup
+        // witness, it creats an exception that IRThisTypeWitness, a non-concrete
+        // witness table, also need to be specialized. Otherwise, there is
+        // no logic in the later passes can handle it. However, we know for
+        // sure that the IRThisTypeWitness is only used to wrap an interface
+        // type, therefore, it must be only used to exact the interface requirement.
         auto witnessTable = as<IRWitnessTable>(lookupInst->getWitnessTable());
+        IRInterfaceType* interfaceType = nullptr;
         if (!witnessTable)
         {
-            return false;
+            if (auto thisTypeWitness = as<IRThisTypeWitness>(lookupInst->getWitnessTable()))
+            {
+                if (auto witnessTableType =
+                        as<IRWitnessTableTypeBase>(thisTypeWitness->getDataType()))
+                {
+                    if (!areAllOperandsFullySpecialized(witnessTableType))
+                        return false;
+
+                    interfaceType = as<IRInterfaceType>(witnessTableType->getConformanceType());
+                }
+            }
+
+            if (!interfaceType)
+                return false;
         }
 
         // Because we have a concrete witness table, we can
@@ -856,7 +877,26 @@ struct SpecializationContext
         // the given interface requirement.
         //
         auto requirementKey = lookupInst->getRequirementKey();
-        auto satisfyingVal = findWitnessVal(witnessTable, requirementKey);
+        IRInst* satisfyingVal = nullptr;
+
+        if (witnessTable)
+            satisfyingVal = findWitnessVal(witnessTable, requirementKey);
+        else
+        {
+            // If we are specializing ThisTypeWitness, the result of the specialization
+            // could be a WitnessTabelType, in such case, in order to not break the generality
+            // the specialziation (we don't specialize WitnessTableType here), we will wrap it
+            // into another ThisTypeWitnes and handle it later.
+            satisfyingVal = findInterfaceRequirement(interfaceType, requirementKey);
+            if (auto witnessTableType = as<IRWitnessTableType>(satisfyingVal))
+            {
+                auto newInterfaceType = as<IRInterfaceType>(witnessTableType->getConformanceType());
+                IRBuilder builderStorage(module);
+                IRBuilder* builder = &builderStorage;
+                builder->setInsertBefore(lookupInst);
+                satisfyingVal = builder->createThisTypeWitness(newInterfaceType);
+            }
+        }
 
         // We expect to always find a satisfying value, but
         // we will go ahead and code defensively so that
@@ -915,6 +955,7 @@ struct SpecializationContext
         }
         return nullptr;
     }
+
     template<typename TDict>
     void _readSpecializationDictionaryImpl(TDict& dict, IRInst* dictInst)
     {
diff --git a/source/slang/slang-lower-to-ir.cpp b/source/slang/slang-lower-to-ir.cpp
index 3cf5d7803..ca32b0f2d 100644
--- a/source/slang/slang-lower-to-ir.cpp
+++ b/source/slang/slang-lower-to-ir.cpp
@@ -11749,6 +11749,33 @@ bool isAbstractWitnessTable(IRInst* inst)
     return false;
 }
 
+static IRInst* maybeCloneThisTypeWitness(
+    IRGenContext* context,
+    IRInst* thisTypeWitness,
+    Type* thisType)
+{
+    auto currentInsertLoc = context->irBuilder->getInsertLoc().getParent();
+    auto parentOfThisTypeWitness = thisTypeWitness->parent;
+
+    while (currentInsertLoc != nullptr)
+    {
+        // If current insert location is same as scope of ThisTypeWitness, don't copy it.
+        if (parentOfThisTypeWitness == currentInsertLoc)
+        {
+            return thisTypeWitness;
+        }
+
+        currentInsertLoc = currentInsertLoc->parent;
+    }
+
+    auto thisTypeIR = as<IRThisType>(lowerType(context, thisType));
+    SLANG_RELEASE_ASSERT(thisTypeIR);
+
+    auto newThisTypeWitness =
+        context->irBuilder->createThisTypeWitness((IRType*)thisTypeIR->getConstraintType());
+    return newThisTypeWitness;
+}
+
 LoweredValInfo emitDeclRef(IRGenContext* context, Decl* decl, DeclRefBase* subst, IRType* type)
 {
     const auto initialSubst = subst;
@@ -11898,40 +11925,23 @@ LoweredValInfo emitDeclRef(IRGenContext* context, Decl* decl, DeclRefBase* subst
             // witness table for the concrete type that conforms to `ISomething<Foo>`.
             //
             auto irWitnessTable = lowerSimpleVal(context, thisTypeSubst->getWitness());
+            SLANG_RELEASE_ASSERT(irWitnessTable);
+
             if (isAbstractWitnessTable(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 unifying our IR representation to
-                // represent associated types with lookupWitness inst even inside
-                // interface definitions.
-                return emitDeclRef(
+                // Copy ThisTypeWitness locally if necessary
+                irWitnessTable = maybeCloneThisTypeWitness(
                     context,
-                    decl->getDefaultDeclRef(),
-                    context->irBuilder->getTypeKind());
+                    irWitnessTable,
+                    thisTypeSubst->getLookupSource());
             }
 
-            SLANG_RELEASE_ASSERT(irWitnessTable);
-
-            //
-            // The key to use for looking up the interface member is
-            // derived from the declaration.
-            //
             auto irRequirementKey = getInterfaceRequirementKey(context, decl);
-            //
-            // Those two pieces of information tell us what we need to
-            // do in order to look up the value that satisfied the requirement.
-            //
-            auto irSatisfyingVal = context->irBuilder->emitLookupInterfaceMethodInst(
+            auto irLookupWitness = context->irBuilder->emitLookupInterfaceMethodInst(
                 type,
                 irWitnessTable,
                 irRequirementKey);
-            return LoweredValInfo::simple(irSatisfyingVal);
+            return LoweredValInfo::simple(irLookupWitness);
         }
         else
         {