Allow plugging in types with resources for interface parameters (#913)

* Allow plugging in types with resources for interface parameters

The key feature enabled by this change is that you can take a shader declared with interface-type parameters:

```hlsl
ConstantBuffer<ILight> gLight;
float4 myShader(IMaterial material, ...)
{ ... }
```

and specialize its interface-type parameters to concrete type that can contain resources like textures, samplers, etc.

The hard part of doing this layout is that we need to support signatures that include a mix of interface and non-interface types. Imagine this contrived example:

```hlsl
float4 myShader(
    Texture2D diffuseMap,
    ILight        light,
    Texture2D specularMap)
{ ... }
```

We end up wanting `diffuseMap` to get `register(t0)` and `specularMap` to get `register(t1)`, so that they have the same location no matter what we plug in for `light`.
But if we plug in a concrete type for `light` that needs a texture register, we need to allocate it *somewhere*.
We handle this by having the `TypeLayout` for `light` come back with a "primary" type layout that doesn't have any texture registers, but with a "pending" type layout that includes the texture register requirements of whatever concrete type we plug in.

This split between "primary" and "pending" layout then needs to work its way up the hierarchy, so that an aggregate `struct` type with a mix of interface and non-interface fields (recursively), needs to compute an aggregate "primary type layout" and an aggregate "pending type layout," and then each field needs to be able to compute its offset in the primary/pending layout of the aggregate.

A large chunk of the work in this PR is then just implementing the split between primary and pending data, and ensuring that layouts are computed appropriately.

The next catch is that when a "parameter group" (either a parameter block or constant buffer) contains one or more values of interface type, then we can allow the parameter group to "mask" some of the resource usage of the concrete types we plug in, but others "bleed through."

For example, if we have:

```hlsl
struct MyStuff { float3 color; ILight light; }
ConstantBuffer<MyStuff> myStuff;

struct SpotLight { float3 position; Texture2D shadowMap; }
``

If we plug in the `SpotLight` type for `myStuff.light`, then the `float3` data for the light can be "masked" by the fact that we have a constant buffer (we can just allocate the `float3` `position` right after `color`), but the `Texture2D` needed for `shadowMap` needs to "bleed through" and become "pending" data for the `myStuff` shader parameter.

Adding support for that detail more or less required a full rewrite of the logic for allocating parameter group type layouts.

The next detail is that when we go to legalize a declaration like the `myStuff` buffer, we will end up with something like:

```hlsl
struct MyStuff_stripped { float3 color; }
struct Wrapped
{
    MyStuff_stripped primary;
    SpotLight pending;
}
ConstantBuffer<Wrapped> myStuff;
```

This "wrapped" version of the buffer type more accurately reflects the layout we need/want for the uniform/ordinary data, but in order to further legalize it and pull out the resource-type fields like `shadowMap` we need to have accurate layout information, and the problem is that layout information for the original buffer can't apply to this new "wrapped" buffer.

The last major piece of this change is logic that runs during existential type legalization to compute new layouts for "wrapped" buffers like these that embeds correct offset/binding/register information for any resources nested inside them. A key challenge in that code is that existential legalization needs to erase any "pending" data from the program entirely, so that offset information that used to be relatie to the "pending" part of a surrounding type now needs to be relative to the primary part.

The work here may not be 100% complete for all scenarios, but it does well enough on the new and existing tests that I want to checkpoint it. Note that a few other tests have had their output changed, but in all cases I've reviewed the diffs and determined that the change in observable behavior is consistent with what we intened Slang's behavior to be.

Note that there is still one major piece of support for interface-type parameters that is missing here, and which might force us to revisit some of the decisions in this code: we don't properly support user-defined `struct` types with interface-type fields.

* fixup: typos

1 files changed, 73 insertions, 6 deletions

diff --git a/source/slang/legalize-types.cpp b/source/slang/legalize-types.cpp
index 6e459ecf1..995b797e4 100644
--- a/source/slang/legalize-types.cpp
+++ b/source/slang/legalize-types.cpp
@@ -591,6 +591,18 @@ LegalType createLegalUniformBufferTypeForResources(
 {
     switch (legalElementType.flavor)
     {
+    case LegalType::Flavor::simple:
+        {
+            // Seeing a simple type here means that it must be a
+            // "special" type (a resource type or array thereof)
+            // because otherwise the catch-all behavior in
+            // `createLegalUniformBufferType()` would have handled it.
+            //
+            // This case is the same as what we do for tuple elements below.
+            //
+            return LegalType::implicitDeref(legalElementType);
+        }
+
     case LegalType::Flavor::pair:
         {
             auto pairType = legalElementType.getPair();
@@ -1080,6 +1092,10 @@ LegalType legalizeTypeImpl(
     TypeLegalizationContext*    context,
     IRType*                     type)
 {
+    if(!type)
+        return LegalType::simple(nullptr);
+
+    context->builder->setInsertBefore(type);
 
     if (auto uniformBufferType = as<IRUniformParameterGroupType>(type))
     {
@@ -1097,9 +1113,27 @@ LegalType legalizeTypeImpl(
         // we'll want to completely eliminate the uniform/ordinary
         // part.
 
+        auto originalElementType = uniformBufferType->getElementType();
+
         // Legalize the element type to see what we are working with.
         auto legalElementType = legalizeType(context,
-            uniformBufferType->getElementType());
+            originalElementType);
+
+        // As a bit of a corner case, if the user requested something
+        // like `ConstantBuffer<Texture2D>` the element type would
+        // legalize to a "simple" type, and that would be interpreted
+        // as an *ordinary* type, but we really need to notice the
+        // case when the element type is simple, but *special*.
+        //
+        if( context->isSpecialType(originalElementType) )
+        {
+            // Anything that has a special element type needs to
+            // be handled by the pass-specific logic in the context.
+            //
+            return context->createLegalUniformBufferType(
+                uniformBufferType->op,
+                legalElementType);
+        }
 
         // Note that even when legalElementType.flavor == Simple
         // we still need to create a new uniform buffer type
@@ -1270,7 +1304,7 @@ RefPtr<TypeLayout> getDerefTypeLayout(
 
 RefPtr<VarLayout> getFieldLayout(
     TypeLayout*     typeLayout,
-    String const&   mangledFieldName)
+    IRInst*         fieldKey)
 {
     if (!typeLayout)
         return nullptr;
@@ -1294,9 +1328,26 @@ RefPtr<VarLayout> getFieldLayout(
 
     if (auto structTypeLayout = as<StructTypeLayout>(typeLayout))
     {
+        // First, let's see if the field had a layout registered
+        // directly using its IR key.
+        //
+        RefPtr<VarLayout> fieldLayout;
+        if(structTypeLayout->mapKeyToLayout.TryGetValue(fieldKey, fieldLayout))
+            return fieldLayout;
+
+        // Otherwise, fall back to doing lookup using the linkage
+        // attached to the key, and its mangled name.
+        //
+        auto fieldLinkage = fieldKey->findDecoration<IRLinkageDecoration>();
+        if(!fieldLinkage)
+            return nullptr;
+        auto mangledFieldName = fieldLinkage->getMangledName();
+
+        // In this case we fall back to a linear search over the fields.
+        //
         for(auto ff : structTypeLayout->fields)
         {
-            if(mangledFieldName == getMangledName(ff->varDecl.getDecl()) )
+            if(mangledFieldName == getMangledName(ff->varDecl.getDecl()).getUnownedSlice() )
             {
                 return ff;
             }
@@ -1306,9 +1357,9 @@ RefPtr<VarLayout> getFieldLayout(
     return nullptr;
 }
 
-RefPtr<VarLayout> createVarLayout(
-    LegalVarChain*  varChain,
-    TypeLayout*     typeLayout)
+RefPtr<VarLayout> createSimpleVarLayout(
+    SimpleLegalVarChain*    varChain,
+    TypeLayout*             typeLayout)
 {
     if (!typeLayout)
         return nullptr;
@@ -1372,7 +1423,23 @@ RefPtr<VarLayout> createVarLayout(
         }
     }
 
+    return varLayout;
+}
+
 
+RefPtr<VarLayout> createVarLayout(
+    LegalVarChain const&    varChain,
+    TypeLayout*             typeLayout)
+{
+    if(!typeLayout)
+        return nullptr;
+
+    auto varLayout = createSimpleVarLayout(varChain.primaryChain, typeLayout);
+
+    if(auto pendingDataTypeLayout = typeLayout->pendingDataTypeLayout)
+    {
+        varLayout->pendingVarLayout = createSimpleVarLayout(varChain.pendingChain, typeLayout);
+    }
 
     return varLayout;
 }