Add d.location attribute support in tgpu-gen (#530)

software-mansion · Nov 5, 2024 · e3fe10b · e3fe10b
1 parent c03f370
commit e3fe10b
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git 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
@@ -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;
 }\`);`);