Skip to content

Commit 4b4d804

Browse files
aleksanderkatanaleksanderkatan
authored
fix: TypeGPU functions using TGSL dependencies declaration order (#1522)
1 parent 83d1428 commit 4b4d804

File tree

8 files changed

+159
-81
lines changed

8 files changed

+159
-81
lines changed

‎packages/typegpu/tests/tgslFn.test.ts

Lines changed: 58 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -574,36 +574,6 @@ describe('TGSL tgpu.fn function', () => {
574574
);
575575
});
576576

577-
it('(when using plugin) can be invoked for a constant with "kernel" directive', () => {
578-
const addKernelJs = (x: number, y: number) => {
579-
'kernel';
580-
return x + y;
581-
};
582-
583-
const add = tgpu.fn([d.u32, d.u32])(addKernelJs);
584-
585-
expect(addKernelJs(2, 3)).toBe(5);
586-
expect(add(2, 3)).toBe(5);
587-
expect(parseResolved({ add })).toBe(
588-
parse(`fn add(x: u32, y: u32){
589-
return (x + y);
590-
}`),
591-
);
592-
});
593-
594-
it('(when using plugin) can be invoked for inline function with no directive', () => {
595-
const add = tgpu.fn([d.u32, d.u32])(
596-
(x, y) => x + y,
597-
);
598-
599-
expect(add(2, 3)).toBe(5);
600-
expect(parseResolved({ add })).toBe(
601-
parse(`fn add(x: u32, y: u32){
602-
return (x + y);
603-
}`),
604-
);
605-
});
606-
607577
it('resolves a function with a pointer parameter', () => {
608578
const addOnes = tgpu.fn([d.ptrStorage(d.vec3f, 'read-write')])((ptr) => {
609579
ptr.x += 1;
@@ -879,3 +849,61 @@ describe('tgpu.fn arguments', () => {
879849
expect(vec).toStrictEqual(d.vec3f());
880850
});
881851
});
852+
853+
describe('tgsl fn when using plugin', () => {
854+
it('can be invoked for a constant with "kernel" directive', () => {
855+
const addKernelJs = (x: number, y: number) => {
856+
'kernel';
857+
return x + y;
858+
};
859+
860+
const add = tgpu.fn([d.u32, d.u32])(addKernelJs);
861+
862+
expect(addKernelJs(2, 3)).toBe(5);
863+
expect(add(2, 3)).toBe(5);
864+
expect(parseResolved({ add })).toBe(
865+
parse(`fn add(x: u32, y: u32){
866+
return (x + y);
867+
}`),
868+
);
869+
});
870+
871+
it('can be invoked for inline function with no directive', () => {
872+
const add = tgpu.fn([d.u32, d.u32])(
873+
(x, y) => x + y,
874+
);
875+
876+
expect(add(2, 3)).toBe(5);
877+
expect(parseResolved({ add })).toBe(
878+
parse(`fn add(x: u32, y: u32){
879+
return (x + y);
880+
}`),
881+
);
882+
});
883+
884+
it('can reference function defined below', () => {
885+
const bar = tgpu.fn([], d.f32)(() => foo() + 2);
886+
const foo = tgpu.fn([], d.f32)(() => 1);
887+
888+
expect(parseResolved({ bar })).toBe(
889+
parse(`
890+
fn foo() -> f32 {
891+
return 1;
892+
}
893+
894+
fn bar() -> f32 {
895+
return (foo() + 2);
896+
}`),
897+
);
898+
});
899+
900+
// TODO: throw an error when cyclic dependency is detected
901+
// it('throws when it detects a cyclic dependency', () => {
902+
// let bar: TgpuFn;
903+
// let foo: TgpuFn;
904+
// bar = tgpu.fn([], d.f32)(() => foo() + 2);
905+
// foo = tgpu.fn([], d.f32)(() => bar() - 2);
906+
907+
// expect(() => parseResolved({ bar })).toThrowErrorMatchingInlineSnapshot(``);
908+
// });
909+
});

‎packages/unplugin-typegpu/src/babel.ts

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -54,13 +54,19 @@ function functionToTranspiled(
5454
i('ast'),
5555
template.expression`${embedJSON({ params, body, externalNames })}`(),
5656
),
57-
types.objectProperty(
57+
types.objectMethod(
58+
'get',
5859
i('externals'),
59-
types.objectExpression(
60-
externalNames.map((name) =>
61-
types.objectProperty(i(name), i(name), false, true)
60+
[],
61+
types.blockStatement([
62+
types.returnStatement(
63+
types.objectExpression(
64+
externalNames.map((name) =>
65+
types.objectProperty(i(name), i(name), false, true)
66+
),
67+
),
6268
),
63-
),
69+
]),
6470
),
6571
]);
6672

‎packages/unplugin-typegpu/src/index.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -187,7 +187,7 @@ const typegpu: UnpluginInstance<Options, false> = createUnplugin(
187187
const metadata = `{
188188
v: ${FORMAT_VERSION},
189189
ast: ${embedJSON({ params, body, externalNames })},
190-
externals: {${externalNames.join(', ')}},
190+
get externals() { return {${externalNames.join(', ')}}; },
191191
}`;
192192

193193
assignMetadata(magicString, def, metadata);

‎packages/unplugin-typegpu/test/aliasing.test.ts

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,9 @@ describe('[BABEL] tgpu alias gathering', () => {
1818
}, {
1919
v: 1,
2020
ast: {"params":[],"body":[0,[[13,"x",[1,[5,"2"],"+",[5,"2"]]]]],"externalNames":[]},
21-
externals: {}
21+
get externals() {
22+
return {};
23+
}
2224
}) && $.f)({}));"
2325
`);
2426
});
@@ -39,7 +41,9 @@ describe('[BABEL] tgpu alias gathering', () => {
3941
}, {
4042
v: 1,
4143
ast: {"params":[],"body":[0,[[13,"x",[1,[5,"2"],"+",[5,"2"]]]]],"externalNames":[]},
42-
externals: {}
44+
get externals() {
45+
return {};
46+
}
4347
}) && $.f)({}));"
4448
`);
4549
});
@@ -60,7 +64,9 @@ describe('[BABEL] tgpu alias gathering', () => {
6064
}, {
6165
v: 1,
6266
ast: {"params":[],"body":[0,[[13,"x",[1,[5,"2"],"+",[5,"2"]]]]],"externalNames":[]},
63-
externals: {}
67+
get externals() {
68+
return {};
69+
}
6470
}) && $.f)({}));"
6571
`);
6672
});
@@ -82,7 +88,7 @@ describe('[ROLLUP] tgpu alias gathering', () => {
8288
}), {
8389
v: 1,
8490
ast: {"params":[],"body":[0,[]],"externalNames":[]},
85-
externals: {},
91+
get externals() { return {}; },
8692
}) && $.f)({})));
8793
"
8894
`);
@@ -104,7 +110,7 @@ describe('[ROLLUP] tgpu alias gathering', () => {
104110
}), {
105111
v: 1,
106112
ast: {"params":[],"body":[0,[]],"externalNames":[]},
107-
externals: {},
113+
get externals() { return {}; },
108114
}) && $.f)({})));
109115
"
110116
`);
@@ -126,7 +132,7 @@ describe('[ROLLUP] tgpu alias gathering', () => {
126132
}), {
127133
v: 1,
128134
ast: {"params":[],"body":[0,[]],"externalNames":[]},
129-
externals: {},
135+
get externals() { return {}; },
130136
}) && $.f)({})));
131137
"
132138
`);

‎packages/unplugin-typegpu/test/auto-naming.test.ts

Lines changed: 21 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,9 @@ describe('[BABEL] auto naming', () => {
2424
var fn = (globalThis.__TYPEGPU_AUTONAME__ ?? (a => a))(tgpu.fn([])(($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = () => {}, {
2525
v: 1,
2626
ast: {"params":[],"body":[0,[]],"externalNames":[]},
27-
externals: {}
27+
get externals() {
28+
return {};
29+
}
2830
}) && $.f)({})), "fn");
2931
let shell = (globalThis.__TYPEGPU_AUTONAME__ ?? (a => a))(tgpu.fn([]), "shell");
3032
console.log(bindGroupLayout, vertexLayout);"
@@ -126,14 +128,18 @@ describe('[BABEL] auto naming', () => {
126128
const myFunction = (globalThis.__TYPEGPU_AUTONAME__ ?? (a => a))(tgpu.fn([])(($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = () => 0, {
127129
v: 1,
128130
ast: {"params":[],"body":[0,[[10,[5,"0"]]]],"externalNames":[]},
129-
externals: {}
131+
get externals() {
132+
return {};
133+
}
130134
}) && $.f)({})), "myFunction");
131135
const myComputeFn = (globalThis.__TYPEGPU_AUTONAME__ ?? (a => a))(tgpu['~unstable'].computeFn({
132136
workgroupSize: [1]
133137
})(($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = () => {}, {
134138
v: 1,
135139
ast: {"params":[],"body":[0,[]],"externalNames":[]},
136-
externals: {}
140+
get externals() {
141+
return {};
142+
}
137143
}) && $.f)({})), "myComputeFn");
138144
const myVertexFn = (globalThis.__TYPEGPU_AUTONAME__ ?? (a => a))(tgpu['~unstable'].vertexFn({
139145
out: {
@@ -144,7 +150,9 @@ describe('[BABEL] auto naming', () => {
144150
}), {
145151
v: 1,
146152
ast: {"params":[],"body":[0,[[10,[104,{"ret":[5,"0"]}]]]],"externalNames":[]},
147-
externals: {}
153+
get externals() {
154+
return {};
155+
}
148156
}) && $.f)({})), "myVertexFn");
149157
const myFragmentFn = (globalThis.__TYPEGPU_AUTONAME__ ?? (a => a))(tgpu['~unstable'].fragmentFn({
150158
in: {
@@ -154,8 +162,10 @@ describe('[BABEL] auto naming', () => {
154162
})(($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = () => d.vec4f(), {
155163
v: 1,
156164
ast: {"params":[],"body":[0,[[10,[6,[7,"d","vec4f"],[]]]]],"externalNames":["d"]},
157-
externals: {
158-
d
165+
get externals() {
166+
return {
167+
d
168+
};
159169
}
160170
}) && $.f)({})), "myFragmentFn");"
161171
`);
@@ -304,7 +314,7 @@ describe('[ROLLUP] auto naming', () => {
304314
((globalThis.__TYPEGPU_AUTONAME__ ?? (a => a))(tgpu.fn([])((($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (() => {}), {
305315
v: 1,
306316
ast: {"params":[],"body":[0,[]],"externalNames":[]},
307-
externals: {},
317+
get externals() { return {}; },
308318
}) && $.f)({}))), "fn"));
309319
310320
console.log(bindGroupLayout, vertexLayout);
@@ -412,20 +422,20 @@ describe('[ROLLUP] auto naming', () => {
412422
((globalThis.__TYPEGPU_AUTONAME__ ?? (a => a))(tgpu.fn([])((($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (() => 0), {
413423
v: 1,
414424
ast: {"params":[],"body":[0,[[10,[5,"0"]]]],"externalNames":[]},
415-
externals: {},
425+
get externals() { return {}; },
416426
}) && $.f)({}))), "myFunction"));
417427
((globalThis.__TYPEGPU_AUTONAME__ ?? (a => a))(tgpu['~unstable'].computeFn({ workgroupSize: [1] })(
418428
(($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (() => {}), {
419429
v: 1,
420430
ast: {"params":[],"body":[0,[]],"externalNames":[]},
421-
externals: {},
431+
get externals() { return {}; },
422432
}) && $.f)({})),
423433
), "myComputeFn"));
424434
((globalThis.__TYPEGPU_AUTONAME__ ?? (a => a))(tgpu['~unstable'].vertexFn({ out: { ret: d.i32 } })(
425435
(($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (() => ({ ret: 0 })), {
426436
v: 1,
427437
ast: {"params":[],"body":[0,[[10,[104,{"ret":[5,"0"]}]]]],"externalNames":[]},
428-
externals: {},
438+
get externals() { return {}; },
429439
}) && $.f)({})),
430440
), "myVertexFn"));
431441
((globalThis.__TYPEGPU_AUTONAME__ ?? (a => a))(tgpu['~unstable'].fragmentFn({
@@ -435,7 +445,7 @@ describe('[ROLLUP] auto naming', () => {
435445
(($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (() => d.vec4f()), {
436446
v: 1,
437447
ast: {"params":[],"body":[0,[[10,[6,[7,"d","vec4f"],[]]]]],"externalNames":["d"]},
438-
externals: {d},
448+
get externals() { return {d}; },
439449
}) && $.f)({})),
440450
), "myFragmentFn"));
441451
"

0 commit comments

Comments
 (0)