From 075e6c25826cd0a1df26bd36af4afe865673a790 Mon Sep 17 00:00:00 2001
From: Yong He <yonghe@outlook.com>
Date: Tue, 31 Dec 2024 17:47:20 -0800
Subject: Fix reflection for metal vector [[id]] location. (#5943)

---
 source/slang/slang-type-layout.cpp              | 13 +++++++++++++
 tests/metal/vector-argument-buffer-layout.slang | 18 ++++++++++++++++++
 2 files changed, 31 insertions(+)
 create mode 100644 tests/metal/vector-argument-buffer-layout.slang

diff --git a/source/slang/slang-type-layout.cpp b/source/slang/slang-type-layout.cpp
index c6acf5250..98a324d96 100644
--- a/source/slang/slang-type-layout.cpp
+++ b/source/slang/slang-type-layout.cpp
@@ -1898,6 +1898,19 @@ struct MetalArgumentBufferElementLayoutRulesImpl : ObjectLayoutRulesImpl, Defaul
         return SimpleLayoutInfo(LayoutResourceKind::MetalArgumentBufferElement, 1);
     }
 
+    SimpleLayoutInfo GetVectorLayout(
+        BaseType elementType,
+        SimpleLayoutInfo elementInfo,
+        size_t elementCount) override
+    {
+        SLANG_UNUSED(elementType);
+        SLANG_UNUSED(elementInfo);
+        SLANG_UNUSED(elementCount);
+
+        // A vector occupies one [[id]] slot in a metal argument buffer.
+        return SimpleLayoutInfo(LayoutResourceKind::MetalArgumentBufferElement, 1);
+    }
+
     virtual ObjectLayoutInfo GetObjectLayout(ShaderParameterKind kind, const Options& /* options */)
         override
     {
diff --git a/tests/metal/vector-argument-buffer-layout.slang b/tests/metal/vector-argument-buffer-layout.slang
new file mode 100644
index 000000000..698fd30d8
--- /dev/null
+++ b/tests/metal/vector-argument-buffer-layout.slang
@@ -0,0 +1,18 @@
+//TEST:REFLECTION(filecheck=CHECK): -target metal
+//CHECK: "binding": {"kind": "metalArgumentBufferElement", "index": 0}
+//CHECK: "binding": {"kind": "metalArgumentBufferElement", "index": 1}
+
+struct Params
+{
+    float3 dir;
+    float4 color;
+}
+
+RWStructuredBuffer<float> o;
+ParameterBlock<Params> gParams;
+
+[numthreads(1,1,1)]
+void computeMain()
+{
+    o[0] = gParams.dir.x;
+}
\ No newline at end of file
-- 
cgit v1.2.3