feat: Allow destructuring arguments in TGSL function implementations

packages/tgpu-jit/src/index.ts

@@ -1,12 +1,12 @@
 import { parse } from 'acorn';
-import type { AnyNode, ArgNames, Block } from 'tinyest';
+import type { AnyNode, Block, FuncParameter } from 'tinyest';
 import { transpileFn, transpileNode } from 'tinyest-for-wgsl';
 
 /**
  * Information extracted from transpiling a JS function.
  */
 type TranspilationResult = {
-  argNames: ArgNames;
+  params: FuncParameter[];
   body: Block;
   /**
    * All identifiers found in the function code that are not declared in the

packages/tinyest-for-wgsl/src/parsers.ts

@@ -1,6 +1,7 @@
 import type * as babel from '@babel/types';
 import type * as acorn from 'acorn';
 import * as tinyest from 'tinyest';
+import { FuncParameterType } from 'tinyest';
 
 const { NodeTypeCatalog: NODE } = tinyest;
 
@@ -378,7 +379,7 @@ function transpile(ctx: Context, node: JsNode): tinyest.AnyNode {
 }
 
 export type TranspilationResult = {
-  argNames: tinyest.ArgNames;
+  params: tinyest.FuncParameter[];
   body: tinyest.Block;
   /**
    * All identifiers found in the function code that are not declared in the
@@ -388,7 +389,7 @@ export type TranspilationResult = {
 };
 
 export function extractFunctionParts(rootNode: JsNode): {
-  params: tinyest.ArgNames;
+  params: tinyest.FuncParameter[];
   body:
     | acorn.BlockStatement
     | acorn.Expression
@@ -452,48 +453,60 @@ export function extractFunctionParts(rootNode: JsNode): {
     throw new Error('tgpu.fn cannot be a generator');
   }
 
-  // destructured object argument
-  if (
-    functionNode.params[0] &&
-    functionNode.params[0].type === 'ObjectPattern'
-  ) {
-    return {
-      params: {
-        type: 'destructured-object',
-        props: functionNode.params[0].properties.flatMap((prop) =>
-          (prop.type === 'Property' || prop.type === 'ObjectProperty') &&
-            prop.key.type === 'Identifier' &&
-            prop.value.type === 'Identifier'
-            ? [{ prop: prop.key.name, alias: prop.value.name }]
-            : []
-        ),
-      },
-      body: functionNode.body,
-    };
+  const unsupportedTypes = new Set(
+    functionNode.params.flatMap((param) =>
+      param.type === 'ObjectPattern' || param.type === 'Identifier'
+        ? []
+        : [param.type]
+    ),
+  );
+  if (unsupportedTypes.size > 0) {
+    throw new Error(
+      `Unsupported function parameter type(s): ${[...unsupportedTypes]}`,
+    );
   }
 
   return {
-    params: {
-      type: 'identifiers',
-      names: functionNode.params.flatMap((x) =>
-        x.type === 'Identifier' ? [x.name] : []
+    params: (functionNode
+      .params as (
+        | babel.Identifier
+        | acorn.Identifier
+        | babel.ObjectPattern
+        | acorn.ObjectPattern
+      )[]).map((param) =>
+        param.type === 'ObjectPattern'
+          ? {
+            type: FuncParameterType.destructuredObject,
+            props: param.properties.flatMap((prop) =>
+              (prop.type === 'Property' || prop.type === 'ObjectProperty') &&
+                prop.key.type === 'Identifier' &&
+                prop.value.type === 'Identifier'
+                ? [{ name: prop.key.name, alias: prop.value.name }]
+                : []
+            ),
+          }
+          : {
+            type: FuncParameterType.identifier,
+            name: param.name,
+          }
       ),
-    },
     body: functionNode.body,
   };
 }
 
 export function transpileFn(rootNode: JsNode): TranspilationResult {
-  const { params: argNames, body } = extractFunctionParts(rootNode);
+  const { params, body } = extractFunctionParts(rootNode);
 
   const ctx: Context = {
     externalNames: new Set(),
     ignoreExternalDepth: 0,
     stack: [
       {
-        declaredNames: argNames.type === 'identifiers'
-          ? argNames.names
-          : argNames.props.map((prop) => prop.alias),
+        declaredNames: params.flatMap((param) =>
+          param.type === FuncParameterType.identifier
+            ? param.name
+            : param.props.map((prop) => prop.alias)
+        ),
       },
     ],
   };
@@ -503,14 +516,14 @@ export function transpileFn(rootNode: JsNode): TranspilationResult {
 
   if (body.type === 'BlockStatement') {
     return {
-      argNames,
+      params,
       body: tinyestBody as tinyest.Block,
       externalNames,
     };
   }
 
   return {
-    argNames,
+    params,
     body: [NODE.block, [[NODE.return, tinyestBody as tinyest.Expression]]],
     externalNames,
   };

packages/tinyest-for-wgsl/tests/parsers.test.ts

@@ -1,7 +1,6 @@
 import babel from '@babel/parser';
 import type { Node } from '@babel/types';
 import * as acorn from 'acorn';
-import type { ArgNames } from 'tinyest';
 import { describe, expect, it } from 'vitest';
 import { transpileFn } from '../src/parsers.ts';
 
@@ -28,12 +27,9 @@ describe('transpileFn', () => {
   it(
     'parses an empty arrow function',
     dualTest((p) => {
-      const { argNames, body, externalNames } = transpileFn(p('() => {}'));
+      const { params, body, externalNames } = transpileFn(p('() => {}'));
 
-      expect(argNames).toStrictEqual({
-        type: 'identifiers',
-        names: [],
-      });
+      expect(params).toStrictEqual([]);
       expect(JSON.stringify(body)).toMatchInlineSnapshot(`"[0,[]]"`);
       expect(externalNames).toStrictEqual([]);
     }),
@@ -42,14 +38,11 @@ describe('transpileFn', () => {
   it(
     'parses an empty named function',
     dualTest((p) => {
-      const { argNames, body, externalNames } = transpileFn(
+      const { params, body, externalNames } = transpileFn(
         p('function example() {}'),
       );
 
-      expect(argNames).toStrictEqual({
-        type: 'identifiers',
-        names: [],
-      });
+      expect(params).toStrictEqual([]);
       expect(JSON.stringify(body)).toMatchInlineSnapshot(`"[0,[]]"`);
       expect(externalNames).toStrictEqual([]);
     }),
@@ -58,14 +51,14 @@ describe('transpileFn', () => {
   it(
     'gathers external names',
     dualTest((p) => {
-      const { argNames, body, externalNames } = transpileFn(
+      const { params, body, externalNames } = transpileFn(
         p('(a, b) => a + b - c'),
       );
 
-      expect(argNames).toStrictEqual({
-        type: 'identifiers',
-        names: ['a', 'b'],
-      });
+      expect(params).toStrictEqual([
+        { type: 'i', name: 'a' },
+        { type: 'i', name: 'b' },
+      ]);
       expect(JSON.stringify(body)).toMatchInlineSnapshot(
         `"[0,[[10,[1,[1,"a","+","b"],"-","c"]]]]"`,
       );
@@ -76,17 +69,14 @@ describe('transpileFn', () => {
   it(
     'respects local declarations when gathering external names',
     dualTest((p) => {
-      const { argNames, body, externalNames } = transpileFn(
+      const { params, body, externalNames } = transpileFn(
         p(`() => {
         const a = 0;
         c = a + 2;
       }`),
       );
 
-      expect(argNames).toStrictEqual({
-        type: 'identifiers',
-        names: [],
-      });
+      expect(params).toStrictEqual([]);
       expect(JSON.stringify(body)).toMatchInlineSnapshot(
         `"[0,[[13,"a",[5,"0"]],[2,"c","=",[1,"a","+",[5,"2"]]]]]"`,
       );
@@ -98,7 +88,7 @@ describe('transpileFn', () => {
   it(
     'respects outer scope when gathering external names',
     dualTest((p) => {
-      const { argNames, body, externalNames } = transpileFn(
+      const { params, body, externalNames } = transpileFn(
         p(`() => {
         const a = 0;
         {
@@ -107,10 +97,7 @@ describe('transpileFn', () => {
       }`),
       );
 
-      expect(argNames).toStrictEqual({
-        type: 'identifiers',
-        names: [],
-      });
+      expect(params).toStrictEqual([]);
       expect(JSON.stringify(body)).toMatchInlineSnapshot(
         `"[0,[[13,"a",[5,"0"]],[0,[[2,"c","=",[1,"a","+",[5,"2"]]]]]]]"`,
       );
@@ -122,14 +109,11 @@ describe('transpileFn', () => {
   it(
     'treats the object as a possible external value when accessing a member',
     dualTest((p) => {
-      const { argNames, body, externalNames } = transpileFn(
+      const { params, body, externalNames } = transpileFn(
         p('() => external.outside.prop'),
       );
 
-      expect(argNames).toStrictEqual({
-        type: 'identifiers',
-        names: [],
-      });
+      expect(params).toStrictEqual([]);
       expect(JSON.stringify(body)).toMatchInlineSnapshot(
         `"[0,[[10,[7,[7,"external","outside"],"prop"]]]]"`,
       );
@@ -141,29 +125,75 @@ describe('transpileFn', () => {
   it(
     'handles destructured args',
     dualTest((p) => {
-      const { argNames, externalNames } = transpileFn(
+      const { params, externalNames } = transpileFn(
         p(`({ pos, a: b }) => {
           const x = pos.x;
         }`),
       );
 
-      expect(argNames).toStrictEqual(
-        {
-          type: 'destructured-object',
+      expect(params).toStrictEqual(
+        [{
+          type: 'd',
           props: [
             {
               alias: 'pos',
-              prop: 'pos',
+              name: 'pos',
             },
             {
               alias: 'b',
-              prop: 'a',
+              name: 'a',
             },
           ],
-        } satisfies ArgNames,
+        }],
       );
 
       expect(externalNames).toStrictEqual([]);
     }),
   );
+
+  it(
+    'handles mixed type parameters',
+    dualTest((p) => {
+      const { params, externalNames } = transpileFn(
+        p(`(y, { pos, a: b }, {c, d}) => {
+          const x = pos.x;
+        }`),
+      );
+
+      expect(params).toStrictEqual([
+        {
+          type: 'i',
+          name: 'y',
+        },
+        {
+          type: 'd',
+          props: [
+            {
+              alias: 'pos',
+              name: 'pos',
+            },
+            {
+              alias: 'b',
+              name: 'a',
+            },
+          ],
+        },
+        {
+          type: 'd',
+          props: [
+            {
+              alias: 'c',
+              name: 'c',
+            },
+            {
+              alias: 'd',
+              name: 'd',
+            },
+          ],
+        },
+      ]);
+
+      expect(externalNames).toStrictEqual([]);
+    }),
+  );
 });

packages/tinyest/src/nodes.ts

@@ -253,15 +253,20 @@ export type Expression =
 
 export type AnyNode = Statement | Expression;
 
-export type ArgNames =
+export const FuncParameterType = {
+  identifier: 'i',
+  destructuredObject: 'd',
+} as const;
+
+export type FuncParameter =
   | {
-    type: 'identifiers';
-    names: string[];
+    type: typeof FuncParameterType.identifier;
+    name: string;
   }
   | {
-    type: 'destructured-object';
+    type: typeof FuncParameterType.destructuredObject;
     props: {
-      prop: string;
+      name: string;
       alias: string;
     }[];
   };

packages/typegpu/src/core/function/astUtils.ts

@@ -1,7 +1,7 @@
-import type { ArgNames, Block } from 'tinyest';
+import type { Block, FuncParameter } from 'tinyest';
 
 export type Ast = {
-  argNames: ArgNames;
+  params: FuncParameter[];
   body: Block;
   externalNames: string[];
 };