From e3fe10b2666d4ba637e8a5dad1b4e2674b51ff0a Mon Sep 17 00:00:00 2001
From: Marcin Hawryluk <70582973+mhawryluk@users.noreply.github.com>
Date: Tue, 5 Nov 2024 15:54:28 +0100
Subject: [PATCH] Add d.location attribute support in tgpu-gen (#530)

---
 packages/tgpu-gen/gen.mjs                 | 4 ++--
 packages/tgpu-gen/tests/functions.test.ts | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/packages/tgpu-gen/gen.mjs b/packages/tgpu-gen/gen.mjs
index 8ef9d1c6..b7825081 100644
--- a/packages/tgpu-gen/gen.mjs
+++ b/packages/tgpu-gen/gen.mjs
@@ -226,7 +226,7 @@ function generateType(type_, options) {
   const result =
     type_.attributes?.reduce(
       (acc, attribute) =>
-        ['align', 'size'].includes(attribute.name)
+        ['align', 'size', 'location'].includes(attribute.name)
           ? `d.${attribute.name}(${attribute.value}, ${acc})`
           : acc,
       tgpuType,
@@ -472,7 +472,7 @@ function generateFunction(func, wgsl, options) {
         : 'fn';
 
   const inputs = `[${func.arguments
-    .flatMap((input) => (input.type ? [generateType(input.type, options)] : []))
+    .flatMap((arg) => (arg.type ? [generateType(arg.type, options)] : []))
     .join(', ')}]`;
 
   const output = func.returnType
diff --git a/packages/tgpu-gen/tests/functions.test.ts b/packages/tgpu-gen/tests/functions.test.ts
index 8f0f2d3b..8fe7efab 100644
--- a/packages/tgpu-gen/tests/functions.test.ts
+++ b/packages/tgpu-gen/tests/functions.test.ts
@@ -67,7 +67,7 @@ fn mainVert(@builtin(instance_index) ii: u32, @location(0) v: vec2f) -> VertexOu
 
     expect(generate(wgsl)).toContain(`\
 export const mainVert = tgpu
-  .vertexFn([d.u32, d.vec2f], VertexOutput)
+  .vertexFn([d.u32, d.location(0, d.vec2f)], VertexOutput)
   .does(/* wgsl */ \`(@builtin(instance_index) ii: u32, @location(0) v: vec2f) -> VertexOutput {
     let instanceInfo = trianglePos[ii];
 
@@ -97,7 +97,7 @@ fn mainFrag(@location(1) color: vec4f) -> @location(0) vec4f {
 
     expect(generate(wgsl)).toContain(`\
 export const mainFrag = tgpu
-  .fragmentFn([d.vec4f], d.vec4f)
+  .fragmentFn([d.location(1, d.vec4f)], d.location(0, d.vec4f))
   .does(/* wgsl */ \`(@location(1) color: vec4f) -> @location(0) vec4f {
     return color;
 }\`);`);