Allow interface types to be used inside of structs (#966)

Previously, interface types were allowed to be used directly as function parameters, local variables, and global shader parameters.
Using an interface type as a field of a `struct` type or a `cbuffer` declaration was not implemented.
This change adds that support, and fixes several unrelated issues that caused problems in doing so.

* The most important work here was adding a case for `IRStructType` to `maybeSpecializeBindExistentialsType` that creates a specialized variant of a `struct` type on-demand based on specialization operands. This logic loops over the fields of the original struct, and creates new fields by binding the existentials/interfaces in the type of each field. Caching is used to ensure that the same `struct` type specialized to the same operands should yield the same result.

* To allow subsequent specialization to occur when a `struct` with interface-type fields is used, it was also necessary to specialize field-address and field-extract instructions in cases where the value that the field is being extracted from is a `wrapExistential`.

* Similarly, we neede to make sure that the logic for specializing called functions based on the concrete types for interfaces in the argument list would also take into account `struct` types with existential-type fields inside of them.

* Doing the above changes revealed some serious flaws in how the `ir-specialize.cpp` logic was tracking which instructions still needed to be processed. It had previously been assuming that it could assume any relevant instructions were on its work list, and when the work list went empty it could exit. This runs into two problems: (1) sometimes we create new instructions when specializing, and it may be impossible to ensure that all the new instructions (e.g., those created by utility routines in other files) get added to the work list, and (2) sometimes the instruction(s) that need to be re-visited when we specialize something aren't its direct users, but instead somethign that transitively depends on the instruction.

  These issues were fixed by two changes to the pass: (1) we now maintain a list of known "clean" instructions instead of implicitly using the work-list as a list of "dirty" instructions (so that implicitly any new instruction is dirty), and periodically iterating over all instructions to add the non-clean ones to the work list for processing, and (2) when an instruction is specialized/replaced we mark everything that transitively depends on it "dirty" (by removing it from the "clean" list).

* Added some logic to "fix up" the type of an IR function after changes that might modify its parameter list. Failing to have this logic meant that certain types were still live (because they were referenced by a function type) that couldn't actually be emitted as legal HLSL/GLSL.

* Added some special cases to IR instruction creation for `wrapExistential` and `BindExistentialsType` so that they act as no-ops when there are no "slots" providing specialization information. This helps avoid some special cases when specializing structure fields (since some fields specialization and others don't, so in general there are zero or more operands specific to each field).

* Added a test case that uses an interface type in a `cbuffer`, as well as an interface type in a `struct` passed as an entry-point `uniform` parameter.

* Fixed up some parts of the `.natvis` files to reflect naming changes from a previous PR and thus restore some of the useful Visual Studio debugging experience for Slang.

1 files changed, 431 insertions, 42 deletions

diff --git a/source/slang/ir-specialize.cpp b/source/slang/ir-specialize.cpp
index 57d14c11a..b57d2b58f 100644
--- a/source/slang/ir-specialize.cpp
+++ b/source/slang/ir-specialize.cpp
@@ -98,6 +98,9 @@ struct SpecializationContext
     // whether generic, existential, etc.
     //
     List<IRInst*> workList;
+    HashSet<IRInst*> workListSet;
+
+    HashSet<IRInst*> cleanInsts;
 
     void addToWorkList(
         IRInst* inst)
@@ -112,7 +115,14 @@ struct SpecializationContext
                 return;
         }
 
+        if(workListSet.Contains(inst))
+            return;
+
         workList.add(inst);
+        workListSet.Add(inst);
+        cleanInsts.Remove(inst);
+
+        addUsersToWorkList(inst);
     }
 
     // When a transformation makes a change to an instruction,
@@ -367,6 +377,7 @@ struct SpecializationContext
         case kIROp_Specialize:
         case kIROp_lookup_interface_method:
         case kIROp_ExtractExistentialType:
+        case kIROp_BindExistentialsType:
             break;
         }
     }
@@ -431,6 +442,13 @@ struct SpecializationContext
             maybeSpecializeLoad(as<IRLoad>(inst));
             break;
 
+        case kIROp_FieldExtract:
+            maybeSpecializeFieldExtract(as<IRFieldExtract>(inst));
+            break;
+        case kIROp_FieldAddress:
+            maybeSpecializeFieldAddress(as<IRFieldAddress>(inst));
+            break;
+
         case kIROp_BindExistentialsType:
             maybeSpecializeBindExistentialsType(as<IRBindExistentialsType>(inst));
             break;
@@ -566,12 +584,18 @@ struct SpecializationContext
         //
         addToWorkList(module->getModuleInst());
 
+        while(workList.getCount() != 0)
+        {
+
         // We will then iterate until our work list goes dry.
         //
         while(workList.getCount() != 0)
         {
             IRInst* inst = workList.getLast();
+
             workList.removeLast();
+            workListSet.Remove(inst);
+            cleanInsts.Add(inst);
 
             // For each instruction we process, we want to perform
             // a few steps.
@@ -610,6 +634,10 @@ struct SpecializationContext
             }
         }
 
+        addDirtyInstsToWorkListRec(module->getModuleInst());
+
+        }
+
         // Once the work list has gone dry, we should have the invariant
         // that there are no `specialize` instructions inside of non-generic
         // functions that in turn reference a generic type/function, *except*
@@ -617,6 +645,19 @@ struct SpecializationContext
         // which case we wouldn't want to specialize it anyway.
     }
 
+    void addDirtyInstsToWorkListRec(IRInst* inst)
+    {
+        if( !cleanInsts.Contains(inst) )
+        {
+            addToWorkList(inst);
+        }
+
+        for(auto child = inst->getLastChild(); child; child = child->getPrevInst())
+        {
+            addDirtyInstsToWorkListRec(child);
+        }
+    }
+
     // Given a `call` instruction in the IR, we need to detect the case
     // where the callee has some interface-type parameter(s) and at the
     // call site it is statically clear what concrete type(s) the arguments
@@ -719,6 +760,19 @@ struct SpecializationContext
                 auto witnessTable = makeExistential->getWitnessTable();
                 key.vals.add(witnessTable);
             }
+            else if( auto wrapExistential = as<IRWrapExistential>(arg) )
+            {
+                auto val = wrapExistential->getWrappedValue();
+                auto valType = val->getFullType();
+                key.vals.add(valType);
+
+                UInt slotOperandCount = wrapExistential->getSlotOperandCount();
+                for( UInt ii = 0; ii < slotOperandCount; ++ii )
+                {
+                    auto slotOperand = wrapExistential->getSlotOperand(ii);
+                    key.vals.add(slotOperand);
+                }
+            }
             else
             {
                 SLANG_UNEXPECTED("missing case for existential argument");
@@ -770,6 +824,11 @@ struct SpecializationContext
                     auto val = makeExistential->getWrappedValue();
                     newArgs.add(val);
                 }
+                else if( auto wrapExistential = as<IRWrapExistential>(arg) )
+                {
+                    auto val = wrapExistential->getWrappedValue();
+                    newArgs.add(val);
+                }
                 else
                 {
                     SLANG_UNEXPECTED("missing case for existential argument");
@@ -838,6 +897,13 @@ struct SpecializationContext
         if(as<IRMakeExistential>(inst))
             return true;
 
+        // A `wrapExistential(v, T0,w0, T1, w1, ...)` instruction
+        // is just a generalization of `makeExistential`, so it
+        // can apply in the same cases.
+        //
+        if(as<IRWrapExistential>(inst))
+            return true;
+
         // If we start to specialize functions that take arrays
         // of existentials as input, we will need a strategy to
         // determine arguments suitable for use in specializing
@@ -904,7 +970,66 @@ struct SpecializationContext
             //
             IRInst* replacementVal = nullptr;
 
-            if( !isExistentialType(oldParam->getDataType()) )
+            // The trickier case is when we have an existential-type
+            // parameter, because we need to extract out the concrete
+            // type that is coming from the call site.
+            //
+            if( auto oldMakeExistential = as<IRMakeExistential>(arg) )
+            {
+                // In this case, the `arg` is `makeExistential(val, witnessTable)`
+                // and we know that the specialized call site will just be
+                // passing in `val`.
+                //
+                auto val = oldMakeExistential->getWrappedValue();
+                auto witnessTable = oldMakeExistential->getWitnessTable();
+
+                // Our specialized function needs to take a parameter with the
+                // same type as `val`, to match the call site(s) that will be
+                // created.
+                //
+                auto valType = val->getFullType();
+                auto newParam = builder->createParam(valType);
+                newParams.add(newParam);
+
+                // Within the body of the function we cannot just use `val`
+                // directly, because the existing code expects an existential
+                // value, including its witness table.
+                //
+                // Therefore we will create a `makeExistential(newParam, witnessTable)`
+                // in the body of the new function and use *that* as the replacement
+                // value for the original parameter (since it will have the
+                // correct existential type, and stores the right witness table).
+                //
+                auto newMakeExistential = builder->emitMakeExistential(oldParam->getFullType(), newParam, witnessTable);
+                newBodyInsts.add(newMakeExistential);
+                replacementVal = newMakeExistential;
+            }
+            else if( auto oldWrapExistential = as<IRWrapExistential>(arg) )
+            {
+                auto val = oldWrapExistential->getWrappedValue();
+                auto valType = val->getFullType();
+
+                auto newParam = builder->createParam(valType);
+                newParams.add(newParam);
+
+                // Within the body of the function we cannot just use `val`
+                // directly, because the existing code expects an existential
+                // value, including its witness table.
+                //
+                // Therefore we will create a `makeExistential(newParam, witnessTable)`
+                // in the body of the new function and use *that* as the replacement
+                // value for the original parameter (since it will have the
+                // correct existential type, and stores the right witness table).
+                //
+                auto newWrapExistential = builder->emitWrapExistential(
+                    oldParam->getFullType(),
+                    newParam,
+                    oldWrapExistential->getSlotOperandCount(),
+                    oldWrapExistential->getSlotOperands());
+                newBodyInsts.add(newWrapExistential);
+                replacementVal = newWrapExistential;
+            }
+            else
             {
                 // For parameters that don't have an existential type,
                 // there is nothing interesting to do. The new function
@@ -915,47 +1040,6 @@ struct SpecializationContext
                 newParams.add(newParam);
                 replacementVal = newParam;
             }
-            else
-            {
-                // The trickier case is when we have an existential-type
-                // parameter, because we need to extract out the concrete
-                // type that is coming from the call site.
-                //
-                if( auto oldMakeExistential = as<IRMakeExistential>(arg) )
-                {
-                    // In this case, the `arg` is `makeExistential(val, witnessTable)`
-                    // and we know that the specialized call site will just be
-                    // passing in `val`.
-                    //
-                    auto val = oldMakeExistential->getWrappedValue();
-                    auto witnessTable = oldMakeExistential->getWitnessTable();
-
-                    // Our specialized function needs to take a parameter with the
-                    // same type as `val`, to match the call site(s) that will be
-                    // created.
-                    //
-                    auto valType = val->getFullType();
-                    auto newParam = builder->createParam(valType);
-                    newParams.add(newParam);
-
-                    // Within the body of the function we cannot just use `val`
-                    // directly, because the existing code expects an existential
-                    // value, including its witness table.
-                    //
-                    // Therefore we will create a `makeExistential(newParam, witnessTable)`
-                    // in the body of the new function and use *that* as the replacement
-                    // value for the original parameter (since it will have the
-                    // correct existential type, and stores the right witness table).
-                    //
-                    auto newMakeExistential = builder->emitMakeExistential(oldParam->getFullType(), newParam, witnessTable);
-                    newBodyInsts.add(newMakeExistential);
-                    replacementVal = newMakeExistential;
-                }
-                else
-                {
-                    SLANG_UNEXPECTED("missing case for existential argument");
-                }
-            }
 
             // Whatever replacement value was constructed, we need to
             // register it as the replacement for the original parameter.
@@ -1207,6 +1291,246 @@ struct SpecializationContext
         }
     }
 
+    UInt calcExistentialBoxSlotCount(IRType* type)
+    {
+    top:
+        if( as<IRExistentialBoxType>(type) )
+        {
+            return 2;
+        }
+        else if( auto ptrType = as<IRPtrTypeBase>(type) )
+        {
+            type = ptrType->getValueType();
+            goto top;
+        }
+        else if( auto ptrLikeType = as<IRPointerLikeType>(type) )
+        {
+            type = ptrLikeType->getElementType();
+            goto top;
+        }
+        else if( auto structType = as<IRStructType>(type) )
+        {
+            UInt count = 0;
+            for( auto field : structType->getFields() )
+            {
+                count += calcExistentialBoxSlotCount(field->getFieldType());
+            }
+            return count;
+        }
+        else
+        {
+            return 0;
+        }
+    }
+
+    void maybeSpecializeFieldExtract(IRFieldExtract* inst)
+    {
+        auto baseArg = inst->getBase();
+        auto fieldKey = inst->getField();
+
+        if( auto wrapInst = as<IRWrapExistential>(baseArg) )
+        {
+            // We have `getField(wrapExistential(val, ...), fieldKey)`
+            //
+            auto val = wrapInst->getWrappedValue();
+
+            // We know what type we are expected to produce.
+            //
+            auto resultType = inst->getFullType();
+
+            IRBuilder builder;
+            builder.sharedBuilder = &sharedBuilderStorage;
+            builder.setInsertBefore(inst);
+
+            // We'd *like* to replace this instruction with
+            // `wrapExistential(getField(val, fieldKey), ...)` instead, since that
+            // will enable subsequent specializations.
+            //
+            // To do that, we need to figure out:
+            //
+            // 1. What type that inner `getField` would return (what
+            // is the type of the `fieldKey` field in `val`?)
+            //
+            // 2. Which of the existential slot operands in `...` there
+            // actually apply to the given field.
+            //
+
+            // To determine these things, we need the type of
+            // `val` to be a structure type so that we can look
+            // up the field corresponding to `fieldKey`.
+            //
+            auto valType = val->getDataType();
+            auto valStructType = as<IRStructType>(valType);
+            if(!valStructType)
+                return;
+
+            UInt slotOperandOffset = 0;
+
+            IRStructField* foundField = nullptr;
+            for( auto valField : valStructType->getFields() )
+            {
+                if( valField->getKey() == fieldKey )
+                {
+                    foundField = valField;
+                    break;
+                }
+
+                slotOperandOffset += calcExistentialBoxSlotCount(valField->getFieldType());
+            }
+
+            if(!foundField)
+                return;
+
+            auto foundFieldType = foundField->getFieldType();
+
+            List<IRInst*> slotOperands;
+            UInt slotOperandCount = calcExistentialBoxSlotCount(foundFieldType);
+
+            for( UInt ii = 0; ii < slotOperandCount; ++ii )
+            {
+                slotOperands.add(wrapInst->getSlotOperand(slotOperandOffset + ii));
+            }
+
+            auto newGetField = builder.emitFieldExtract(
+                foundFieldType,
+                val,
+                fieldKey);
+
+            auto newWrapExistentialInst = builder.emitWrapExistential(
+                resultType,
+                newGetField,
+                slotOperandCount,
+                slotOperands.getBuffer());
+
+            addUsersToWorkList(inst);
+            inst->replaceUsesWith(newWrapExistentialInst);
+            inst->removeAndDeallocate();
+        }
+    }
+
+
+    void maybeSpecializeFieldAddress(IRFieldAddress* inst)
+    {
+        auto baseArg = inst->getBase();
+        auto fieldKey = inst->getField();
+
+        if( auto wrapInst = as<IRWrapExistential>(baseArg) )
+        {
+            // We have `getFieldAddr(wrapExistential(val, ...), fieldKey)`
+            //
+            auto val = wrapInst->getWrappedValue();
+
+            // We know what type we are expected to produce.
+            //
+            auto resultType = inst->getFullType();
+
+            IRBuilder builder;
+            builder.sharedBuilder = &sharedBuilderStorage;
+            builder.setInsertBefore(inst);
+
+            // We'd *like* to replace this instruction with
+            // `wrapExistential(getFieldAddr(val, fieldKey), ...)` instead, since that
+            // will enable subsequent specializations.
+            //
+            // To do that, we need to figure out:
+            //
+            // 1. What type that inner `getFieldAddr` would return (what
+            // is the type of the `fieldKey` field in `val`?)
+            //
+            // 2. Which of the existential slot operands in `...` there
+            // actually apply to the given field.
+            //
+
+            // To determine these things, we need the type of
+            // `val` to be a (pointer to a) structure type so that we can look
+            // up the field corresponding to `fieldKey`.
+            //
+            auto valType = tryGetPointedToType(&builder, val->getDataType());
+            if(!valType)
+                return;
+
+            auto valStructType = as<IRStructType>(valType);
+            if(!valStructType)
+                return;
+
+            UInt slotOperandOffset = 0;
+
+            IRStructField* foundField = nullptr;
+            for( auto valField : valStructType->getFields() )
+            {
+                if( valField->getKey() == fieldKey )
+                {
+                    foundField = valField;
+                    break;
+                }
+
+                slotOperandOffset += calcExistentialBoxSlotCount(valField->getFieldType());
+            }
+
+            if(!foundField)
+                return;
+
+            auto foundFieldType = foundField->getFieldType();
+
+            List<IRInst*> slotOperands;
+            UInt slotOperandCount = calcExistentialBoxSlotCount(foundFieldType);
+
+            for( UInt ii = 0; ii < slotOperandCount; ++ii )
+            {
+                slotOperands.add(wrapInst->getSlotOperand(slotOperandOffset + ii));
+            }
+
+            auto newGetFieldAddr = builder.emitFieldAddress(
+                builder.getPtrType(foundFieldType),
+                val,
+                fieldKey);
+
+            auto newWrapExistentialInst = builder.emitWrapExistential(
+                resultType,
+                newGetFieldAddr,
+                slotOperandCount,
+                slotOperands.getBuffer());
+
+            addUsersToWorkList(inst);
+            inst->replaceUsesWith(newWrapExistentialInst);
+            inst->removeAndDeallocate();
+        }
+    }
+
+    UInt calcExistentialTypeParamSlotCount(IRType* type)
+    {
+    top:
+        if( as<IRInterfaceType>(type) )
+        {
+            return 2;
+        }
+        else if( auto ptrType = as<IRPtrTypeBase>(type) )
+        {
+            type = ptrType->getValueType();
+            goto top;
+        }
+        else if( auto ptrLikeType = as<IRPointerLikeType>(type) )
+        {
+            type = ptrLikeType->getElementType();
+            goto top;
+        }
+        else if( auto structType = as<IRStructType>(type) )
+        {
+            UInt count = 0;
+            for( auto field : structType->getFields() )
+            {
+                count += calcExistentialTypeParamSlotCount(field->getFieldType());
+            }
+            return count;
+        }
+        else
+        {
+            return 0;
+        }
+    }
+
+    Dictionary<IRSimpleSpecializationKey, IRStructType*> existentialSpecializedStructs;
+
     void maybeSpecializeBindExistentialsType(IRBindExistentialsType* type)
     {
         auto baseType = type->getBaseType();
@@ -1253,17 +1577,82 @@ struct SpecializationContext
                 baseElementType,
                 slotOperandCount,
                 type->getExistentialArgs());
+            addToWorkList(wrappedElementType);
 
             auto newPtrLikeType = builder.getType(
                 basePtrLikeType->op,
                 1,
                 &wrappedElementType);
+            addToWorkList(newPtrLikeType);
 
             addUsersToWorkList(type);
             type->replaceUsesWith(newPtrLikeType);
             type->removeAndDeallocate();
             return;
         }
+        else if( auto baseStructType = as<IRStructType>(baseType) )
+        {
+            // In order to bind a `struct` type we will generate
+            // a new specialized `struct` type on demand and then
+            // cache and re-use it.
+            //
+            // We don't want to start specializing here unless
+            // all the operand types (and witness tables) we
+            // will be specializing to are themselves fully
+            // specialized, so that we can be sure that we
+            // have a unique type.
+            //
+            if( !areAllOperandsFullySpecialized(type) )
+                return;
+
+            // Now we we check to see if we've already created
+            // a specialized struct type or not.
+            //
+            IRSimpleSpecializationKey key;
+            key.vals.add(baseStructType);
+            for( UInt ii = 0; ii < slotOperandCount; ++ii )
+            {
+                key.vals.add(type->getExistentialArg(ii));
+            }
+
+            IRStructType* newStructType = nullptr;
+            if( !existentialSpecializedStructs.TryGetValue(key, newStructType) )
+            {
+                builder.setInsertBefore(baseStructType);
+                newStructType = builder.createStructType();
+
+                auto fieldSlotArgs = type->getExistentialArgs();
+
+                for( auto oldField : baseStructType->getFields() )
+                {
+                    // TODO: we need to figure out which of the specialization arguments
+                    // apply to this field...
+
+                    auto oldFieldType = oldField->getFieldType();
+                    auto fieldSlotArgCount = calcExistentialTypeParamSlotCount(oldFieldType);
+
+                    auto newFieldType = builder.getBindExistentialsType(
+                        oldFieldType,
+                        fieldSlotArgCount,
+                        fieldSlotArgs);
+
+                    addToWorkList(newFieldType);
+
+                    fieldSlotArgs += fieldSlotArgCount;
+
+                    builder.createStructField(newStructType, oldField->getKey(), newFieldType);
+                }
+
+                existentialSpecializedStructs.Add(key, newStructType);
+                addToWorkList(newStructType);
+            }
+
+            addUsersToWorkList(type);
+            type->replaceUsesWith(newStructType);
+            type->removeAndDeallocate();
+            return;
+
+        }
     }
 
     // The handling of specialization for global generic type