From 2f74d793dce71a011cf514b55df755b39d102a7c Mon Sep 17 00:00:00 2001 From: Aleksander Katan <56294622+aleksanderkatan@users.noreply.github.com> Date: Thu, 9 Jul 2026 11:32:10 +0200 Subject: [PATCH 1/6] Add tests --- .../test/use-gpu-directive.test.ts | 46 +++++++++++++++++++ 1 file changed, 46 insertions(+) diff --git a/packages/unplugin-typegpu/test/use-gpu-directive.test.ts b/packages/unplugin-typegpu/test/use-gpu-directive.test.ts index 9e8dea4ea2..774f452480 100644 --- a/packages/unplugin-typegpu/test/use-gpu-directive.test.ts +++ b/packages/unplugin-typegpu/test/use-gpu-directive.test.ts @@ -1,6 +1,52 @@ import { describe, expect, test } from 'vitest'; import { babelTransform, rollupTransform } from './transform.ts'; +describe('"use gpu" is removed after transform', () => { + const code = `\ + const fn = () => { + 'use gpu'; + }; + + console.log(fn); + `; + + test('babel', () => { + expect(babelTransform(code)).toMatchInlineSnapshot(` + "const fn = /*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = () => { + 'use gpu'; + }, { + v: 1, + name: "fn", + ast: { + params: [], + body: [0, []], + externalNames: [] + }, + externals: () => { + return {}; + } + }) && $.f)({}); + console.log(fn);" + `); + }); + + test('rollup', async () => { + expect(await rollupTransform(code)).toMatchInlineSnapshot(` + "const fn = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (() => { + 'use gpu'; + }), { + v: 1, + name: "fn", + ast: {"params":[],"body":[0,[]],"externalNames":[]}, + externals: () => ({}), + }) && $.f)({})); + + console.log(fn); + " + `); + }); +}); + describe('"use gpu" marked arrow function, assigned to a const', () => { const code = `\ /** ADD */ From 07e2bab33df80860ff99a8eb3fe011af1eec4ce1 Mon Sep 17 00:00:00 2001 From: Aleksander Katan <56294622+aleksanderkatan@users.noreply.github.com> Date: Thu, 9 Jul 2026 12:01:50 +0200 Subject: [PATCH 2/6] Add method stubs --- packages/unplugin-typegpu/src/babel.ts | 4 ++++ packages/unplugin-typegpu/src/core/common.ts | 5 ++++- packages/unplugin-typegpu/src/core/factory.ts | 3 +++ 3 files changed, 11 insertions(+), 1 deletion(-) diff --git a/packages/unplugin-typegpu/src/babel.ts b/packages/unplugin-typegpu/src/babel.ts index 371323aba7..6b9585ec46 100644 --- a/packages/unplugin-typegpu/src/babel.ts +++ b/packages/unplugin-typegpu/src/babel.ts @@ -4,6 +4,7 @@ import { transpileFn } from 'tinyest-for-wgsl'; import * as t from '@babel/types'; import { METADATA_FORMAT_VERSION, + type MetadatableFunction, type PluginState, defaultOptions, functionVisitor, @@ -148,6 +149,8 @@ function replaceWithAssignmentOverload( ); } +function removeUseGpuDirective(path: NodePath) {} + function replaceWithBinaryOverload(path: NodePath, runtimeFn: string): void { path.replaceWith( t.callExpression(i(runtimeFn), [path.node.left as t.Expression, path.node.right]), @@ -165,6 +168,7 @@ export default function TypeGPUPlugin() { wrapInAutoName, replaceWithAssignmentOverload, replaceWithBinaryOverload, + removeUseGpuDirective, }); }, visitor: { diff --git a/packages/unplugin-typegpu/src/core/common.ts b/packages/unplugin-typegpu/src/core/common.ts index 2a2de42f9d..12f47c57ce 100644 --- a/packages/unplugin-typegpu/src/core/common.ts +++ b/packages/unplugin-typegpu/src/core/common.ts @@ -85,6 +85,8 @@ export interface TransformMethods { replaceWithAssignmentOverload(path: NodePath, runtimeFn: string): void; replaceWithBinaryOverload(path: NodePath, runtimeFn: string): void; + + removeUseGpuDirective(path: NodePath): void; } export interface PluginState extends TransformMethods { @@ -108,7 +110,7 @@ export interface PluginState extends TransformMethods { inUseGpuScope: boolean; - alreadyTransformed: WeakSet; + alreadyTransformed: WeakSet; // TODO: take a look } export interface NodeLocation { @@ -450,6 +452,7 @@ function functionOnExit( if (!containsUseGpuDirective(node)) { return; } + state.removeUseGpuDirective(path); state.inUseGpuScope = false; diff --git a/packages/unplugin-typegpu/src/core/factory.ts b/packages/unplugin-typegpu/src/core/factory.ts index 3005352573..d2e51adb94 100644 --- a/packages/unplugin-typegpu/src/core/factory.ts +++ b/packages/unplugin-typegpu/src/core/factory.ts @@ -127,6 +127,8 @@ function replaceWithBinaryOverload( this.overwrite(path.node, `${runtimeFn}(${lhs}, ${rhs})`); } +function removeUseGpuDirective(path: NodePath) {} + const NodeUtils = { slice(this: UnpluginPluginState, node: NodeLocation): string { return this.magicString.slice(node.start ?? 0, node.end ?? 0); @@ -188,6 +190,7 @@ export const unpluginFactory = ((rawOptions, _meta) => { wrapInAutoName, replaceWithAssignmentOverload, replaceWithBinaryOverload, + removeUseGpuDirective, }); traverse(ast, functionVisitor, undefined, state); From 0f955bdd152d952e7e822f8d8468f48f67ce27d8 Mon Sep 17 00:00:00 2001 From: Aleksander Katan <56294622+aleksanderkatan@users.noreply.github.com> Date: Thu, 9 Jul 2026 12:24:01 +0200 Subject: [PATCH 3/6] remove directive in rollup --- packages/unplugin-typegpu/src/babel.ts | 2 +- packages/unplugin-typegpu/src/core/common.ts | 2 +- packages/unplugin-typegpu/src/core/factory.ts | 9 +++- .../unplugin-typegpu/test/auto-naming.test.ts | 10 ++-- .../test/use-gpu-directive.test.ts | 54 +++++++++---------- 5 files changed, 42 insertions(+), 35 deletions(-) diff --git a/packages/unplugin-typegpu/src/babel.ts b/packages/unplugin-typegpu/src/babel.ts index 6b9585ec46..8ff3446e9f 100644 --- a/packages/unplugin-typegpu/src/babel.ts +++ b/packages/unplugin-typegpu/src/babel.ts @@ -149,7 +149,7 @@ function replaceWithAssignmentOverload( ); } -function removeUseGpuDirective(path: NodePath) {} +function removeUseGpuDirective(this: PluginState, path: NodePath) {} function replaceWithBinaryOverload(path: NodePath, runtimeFn: string): void { path.replaceWith( diff --git a/packages/unplugin-typegpu/src/core/common.ts b/packages/unplugin-typegpu/src/core/common.ts index 12f47c57ce..a403113e51 100644 --- a/packages/unplugin-typegpu/src/core/common.ts +++ b/packages/unplugin-typegpu/src/core/common.ts @@ -86,7 +86,7 @@ export interface TransformMethods { replaceWithBinaryOverload(path: NodePath, runtimeFn: string): void; - removeUseGpuDirective(path: NodePath): void; + removeUseGpuDirective(this: PluginState, path: NodePath): void; } export interface PluginState extends TransformMethods { diff --git a/packages/unplugin-typegpu/src/core/factory.ts b/packages/unplugin-typegpu/src/core/factory.ts index d2e51adb94..4c438a6c1c 100644 --- a/packages/unplugin-typegpu/src/core/factory.ts +++ b/packages/unplugin-typegpu/src/core/factory.ts @@ -127,7 +127,14 @@ function replaceWithBinaryOverload( this.overwrite(path.node, `${runtimeFn}(${lhs}, ${rhs})`); } -function removeUseGpuDirective(path: NodePath) {} +function removeUseGpuDirective(this: UnpluginPluginState, path: NodePath) { + const directives = 'directives' in path.node.body ? (path.node.body?.directives ?? []) : []; + for (const directive of directives) { + if (directive.value.value === 'use gpu') { + this.overwrite(directive, ''); + } + } +} const NodeUtils = { slice(this: UnpluginPluginState, node: NodeLocation): string { diff --git a/packages/unplugin-typegpu/test/auto-naming.test.ts b/packages/unplugin-typegpu/test/auto-naming.test.ts index 9c1fc89193..dc5c4802d9 100644 --- a/packages/unplugin-typegpu/test/auto-naming.test.ts +++ b/packages/unplugin-typegpu/test/auto-naming.test.ts @@ -829,7 +829,7 @@ describe('[ROLLUP] auto naming', () => { expect(await rollupTransform(code, { autoNamingEnabled: true })).toMatchInlineSnapshot(` "const myFun3 = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (function myFun3() { - 'use gpu'; + return 0; }), { v: 1, @@ -839,7 +839,7 @@ describe('[ROLLUP] auto naming', () => { }) && $.f)({})); const myFun1 = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (() => { - 'use gpu'; + return 0; }), { v: 1, @@ -849,7 +849,7 @@ describe('[ROLLUP] auto naming', () => { }) && $.f)({})); const myFun2 = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (function () { - 'use gpu'; + return 0; }), { v: 1, @@ -1027,7 +1027,7 @@ describe('[ROLLUP] auto naming', () => { const root = await tgpu.init(); const myGuardedPipeline = (/*#__PURE__*/(globalThis.__TYPEGPU_AUTONAME__ ?? (a => a))(root.createGuardedComputePipeline((/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (() => { - 'use gpu'; + }), { v: 1, name: undefined, @@ -1037,7 +1037,7 @@ describe('[ROLLUP] auto naming', () => { const anotherGuardedPipeline = (/*#__PURE__*/(globalThis.__TYPEGPU_AUTONAME__ ?? (a => a))(root .createGuardedComputePipeline((/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (() => { - 'use gpu'; + }), { v: 1, name: undefined, diff --git a/packages/unplugin-typegpu/test/use-gpu-directive.test.ts b/packages/unplugin-typegpu/test/use-gpu-directive.test.ts index 774f452480..b84394bd30 100644 --- a/packages/unplugin-typegpu/test/use-gpu-directive.test.ts +++ b/packages/unplugin-typegpu/test/use-gpu-directive.test.ts @@ -33,7 +33,7 @@ describe('"use gpu" is removed after transform', () => { test('rollup', async () => { expect(await rollupTransform(code)).toMatchInlineSnapshot(` "const fn = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (() => { - 'use gpu'; + }), { v: 1, name: "fn", @@ -101,7 +101,7 @@ describe('"use gpu" marked arrow function, assigned to a const', () => { "/** ADD */ // another comment const addGPU = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = ((a, b) => { - 'use gpu'; + return __tsover_add(a, b); }), { v: 1, @@ -175,7 +175,7 @@ describe('marked arrow functions passed to shells', () => { const shell = tgpu.fn([]); shell((/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = ((a, b) => { - 'use gpu'; + return __tsover_add(a, b); }), { v: 1, @@ -247,7 +247,7 @@ describe('marked anonymous function expressions passed to shells', () => { const shell = tgpu.fn([]); shell((/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (function(a, b){ - 'use gpu'; + return __tsover_add(a, b); }), { v: 1, @@ -319,7 +319,7 @@ describe('marked named function expressions passed to shells', () => { const shell = tgpu.fn([]); shell((/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (function addGPU(a, b){ - 'use gpu'; + return __tsover_add(a, b); }), { v: 1, @@ -390,7 +390,7 @@ describe('marked function statements', () => { expect(await rollupTransform(code)).toMatchInlineSnapshot(` "/** ADD */ const addGPU = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (function addGPU(a, b) { - 'use gpu'; + // hello there return __tsover_add(a, b); }), { @@ -508,7 +508,7 @@ describe('marked object methods', () => { "const obj = { /** MOD */ mod: (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = ((a, b) => { - 'use gpu'; + return __tsover_mod(a, b); }), { v: 1, @@ -520,7 +520,7 @@ describe('marked object methods', () => { /** PRIME */ const isPrime = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = ((n) => { - 'use gpu'; + if (n <= 1) { return false; } @@ -611,7 +611,7 @@ describe('transforms numeric operations', () => { /** the main function */ // another comment const main = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = ((a, b) => { - 'use gpu'; + let c = __tsover_add(__tsover_add(a, b), 2); c = __tsover_add(c, __tsover_mul(2, b)); countMutable.$ = __tsover_add(countMutable.$, 3); @@ -707,7 +707,7 @@ describe('hoists global function statements marked with "use gpu"', () => { "/** MUL */ // another comment const mul = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (function mul(a, b) { - 'use gpu'; + return __tsover_mul(a, b); }), { v: 1, @@ -719,7 +719,7 @@ describe('hoists global function statements marked with "use gpu"', () => { /** ADD */ // another comment const add = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (function add(a, b) { - 'use gpu'; + return __tsover_add(a, b); }), { v: 1, @@ -817,7 +817,7 @@ describe('hoists function statements marked with "use gpu", scoped inside anothe /** MUL */ // another comment const mul = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (function mul(a, b) { - 'use gpu'; + return __tsover_mul(a, b); }), { v: 1, @@ -829,7 +829,7 @@ describe('hoists function statements marked with "use gpu", scoped inside anothe /** ADD */ // another comment const add = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (function add(a, b) { - 'use gpu'; + return __tsover_add(a, b); }), { v: 1, @@ -938,7 +938,7 @@ describe('hoists function statements marked with "use gpu", scoped inside an arr /** MUL */ // another comment const mul = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (function mul(a, b) { - 'use gpu'; + return __tsover_mul(a, b); }), { v: 1, @@ -950,7 +950,7 @@ describe('hoists function statements marked with "use gpu", scoped inside an arr /** ADD */ // another comment const add = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (function add(a, b) { - 'use gpu'; + return __tsover_add(a, b); }), { v: 1, @@ -1066,7 +1066,7 @@ describe('hoists function statements marked with "use gpu", scoped inside an if /** MUL */ // another comment const mul = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (function mul(a, b) { - 'use gpu'; + return __tsover_mul(__tsover_mul(a, b), c); }), { v: 1, @@ -1078,7 +1078,7 @@ describe('hoists function statements marked with "use gpu", scoped inside an if /** ADD */ // another comment const add = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (function add(a, b) { - 'use gpu'; + return __tsover_add(__tsover_add(a, b), c); }), { v: 1, @@ -1203,7 +1203,7 @@ describe('replaces function statements marked with "use gpu" in place when condi /** ADD */ // another comment const add = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (function add(a, b) { - 'use gpu'; + return __tsover_add(__tsover_add(a, b), c); }), { v: 1, @@ -1218,7 +1218,7 @@ describe('replaces function statements marked with "use gpu" in place when condi /** MUL */ // another comment const mul = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (function mul(a, b) { - 'use gpu'; + return __tsover_mul(__tsover_mul(a, b), c); }), { v: 1, @@ -1312,7 +1312,7 @@ describe('hoists exported marked function statements', () => { expect(await rollupTransform(code)).toMatchInlineSnapshot(` "/** MUL */ const mul = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (function mul(a, b) { - 'use gpu'; + return __tsover_mul(a, b); }), { v: 1, @@ -1323,7 +1323,7 @@ describe('hoists exported marked function statements', () => { /** ADD */ const add = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (function add(a, b) { - 'use gpu'; + return __tsover_add(a, b); }), { v: 1, @@ -1386,7 +1386,7 @@ describe('hoists default exported marked function statement', () => { expect(await rollupTransform(code)).toMatchInlineSnapshot(` "/** ADD */ const add = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (function add(a, b) { - 'use gpu'; + return __tsover_add(a, b); }), { v: 1, @@ -1497,7 +1497,7 @@ describe('export marked arrow function', () => { test('rollup', async () => { expect(await rollupTransform(code)).toMatchInlineSnapshot(` "const add = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = ((a, b) => { - 'use gpu'; + return __tsover_add(a, b); }), { v: 1, @@ -1507,7 +1507,7 @@ describe('export marked arrow function', () => { }) && $.f)({})); const increment = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = ((n) => { - 'use gpu'; + return __tsover_add(n, 1); }), { v: 1, @@ -1517,7 +1517,7 @@ describe('export marked arrow function', () => { }) && $.f)({})); const mul = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = ((a, b) => { - 'use gpu'; + return __tsover_mul(a, b); }), { v: 1, @@ -1570,7 +1570,7 @@ describe('anonymous default export marked function statement', () => { test('rollup', async () => { expect(await rollupTransform(code)).toMatchInlineSnapshot(` "var _virtual_code = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (function (a, b) { - 'use gpu'; + return __tsover_add(a, b); }), { v: 1, @@ -1623,7 +1623,7 @@ describe('anonymous default export marked arrow function', () => { test('rollup', async () => { expect(await rollupTransform(code)).toMatchInlineSnapshot(` "var _virtual_code = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = ((a, b) => { - 'use gpu'; + return __tsover_add(a, b); }), { v: 1, From c03764a203dec1f548dc7999ffcbaa54820d1b41 Mon Sep 17 00:00:00 2001 From: Aleksander Katan <56294622+aleksanderkatan@users.noreply.github.com> Date: Fri, 10 Jul 2026 11:06:56 +0200 Subject: [PATCH 4/6] Remove directive in babel --- .../typegpu/tests/internal/autoname.test.ts | 2 +- packages/unplugin-typegpu/src/babel.ts | 6 +- .../unplugin-typegpu/test/auto-naming.test.ts | 14 +---- .../test/typescript-syntax.test.ts | 2 - .../test/use-gpu-directive.test.ts | 56 +------------------ 5 files changed, 9 insertions(+), 71 deletions(-) diff --git a/packages/typegpu/tests/internal/autoname.test.ts b/packages/typegpu/tests/internal/autoname.test.ts index de3b55d8cc..d66680aa95 100644 --- a/packages/typegpu/tests/internal/autoname.test.ts +++ b/packages/typegpu/tests/internal/autoname.test.ts @@ -195,7 +195,7 @@ describe('autonaming', () => { expect(scope.toString()).toMatchInlineSnapshot(` "() => { const myFun = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (function myFun() { - "use gpu"; + return 0; }), { v: 1, diff --git a/packages/unplugin-typegpu/src/babel.ts b/packages/unplugin-typegpu/src/babel.ts index 8ff3446e9f..588dbf3831 100644 --- a/packages/unplugin-typegpu/src/babel.ts +++ b/packages/unplugin-typegpu/src/babel.ts @@ -149,7 +149,11 @@ function replaceWithAssignmentOverload( ); } -function removeUseGpuDirective(this: PluginState, path: NodePath) {} +function removeUseGpuDirective(this: PluginState, path: NodePath) { + const directives = path.get('body').get('directives'); + const maybeUseGpu = directives.find((directive) => directive.node.value.value === 'use gpu'); + maybeUseGpu?.remove(); +} function replaceWithBinaryOverload(path: NodePath, runtimeFn: string): void { path.replaceWith( diff --git a/packages/unplugin-typegpu/test/auto-naming.test.ts b/packages/unplugin-typegpu/test/auto-naming.test.ts index dc5c4802d9..880e335f42 100644 --- a/packages/unplugin-typegpu/test/auto-naming.test.ts +++ b/packages/unplugin-typegpu/test/auto-naming.test.ts @@ -327,8 +327,6 @@ describe('[BABEL] auto naming', () => { expect(babelTransform(code, { autoNamingEnabled: true })).toMatchInlineSnapshot(` "const myFun3 = /*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = function myFun3() { - 'use gpu'; - return 0; }, { v: 1, @@ -343,8 +341,6 @@ describe('[BABEL] auto naming', () => { } }) && $.f)({}); const myFun1 = /*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = () => { - 'use gpu'; - return 0; }, { v: 1, @@ -359,8 +355,6 @@ describe('[BABEL] auto naming', () => { } }) && $.f)({}); const myFun2 = /*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = function () { - 'use gpu'; - return 0; }, { v: 1, @@ -510,9 +504,7 @@ describe('[BABEL] auto naming', () => { expect(babelTransform(code, { autoNamingEnabled: true })).toMatchInlineSnapshot(` "import { tgpu } from 'typegpu'; const root = await tgpu.init(); - const myGuardedPipeline = /*#__PURE__*/(globalThis.__TYPEGPU_AUTONAME__ ?? (a => a))(root.createGuardedComputePipeline(/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = () => { - 'use gpu'; - }, { + const myGuardedPipeline = /*#__PURE__*/(globalThis.__TYPEGPU_AUTONAME__ ?? (a => a))(root.createGuardedComputePipeline(/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = () => {}, { v: 1, name: undefined, ast: { @@ -524,9 +516,7 @@ describe('[BABEL] auto naming', () => { return {}; } }) && $.f)({})), "myGuardedPipeline"); - const anotherGuardedPipeline = /*#__PURE__*/(globalThis.__TYPEGPU_AUTONAME__ ?? (a => a))(root.createGuardedComputePipeline(/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = () => { - 'use gpu'; - }, { + const anotherGuardedPipeline = /*#__PURE__*/(globalThis.__TYPEGPU_AUTONAME__ ?? (a => a))(root.createGuardedComputePipeline(/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = () => {}, { v: 1, name: undefined, ast: { diff --git a/packages/unplugin-typegpu/test/typescript-syntax.test.ts b/packages/unplugin-typegpu/test/typescript-syntax.test.ts index b6af6fbff9..d539afa3cb 100644 --- a/packages/unplugin-typegpu/test/typescript-syntax.test.ts +++ b/packages/unplugin-typegpu/test/typescript-syntax.test.ts @@ -14,8 +14,6 @@ describe('as type', () => { expect(babelTransform(code)).toMatchInlineSnapshot(` "const hello = /*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (a: number, b: number | undefined) => { - 'use gpu'; - return __tsover_add(a, b as number); }, { v: 1, diff --git a/packages/unplugin-typegpu/test/use-gpu-directive.test.ts b/packages/unplugin-typegpu/test/use-gpu-directive.test.ts index b84394bd30..a5bdccbcb6 100644 --- a/packages/unplugin-typegpu/test/use-gpu-directive.test.ts +++ b/packages/unplugin-typegpu/test/use-gpu-directive.test.ts @@ -12,9 +12,7 @@ describe('"use gpu" is removed after transform', () => { test('babel', () => { expect(babelTransform(code)).toMatchInlineSnapshot(` - "const fn = /*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = () => { - 'use gpu'; - }, { + "const fn = /*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = () => {}, { v: 1, name: "fn", ast: { @@ -68,8 +66,6 @@ describe('"use gpu" marked arrow function, assigned to a const', () => { "/** ADD */ // another comment const addGPU = /*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (a, b) => { - 'use gpu'; - return __tsover_add(a, b); }, { v: 1, @@ -141,8 +137,6 @@ describe('marked arrow functions passed to shells', () => { "import { tgpu } from 'typegpu'; const shell = tgpu.fn([]); shell(/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (a, b) => { - 'use gpu'; - return __tsover_add(a, b); }, { v: 1, @@ -213,8 +207,6 @@ describe('marked anonymous function expressions passed to shells', () => { "import { tgpu } from 'typegpu'; const shell = tgpu.fn([]); shell(/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = function (a, b) { - 'use gpu'; - return __tsover_add(a, b); }, { v: 1, @@ -285,8 +277,6 @@ describe('marked named function expressions passed to shells', () => { "import { tgpu } from 'typegpu'; const shell = tgpu.fn([]); shell(/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = function addGPU(a, b) { - 'use gpu'; - return __tsover_add(a, b); }, { v: 1, @@ -357,8 +347,6 @@ describe('marked function statements', () => { expect(babelTransform(code)).toMatchInlineSnapshot(` "/** ADD */ const addGPU = /*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = function addGPU(a, b) { - 'use gpu'; - // hello there return __tsover_add(a, b); }, { @@ -446,8 +434,6 @@ describe('marked object methods', () => { "const obj = { /** MOD */ mod: /*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (a, b) => { - 'use gpu'; - return __tsover_mod(a, b); }, { v: 1, @@ -471,8 +457,6 @@ describe('marked object methods', () => { /** PRIME */ const isPrime = /*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = n => { - 'use gpu'; - if (n <= 1) { return false; } @@ -572,8 +556,6 @@ describe('transforms numeric operations', () => { /** the main function */ // another comment const main = /*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (a, b) => { - 'use gpu'; - let c = __tsover_add(__tsover_add(a, b), 2); c = __tsover_add(c, __tsover_mul(2, b)); countMutable.$ = __tsover_add(countMutable.$, 3); @@ -653,8 +635,6 @@ describe('hoists global function statements marked with "use gpu"', () => { "/** MUL */ // another comment const mul = /*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = function mul(a, b) { - 'use gpu'; - return __tsover_mul(a, b); }, { v: 1, @@ -677,8 +657,6 @@ describe('hoists global function statements marked with "use gpu"', () => { /** ADD */ // another comment const add = /*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = function add(a, b) { - 'use gpu'; - return __tsover_add(a, b); }, { v: 1, @@ -761,8 +739,6 @@ describe('hoists function statements marked with "use gpu", scoped inside anothe /** MUL */ // another comment const mul = /*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = function mul(a, b) { - 'use gpu'; - return __tsover_mul(a, b); }, { v: 1, @@ -785,8 +761,6 @@ describe('hoists function statements marked with "use gpu", scoped inside anothe /** ADD */ // another comment const add = /*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = function add(a, b) { - 'use gpu'; - return __tsover_add(a, b); }, { v: 1, @@ -882,8 +856,6 @@ describe('hoists function statements marked with "use gpu", scoped inside an arr /** MUL */ // another comment const mul = /*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = function mul(a, b) { - 'use gpu'; - return __tsover_mul(a, b); }, { v: 1, @@ -906,8 +878,6 @@ describe('hoists function statements marked with "use gpu", scoped inside an arr /** ADD */ // another comment const add = /*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = function add(a, b) { - 'use gpu'; - return __tsover_add(a, b); }, { v: 1, @@ -1005,8 +975,6 @@ describe('hoists function statements marked with "use gpu", scoped inside an if /** MUL */ // another comment const mul = /*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = function mul(a, b) { - 'use gpu'; - return __tsover_mul(__tsover_mul(a, b), c); }, { v: 1, @@ -1031,8 +999,6 @@ describe('hoists function statements marked with "use gpu", scoped inside an if /** ADD */ // another comment const add = /*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = function add(a, b) { - 'use gpu'; - return __tsover_add(__tsover_add(a, b), c); }, { v: 1, @@ -1137,8 +1103,6 @@ describe('replaces function statements marked with "use gpu" in place when condi /** ADD */ // another comment const add = /*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = function add(a, b) { - 'use gpu'; - return __tsover_add(__tsover_add(a, b), c); }, { v: 1, @@ -1165,8 +1129,6 @@ describe('replaces function statements marked with "use gpu" in place when condi /** MUL */ // another comment const mul = /*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = function mul(a, b) { - 'use gpu'; - return __tsover_mul(__tsover_mul(a, b), c); }, { v: 1, @@ -1257,8 +1219,6 @@ describe('hoists exported marked function statements', () => { expect(babelTransform(code)).toMatchInlineSnapshot(` "/** MUL */ const mul = /*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = function mul(a, b) { - 'use gpu'; - return __tsover_mul(a, b); }, { v: 1, @@ -1280,8 +1240,6 @@ describe('hoists exported marked function statements', () => { }) && $.f)({}); /** ADD */ const add = /*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = function add(a, b) { - 'use gpu'; - return __tsover_add(a, b); }, { v: 1, @@ -1356,8 +1314,6 @@ describe('hoists default exported marked function statement', () => { expect(babelTransform(code)).toMatchInlineSnapshot(` "/** ADD */ const add = /*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = function add(a, b) { - 'use gpu'; - return __tsover_add(a, b); }, { v: 1, @@ -1427,8 +1383,6 @@ describe('export marked arrow function', () => { test('babel', () => { expect(babelTransform(code)).toMatchInlineSnapshot(` "export const add = /*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (a, b) => { - 'use gpu'; - return __tsover_add(a, b); }, { v: 1, @@ -1449,8 +1403,6 @@ describe('export marked arrow function', () => { } }) && $.f)({}); const increment = /*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = n => { - 'use gpu'; - return __tsover_add(n, 1); }, { v: 1, @@ -1469,8 +1421,6 @@ describe('export marked arrow function', () => { }) && $.f)({}); export { increment }; const mul = /*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (a, b) => { - 'use gpu'; - return __tsover_mul(a, b); }, { v: 1, @@ -1543,8 +1493,6 @@ describe('anonymous default export marked function statement', () => { test('babel', () => { expect(babelTransform(code)).toMatchInlineSnapshot(` "export default /*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = function (a, b) { - 'use gpu'; - return __tsover_add(a, b); }, { v: 1, @@ -1596,8 +1544,6 @@ describe('anonymous default export marked arrow function', () => { test('babel', () => { expect(babelTransform(code)).toMatchInlineSnapshot(` "export default /*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (a, b) => { - 'use gpu'; - return __tsover_add(a, b); }, { v: 1, From 5cdb85918b8a1de16124e2ccfe00ad89012331cb Mon Sep 17 00:00:00 2001 From: Aleksander Katan <56294622+aleksanderkatan@users.noreply.github.com> Date: Fri, 10 Jul 2026 11:08:36 +0200 Subject: [PATCH 5/6] Update tests --- .../test/use-gpu-directive.test.ts | 74 +++++++++++++++++-- 1 file changed, 66 insertions(+), 8 deletions(-) diff --git a/packages/unplugin-typegpu/test/use-gpu-directive.test.ts b/packages/unplugin-typegpu/test/use-gpu-directive.test.ts index a5bdccbcb6..758f960bf2 100644 --- a/packages/unplugin-typegpu/test/use-gpu-directive.test.ts +++ b/packages/unplugin-typegpu/test/use-gpu-directive.test.ts @@ -3,18 +3,56 @@ import { babelTransform, rollupTransform } from './transform.ts'; describe('"use gpu" is removed after transform', () => { const code = `\ - const fn = () => { + const fn1 = () => { 'use gpu'; }; - console.log(fn); + const fn2 = () => { + 'use gpu'; + 'worklet'; + }; + + const fn3 = () => { + 'worklet'; + 'use gpu'; + }; + + console.log(fn1, fn2, fn3); `; test('babel', () => { expect(babelTransform(code)).toMatchInlineSnapshot(` - "const fn = /*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = () => {}, { + "const fn1 = /*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = () => {}, { + v: 1, + name: "fn1", + ast: { + params: [], + body: [0, []], + externalNames: [] + }, + externals: () => { + return {}; + } + }) && $.f)({}); + const fn2 = /*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = () => { + 'worklet'; + }, { + v: 1, + name: "fn2", + ast: { + params: [], + body: [0, []], + externalNames: [] + }, + externals: () => { + return {}; + } + }) && $.f)({}); + const fn3 = /*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = () => { + 'worklet'; + }, { v: 1, - name: "fn", + name: "fn3", ast: { params: [], body: [0, []], @@ -24,22 +62,42 @@ describe('"use gpu" is removed after transform', () => { return {}; } }) && $.f)({}); - console.log(fn);" + console.log(fn1, fn2, fn3);" `); }); test('rollup', async () => { expect(await rollupTransform(code)).toMatchInlineSnapshot(` - "const fn = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (() => { + "const fn1 = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (() => { + + }), { + v: 1, + name: "fn1", + ast: {"params":[],"body":[0,[]],"externalNames":[]}, + externals: () => ({}), + }) && $.f)({})); + + const fn2 = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (() => { + + 'worklet'; + }), { + v: 1, + name: "fn2", + ast: {"params":[],"body":[0,[]],"externalNames":[]}, + externals: () => ({}), + }) && $.f)({})); + + const fn3 = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (() => { + 'worklet'; }), { v: 1, - name: "fn", + name: "fn3", ast: {"params":[],"body":[0,[]],"externalNames":[]}, externals: () => ({}), }) && $.f)({})); - console.log(fn); + console.log(fn1, fn2, fn3); " `); }); From 7fcb443760f404fbc2a82f4859e0c0a8cee4d850 Mon Sep 17 00:00:00 2001 From: Aleksander Katan <56294622+aleksanderkatan@users.noreply.github.com> Date: Fri, 10 Jul 2026 11:37:05 +0200 Subject: [PATCH 6/6] Remove `alreadyTransformed`, add new test --- packages/unplugin-typegpu/src/babel.ts | 1 - packages/unplugin-typegpu/src/core/common.ts | 8 ---- packages/unplugin-typegpu/src/core/factory.ts | 2 +- .../test/use-gpu-directive.test.ts | 44 +++++++++++++++++++ 4 files changed, 45 insertions(+), 10 deletions(-) diff --git a/packages/unplugin-typegpu/src/babel.ts b/packages/unplugin-typegpu/src/babel.ts index 588dbf3831..ae10acab24 100644 --- a/packages/unplugin-typegpu/src/babel.ts +++ b/packages/unplugin-typegpu/src/babel.ts @@ -102,7 +102,6 @@ function assignMetadata( if (visibility) { // Hoisting the declaration to the top of the scope visibility.unshiftContainer('body', replacement as t.Statement); - this.alreadyTransformed.add(expression); const id = t.isFunctionDeclaration(path.node) ? path.node.id : undefined; if (id && path.parentPath.isExportNamedDeclaration()) { diff --git a/packages/unplugin-typegpu/src/core/common.ts b/packages/unplugin-typegpu/src/core/common.ts index a403113e51..696e5ed996 100644 --- a/packages/unplugin-typegpu/src/core/common.ts +++ b/packages/unplugin-typegpu/src/core/common.ts @@ -109,8 +109,6 @@ export interface PluginState extends TransformMethods { opts: Options; inUseGpuScope: boolean; - - alreadyTransformed: WeakSet; // TODO: take a look } export interface NodeLocation { @@ -129,7 +127,6 @@ export function initPluginState(state: PluginState, methods: TransformMethods): state.tgpuAliases = new Set(state.opts.forceTgpuAlias ? [state.opts.forceTgpuAlias] : []); state.autoNamingEnabled = state.opts.autoNamingEnabled ?? true; state.inUseGpuScope = false; - state.alreadyTransformed = new WeakSet(); Object.assign(state, methods); } @@ -456,17 +453,12 @@ function functionOnExit( state.inUseGpuScope = false; - if (state.alreadyTransformed.has(node)) { - return; - } - const ast = fnNodeToTranspiledMap.get(path.node); const maybeName = getFunctionName(path); if (!ast) { throw new Error(`No metadata found for function ${maybeName ?? ''}`); } state.assignMetadata(path, maybeName, ast); - state.alreadyTransformed.add(node); path.skip(); } diff --git a/packages/unplugin-typegpu/src/core/factory.ts b/packages/unplugin-typegpu/src/core/factory.ts index 4c438a6c1c..27e1d11401 100644 --- a/packages/unplugin-typegpu/src/core/factory.ts +++ b/packages/unplugin-typegpu/src/core/factory.ts @@ -131,7 +131,7 @@ function removeUseGpuDirective(this: UnpluginPluginState, path: NodePath { }); }); +describe('double transformation', () => { + const code = `\ + const fn = () => { + 'use gpu'; + }; + + console.log(fn); + `; + + test('babel', () => { + expect(babelTransform(babelTransform(code) ?? '')).toMatchInlineSnapshot(` + "const fn = /*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = () => {}, { + v: 1, + name: "fn", + ast: { + params: [], + body: [0, []], + externalNames: [] + }, + externals: () => { + return {}; + } + }) && $.f)({}); + console.log(fn);" + `); + }); + + test('rollup', async () => { + expect(await rollupTransform(await rollupTransform(code))).toMatchInlineSnapshot(` + "const fn = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (() => { + + }), { + v: 1, + name: "fn", + ast: {"params":[],"body":[0,[]],"externalNames":[]}, + externals: () => ({}), + }) && $.f)({})); + + console.log(fn); + " + `); + }); +}); + describe('"use gpu" marked arrow function, assigned to a const', () => { const code = `\ /** ADD */