Fix reflection to json issue (#7379)

Apply argument buffer tier2 rule when using parameter block for Metal target.

Close #6803.

4 files changed, 80 insertions, 11 deletions

diff --git a/source/slang/slang-reflection-json.cpp b/source/slang/slang-reflection-json.cpp
index a84536f20..72855a9b1 100644
--- a/source/slang/slang-reflection-json.cpp
+++ b/source/slang/slang-reflection-json.cpp
@@ -3,6 +3,7 @@
 
 #include "../core/slang-blob.h"
 #include "slang-ast-support-types.h"
+#include "slang.h"
 
 template<typename T>
 struct Range
@@ -56,6 +57,7 @@ namespace Slang
 static void emitReflectionVarInfoJSON(PrettyWriter& writer, slang::VariableReflection* var);
 static void emitReflectionTypeLayoutJSON(PrettyWriter& writer, slang::TypeLayoutReflection* type);
 static void emitReflectionTypeJSON(PrettyWriter& writer, slang::TypeReflection* type);
+static slang::ShaderReflection* g_inProgramLayout = nullptr;
 
 static void emitReflectionVarBindingInfoJSON(
     PrettyWriter& writer,
@@ -361,6 +363,26 @@ static void emitUserAttributes(PrettyWriter& writer, slang::FunctionReflection*
     }
 }
 
+static slang::TypeLayoutReflection* maybeChangeTypeLayoutToAgumentBufferTier2(
+    slang::VariableLayoutReflection* varLayout)
+{
+    if (varLayout->getCategoryCount() != 0)
+    {
+        for (unsigned int categoryIdx = 0; categoryIdx < varLayout->getCategoryCount();
+             categoryIdx++)
+        {
+            auto category = varLayout->getCategoryByIndex(categoryIdx);
+            if (category == slang::MetalArgumentBufferElement)
+            {
+                return g_inProgramLayout->getTypeLayout(
+                    varLayout->getTypeLayout()->getType(),
+                    slang::LayoutRules::MetalArgumentBufferTier2);
+            }
+        }
+    }
+    return nullptr;
+}
+
 static void emitReflectionVarLayoutJSON(PrettyWriter& writer, slang::VariableLayoutReflection* var)
 {
     writer << "{\n";
@@ -376,7 +398,14 @@ static void emitReflectionVarLayoutJSON(PrettyWriter& writer, slang::VariableLay
 
     writer.maybeComma();
     writer << "\"type\": ";
-    emitReflectionTypeLayoutJSON(writer, var->getTypeLayout());
+    if (auto newTypeLayout = maybeChangeTypeLayoutToAgumentBufferTier2(var))
+    {
+        emitReflectionTypeLayoutJSON(writer, newTypeLayout);
+    }
+    else
+    {
+        emitReflectionTypeLayoutJSON(writer, var->getTypeLayout());
+    }
 
     emitReflectionModifierInfoJSON(writer, var->getVariable());
 
@@ -693,7 +722,18 @@ static void emitReflectionParameterGroupTypeLayoutInfoJSON(
     writer << "\"";
 
     writer << ",\n\"elementType\": ";
-    emitReflectionTypeLayoutJSON(writer, typeLayout->getElementTypeLayout());
+
+    if (auto newElementTypeLayout =
+            maybeChangeTypeLayoutToAgumentBufferTier2(typeLayout->getElementVarLayout()))
+    {
+        // If we are in argument buffer tier 2, we need to use the new type layout
+        // that has the correct binding information.
+        emitReflectionTypeLayoutJSON(writer, newElementTypeLayout);
+    }
+    else
+    {
+        emitReflectionTypeLayoutJSON(writer, typeLayout->getElementTypeLayout());
+    }
 
     // Note: There is a subtle detail below when it comes to the
     // container/element variable layouts that get nested inside
@@ -1231,6 +1271,7 @@ void emitReflectionJSON(
     PrettyWriter& writer)
 {
     auto programReflection = (slang::ShaderReflection*)reflection;
+    g_inProgramLayout = programReflection;
     emitReflectionJSON(writer, request, programReflection);
 }
 
diff --git a/tests/metal/nested-parameter-block-reflection.slang b/tests/metal/nested-parameter-block-reflection.slang
index 8282f5573..a294755da 100644
--- a/tests/metal/nested-parameter-block-reflection.slang
+++ b/tests/metal/nested-parameter-block-reflection.slang
@@ -1,10 +1,27 @@
 //TEST:REFLECTION(filecheck=CHECK): -target metal
 
+
+
+// CHECK: "name": "gParams",
+// CHECK: "binding": {"kind": "constantBuffer", "index": 0},
+
+// CHECK:"name": "pdata",
+// CHECK:"binding": {"kind": "uniform", "offset": 0, "size": 16}
 // CHECK:"name": "tex",
-// CHECK:"binding": {"kind": "metalArgumentBufferElement", "index": 1}
 
+// Since we will apply MetalArgumentBufferTier2, 'tex' here will just be a uniform.
+// The pdata is a nested parameter block, so it will be a 64-bit device pointer which take
+// 8 bytes. So the offset of `tex` will be 8 bytes.
+// CHECK:"binding": {"kind": "uniform", "offset": 8, "size": 8}
+
+
+// Check that there will be only two bindings.
+//
+// CHECK: "name": "output",
+// CHECK: "binding": {"kind": "constantBuffer", "index": 1},
+//
 struct Data { int3 content; }
-struct Params 
+struct Params
 {
     ParameterBlock<Data> pdata;
     Texture2D tex;
@@ -15,4 +32,4 @@ RWStructuredBuffer<float4> output;
 void computeMain()
 {
     output[0] = gParams.tex.Load(gParams.pdata.content);
-}
\ No newline at end of file
+}
diff --git a/tests/metal/simple-compute.slang b/tests/metal/simple-compute.slang
index 85abaffb4..d70d37f5c 100644
--- a/tests/metal/simple-compute.slang
+++ b/tests/metal/simple-compute.slang
@@ -18,10 +18,10 @@ ParameterBlock<MyBlock> block2;
 
 // REFLECT: "elementVarLayout": {
 // REFLECT: "name": "b1",
-// REFLECT: "binding": {"kind": "metalArgumentBufferElement", "index": 0}
+// REFLECT: "binding": {"kind": "uniform", "offset": 0, "size": 8}
 
 // REFLECT: "name": "b2",
-// REFLECT: "binding": {"kind": "metalArgumentBufferElement", "index": 1}
+// REFLECT: "binding": {"kind": "uniform", "offset": 8, "size": 8}
 
 // REFLECT: "binding": {"kind": "metalArgumentBufferElement", "index": 0, "count": 2}
 
@@ -42,4 +42,4 @@ void func(float v)
 void main_kernel()
 {
     func(3.0f);
-}
\ No newline at end of file
+}
diff --git a/tests/metal/vector-argument-buffer-layout.slang b/tests/metal/vector-argument-buffer-layout.slang
index 698fd30d8..8b74e90d3 100644
--- a/tests/metal/vector-argument-buffer-layout.slang
+++ b/tests/metal/vector-argument-buffer-layout.slang
@@ -1,6 +1,17 @@
 //TEST:REFLECTION(filecheck=CHECK): -target metal
-//CHECK: "binding": {"kind": "metalArgumentBufferElement", "index": 0}
-//CHECK: "binding": {"kind": "metalArgumentBufferElement", "index": 1}
+
+
+//CHECK: "name": "o",
+//CHECK: "binding": {"kind": "constantBuffer", "index": 0},
+
+//CHECK: "name": "gParams",
+//CHECK: "binding": {"kind": "constantBuffer", "index": 1},
+
+
+//CHECK: "dir"
+//CHECK: "binding": {"kind": "uniform", "offset": 0, "size": 16}
+//CHECK: "color"
+//CHECK: "binding": {"kind": "uniform", "offset": 16, "size": 16}
 
 struct Params
 {
@@ -15,4 +26,4 @@ ParameterBlock<Params> gParams;
 void computeMain()
 {
     o[0] = gParams.dir.x;
-}
\ No newline at end of file
+}