Fix a bug in handling explicit register space bindings (#1175)

The logic for handling explicit `space`/`set` bindings on shader parameters for parameter blocks was not correctly marking the `space`/`set` that gets grabbed as used, and as a result it was possible for another parameter block that relies on implicit assignment to end up with a conflicting space.

This change fixes the original oversight, and adds a test case to prevent against regression.

3 files changed, 198 insertions, 17 deletions

diff --git a/source/slang/slang-parameter-binding.cpp b/source/slang/slang-parameter-binding.cpp
index 8ebc4f420..a6bd52122 100644
--- a/source/slang/slang-parameter-binding.cpp
+++ b/source/slang/slang-parameter-binding.cpp
@@ -757,11 +757,12 @@ static RefPtr<UsedRangeSet> findUsedRangeSetForSpace(
 // has been used in at least one binding, and so it should not
 // be used by auto-generated bindings that need to claim entire
 // spaces.
-static void markSpaceUsed(
+static VarLayout* markSpaceUsed(
     ParameterBindingContext*    context,
+    VarLayout*                  varLayout,
     UInt                        space)
 {
-    context->shared->usedSpaces.Add(nullptr, space, space+1);
+    return context->shared->usedSpaces.Add(varLayout, space, space+1);
 }
 
 static UInt allocateUnusedSpaces(
@@ -795,8 +796,7 @@ static void addExplicitParameterBinding(
     RefPtr<ParameterInfo>       parameterInfo,
     VarDeclBase*                varDecl,
     LayoutSemanticInfo const&   semanticInfo,
-    LayoutSize                  count,
-    RefPtr<UsedRangeSet>        usedRangeSet = nullptr)
+    LayoutSize                  count)
 {
     auto kind = semanticInfo.kind;
 
@@ -822,20 +822,31 @@ static void addExplicitParameterBinding(
         bindingInfo.index = semanticInfo.index;
         bindingInfo.space = semanticInfo.space;
 
-        if (!usedRangeSet)
+        VarLayout* overlappedVarLayout = nullptr;
+        if( kind == LayoutResourceKind::RegisterSpace )
         {
-            usedRangeSet = findUsedRangeSetForSpace(context, semanticInfo.space);
+            // Parameter is being bound to an entire space, so we
+            // need to mark the given space as used and report
+            // an error if another parameter was already allocated
+            // there.
+            //
+            overlappedVarLayout = markSpaceUsed(context, parameterInfo->varLayout, semanticInfo.index);
+        }
+        else
+        {
+            auto usedRangeSet = findUsedRangeSetForSpace(context, semanticInfo.space);
 
             // Record that the particular binding space was
             // used by an explicit binding, so that we don't
             // claim it for auto-generated bindings that
             // need to grab a full space
-            markSpaceUsed(context, semanticInfo.space);
+            markSpaceUsed(context, parameterInfo->varLayout, semanticInfo.space);
+
+            overlappedVarLayout = usedRangeSet->usedResourceRanges[(int)semanticInfo.kind].Add(
+                parameterInfo->varLayout,
+                semanticInfo.index,
+                semanticInfo.index + count);
         }
-        auto overlappedVarLayout = usedRangeSet->usedResourceRanges[(int)semanticInfo.kind].Add(
-            parameterInfo->varLayout,
-            semanticInfo.index,
-            semanticInfo.index + count);
 
         if (overlappedVarLayout)
         {
@@ -965,11 +976,6 @@ static void addExplicitParameterBindings_GLSL(
     // the index/offset/etc.
     //
 
-    // We also may need to store explicit binding info in a different place,
-    // in the case of varying input/output, since we don't want to collect
-    // things globally;
-    RefPtr<UsedRangeSet> usedRangeSet;
-
     TypeLayout::ResourceInfo* resInfo = nullptr;
     LayoutSemanticInfo semanticInfo;
     semanticInfo.index = 0;
@@ -1017,7 +1023,7 @@ static void addExplicitParameterBindings_GLSL(
     auto count = resInfo->count;
     semanticInfo.kind = kind;
 
-    addExplicitParameterBinding(context, parameterInfo, varDecl, semanticInfo, count, usedRangeSet);
+    addExplicitParameterBinding(context, parameterInfo, varDecl, semanticInfo, count);
 }
 
 // Given a single parameter, collect whatever information we have on
diff --git a/tests/reflection/mix-explicit-and-implicit-spaces.slang b/tests/reflection/mix-explicit-and-implicit-spaces.slang
new file mode 100644
index 000000000..52dd4321e
--- /dev/null
+++ b/tests/reflection/mix-explicit-and-implicit-spaces.slang
@@ -0,0 +1,18 @@
+// mix-explicit-and-implicit-spaces.slang
+//TEST:REFLECTION:-profile cs_5_1 -target hlsl
+
+// Ensure that correct layout/reflection is computed
+// when mixing implicit and explicit spaces for
+// parameter blocks.
+
+struct A { float x; }
+ParameterBlock<A> a : register(space0);
+
+struct B { float y; }
+ParameterBlock<B> b : register(space1);
+
+struct C { float z; }
+ParameterBlock<C> c;
+
+void main()
+{}
diff --git a/tests/reflection/mix-explicit-and-implicit-spaces.slang.expected b/tests/reflection/mix-explicit-and-implicit-spaces.slang.expected
new file mode 100644
index 000000000..2e4ead342
--- /dev/null
+++ b/tests/reflection/mix-explicit-and-implicit-spaces.slang.expected
@@ -0,0 +1,157 @@
+result code = 0
+standard error = {
+}
+standard output = {
+{
+    "parameters": [
+        {
+            "name": "a",
+            "bindings": [
+                {"kind": "constantBuffer", "index": 0, "count": 0},
+                {"kind": "registerSpace", "index": 0}
+            ],
+            "type": {
+                "kind": "parameterBlock",
+                "elementType": {
+                    "kind": "struct",
+                    "name": "A",
+                    "fields": [
+                        {
+                            "name": "x",
+                            "type": {
+                                "kind": "scalar",
+                                "scalarType": "float32"
+                            },
+                            "binding": {"kind": "uniform", "offset": 0, "size": 4}
+                        }
+                    ]
+                },
+                "containerVarLayout": {
+                    "bindings": [
+                        {"kind": "constantBuffer", "index": 0},
+                        {"kind": "registerSpace", "index": 0}
+                    ]
+                },
+                "elementVarLayout": {
+                    "type": {
+                        "kind": "struct",
+                        "name": "A",
+                        "fields": [
+                            {
+                                "name": "x",
+                                "type": {
+                                    "kind": "scalar",
+                                    "scalarType": "float32"
+                                },
+                                "binding": {"kind": "uniform", "offset": 0, "size": 4}
+                            }
+                        ]
+                    },
+                    "binding": {"kind": "uniform", "offset": 0, "size": 4}
+                }
+            }
+        },
+        {
+            "name": "b",
+            "bindings": [
+                {"kind": "constantBuffer", "index": 0, "count": 0},
+                {"kind": "registerSpace", "index": 1}
+            ],
+            "type": {
+                "kind": "parameterBlock",
+                "elementType": {
+                    "kind": "struct",
+                    "name": "B",
+                    "fields": [
+                        {
+                            "name": "y",
+                            "type": {
+                                "kind": "scalar",
+                                "scalarType": "float32"
+                            },
+                            "binding": {"kind": "uniform", "offset": 0, "size": 4}
+                        }
+                    ]
+                },
+                "containerVarLayout": {
+                    "bindings": [
+                        {"kind": "constantBuffer", "index": 0},
+                        {"kind": "registerSpace", "index": 0}
+                    ]
+                },
+                "elementVarLayout": {
+                    "type": {
+                        "kind": "struct",
+                        "name": "B",
+                        "fields": [
+                            {
+                                "name": "y",
+                                "type": {
+                                    "kind": "scalar",
+                                    "scalarType": "float32"
+                                },
+                                "binding": {"kind": "uniform", "offset": 0, "size": 4}
+                            }
+                        ]
+                    },
+                    "binding": {"kind": "uniform", "offset": 0, "size": 4}
+                }
+            }
+        },
+        {
+            "name": "c",
+            "bindings": [
+                {"kind": "constantBuffer", "index": 0, "count": 0},
+                {"kind": "registerSpace", "index": 2}
+            ],
+            "type": {
+                "kind": "parameterBlock",
+                "elementType": {
+                    "kind": "struct",
+                    "name": "C",
+                    "fields": [
+                        {
+                            "name": "z",
+                            "type": {
+                                "kind": "scalar",
+                                "scalarType": "float32"
+                            },
+                            "binding": {"kind": "uniform", "offset": 0, "size": 4}
+                        }
+                    ]
+                },
+                "containerVarLayout": {
+                    "bindings": [
+                        {"kind": "constantBuffer", "index": 0},
+                        {"kind": "registerSpace", "index": 0}
+                    ]
+                },
+                "elementVarLayout": {
+                    "type": {
+                        "kind": "struct",
+                        "name": "C",
+                        "fields": [
+                            {
+                                "name": "z",
+                                "type": {
+                                    "kind": "scalar",
+                                    "scalarType": "float32"
+                                },
+                                "binding": {"kind": "uniform", "offset": 0, "size": 4}
+                            }
+                        ]
+                    },
+                    "binding": {"kind": "uniform", "offset": 0, "size": 4}
+                }
+            }
+        }
+    ],
+    "entryPoints": [
+        {
+            "name": "main",
+            "stage:": "compute",
+            "threadGroupSize": [1, 1, 1]
+        }
+    ]
+}
+}