Skip to content

Commit 68fd55c

Browse files
mhawrylukmhawryluk
authored
feat: Allow destructuring arguments in TGSL function implementations (#1257)
1 parent 51354a8 commit 68fd55c

20 files changed

+477
-184
lines changed

‎packages/tgpu-jit/src/index.ts

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,12 @@
11
import { parse } from 'acorn';
2-
import type { AnyNode, ArgNames, Block } from 'tinyest';
2+
import type { AnyNode, Block, FuncParameter } from 'tinyest';
33
import { transpileFn, transpileNode } from 'tinyest-for-wgsl';
44

55
/**
66
* Information extracted from transpiling a JS function.
77
*/
88
type TranspilationResult = {
9-
argNames: ArgNames;
9+
params: FuncParameter[];
1010
body: Block;
1111
/**
1212
* All identifiers found in the function code that are not declared in the

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

Lines changed: 44 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import type * as babel from '@babel/types';
22
import type * as acorn from 'acorn';
33
import * as tinyest from 'tinyest';
4+
import { FuncParameterType } from 'tinyest';
45

56
const { NodeTypeCatalog: NODE } = tinyest;
67

@@ -378,7 +379,7 @@ function transpile(ctx: Context, node: JsNode): tinyest.AnyNode {
378379
}
379380

380381
export type TranspilationResult = {
381-
argNames: tinyest.ArgNames;
382+
params: tinyest.FuncParameter[];
382383
body: tinyest.Block;
383384
/**
384385
* All identifiers found in the function code that are not declared in the
@@ -388,7 +389,7 @@ export type TranspilationResult = {
388389
};
389390

390391
export function extractFunctionParts(rootNode: JsNode): {
391-
params: tinyest.ArgNames;
392+
params: tinyest.FuncParameter[];
392393
body:
393394
| acorn.BlockStatement
394395
| acorn.Expression
@@ -452,48 +453,60 @@ export function extractFunctionParts(rootNode: JsNode): {
452453
throw new Error('tgpu.fn cannot be a generator');
453454
}
454455

455-
// destructured object argument
456-
if (
457-
functionNode.params[0] &&
458-
functionNode.params[0].type === 'ObjectPattern'
459-
) {
460-
return {
461-
params: {
462-
type: 'destructured-object',
463-
props: functionNode.params[0].properties.flatMap((prop) =>
464-
(prop.type === 'Property' || prop.type === 'ObjectProperty') &&
465-
prop.key.type === 'Identifier' &&
466-
prop.value.type === 'Identifier'
467-
? [{ prop: prop.key.name, alias: prop.value.name }]
468-
: []
469-
),
470-
},
471-
body: functionNode.body,
472-
};
456+
const unsupportedTypes = new Set(
457+
functionNode.params.flatMap((param) =>
458+
param.type === 'ObjectPattern' || param.type === 'Identifier'
459+
? []
460+
: [param.type]
461+
),
462+
);
463+
if (unsupportedTypes.size > 0) {
464+
throw new Error(
465+
`Unsupported function parameter type(s): ${[...unsupportedTypes]}`,
466+
);
473467
}
474468

475469
return {
476-
params: {
477-
type: 'identifiers',
478-
names: functionNode.params.flatMap((x) =>
479-
x.type === 'Identifier' ? [x.name] : []
470+
params: (functionNode
471+
.params as (
472+
| babel.Identifier
473+
| acorn.Identifier
474+
| babel.ObjectPattern
475+
| acorn.ObjectPattern
476+
)[]).map((param) =>
477+
param.type === 'ObjectPattern'
478+
? {
479+
type: FuncParameterType.destructuredObject,
480+
props: param.properties.flatMap((prop) =>
481+
(prop.type === 'Property' || prop.type === 'ObjectProperty') &&
482+
prop.key.type === 'Identifier' &&
483+
prop.value.type === 'Identifier'
484+
? [{ name: prop.key.name, alias: prop.value.name }]
485+
: []
486+
),
487+
}
488+
: {
489+
type: FuncParameterType.identifier,
490+
name: param.name,
491+
}
480492
),
481-
},
482493
body: functionNode.body,
483494
};
484495
}
485496

486497
export function transpileFn(rootNode: JsNode): TranspilationResult {
487-
const { params: argNames, body } = extractFunctionParts(rootNode);
498+
const { params, body } = extractFunctionParts(rootNode);
488499

489500
const ctx: Context = {
490501
externalNames: new Set(),
491502
ignoreExternalDepth: 0,
492503
stack: [
493504
{
494-
declaredNames: argNames.type === 'identifiers'
495-
? argNames.names
496-
: argNames.props.map((prop) => prop.alias),
505+
declaredNames: params.flatMap((param) =>
506+
param.type === FuncParameterType.identifier
507+
? param.name
508+
: param.props.map((prop) => prop.alias)
509+
),
497510
},
498511
],
499512
};
@@ -503,14 +516,14 @@ export function transpileFn(rootNode: JsNode): TranspilationResult {
503516

504517
if (body.type === 'BlockStatement') {
505518
return {
506-
argNames,
519+
params,
507520
body: tinyestBody as tinyest.Block,
508521
externalNames,
509522
};
510523
}
511524

512525
return {
513-
argNames,
526+
params,
514527
body: [NODE.block, [[NODE.return, tinyestBody as tinyest.Expression]]],
515528
externalNames,
516529
};

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

Lines changed: 68 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
import babel from '@babel/parser';
22
import type { Node } from '@babel/types';
33
import * as acorn from 'acorn';
4-
import type { ArgNames } from 'tinyest';
54
import { describe, expect, it } from 'vitest';
65
import { transpileFn } from '../src/parsers.ts';
76

@@ -28,12 +27,9 @@ describe('transpileFn', () => {
2827
it(
2928
'parses an empty arrow function',
3029
dualTest((p) => {
31-
const { argNames, body, externalNames } = transpileFn(p('() => {}'));
30+
const { params, body, externalNames } = transpileFn(p('() => {}'));
3231

33-
expect(argNames).toStrictEqual({
34-
type: 'identifiers',
35-
names: [],
36-
});
32+
expect(params).toStrictEqual([]);
3733
expect(JSON.stringify(body)).toMatchInlineSnapshot(`"[0,[]]"`);
3834
expect(externalNames).toStrictEqual([]);
3935
}),
@@ -42,14 +38,11 @@ describe('transpileFn', () => {
4238
it(
4339
'parses an empty named function',
4440
dualTest((p) => {
45-
const { argNames, body, externalNames } = transpileFn(
41+
const { params, body, externalNames } = transpileFn(
4642
p('function example() {}'),
4743
);
4844

49-
expect(argNames).toStrictEqual({
50-
type: 'identifiers',
51-
names: [],
52-
});
45+
expect(params).toStrictEqual([]);
5346
expect(JSON.stringify(body)).toMatchInlineSnapshot(`"[0,[]]"`);
5447
expect(externalNames).toStrictEqual([]);
5548
}),
@@ -58,14 +51,14 @@ describe('transpileFn', () => {
5851
it(
5952
'gathers external names',
6053
dualTest((p) => {
61-
const { argNames, body, externalNames } = transpileFn(
54+
const { params, body, externalNames } = transpileFn(
6255
p('(a, b) => a + b - c'),
6356
);
6457

65-
expect(argNames).toStrictEqual({
66-
type: 'identifiers',
67-
names: ['a', 'b'],
68-
});
58+
expect(params).toStrictEqual([
59+
{ type: 'i', name: 'a' },
60+
{ type: 'i', name: 'b' },
61+
]);
6962
expect(JSON.stringify(body)).toMatchInlineSnapshot(
7063
`"[0,[[10,[1,[1,"a","+","b"],"-","c"]]]]"`,
7164
);
@@ -76,17 +69,14 @@ describe('transpileFn', () => {
7669
it(
7770
'respects local declarations when gathering external names',
7871
dualTest((p) => {
79-
const { argNames, body, externalNames } = transpileFn(
72+
const { params, body, externalNames } = transpileFn(
8073
p(`() => {
8174
const a = 0;
8275
c = a + 2;
8376
}`),
8477
);
8578

86-
expect(argNames).toStrictEqual({
87-
type: 'identifiers',
88-
names: [],
89-
});
79+
expect(params).toStrictEqual([]);
9080
expect(JSON.stringify(body)).toMatchInlineSnapshot(
9181
`"[0,[[13,"a",[5,"0"]],[2,"c","=",[1,"a","+",[5,"2"]]]]]"`,
9282
);
@@ -98,7 +88,7 @@ describe('transpileFn', () => {
9888
it(
9989
'respects outer scope when gathering external names',
10090
dualTest((p) => {
101-
const { argNames, body, externalNames } = transpileFn(
91+
const { params, body, externalNames } = transpileFn(
10292
p(`() => {
10393
const a = 0;
10494
{
@@ -107,10 +97,7 @@ describe('transpileFn', () => {
10797
}`),
10898
);
10999

110-
expect(argNames).toStrictEqual({
111-
type: 'identifiers',
112-
names: [],
113-
});
100+
expect(params).toStrictEqual([]);
114101
expect(JSON.stringify(body)).toMatchInlineSnapshot(
115102
`"[0,[[13,"a",[5,"0"]],[0,[[2,"c","=",[1,"a","+",[5,"2"]]]]]]]"`,
116103
);
@@ -122,14 +109,11 @@ describe('transpileFn', () => {
122109
it(
123110
'treats the object as a possible external value when accessing a member',
124111
dualTest((p) => {
125-
const { argNames, body, externalNames } = transpileFn(
112+
const { params, body, externalNames } = transpileFn(
126113
p('() => external.outside.prop'),
127114
);
128115

129-
expect(argNames).toStrictEqual({
130-
type: 'identifiers',
131-
names: [],
132-
});
116+
expect(params).toStrictEqual([]);
133117
expect(JSON.stringify(body)).toMatchInlineSnapshot(
134118
`"[0,[[10,[7,[7,"external","outside"],"prop"]]]]"`,
135119
);
@@ -141,29 +125,75 @@ describe('transpileFn', () => {
141125
it(
142126
'handles destructured args',
143127
dualTest((p) => {
144-
const { argNames, externalNames } = transpileFn(
128+
const { params, externalNames } = transpileFn(
145129
p(`({ pos, a: b }) => {
146130
const x = pos.x;
147131
}`),
148132
);
149133

150-
expect(argNames).toStrictEqual(
151-
{
152-
type: 'destructured-object',
134+
expect(params).toStrictEqual(
135+
[{
136+
type: 'd',
153137
props: [
154138
{
155139
alias: 'pos',
156-
prop: 'pos',
140+
name: 'pos',
157141
},
158142
{
159143
alias: 'b',
160-
prop: 'a',
144+
name: 'a',
161145
},
162146
],
163-
} satisfies ArgNames,
147+
}],
164148
);
165149

166150
expect(externalNames).toStrictEqual([]);
167151
}),
168152
);
153+
154+
it(
155+
'handles mixed type parameters',
156+
dualTest((p) => {
157+
const { params, externalNames } = transpileFn(
158+
p(`(y, { pos, a: b }, {c, d}) => {
159+
const x = pos.x;
160+
}`),
161+
);
162+
163+
expect(params).toStrictEqual([
164+
{
165+
type: 'i',
166+
name: 'y',
167+
},
168+
{
169+
type: 'd',
170+
props: [
171+
{
172+
alias: 'pos',
173+
name: 'pos',
174+
},
175+
{
176+
alias: 'b',
177+
name: 'a',
178+
},
179+
],
180+
},
181+
{
182+
type: 'd',
183+
props: [
184+
{
185+
alias: 'c',
186+
name: 'c',
187+
},
188+
{
189+
alias: 'd',
190+
name: 'd',
191+
},
192+
],
193+
},
194+
]);
195+
196+
expect(externalNames).toStrictEqual([]);
197+
}),
198+
);
169199
});

‎packages/tinyest/src/nodes.ts

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -253,15 +253,20 @@ export type Expression =
253253

254254
export type AnyNode = Statement | Expression;
255255

256-
export type ArgNames =
256+
export const FuncParameterType = {
257+
identifier: 'i',
258+
destructuredObject: 'd',
259+
} as const;
260+
261+
export type FuncParameter =
257262
| {
258-
type: 'identifiers';
259-
names: string[];
263+
type: typeof FuncParameterType.identifier;
264+
name: string;
260265
}
261266
| {
262-
type: 'destructured-object';
267+
type: typeof FuncParameterType.destructuredObject;
263268
props: {
264-
prop: string;
269+
name: string;
265270
alias: string;
266271
}[];
267272
};

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
1-
import type { ArgNames, Block } from 'tinyest';
1+
import type { Block, FuncParameter } from 'tinyest';
22

33
export type Ast = {
4-
argNames: ArgNames;
4+
params: FuncParameter[];
55
body: Block;
66
externalNames: string[];
77
};

0 commit comments

Comments
 (0)