From 0b5f00585d296922e7fe21035e0b9e10b060688c Mon Sep 17 00:00:00 2001 From: Szymon Szulc Date: Fri, 14 Aug 2026 13:21:19 +0200 Subject: [PATCH 01/18] ast parser --- .../src/core/embeddedMetadata.ts | 261 +++++++++++++++--- .../test/embedded-metadata.test.ts | 11 +- 2 files changed, 222 insertions(+), 50 deletions(-) diff --git a/packages/unplugin-typegpu/src/core/embeddedMetadata.ts b/packages/unplugin-typegpu/src/core/embeddedMetadata.ts index 8da8645d7a..251e7ef663 100644 --- a/packages/unplugin-typegpu/src/core/embeddedMetadata.ts +++ b/packages/unplugin-typegpu/src/core/embeddedMetadata.ts @@ -1,10 +1,12 @@ import type { NodePath } from '@babel/traverse'; import * as t from '@babel/types'; import { type MetadatableFunction } from './common.ts'; +import type { Block, FuncParameter } from 'tinyest'; export interface EmbeddedTypegpuMetadata { v: number; name: string | undefined; + ast?: { params: FuncParameter[]; body: Block }; // TODO: parse AST and externals } @@ -15,7 +17,11 @@ const embeddedTypegpuMetadataCache = new WeakMap< /** * Returns the node after unwrapping any parenthesized expressions. + * + * Example: `(a + b)` returns `a + b`. */ +function unwrapParentheses(node: t.Expression): t.Expression; +function unwrapParentheses(node: t.Node): t.Node; function unwrapParentheses(node: t.Node): t.Node { let current = node; @@ -26,41 +32,56 @@ function unwrapParentheses(node: t.Node): t.Node { return current; } +/** + * Returns the first parent path that is not a parenthesized expression. + */ +function parentPathSkippingParentheses(path: NodePath): NodePath | null { + let parentPath = path.parentPath; + + while (parentPath?.isParenthesizedExpression()) { + parentPath = parentPath.parentPath; + } + + return parentPath; +} + /** * Returns the property name of a member expression. */ function memberPropertyName(node: t.MemberExpression): string | undefined { - if (!node.computed && t.isIdentifier(node.property)) { - return node.property.name; + const property = unwrapParentheses(node.property); // foo[('bar')] + + if (!node.computed && t.isIdentifier(property)) { + return property.name; } - if (node.computed && t.isStringLiteral(node.property)) { - return node.property.value; + if (node.computed && t.isStringLiteral(property)) { + return property.value; } return undefined; } /** - * Returns whether the node is a global TypeGPU metadata expression (`globalThis.__TYPEGPU_META__`). + * Returns whether the node is a global TypeGPU metadata expression `globalThis.__TYPEGPU_META__`. */ function isGlobalTypegpuMetadata(node: t.Node): boolean { - const unwrapped = unwrapParentheses(node); + const expression = unwrapParentheses(node); return ( - t.isMemberExpression(unwrapped) && - t.isIdentifier(unwrapped.object, { name: 'globalThis' }) && - memberPropertyName(unwrapped) === '__TYPEGPU_META__' + t.isMemberExpression(expression) && + t.isIdentifier(unwrapParentheses(expression.object), { name: 'globalThis' }) && + memberPropertyName(expression) === '__TYPEGPU_META__' ); } /** - * Returns whether the node is a TypeGPU metadata set call (`(globalThis.__TYPEGPU_META__ ??= new WeakMap()).set`) + * Returns whether the node is a TypeGPU metadata set call `(globalThis.__TYPEGPU_META__ ??= new WeakMap()).set(...)`. */ function isTypegpuMetadataSetCall(node: t.CallExpression): boolean { const callee = unwrapParentheses(node.callee); - if (!t.isMemberExpression(callee) || memberPropertyName(callee) !== 'set') { + if (!(t.isMemberExpression(callee) && memberPropertyName(callee) === 'set')) { return false; } @@ -73,31 +94,160 @@ function isTypegpuMetadataSetCall(node: t.CallExpression): boolean { } /** - * Returns the node of the property with the given name in the object, if it exists. + * Returns the value node of the property with the given name in the object, if it exists. */ -function objectPropertyValue(object: t.ObjectExpression, expectedName: string): t.Node | undefined { +function objectPropertyValue( + object: t.ObjectExpression, + expectedName: string, +): t.Expression | undefined { for (const property of object.properties) { - if (!t.isObjectProperty(property) || property.computed) { + if (!t.isObjectProperty(property)) { continue; } - const name = t.isIdentifier(property.key) - ? property.key.name - : t.isStringLiteral(property.key) - ? property.key.value - : undefined; + const key = unwrapParentheses(property.key); // { foo: (2) } + const name = + !property.computed && t.isIdentifier(key) + ? key.name + : t.isStringLiteral(key) + ? key.value + : undefined; + + if (name === expectedName && t.isExpression(property.value)) { + return unwrapParentheses(property.value); + } + } + + return undefined; +} + +type EncodedTinyestValue = + | number + | string + | boolean + | null + | EncodedTinyestValue[] + | { [key: string]: EncodedTinyestValue }; + +function parseTinyestValue(node: t.Node): EncodedTinyestValue | undefined { + if (t.isNumericLiteral(node) || t.isStringLiteral(node) || t.isBooleanLiteral(node)) { + return node.value; + } + + if (t.isNullLiteral(node)) { + return null; + } + + if (t.isArrayExpression(node)) { + const result: EncodedTinyestValue[] = []; + + for (const element of node.elements) { + if (element === null || !t.isExpression(element)) { + return undefined; + } + + const parsed = parseTinyestValue(element); + if (parsed === undefined) { + return undefined; + } + + result.push(parsed); + } + + return result; + } + + if (t.isObjectExpression(node)) { + const result: Record = {}; + + for (const property of node.properties) { + if (!t.isObjectProperty(property) || property.computed || !t.isExpression(property.value)) { + return undefined; + } - if (name === expectedName) { - return property.value; + const keyNode = property.key; + const key = t.isIdentifier(keyNode) + ? keyNode.name + : t.isStringLiteral(keyNode) || t.isNumericLiteral(keyNode) + ? String(keyNode.value) + : undefined; + + if (key === undefined) { + return undefined; + } + + const value = parseTinyestValue(property.value); + if (value === undefined) { + return undefined; + } + + result[key] = value; } + + return result; } return undefined; } +/** + * Given AST of a function's parameters, returns the parsed parameters. + */ +function parseFuncParameters(paramsNode: t.ArrayExpression): FuncParameter[] | undefined { + return []; +} + +/** + * Given AST of a function's body in tinyest encoding, returns the parsed body in tinyest encoding. + */ +function parseBody(bodyNode: t.ArrayExpression): Block { + const parsed = parseTinyestValue(bodyNode); + + if ( + !( + parsed !== undefined && + Array.isArray(parsed) && + parsed.length === 2 && + parsed[0] === 0 && + Array.isArray(parsed[1]) + ) + ) { + throw new Error( + 'unplugin-typegpu: While parsing metadata encountered an invalid ast.body node.', + ); + } + + return parsed as unknown as Block; +} + +function parseAstNode(astNode: t.ObjectExpression): EmbeddedTypegpuMetadata['ast'] | undefined { + const paramsNode = objectPropertyValue(astNode, 'params'); + const bodyNode = objectPropertyValue(astNode, 'body'); + + if ( + !paramsNode || + !bodyNode || + !t.isArrayExpression(paramsNode) || + !t.isArrayExpression(bodyNode) + ) { + return undefined; + } + + const params = parseFuncParameters(paramsNode); + const body = parseBody(bodyNode); + + if (!params) { + return undefined; + } + + return { + params, + body, + }; +} + /** * Returns metadata embedded by unplugin-typegpu for this exact function. - * Externals are required as part of the emitted shape, but are intentionally not evaluated. * * Consider: * ```ts @@ -113,6 +263,17 @@ function objectPropertyValue(object: t.ObjectExpression, expectedName: string): * externals: {} * }) && $.f)({}); * ``` + * + * Function will return: + * { + * v: 2, + * name: "f", + * ast: { + * params: [], + * body: [0, []] + * }, + * externals: {} + * } */ export function getEmbeddedTypegpuMetadata( path: NodePath, @@ -122,34 +283,25 @@ export function getEmbeddedTypegpuMetadata( return cached; } - // we start with () => { 'use gpu'; ... } - let expressionPath: NodePath = path; - - // rollup puts it in parentheses - while (expressionPath.parentPath?.isParenthesizedExpression()) { - expressionPath = expressionPath.parentPath; - } - // we check for f.$ = () => { 'use gpu'; ... } - const assignmentPath = expressionPath.parentPath; + const assignmentPath = parentPathSkippingParentheses(path); if (!assignmentPath?.isAssignmentExpression({ operator: '=' })) { return undefined; } - // get rid of parentheses - let assignedPath: NodePath = assignmentPath; - while (assignedPath.parentPath?.isParenthesizedExpression()) { - assignedPath = assignedPath.parentPath; - } - // we check for `(globalThis.__TYPEGPU_META__ ??= new WeakMap()).set()` - const callPath = assignedPath.parentPath; + const callPath = parentPathSkippingParentheses(assignmentPath); if (!(callPath?.isCallExpression() && isTypegpuMetadataSetCall(callPath.node))) { return undefined; } // we check for the metadata object - const metadataNode = callPath.node.arguments[1]; + const metadataArgument = callPath.node.arguments[1]; + if (metadataArgument === undefined) { + return undefined; + } + + const metadataNode = unwrapParentheses(metadataArgument); if (!t.isObjectExpression(metadataNode)) { return undefined; } @@ -157,23 +309,40 @@ export function getEmbeddedTypegpuMetadata( // get the metadata properties const versionNode = objectPropertyValue(metadataNode, 'v'); const nameNode = objectPropertyValue(metadataNode, 'name'); + + if (!t.isNumericLiteral(versionNode) || nameNode === undefined) { + return undefined; + } + + const name = t.isStringLiteral(nameNode) ? nameNode.value : undefined; + + if (versionNode.value == 1) { + return { + v: versionNode.value, + name, + }; + } + const astNode = objectPropertyValue(metadataNode, 'ast'); const externalsNode = objectPropertyValue(metadataNode, 'externals'); - if ( - !t.isNumericLiteral(versionNode) || - nameNode === undefined || - astNode === undefined || - externalsNode === undefined - ) { + if (astNode === undefined || externalsNode === undefined) { return undefined; } - const name = t.isStringLiteral(nameNode) ? nameNode.value : undefined; + if (!t.isObjectExpression(astNode)) { + return undefined; + } + + const ast = parseAstNode(astNode); + if (ast === undefined) { + return undefined; + } const embeddedTypegpuMetadata = { v: versionNode.value, name, + ast, }; embeddedTypegpuMetadataCache.set(path, embeddedTypegpuMetadata); diff --git a/packages/unplugin-typegpu/test/embedded-metadata.test.ts b/packages/unplugin-typegpu/test/embedded-metadata.test.ts index bdec719bff..1b3d52bb01 100644 --- a/packages/unplugin-typegpu/test/embedded-metadata.test.ts +++ b/packages/unplugin-typegpu/test/embedded-metadata.test.ts @@ -45,7 +45,10 @@ function createRollupMetadataCollector(metadata: EmbeddedTypegpuMetadata[]): Plu return { name: 'collect-typegpu-metadata', transform(code) { - const ast = parser.parse(code, { sourceType: 'module' }); + const ast = parser.parse(code, { + sourceType: 'module', + createParenthesizedExpressions: true, + }); traverse(ast, { ArrowFunctionExpression(path) { @@ -86,10 +89,10 @@ describe('collects embedded TypeGPU metadata', () => { test('babel', () => { const metadata: EmbeddedTypegpuMetadata[] = []; - babelTransform(code, undefined, [createBabelMetadataCollector(metadata)]); + babelTransform(code, {}, [createBabelMetadataCollector(metadata)]); expect(JSON.stringify(metadata)).toMatchInlineSnapshot( - `"[{"v":2,"name":"fn1"},{"v":2,"name":"fn2"},{"v":2,"name":"fn3"}]"`, + `"[{"v":2,"name":"fn1","ast":{"params":[],"body":[0,[]]}},{"v":2,"name":"fn2","ast":{"params":[],"body":[0,[]]}},{"v":2,"name":"fn3","ast":{"params":[],"body":[0,[]]}}]"`, ); }); @@ -99,7 +102,7 @@ describe('collects embedded TypeGPU metadata', () => { await rollupTransform(code, undefined, [createRollupMetadataCollector(metadata)]); expect(JSON.stringify(metadata)).toMatchInlineSnapshot( - `"[{"v":2,"name":"fn1"},{"v":2,"name":"fn2"},{"v":2,"name":"fn3"}]"`, + `"[{"v":2,"name":"fn1","ast":{"params":[],"body":[0,[]]}},{"v":2,"name":"fn2","ast":{"params":[],"body":[0,[]]}},{"v":2,"name":"fn3","ast":{"params":[],"body":[0,[]]}}]"`, ); }); }); From b8432e68fcca6760c0b650695d275801f96fe180 Mon Sep 17 00:00:00 2001 From: Szymon Szulc Date: Mon, 7 Sep 2026 17:40:28 +0200 Subject: [PATCH 02/18] minor cleanup --- .../src/core/embeddedMetadata.ts | 29 ++++++++++--------- 1 file changed, 16 insertions(+), 13 deletions(-) diff --git a/packages/unplugin-typegpu/src/core/embeddedMetadata.ts b/packages/unplugin-typegpu/src/core/embeddedMetadata.ts index 251e7ef663..b2631b6657 100644 --- a/packages/unplugin-typegpu/src/core/embeddedMetadata.ts +++ b/packages/unplugin-typegpu/src/core/embeddedMetadata.ts @@ -85,11 +85,10 @@ function isTypegpuMetadataSetCall(node: t.CallExpression): boolean { return false; } - const receiver = unwrapParentheses(callee.object); + const inner = unwrapParentheses(callee.object); return ( - t.isAssignmentExpression(receiver, { operator: '??=' }) && - isGlobalTypegpuMetadata(receiver.left) + t.isAssignmentExpression(inner, { operator: '??=' }) && isGlobalTypegpuMetadata(inner.left) ); } @@ -105,7 +104,7 @@ function objectPropertyValue( continue; } - const key = unwrapParentheses(property.key); // { foo: (2) } + const key = property.key; const name = !property.computed && t.isIdentifier(key) ? key.name @@ -190,13 +189,6 @@ function parseTinyestValue(node: t.Node): EncodedTinyestValue | undefined { return undefined; } -/** - * Given AST of a function's parameters, returns the parsed parameters. - */ -function parseFuncParameters(paramsNode: t.ArrayExpression): FuncParameter[] | undefined { - return []; -} - /** * Given AST of a function's body in tinyest encoding, returns the parsed body in tinyest encoding. */ @@ -220,6 +212,13 @@ function parseBody(bodyNode: t.ArrayExpression): Block { return parsed as unknown as Block; } +/** + * Given AST of a function's parameters, returns the parsed parameters. + */ +function parseFuncParameters(paramsNode: t.ArrayExpression): FuncParameter[] | undefined { + return []; +} + function parseAstNode(astNode: t.ObjectExpression): EmbeddedTypegpuMetadata['ast'] | undefined { const paramsNode = objectPropertyValue(astNode, 'params'); const bodyNode = objectPropertyValue(astNode, 'body'); @@ -236,7 +235,7 @@ function parseAstNode(astNode: t.ObjectExpression): EmbeddedTypegpuMetadata['ast const params = parseFuncParameters(paramsNode); const body = parseBody(bodyNode); - if (!params) { + if (!params || !body) { return undefined; } @@ -247,7 +246,9 @@ function parseAstNode(astNode: t.ObjectExpression): EmbeddedTypegpuMetadata['ast } /** - * Returns metadata embedded by unplugin-typegpu for this exact function. + * Returns metadata embedded by unplugin-typegpu for this exact function: + * + * @note metadata v1 support is limited. Only the version and name are parsed. * * Consider: * ```ts @@ -274,6 +275,7 @@ function parseAstNode(astNode: t.ObjectExpression): EmbeddedTypegpuMetadata['ast * }, * externals: {} * } + * */ export function getEmbeddedTypegpuMetadata( path: NodePath, @@ -316,6 +318,7 @@ export function getEmbeddedTypegpuMetadata( const name = t.isStringLiteral(nameNode) ? nameNode.value : undefined; + // metadata v1 support is limited if (versionNode.value == 1) { return { v: versionNode.value, From ccc9ce0891a41cd51f0dad92f113d42bf6a63cbf Mon Sep 17 00:00:00 2001 From: Szymon Szulc Date: Mon, 14 Sep 2026 18:27:54 +0200 Subject: [PATCH 03/18] func body and params tests --- .../src/core/embeddedMetadata.ts | 19 +- .../test/embedded-metadata.test.ts | 487 +++++++++++++++++- 2 files changed, 471 insertions(+), 35 deletions(-) diff --git a/packages/unplugin-typegpu/src/core/embeddedMetadata.ts b/packages/unplugin-typegpu/src/core/embeddedMetadata.ts index b2631b6657..29a9a0013d 100644 --- a/packages/unplugin-typegpu/src/core/embeddedMetadata.ts +++ b/packages/unplugin-typegpu/src/core/embeddedMetadata.ts @@ -1,13 +1,16 @@ import type { NodePath } from '@babel/traverse'; import * as t from '@babel/types'; -import { type MetadatableFunction } from './common.ts'; -import type { Block, FuncParameter } from 'tinyest'; +import type { MetadatableFunction } from './common.ts'; +import type { TranspilationResult } from 'tinyest-for-wgsl'; export interface EmbeddedTypegpuMetadata { v: number; name: string | undefined; - ast?: { params: FuncParameter[]; body: Block }; - // TODO: parse AST and externals + ast?: { + params: TranspilationResult['params']; + body: TranspilationResult['body']; + }; + externals?: TranspilationResult['externalNames']; } const embeddedTypegpuMetadataCache = new WeakMap< @@ -192,7 +195,7 @@ function parseTinyestValue(node: t.Node): EncodedTinyestValue | undefined { /** * Given AST of a function's body in tinyest encoding, returns the parsed body in tinyest encoding. */ -function parseBody(bodyNode: t.ArrayExpression): Block { +function parseBody(bodyNode: t.ArrayExpression): TranspilationResult['body'] { const parsed = parseTinyestValue(bodyNode); if ( @@ -209,13 +212,15 @@ function parseBody(bodyNode: t.ArrayExpression): Block { ); } - return parsed as unknown as Block; + return parsed as unknown as TranspilationResult['body']; } /** * Given AST of a function's parameters, returns the parsed parameters. */ -function parseFuncParameters(paramsNode: t.ArrayExpression): FuncParameter[] | undefined { +function parseFuncParameters( + paramsNode: t.ArrayExpression, +): TranspilationResult['params'] | undefined { return []; } diff --git a/packages/unplugin-typegpu/test/embedded-metadata.test.ts b/packages/unplugin-typegpu/test/embedded-metadata.test.ts index 1b3d52bb01..84e6e4503a 100644 --- a/packages/unplugin-typegpu/test/embedded-metadata.test.ts +++ b/packages/unplugin-typegpu/test/embedded-metadata.test.ts @@ -2,8 +2,9 @@ import * as parser from '@babel/parser'; import _traverse, { type NodePath } from '@babel/traverse'; import type { Plugin } from 'rollup'; import { describe, expect, test } from 'vitest'; +import { transpileBabelFn } from 'tinyest-for-wgsl'; import { type BabelTestPlugin, babelTransform, rollupTransform } from './transform.ts'; -import { type MetadatableFunction } from '../src/core/common.ts'; +import type { MetadatableFunction } from '../src/core/common.ts'; import { type EmbeddedTypegpuMetadata, getEmbeddedTypegpuMetadata, @@ -47,7 +48,7 @@ function createRollupMetadataCollector(metadata: EmbeddedTypegpuMetadata[]): Plu transform(code) { const ast = parser.parse(code, { sourceType: 'module', - createParenthesizedExpressions: true, + createParenthesizedExpressions: true, // to catch cases where `unwrapParentheses` or `parentPathSkippingParentheses` are needed }); traverse(ast, { @@ -67,42 +68,472 @@ function createRollupMetadataCollector(metadata: EmbeddedTypegpuMetadata[]): Plu }; } -describe('collects embedded TypeGPU metadata', () => { - const code = `\ - const fn1 = () => { - 'use gpu'; - }; - - const fn2 = () => { - 'use gpu'; - 'worklet'; - }; +/** + * Traverses the program AST to extract the expected ASTs from functions containing 'use gpu' directive. + * + * @note For simplicity, shelled functions are omitted. + */ +function extractExpectedAstsFromSource(code: string) { + const ast = parser.parse(code, { + sourceType: 'module', + plugins: ['typescript'], + }); - const fn3 = () => { - 'worklet'; - 'use gpu'; - }; + const expected: NonNullable[] = []; - console.log(fn1, fn2, fn3); - `; + function collect(path: NodePath) { + const body = path.node.body; - test('babel', () => { - const metadata: EmbeddedTypegpuMetadata[] = []; + if ( + body.type !== 'BlockStatement' || + !body.directives.some((directive) => directive.value.value === 'use gpu') + ) { + return; + } - babelTransform(code, {}, [createBabelMetadataCollector(metadata)]); + const { params, body: transpiledBody } = transpileBabelFn(path.node); + expected.push({ params, body: transpiledBody }); + } - expect(JSON.stringify(metadata)).toMatchInlineSnapshot( - `"[{"v":2,"name":"fn1","ast":{"params":[],"body":[0,[]]}},{"v":2,"name":"fn2","ast":{"params":[],"body":[0,[]]}},{"v":2,"name":"fn3","ast":{"params":[],"body":[0,[]]}}]"`, - ); + traverse(ast, { + ArrowFunctionExpression: collect, + FunctionExpression: collect, + FunctionDeclaration: collect, }); - test('rollup', async () => { + return expected; +} + +function dualTest( + code: string, + check: ( + metadata: EmbeddedTypegpuMetadata[], + expectedAsts: NonNullable[], + ) => void, +) { + test('[BABEL]', () => { + const expectedAsts = extractExpectedAstsFromSource(code); + const metadata: EmbeddedTypegpuMetadata[] = []; + babelTransform(code, {}, [createBabelMetadataCollector(metadata)]); + check(metadata, expectedAsts); + }); + + test('[ROLLUP]', async () => { + const expectedAsts = extractExpectedAstsFromSource(code); + const metadata: EmbeddedTypegpuMetadata[] = []; await rollupTransform(code, undefined, [createRollupMetadataCollector(metadata)]); + check(metadata, expectedAsts); + }); +} + +describe('getEmbeddedTypegpuMetadata', () => { + describe('omits ast and externals fields for metadata v1', () => { + const code = `\ + const fn = ($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = () => { + 'use gpu'; + return 1; + }, { + v: 1, + name: 'fn', + }) && $.f)({}); + + console.log(fn); + `; + + dualTest(code, (metadata) => { + expect(metadata).toStrictEqual([{ v: 1, name: 'fn' }]); + }); + }); + + describe.skip('parsers function parameters', () => { + const code = `\ + const noParams = () => { + 'use gpu'; + }; + + const identifierParams = (a, b) => { + 'use gpu'; + return a + b; + }; + + const destructuredParams = ({ pos, a: b }) => { + 'use gpu'; + return pos + b; + }; + + const mixedParams = (y, { pos, a: b }, { c, d }) => { + 'use gpu'; + return y + pos + b + c + d; + }; + + console.log(noParams, identifierParams, destructuredParams, mixedParams); + `; + + dualTest(code, (metadata, expectedAsts) => { + expect(metadata.map((m) => m.ast?.params)).toStrictEqual( + expectedAsts.map((ast) => ast.params), + ); + }); + }); + + describe('parses literals', () => { + const code = `\ + const fn = () => { + 'use gpu'; + const numericLiteral = 1; + const bigIntLiteral = 2n; + const stringLiteral = 'text'; + const nullLiteral = null; + const boolLiteral = true; + }; + + console.log(fn); + `; + + dualTest(code, (metadata, expectedAsts) => { + expect(metadata.map((m) => m.ast?.body)).toStrictEqual(expectedAsts.map((ast) => ast.body)); + }); + }); + + describe('parses blocks', () => { + const code = `\ + const fn = (a) => { + 'use gpu'; + let value = a; + { + value += 1; + { + value += 2; + } + } + }; + + console.log(fn); + `; + + dualTest(code, (metadata, expectedAsts) => { + expect(metadata.map((m) => m.ast?.body)).toStrictEqual(expectedAsts.map((ast) => ast.body)); + }); + }); + + describe('parses binary expressions', () => { + const code = `\ + const fn = (a, b) => { + 'use gpu'; + const plus = a + b; + const xor = a ^ b; + const eq = a === b; + }; + + console.log(fn); + `; + + dualTest(code, (metadata, expectedAsts) => { + expect(metadata.map((m) => m.ast?.body)).toStrictEqual(expectedAsts.map((ast) => ast.body)); + }); + }); + + describe('parses assignment expressions', () => { + const code = `\ + const fn = (a) => { + 'use gpu'; + let v = 0; + v = a; + v += a; + }; + + console.log(fn); + `; + + dualTest(code, (metadata, expectedAsts) => { + expect(metadata.map((m) => m.ast?.body)).toStrictEqual(expectedAsts.map((ast) => ast.body)); + }); + }); + + describe('parses logical expressions', () => { + const code = `\ + const fn = (a, b) => { + 'use gpu'; + const and = a && b; + const or = a || b; + }; + + console.log(fn); + `; + + dualTest(code, (metadata, expectedAsts) => { + expect(metadata.map((m) => m.ast?.body)).toStrictEqual(expectedAsts.map((ast) => ast.body)); + }); + }); + + describe('parses unary expressions', () => { + const code = `\ + const fn = (a) => { + 'use gpu'; + const negation = !a; + }; + + console.log(fn); + `; + + dualTest(code, (metadata, expectedAsts) => { + expect(metadata.map((m) => m.ast?.body)).toStrictEqual(expectedAsts.map((ast) => ast.body)); + }); + }); + + describe('parses call expressions', () => { + const code = `\ + const fn = (f, a) => { + 'use gpu'; + const noArgs = f(); + const withArgs = f(a, 1); + const nested = f(f(), a) + }; + + console.log(fn); + `; + + dualTest(code, (metadata, expectedAsts) => { + expect(metadata.map((m) => m.ast?.body)).toStrictEqual(expectedAsts.map((ast) => ast.body)); + }); + }); + + describe('parses member and index access', () => { + const code = `\ + const fn = (obj, arr, i) => { + 'use gpu'; + const member = obj.prop; + const index = arr[i]; + const mixed = obj.a[i].b; + }; + + console.log(fn); + `; + + dualTest(code, (metadata, expectedAsts) => { + expect(metadata.map((m) => m.ast?.body)).toStrictEqual(expectedAsts.map((ast) => ast.body)); + }); + }); + + describe('parses return statements', () => { + const code = `\ + const bareReturn = () => { + 'use gpu'; + return; + }; + + const valueReturn = (a) => { + 'use gpu'; + return a + 1; + }; + + console.log(bareReturn, valueReturn); + `; + + dualTest(code, (metadata, expectedAsts) => { + expect(metadata.map((m) => m.ast?.body)).toStrictEqual(expectedAsts.map((ast) => ast.body)); + }); + }); + + describe('parses if statements', () => { + const code = `\ + const bare = (a) => { + 'use gpu'; + if (a > 0) { + return 1; + } + }; + + const alternative = (a) => { + 'use gpu'; + if (a > 0) { + return 1; + } else { + return 2; + } + }; + + const elseIf = (a) => { + 'use gpu'; + if (a > 0) { + return 1; + } else if (a < 0) { + return 2; + } else { + return 3; + } + }; + + const withoutBlocks = (a) => { + 'use gpu'; + if (a > 0) + return 1; + else + return 2; + }; + + console.log(bare, alternative, elseIf, withoutBlocks); + `; + + dualTest(code, (metadata, expectedAsts) => { + expect(metadata.map((m) => m.ast?.body)).toStrictEqual(expectedAsts.map((ast) => ast.body)); + }); + }); + + describe('parses variable declarations', () => { + const code = `\ + const fn = (a) => { + 'use gpu'; + let x = a; + const y = a + 1; + }; + + console.log(fn); + `; + + dualTest(code, (metadata, expectedAsts) => { + expect(metadata.map((m) => m.ast?.body)).toStrictEqual(expectedAsts.map((ast) => ast.body)); + }); + }); + + describe('parses loops', () => { + const code = `\ + const forLoop = (n) => { + 'use gpu'; + let total = 0; + for (let i = 0; i < n; i++) { + total += i; + } + }; + + const whileLoop = (n) => { + 'use gpu'; + let total = 0; + while (total < n) { + total++; + } + }; + + const forOfLoop = (items) => { + 'use gpu'; + let total = 0; + for (const item of items) { + total += item; + } + }; + + console.log(forLoop, whileLoop, forOfLoop); + `; + + dualTest(code, (metadata, expectedAsts) => { + expect(metadata.map((m) => m.ast?.body)).toStrictEqual(expectedAsts.map((ast) => ast.body)); + }); + }); + + describe('parses break and continue', () => { + const code = `\ + const fn = (n) => { + 'use gpu'; + let total = 0; + let i = 1; + while (true) { + if (i % n === 0) { + break; + } + if (i % n === 1) { + continue; + } + total++ + i++; + } + }; + + console.log(fn); + `; + + dualTest(code, (metadata, expectedAsts) => { + expect(metadata.map((m) => m.ast?.body)).toStrictEqual(expectedAsts.map((ast) => ast.body)); + }); + }); + + describe('parses array expressions', () => { + const code = `\ + const fn = (a) => { + 'use gpu'; + const empty = []; + const literals = [1, 2, 3]; + }; + + console.log(fn); + `; + + dualTest(code, (metadata, expectedAsts) => { + expect(metadata.map((m) => m.ast?.body)).toStrictEqual(expectedAsts.map((ast) => ast.body)); + }); + }); + + describe('parses post-update expressions', () => { + const code = `\ + const fn = () => { + 'use gpu'; + let i = 0; + i++; + i--; + }; + + console.log(fn); + `; + + dualTest(code, (metadata, expectedAsts) => { + expect(metadata.map((m) => m.ast?.body)).toStrictEqual(expectedAsts.map((ast) => ast.body)); + }); + }); + + describe('parses object expressions', () => { + const code = `\ + const fn = (a) => { + 'use gpu'; + const empty = {}; + const plain = { x: 1, y: a }; + }; + + console.log(fn); + `; + + dualTest(code, (metadata, expectedAsts) => { + expect(metadata.map((m) => m.ast?.body)).toStrictEqual(expectedAsts.map((ast) => ast.body)); + }); + }); + + describe('parses object expressions with computed properties', () => { + const code = `\ + const fn = (id, f) => { + 'use gpu'; + const computed = { [id]: 1, [f()]: 2 }; + }; + + console.log(fn); + `; + + dualTest(code, (metadata, expectedAsts) => { + expect(metadata.map((m) => m.ast?.body)).toStrictEqual(expectedAsts.map((ast) => ast.body)); + }); + }); + + describe('parses conditional expressions', () => { + const code = `\ + const fn = (a, b) => { + 'use gpu'; + return a ? b : 1; + }; + + console.log(fn); + `; - expect(JSON.stringify(metadata)).toMatchInlineSnapshot( - `"[{"v":2,"name":"fn1","ast":{"params":[],"body":[0,[]]}},{"v":2,"name":"fn2","ast":{"params":[],"body":[0,[]]}},{"v":2,"name":"fn3","ast":{"params":[],"body":[0,[]]}}]"`, - ); + dualTest(code, (metadata, expectedAsts) => { + expect(metadata.map((m) => m.ast?.body)).toStrictEqual(expectedAsts.map((ast) => ast.body)); + }); }); }); From 720005525cee311dc17101a7eee50587939cddff Mon Sep 17 00:00:00 2001 From: Szymon Szulc Date: Tue, 15 Sep 2026 13:48:15 +0200 Subject: [PATCH 04/18] babel evaluate instead of parsing --- .../src/core/embeddedMetadata.ts | 216 +++++------------- .../test/embedded-metadata.test.ts | 2 +- 2 files changed, 63 insertions(+), 155 deletions(-) diff --git a/packages/unplugin-typegpu/src/core/embeddedMetadata.ts b/packages/unplugin-typegpu/src/core/embeddedMetadata.ts index 29a9a0013d..3039384a33 100644 --- a/packages/unplugin-typegpu/src/core/embeddedMetadata.ts +++ b/packages/unplugin-typegpu/src/core/embeddedMetadata.ts @@ -1,7 +1,7 @@ import type { NodePath } from '@babel/traverse'; import * as t from '@babel/types'; -import type { MetadatableFunction } from './common.ts'; import type { TranspilationResult } from 'tinyest-for-wgsl'; +import type { MetadatableFunction } from './common.ts'; export interface EmbeddedTypegpuMetadata { v: number; @@ -35,6 +35,19 @@ function unwrapParentheses(node: t.Node): t.Node { return current; } +/** + * Works like {@link unwrapParentheses}, but for a `NodePath`. + */ +function unwrapParenthesesPath(path: NodePath): NodePath { + let current = path; + + while (current.isParenthesizedExpression()) { + current = current.get('expression'); + } + + return current; +} + /** * Returns the first parent path that is not a parenthesized expression. */ @@ -96,158 +109,45 @@ function isTypegpuMetadataSetCall(node: t.CallExpression): boolean { } /** - * Returns the value node of the property with the given name in the object, if it exists. + * Returns the value path of the first matching property. If there is no match, returns `undefined`. */ -function objectPropertyValue( - object: t.ObjectExpression, - expectedName: string, -): t.Expression | undefined { - for (const property of object.properties) { - if (!t.isObjectProperty(property)) { +function objectPropertyPath( + objectPath: NodePath, + name: string, +): NodePath | undefined { + for (const property of objectPath.get('properties')) { + if (!property.isObjectProperty()) { continue; } - const key = property.key; - const name = - !property.computed && t.isIdentifier(key) - ? key.name - : t.isStringLiteral(key) - ? key.value - : undefined; + const { key, computed } = property.node; - if (name === expectedName && t.isExpression(property.value)) { - return unwrapParentheses(property.value); - } - } + const matches = + (!computed && t.isIdentifier(key, { name })) || t.isStringLiteral(key, { value: name }); - return undefined; -} - -type EncodedTinyestValue = - | number - | string - | boolean - | null - | EncodedTinyestValue[] - | { [key: string]: EncodedTinyestValue }; - -function parseTinyestValue(node: t.Node): EncodedTinyestValue | undefined { - if (t.isNumericLiteral(node) || t.isStringLiteral(node) || t.isBooleanLiteral(node)) { - return node.value; - } - - if (t.isNullLiteral(node)) { - return null; - } - - if (t.isArrayExpression(node)) { - const result: EncodedTinyestValue[] = []; - - for (const element of node.elements) { - if (element === null || !t.isExpression(element)) { - return undefined; - } - - const parsed = parseTinyestValue(element); - if (parsed === undefined) { - return undefined; - } - - result.push(parsed); + if (!matches) { + continue; } - return result; - } - - if (t.isObjectExpression(node)) { - const result: Record = {}; - - for (const property of node.properties) { - if (!t.isObjectProperty(property) || property.computed || !t.isExpression(property.value)) { - return undefined; - } - - const keyNode = property.key; - const key = t.isIdentifier(keyNode) - ? keyNode.name - : t.isStringLiteral(keyNode) || t.isNumericLiteral(keyNode) - ? String(keyNode.value) - : undefined; - - if (key === undefined) { - return undefined; - } - - const value = parseTinyestValue(property.value); - if (value === undefined) { - return undefined; - } - - result[key] = value; + const value = property.get('value'); + if (value.isExpression()) { + return value; } - - return result; } return undefined; } - /** - * Given AST of a function's body in tinyest encoding, returns the parsed body in tinyest encoding. + * Returns the parsed AST */ -function parseBody(bodyNode: t.ArrayExpression): TranspilationResult['body'] { - const parsed = parseTinyestValue(bodyNode); - - if ( - !( - parsed !== undefined && - Array.isArray(parsed) && - parsed.length === 2 && - parsed[0] === 0 && - Array.isArray(parsed[1]) - ) - ) { - throw new Error( - 'unplugin-typegpu: While parsing metadata encountered an invalid ast.body node.', - ); - } - - return parsed as unknown as TranspilationResult['body']; -} - -/** - * Given AST of a function's parameters, returns the parsed parameters. - */ -function parseFuncParameters( - paramsNode: t.ArrayExpression, -): TranspilationResult['params'] | undefined { - return []; -} - -function parseAstNode(astNode: t.ObjectExpression): EmbeddedTypegpuMetadata['ast'] | undefined { - const paramsNode = objectPropertyValue(astNode, 'params'); - const bodyNode = objectPropertyValue(astNode, 'body'); - - if ( - !paramsNode || - !bodyNode || - !t.isArrayExpression(paramsNode) || - !t.isArrayExpression(bodyNode) - ) { - return undefined; - } - - const params = parseFuncParameters(paramsNode); - const body = parseBody(bodyNode); - - if (!params || !body) { - return undefined; - } - - return { - params, - body, - }; +function parseAstPath( + astPath: NodePath, +): EmbeddedTypegpuMetadata['ast'] | undefined { + const evaluated = astPath.evaluate(); + + return !evaluated.confident + ? undefined + : (evaluated.value as NonNullable); } /** @@ -303,52 +203,60 @@ export function getEmbeddedTypegpuMetadata( } // we check for the metadata object - const metadataArgument = callPath.node.arguments[1]; - if (metadataArgument === undefined) { + const metadataPath = callPath.get('arguments')[1]; + if (metadataPath === undefined) { return undefined; } - const metadataNode = unwrapParentheses(metadataArgument); - if (!t.isObjectExpression(metadataNode)) { + const unwrappedMetadataPatah = unwrapParenthesesPath(metadataPath); + if (!unwrappedMetadataPatah.isObjectExpression()) { return undefined; } // get the metadata properties - const versionNode = objectPropertyValue(metadataNode, 'v'); - const nameNode = objectPropertyValue(metadataNode, 'name'); + const versionPath = objectPropertyPath(unwrappedMetadataPatah, 'v'); + const namePath = objectPropertyPath(unwrappedMetadataPatah, 'name'); + + if (versionPath === undefined || namePath === undefined) { + return undefined; + } + + const versionResult = versionPath.evaluate(); + const nameResult = namePath.evaluate(); - if (!t.isNumericLiteral(versionNode) || nameNode === undefined) { + if (!versionResult.confident || !nameResult.confident) { return undefined; } - const name = t.isStringLiteral(nameNode) ? nameNode.value : undefined; + const version = versionResult.value as number; + const name = nameResult.value as string | undefined; // metadata v1 support is limited - if (versionNode.value == 1) { + if (version == 1) { return { - v: versionNode.value, + v: version, name, }; } - const astNode = objectPropertyValue(metadataNode, 'ast'); - const externalsNode = objectPropertyValue(metadataNode, 'externals'); + const astPath = objectPropertyPath(unwrappedMetadataPatah, 'ast'); + const externalsPath = objectPropertyPath(unwrappedMetadataPatah, 'externals'); - if (astNode === undefined || externalsNode === undefined) { + if (astPath === undefined || externalsPath === undefined) { return undefined; } - if (!t.isObjectExpression(astNode)) { + if (!astPath.isObjectExpression()) { return undefined; } - const ast = parseAstNode(astNode); + const ast = parseAstPath(astPath); if (ast === undefined) { return undefined; } const embeddedTypegpuMetadata = { - v: versionNode.value, + v: version, name, ast, }; diff --git a/packages/unplugin-typegpu/test/embedded-metadata.test.ts b/packages/unplugin-typegpu/test/embedded-metadata.test.ts index 84e6e4503a..1c8f578f90 100644 --- a/packages/unplugin-typegpu/test/embedded-metadata.test.ts +++ b/packages/unplugin-typegpu/test/embedded-metadata.test.ts @@ -147,7 +147,7 @@ describe('getEmbeddedTypegpuMetadata', () => { }); }); - describe.skip('parsers function parameters', () => { + describe('parsers function parameters', () => { const code = `\ const noParams = () => { 'use gpu'; From 0e1183c9b376bb79b2c28bfbdefb1db79531dd3b Mon Sep 17 00:00:00 2001 From: Szymon Szulc Date: Tue, 15 Sep 2026 14:05:46 +0200 Subject: [PATCH 05/18] throwing errors instead of returning undefined --- .../src/core/embeddedMetadata.ts | 15 +++++--- .../test/embedded-metadata.test.ts | 37 +++++++++++++++++++ 2 files changed, 46 insertions(+), 6 deletions(-) diff --git a/packages/unplugin-typegpu/src/core/embeddedMetadata.ts b/packages/unplugin-typegpu/src/core/embeddedMetadata.ts index 3039384a33..c515f12e52 100644 --- a/packages/unplugin-typegpu/src/core/embeddedMetadata.ts +++ b/packages/unplugin-typegpu/src/core/embeddedMetadata.ts @@ -3,6 +3,9 @@ import * as t from '@babel/types'; import type { TranspilationResult } from 'tinyest-for-wgsl'; import type { MetadatableFunction } from './common.ts'; +const METADATA_PARSING_ERROR_MESSAGE = + 'unplugin-typegpu: Error when parsing metadata: required fields are missing or could not be evaluated.'; + export interface EmbeddedTypegpuMetadata { v: number; name: string | undefined; @@ -218,21 +221,21 @@ export function getEmbeddedTypegpuMetadata( const namePath = objectPropertyPath(unwrappedMetadataPatah, 'name'); if (versionPath === undefined || namePath === undefined) { - return undefined; + throw new Error(METADATA_PARSING_ERROR_MESSAGE); } const versionResult = versionPath.evaluate(); const nameResult = namePath.evaluate(); if (!versionResult.confident || !nameResult.confident) { - return undefined; + throw new Error(METADATA_PARSING_ERROR_MESSAGE); } const version = versionResult.value as number; const name = nameResult.value as string | undefined; // metadata v1 support is limited - if (version == 1) { + if (version === 1) { return { v: version, name, @@ -243,16 +246,16 @@ export function getEmbeddedTypegpuMetadata( const externalsPath = objectPropertyPath(unwrappedMetadataPatah, 'externals'); if (astPath === undefined || externalsPath === undefined) { - return undefined; + throw new Error(METADATA_PARSING_ERROR_MESSAGE); } if (!astPath.isObjectExpression()) { - return undefined; + throw new Error(METADATA_PARSING_ERROR_MESSAGE); } const ast = parseAstPath(astPath); if (ast === undefined) { - return undefined; + throw new Error(METADATA_PARSING_ERROR_MESSAGE); } const embeddedTypegpuMetadata = { diff --git a/packages/unplugin-typegpu/test/embedded-metadata.test.ts b/packages/unplugin-typegpu/test/embedded-metadata.test.ts index 1c8f578f90..1d6ee2e43a 100644 --- a/packages/unplugin-typegpu/test/embedded-metadata.test.ts +++ b/packages/unplugin-typegpu/test/embedded-metadata.test.ts @@ -129,6 +129,43 @@ function dualTest( } describe('getEmbeddedTypegpuMetadata', () => { + describe.each([ + ['missing version', "{ name: 'fn' }"], + ['missing name', '{ v: 2 }'], + ['unevaluable version', "{ v: unknownVersion, name: 'fn' }"], + ['unevaluable name', '{ v: 2, name: unknownName }'], + ['missing ast', "{ v: 2, name: 'fn', externals: {} }"], + ['missing externals', "{ v: 2, name: 'fn', ast: { params: [], body: [0, []] } }"], + [ + 'unevaluable ast', + "{ v: 2, name: 'fn', ast: { params: [], body: unknownBody }, externals: {} }", + ], + ])('throws for %s', (_, metadata) => { + const code = `\ + const fn = ($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = () => { + 'use gpu'; + }, ${metadata}) && $.f)({}); + + console.log(fn); + `; + + test('[BABEL]', () => { + expect(() => + babelTransform(code, {}, [createBabelMetadataCollector([])]), + ).toThrowErrorMatchingInlineSnapshot( + `[Error: unknown file: unplugin-typegpu: Error when parsing metadata: required fields are missing or could not be evaluated.]`, + ); + }); + + test('[ROLLUP]', async () => { + await expect( + rollupTransform(code, undefined, [createRollupMetadataCollector([])]), + ).rejects.toThrowErrorMatchingInlineSnapshot( + `[Error: unplugin-typegpu: Error when parsing metadata: required fields are missing or could not be evaluated.]`, + ); + }); + }); + describe('omits ast and externals fields for metadata v1', () => { const code = `\ const fn = ($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = () => { From 89e0d468c43878e7c7f7a1a247d087a3b50c4f6b Mon Sep 17 00:00:00 2001 From: Szymon Szulc Date: Mon, 21 Sep 2026 14:33:12 +0200 Subject: [PATCH 06/18] new shape --- .../src/core/embeddedMetadata.ts | 34 ++-- .../test/embedded-metadata.test.ts | 148 ++++++++++++------ 2 files changed, 115 insertions(+), 67 deletions(-) diff --git a/packages/unplugin-typegpu/src/core/embeddedMetadata.ts b/packages/unplugin-typegpu/src/core/embeddedMetadata.ts index c515f12e52..9060d5509b 100644 --- a/packages/unplugin-typegpu/src/core/embeddedMetadata.ts +++ b/packages/unplugin-typegpu/src/core/embeddedMetadata.ts @@ -6,15 +6,17 @@ import type { MetadatableFunction } from './common.ts'; const METADATA_PARSING_ERROR_MESSAGE = 'unplugin-typegpu: Error when parsing metadata: required fields are missing or could not be evaluated.'; -export interface EmbeddedTypegpuMetadata { +type FunctionAst = Pick; + +export type EmbeddedTypegpuMetadata = { v: number; name: string | undefined; - ast?: { - params: TranspilationResult['params']; - body: TranspilationResult['body']; + function?: { + ast: FunctionAst; + astPath: NodePath; + externals: Map; }; - externals?: TranspilationResult['externalNames']; -} +}; const embeddedTypegpuMetadataCache = new WeakMap< NodePath, @@ -143,14 +145,10 @@ function objectPropertyPath( /** * Returns the parsed AST */ -function parseAstPath( - astPath: NodePath, -): EmbeddedTypegpuMetadata['ast'] | undefined { +function parseAstPath(astPath: NodePath): FunctionAst | undefined { const evaluated = astPath.evaluate(); - return !evaluated.confident - ? undefined - : (evaluated.value as NonNullable); + return !evaluated.confident ? undefined : (evaluated.value as FunctionAst); } /** @@ -206,7 +204,7 @@ export function getEmbeddedTypegpuMetadata( } // we check for the metadata object - const metadataPath = callPath.get('arguments')[1]; + const metadataPath = callPath.get('arguments.1'); if (metadataPath === undefined) { return undefined; } @@ -236,10 +234,12 @@ export function getEmbeddedTypegpuMetadata( // metadata v1 support is limited if (version === 1) { - return { + const embeddedTypegpuMetadata = { v: version, name, }; + embeddedTypegpuMetadataCache.set(path, embeddedTypegpuMetadata); + return embeddedTypegpuMetadata; } const astPath = objectPropertyPath(unwrappedMetadataPatah, 'ast'); @@ -261,7 +261,11 @@ export function getEmbeddedTypegpuMetadata( const embeddedTypegpuMetadata = { v: version, name, - ast, + function: { + ast, + astPath, + externals: new Map(), + }, }; embeddedTypegpuMetadataCache.set(path, embeddedTypegpuMetadata); diff --git a/packages/unplugin-typegpu/test/embedded-metadata.test.ts b/packages/unplugin-typegpu/test/embedded-metadata.test.ts index 1d6ee2e43a..13f55ac1eb 100644 --- a/packages/unplugin-typegpu/test/embedded-metadata.test.ts +++ b/packages/unplugin-typegpu/test/embedded-metadata.test.ts @@ -2,7 +2,7 @@ import * as parser from '@babel/parser'; import _traverse, { type NodePath } from '@babel/traverse'; import type { Plugin } from 'rollup'; import { describe, expect, test } from 'vitest'; -import { transpileBabelFn } from 'tinyest-for-wgsl'; +import { transpileBabelFn, type TranspilationResult } from 'tinyest-for-wgsl'; import { type BabelTestPlugin, babelTransform, rollupTransform } from './transform.ts'; import type { MetadatableFunction } from '../src/core/common.ts'; import { @@ -15,6 +15,31 @@ if (typeof (traverse as unknown as { default: typeof traverse }).default === 'fu traverse = (traverse as unknown as { default: typeof traverse }).default; } +function getFirstUseGpuFunctionPath( + ast: ReturnType, +): NodePath { + let functionPath: NodePath | undefined; + + traverse(ast, { + ArrowFunctionExpression(path) { + const body = path.node.body; + if ( + body.type === 'BlockStatement' && + body.directives.some((d) => d.value.value === 'use gpu') + ) { + functionPath = path; + path.stop(); + } + }, + }); + + if (!functionPath) { + throw new Error('No TypeGPU function path found'); + } + + return functionPath; +} + function collectEmbeddedMetadata( path: NodePath, metadata: EmbeddedTypegpuMetadata[], @@ -73,13 +98,13 @@ function createRollupMetadataCollector(metadata: EmbeddedTypegpuMetadata[]): Plu * * @note For simplicity, shelled functions are omitted. */ -function extractExpectedAstsFromSource(code: string) { +function extractTranspilationResultFromSource(code: string) { const ast = parser.parse(code, { sourceType: 'module', plugins: ['typescript'], }); - const expected: NonNullable[] = []; + const expected: TranspilationResult[] = []; function collect(path: NodePath) { const body = path.node.body; @@ -91,8 +116,8 @@ function extractExpectedAstsFromSource(code: string) { return; } - const { params, body: transpiledBody } = transpileBabelFn(path.node); - expected.push({ params, body: transpiledBody }); + const result = transpileBabelFn(path.node); + expected.push(result); } traverse(ast, { @@ -106,29 +131,48 @@ function extractExpectedAstsFromSource(code: string) { function dualTest( code: string, - check: ( - metadata: EmbeddedTypegpuMetadata[], - expectedAsts: NonNullable[], - ) => void, + check: (metadata: EmbeddedTypegpuMetadata[], expected: TranspilationResult[]) => void, ) { test('[BABEL]', () => { - const expectedAsts = extractExpectedAstsFromSource(code); + const expected = extractTranspilationResultFromSource(code); const metadata: EmbeddedTypegpuMetadata[] = []; babelTransform(code, {}, [createBabelMetadataCollector(metadata)]); - check(metadata, expectedAsts); + check(metadata, expected); }); test('[ROLLUP]', async () => { - const expectedAsts = extractExpectedAstsFromSource(code); + const expected = extractTranspilationResultFromSource(code); const metadata: EmbeddedTypegpuMetadata[] = []; await rollupTransform(code, undefined, [createRollupMetadataCollector(metadata)]); - check(metadata, expectedAsts); + check(metadata, expected); }); } describe('getEmbeddedTypegpuMetadata', () => { + test.each([1, 2])('reuses cached v%s metadata', (version) => { + const ast = parser.parse( + ` + const fn = ($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set( + $.f = () => { 'use gpu'; }, + { + v: ${version}, + name: 'fn', + ast: { params: [], body: [0, []] }, + externals: {} + } + ) && $.f)({}); + `, + { sourceType: 'module' }, + ); + + const functionPath = getFirstUseGpuFunctionPath(ast); + const firstRead = getEmbeddedTypegpuMetadata(functionPath); + expect(firstRead).toBeDefined(); + expect(getEmbeddedTypegpuMetadata(functionPath)).toBe(firstRead); + }); + describe.each([ ['missing version', "{ name: 'fn' }"], ['missing name', '{ v: 2 }'], @@ -208,9 +252,9 @@ describe('getEmbeddedTypegpuMetadata', () => { console.log(noParams, identifierParams, destructuredParams, mixedParams); `; - dualTest(code, (metadata, expectedAsts) => { - expect(metadata.map((m) => m.ast?.params)).toStrictEqual( - expectedAsts.map((ast) => ast.params), + dualTest(code, (metadata, expected) => { + expect(metadata.map((m) => m.function?.ast.params)).toStrictEqual( + expected.map((e) => e.params), ); }); }); @@ -229,8 +273,8 @@ describe('getEmbeddedTypegpuMetadata', () => { console.log(fn); `; - dualTest(code, (metadata, expectedAsts) => { - expect(metadata.map((m) => m.ast?.body)).toStrictEqual(expectedAsts.map((ast) => ast.body)); + dualTest(code, (metadata, expected) => { + expect(metadata.map((m) => m.function?.ast.body)).toStrictEqual(expected.map((e) => e.body)); }); }); @@ -250,8 +294,8 @@ describe('getEmbeddedTypegpuMetadata', () => { console.log(fn); `; - dualTest(code, (metadata, expectedAsts) => { - expect(metadata.map((m) => m.ast?.body)).toStrictEqual(expectedAsts.map((ast) => ast.body)); + dualTest(code, (metadata, expected) => { + expect(metadata.map((m) => m.function?.ast.body)).toStrictEqual(expected.map((e) => e.body)); }); }); @@ -267,8 +311,8 @@ describe('getEmbeddedTypegpuMetadata', () => { console.log(fn); `; - dualTest(code, (metadata, expectedAsts) => { - expect(metadata.map((m) => m.ast?.body)).toStrictEqual(expectedAsts.map((ast) => ast.body)); + dualTest(code, (metadata, expected) => { + expect(metadata.map((m) => m.function?.ast.body)).toStrictEqual(expected.map((e) => e.body)); }); }); @@ -284,8 +328,8 @@ describe('getEmbeddedTypegpuMetadata', () => { console.log(fn); `; - dualTest(code, (metadata, expectedAsts) => { - expect(metadata.map((m) => m.ast?.body)).toStrictEqual(expectedAsts.map((ast) => ast.body)); + dualTest(code, (metadata, expected) => { + expect(metadata.map((m) => m.function?.ast.body)).toStrictEqual(expected.map((e) => e.body)); }); }); @@ -300,8 +344,8 @@ describe('getEmbeddedTypegpuMetadata', () => { console.log(fn); `; - dualTest(code, (metadata, expectedAsts) => { - expect(metadata.map((m) => m.ast?.body)).toStrictEqual(expectedAsts.map((ast) => ast.body)); + dualTest(code, (metadata, expected) => { + expect(metadata.map((m) => m.function?.ast.body)).toStrictEqual(expected.map((e) => e.body)); }); }); @@ -315,8 +359,8 @@ describe('getEmbeddedTypegpuMetadata', () => { console.log(fn); `; - dualTest(code, (metadata, expectedAsts) => { - expect(metadata.map((m) => m.ast?.body)).toStrictEqual(expectedAsts.map((ast) => ast.body)); + dualTest(code, (metadata, expected) => { + expect(metadata.map((m) => m.function?.ast.body)).toStrictEqual(expected.map((e) => e.body)); }); }); @@ -332,8 +376,8 @@ describe('getEmbeddedTypegpuMetadata', () => { console.log(fn); `; - dualTest(code, (metadata, expectedAsts) => { - expect(metadata.map((m) => m.ast?.body)).toStrictEqual(expectedAsts.map((ast) => ast.body)); + dualTest(code, (metadata, expected) => { + expect(metadata.map((m) => m.function?.ast.body)).toStrictEqual(expected.map((e) => e.body)); }); }); @@ -349,8 +393,8 @@ describe('getEmbeddedTypegpuMetadata', () => { console.log(fn); `; - dualTest(code, (metadata, expectedAsts) => { - expect(metadata.map((m) => m.ast?.body)).toStrictEqual(expectedAsts.map((ast) => ast.body)); + dualTest(code, (metadata, expected) => { + expect(metadata.map((m) => m.function?.ast.body)).toStrictEqual(expected.map((e) => e.body)); }); }); @@ -369,8 +413,8 @@ describe('getEmbeddedTypegpuMetadata', () => { console.log(bareReturn, valueReturn); `; - dualTest(code, (metadata, expectedAsts) => { - expect(metadata.map((m) => m.ast?.body)).toStrictEqual(expectedAsts.map((ast) => ast.body)); + dualTest(code, (metadata, expected) => { + expect(metadata.map((m) => m.function?.ast.body)).toStrictEqual(expected.map((e) => e.body)); }); }); @@ -414,8 +458,8 @@ describe('getEmbeddedTypegpuMetadata', () => { console.log(bare, alternative, elseIf, withoutBlocks); `; - dualTest(code, (metadata, expectedAsts) => { - expect(metadata.map((m) => m.ast?.body)).toStrictEqual(expectedAsts.map((ast) => ast.body)); + dualTest(code, (metadata, expected) => { + expect(metadata.map((m) => m.function?.ast.body)).toStrictEqual(expected.map((e) => e.body)); }); }); @@ -430,8 +474,8 @@ describe('getEmbeddedTypegpuMetadata', () => { console.log(fn); `; - dualTest(code, (metadata, expectedAsts) => { - expect(metadata.map((m) => m.ast?.body)).toStrictEqual(expectedAsts.map((ast) => ast.body)); + dualTest(code, (metadata, expected) => { + expect(metadata.map((m) => m.function?.ast.body)).toStrictEqual(expected.map((e) => e.body)); }); }); @@ -464,8 +508,8 @@ describe('getEmbeddedTypegpuMetadata', () => { console.log(forLoop, whileLoop, forOfLoop); `; - dualTest(code, (metadata, expectedAsts) => { - expect(metadata.map((m) => m.ast?.body)).toStrictEqual(expectedAsts.map((ast) => ast.body)); + dualTest(code, (metadata, expected) => { + expect(metadata.map((m) => m.function?.ast.body)).toStrictEqual(expected.map((e) => e.body)); }); }); @@ -490,8 +534,8 @@ describe('getEmbeddedTypegpuMetadata', () => { console.log(fn); `; - dualTest(code, (metadata, expectedAsts) => { - expect(metadata.map((m) => m.ast?.body)).toStrictEqual(expectedAsts.map((ast) => ast.body)); + dualTest(code, (metadata, expected) => { + expect(metadata.map((m) => m.function?.ast.body)).toStrictEqual(expected.map((e) => e.body)); }); }); @@ -506,8 +550,8 @@ describe('getEmbeddedTypegpuMetadata', () => { console.log(fn); `; - dualTest(code, (metadata, expectedAsts) => { - expect(metadata.map((m) => m.ast?.body)).toStrictEqual(expectedAsts.map((ast) => ast.body)); + dualTest(code, (metadata, expected) => { + expect(metadata.map((m) => m.function?.ast.body)).toStrictEqual(expected.map((e) => e.body)); }); }); @@ -523,8 +567,8 @@ describe('getEmbeddedTypegpuMetadata', () => { console.log(fn); `; - dualTest(code, (metadata, expectedAsts) => { - expect(metadata.map((m) => m.ast?.body)).toStrictEqual(expectedAsts.map((ast) => ast.body)); + dualTest(code, (metadata, expected) => { + expect(metadata.map((m) => m.function?.ast.body)).toStrictEqual(expected.map((e) => e.body)); }); }); @@ -539,8 +583,8 @@ describe('getEmbeddedTypegpuMetadata', () => { console.log(fn); `; - dualTest(code, (metadata, expectedAsts) => { - expect(metadata.map((m) => m.ast?.body)).toStrictEqual(expectedAsts.map((ast) => ast.body)); + dualTest(code, (metadata, expected) => { + expect(metadata.map((m) => m.function?.ast.body)).toStrictEqual(expected.map((e) => e.body)); }); }); @@ -554,8 +598,8 @@ describe('getEmbeddedTypegpuMetadata', () => { console.log(fn); `; - dualTest(code, (metadata, expectedAsts) => { - expect(metadata.map((m) => m.ast?.body)).toStrictEqual(expectedAsts.map((ast) => ast.body)); + dualTest(code, (metadata, expected) => { + expect(metadata.map((m) => m.function?.ast.body)).toStrictEqual(expected.map((e) => e.body)); }); }); @@ -569,8 +613,8 @@ describe('getEmbeddedTypegpuMetadata', () => { console.log(fn); `; - dualTest(code, (metadata, expectedAsts) => { - expect(metadata.map((m) => m.ast?.body)).toStrictEqual(expectedAsts.map((ast) => ast.body)); + dualTest(code, (metadata, expected) => { + expect(metadata.map((m) => m.function?.ast.body)).toStrictEqual(expected.map((e) => e.body)); }); }); }); From dbbc5d21fcd5ba05ddb16e01f65ada614634507e Mon Sep 17 00:00:00 2001 From: Szymon Szulc Date: Tue, 22 Sep 2026 12:11:27 +0200 Subject: [PATCH 07/18] parsing externals --- packages/unplugin-typegpu/package.json | 2 + .../src/core/embeddedMetadata.ts | 43 ++++++++++++-- .../test/embedded-metadata.test.ts | 59 +++++++++++++++++++ pnpm-lock.yaml | 49 +++++++-------- 4 files changed, 125 insertions(+), 28 deletions(-) diff --git a/packages/unplugin-typegpu/package.json b/packages/unplugin-typegpu/package.json index 9c1c652fcf..6f1f44ffdb 100644 --- a/packages/unplugin-typegpu/package.json +++ b/packages/unplugin-typegpu/package.json @@ -300,9 +300,11 @@ "unplugin": "^3.0.0" }, "devDependencies": { + "@babel/generator": "^7.29.0", "@babel/standalone": "^7.28.6", "@rollup/plugin-virtual": "^3.0.2", "@typegpu/tgpu-dev-cli": "workspace:*", + "@types/babel__generator": "^7.27.0", "@types/babel__standalone": "^7.1.9", "@types/babel__traverse": "^7.20.7", "@types/bun": "^1.3.11", diff --git a/packages/unplugin-typegpu/src/core/embeddedMetadata.ts b/packages/unplugin-typegpu/src/core/embeddedMetadata.ts index 9060d5509b..b17f98e69b 100644 --- a/packages/unplugin-typegpu/src/core/embeddedMetadata.ts +++ b/packages/unplugin-typegpu/src/core/embeddedMetadata.ts @@ -7,6 +7,7 @@ const METADATA_PARSING_ERROR_MESSAGE = 'unplugin-typegpu: Error when parsing metadata: required fields are missing or could not be evaluated.'; type FunctionAst = Pick; +type FunctionExternals = Map>; export type EmbeddedTypegpuMetadata = { v: number; @@ -14,7 +15,7 @@ export type EmbeddedTypegpuMetadata = { function?: { ast: FunctionAst; astPath: NodePath; - externals: Map; + externals: FunctionExternals; }; }; @@ -151,6 +152,38 @@ function parseAstPath(astPath: NodePath): FunctionAst | unde return !evaluated.confident ? undefined : (evaluated.value as FunctionAst); } +/** + * Externals are represented as an object with string keys and external getters. + * This function parses such objects and returns a map keyed by string, with values that point to the corresponding object property node paths. + */ +function parseExternalsPath( + externalsPath: NodePath, +): FunctionExternals | undefined { + const externals: FunctionExternals = new Map(); + + for (const propertyPath of externalsPath.get('properties')) { + if (!propertyPath.isObjectProperty()) { + return undefined; + } + + const { computed, key } = propertyPath.node; + + const name = t.isStringLiteral(key) + ? key.value + : !computed && t.isIdentifier(key) + ? key.name + : undefined; + + if (name === undefined || !propertyPath.get('value').isArrowFunctionExpression()) { + return undefined; + } + + externals.set(name, propertyPath); + } + + return externals; +} + /** * Returns metadata embedded by unplugin-typegpu for this exact function: * @@ -249,12 +282,14 @@ export function getEmbeddedTypegpuMetadata( throw new Error(METADATA_PARSING_ERROR_MESSAGE); } - if (!astPath.isObjectExpression()) { + if (!astPath.isObjectExpression() || !externalsPath.isObjectExpression()) { throw new Error(METADATA_PARSING_ERROR_MESSAGE); } const ast = parseAstPath(astPath); - if (ast === undefined) { + const externals = parseExternalsPath(externalsPath); + + if (ast === undefined || externals === undefined) { throw new Error(METADATA_PARSING_ERROR_MESSAGE); } @@ -264,7 +299,7 @@ export function getEmbeddedTypegpuMetadata( function: { ast, astPath, - externals: new Map(), + externals, }, }; embeddedTypegpuMetadataCache.set(path, embeddedTypegpuMetadata); diff --git a/packages/unplugin-typegpu/test/embedded-metadata.test.ts b/packages/unplugin-typegpu/test/embedded-metadata.test.ts index 13f55ac1eb..c0b2036e67 100644 --- a/packages/unplugin-typegpu/test/embedded-metadata.test.ts +++ b/packages/unplugin-typegpu/test/embedded-metadata.test.ts @@ -1,4 +1,5 @@ import * as parser from '@babel/parser'; +import generator from '@babel/generator'; import _traverse, { type NodePath } from '@babel/traverse'; import type { Plugin } from 'rollup'; import { describe, expect, test } from 'vitest'; @@ -129,6 +130,17 @@ function extractTranspilationResultFromSource(code: string) { return expected; } +/** + * Maps the parsed externals of every function to `[name, getterBody]` pairs. + */ +function extractExternalGetters(metadata: EmbeddedTypegpuMetadata[]) { + return metadata.map((m) => + [...(m.function?.externals ?? [])].map(([name, propertyPath]) => { + return [name, generator(propertyPath.get('value').get('body').node).code]; + }), + ); +} + function dualTest( code: string, check: (metadata: EmbeddedTypegpuMetadata[], expected: TranspilationResult[]) => void, @@ -184,6 +196,10 @@ describe('getEmbeddedTypegpuMetadata', () => { 'unevaluable ast', "{ v: 2, name: 'fn', ast: { params: [], body: unknownBody }, externals: {} }", ], + [ + 'unevaluable externals', + "{ v: 2, name: 'fn', ast: { params: [], body: [0, []] }, externals: unknownExternals }", + ], ])('throws for %s', (_, metadata) => { const code = `\ const fn = ($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = () => { @@ -617,4 +633,47 @@ describe('getEmbeddedTypegpuMetadata', () => { expect(metadata.map((m) => m.function?.ast.body)).toStrictEqual(expected.map((e) => e.body)); }); }); + + describe('parses externals', () => { + const code = `\ + const noExternals = () => { + 'use gpu'; + const a = 1; + }; + + const withExternals = () => { + 'use gpu'; + const a = ext.value; + const b = ext.config.multiplier; + const c = lookup[0]; + }; + + console.log(noExternals, withExternals); + `; + + dualTest(code, (metadata, expected) => { + expect(metadata.map((m) => [...(m.function?.externals.keys() ?? [])])).toStrictEqual( + expected.map((e) => [...e.externalNames.keys()]), + ); + }); + }); + + describe('maps externals to their getters', () => { + const code = `\ + const fn = () => { + 'use gpu'; + const a = ext.value; + const b = ext.config.multiplier; + const c = lookup[0]; + }; + + console.log(fn); + `; + + dualTest(code, (metadata, expected) => { + expect(extractExternalGetters(metadata)).toStrictEqual( + expected.map((e) => [...e.externalNames]), + ); + }); + }); }); diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 2553faf28d..fdfcddd7b8 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -156,7 +156,7 @@ importers: devDependencies: '@types/bun': specifier: latest - version: 1.4.0 + version: 1.4.2 apps/infra-benchmarks: devDependencies: @@ -1067,6 +1067,9 @@ importers: specifier: ^3.0.0 version: 3.0.0 devDependencies: + '@babel/generator': + specifier: ^7.29.0 + version: 7.29.7 '@babel/standalone': specifier: ^7.28.6 version: 7.28.6 @@ -1076,6 +1079,9 @@ importers: '@typegpu/tgpu-dev-cli': specifier: workspace:* version: link:../tgpu-dev-cli + '@types/babel__generator': + specifier: ^7.27.0 + version: 7.27.0 '@types/babel__standalone': specifier: ^7.1.9 version: 7.1.9 @@ -1241,10 +1247,6 @@ packages: resolution: {integrity: sha512-CGOfOJqWjg2qW/Mb6zNsDm+u5vFQ8DxXfbM09z69p5Z6+mE1ikP2jUXw+j42Pf1XTYED2Rni5f95npYeuwMDQA==} engines: {node: '>=6.9.0'} - '@babel/generator@7.29.1': - resolution: {integrity: sha512-qsaF+9Qcm2Qv8SRIMMscAvG4O3lJ0F1GuMo5HR/Bp02LopNgnZBC/EkbevHFeGs4ls/oPz9v+Bsmzbkbe+0dUw==} - engines: {node: '>=6.9.0'} - '@babel/generator@7.29.7': resolution: {integrity: sha512-DkXD5OJQaAQIdZ1bt3UZdEnHAn9Imd3IVBdX03UFe+ony9Ojw5pzr9YVKGDY1jt+Gcn/FnGkNf8r+Vj5NOJWtQ==} engines: {node: '>=6.9.0'} @@ -4533,6 +4535,9 @@ packages: '@types/babel__core@7.20.5': resolution: {integrity: sha512-qoQprZvz5wQFJwMDqeseRXWv3rqMvhgpbXFfVyWhbx9X47POIA6i/+dXefEmZKoAgOaTdaIgNSMqMIU61yRyzA==} + '@types/babel__generator@7.27.0': + resolution: {integrity: sha512-ufFd2Xi92OAVPYsy+P4n7/U7e68fex0+Ee8gSG9KX7eo084CWiQ4sdxktvdl0bOPupXtVJPY19zk6EwWqUQ8lg==} + '@types/babel__generator@7.6.8': resolution: {integrity: sha512-ASsj+tpEDsEiFr1arWrlN6V3mdfjRMZt6LtK/Vp/kreFLnr5QH5+DhvD5nINYZXzwJvXeGq+05iUXcAzVrqWtw==} @@ -4551,8 +4556,8 @@ packages: '@types/bun@1.3.12': resolution: {integrity: sha512-DBv81elK+/VSwXHDlnH3Qduw+KxkTIWi7TXkAeh24zpi5l0B2kUg9Ga3tb4nJaPcOFswflgi/yAvMVBPrxMB+A==} - '@types/bun@1.4.0': - resolution: {integrity: sha512-K+lZULY23vRgK/CfTjFIV+tyifaNdSMlPh9j+6mQ/cLfpOznLyAuzgV/JQysyECpkBQLVMSyvjlr2fBUSA9wFQ==} + '@types/bun@1.4.2': + resolution: {integrity: sha512-GimotNn7+ZV0uVArItBbriZsR1oNf0+WTzPkdcFrzShI7k2norL0uzEaJT8T33dWr7O/c9ZDuAFQrctKCi72oQ==} '@types/chai@5.2.2': resolution: {integrity: sha512-8kB30R7Hwqf40JPiKhVzodJs2Qc1ZJ5zuT3uzw5Hq/dhNCl3G3l83jfpdI1e20BP348+fV7VIL/+FxaXkqBmWg==} @@ -5331,8 +5336,8 @@ packages: bun-types@1.3.12: resolution: {integrity: sha512-HqOLj5PoFajAQciOMRiIZGNoKxDJSr6qigAttOX40vJuSp6DN/CxWp9s3C1Xwm4oH7ybueITwiaOcWXoYVoRkA==} - bun-types@1.4.0: - resolution: {integrity: sha512-iIKw23BspnQQYd3prITOBxeUsxBHnwzX6YJfGMuNOZzeNcMmVqzIIVGRm1l69ogaPQmb4wB6BN8mA5bE9YuC5Q==} + bun-types@1.4.2: + resolution: {integrity: sha512-bxV1FgK7yBIzjRe5zBozIM4Bem11ZJcCXSrjWRG3YWLt8yFDePu4cLjpebO8OvPeIE9trbyPF4fuj3Cia4Fj3w==} bun@1.3.10: resolution: {integrity: sha512-S/CXaXXIyA4CMjdMkYQ4T2YMqnAn4s0ysD3mlsY4bUiOCqGlv28zck4Wd4H4kpvbekx15S9mUeLQ7Uxd0tYTLA==} @@ -10330,7 +10335,7 @@ snapshots: '@babel/helper-compilation-targets': 7.28.6 '@babel/helper-module-transforms': 7.28.6(@babel/core@7.29.0) '@babel/helpers': 7.28.6 - '@babel/parser': 7.29.0 + '@babel/parser': 7.29.7 '@babel/template': 7.28.6 '@babel/traverse': 7.29.7 '@babel/types': 7.29.0 @@ -10343,14 +10348,6 @@ snapshots: transitivePeerDependencies: - supports-color - '@babel/generator@7.29.1': - dependencies: - '@babel/parser': 7.29.0 - '@babel/types': 7.29.0 - '@jridgewell/gen-mapping': 0.3.13 - '@jridgewell/trace-mapping': 0.3.31 - jsesc: 3.1.0 - '@babel/generator@7.29.7': dependencies: '@babel/parser': 7.29.7 @@ -10870,7 +10867,7 @@ snapshots: '@babel/traverse@7.29.0': dependencies: '@babel/code-frame': 7.29.0 - '@babel/generator': 7.29.1 + '@babel/generator': 7.29.7 '@babel/helper-globals': 7.28.0 '@babel/parser': 7.29.0 '@babel/template': 7.28.6 @@ -13099,6 +13096,10 @@ snapshots: '@types/babel__template': 7.4.4 '@types/babel__traverse': 7.20.7 + '@types/babel__generator@7.27.0': + dependencies: + '@babel/types': 7.29.0 + '@types/babel__generator@7.6.8': dependencies: '@babel/types': 7.29.0 @@ -13129,9 +13130,9 @@ snapshots: dependencies: bun-types: 1.3.12 - '@types/bun@1.4.0': + '@types/bun@1.4.2': dependencies: - bun-types: 1.4.0 + bun-types: 1.4.2 '@types/chai@5.2.2': dependencies: @@ -14192,7 +14193,7 @@ snapshots: dependencies: '@types/node': 24.10.0 - bun-types@1.4.0: + bun-types@1.4.2: dependencies: '@types/node': 24.10.0 @@ -16734,7 +16735,7 @@ snapshots: '@babel/code-frame': 7.29.0 '@babel/core': 7.29.0 '@babel/generator': 7.29.7 - '@babel/parser': 7.29.0 + '@babel/parser': 7.29.7 '@babel/template': 7.28.6 '@babel/traverse': 7.29.7 '@babel/types': 7.29.0 @@ -18250,7 +18251,7 @@ snapshots: rolldown-plugin-dts@0.17.8(rolldown@1.0.0-beta.45(@emnapi/core@1.9.2)(@emnapi/runtime@1.9.2))(tsover@6.0.2): dependencies: - '@babel/generator': 7.29.1 + '@babel/generator': 7.29.7 '@babel/parser': 7.29.0 '@babel/types': 7.29.0 ast-kit: 2.2.0 From 792b51934977f239fff9daa6f4a663994256394f Mon Sep 17 00:00:00 2001 From: Szymon Szulc Date: Tue, 22 Sep 2026 12:58:32 +0200 Subject: [PATCH 08/18] nullish operator fallback --- .../src/core/embeddedMetadata.ts | 10 ++++++++- .../test/embedded-metadata.test.ts | 22 +++++++++++++++++++ 2 files changed, 31 insertions(+), 1 deletion(-) diff --git a/packages/unplugin-typegpu/src/core/embeddedMetadata.ts b/packages/unplugin-typegpu/src/core/embeddedMetadata.ts index b17f98e69b..7794842889 100644 --- a/packages/unplugin-typegpu/src/core/embeddedMetadata.ts +++ b/packages/unplugin-typegpu/src/core/embeddedMetadata.ts @@ -109,8 +109,16 @@ function isTypegpuMetadataSetCall(node: t.CallExpression): boolean { const inner = unwrapParentheses(callee.object); + // globalThis.__TYPEGPU_META__ ??= + if (t.isAssignmentExpression(inner, { operator: '??=' }) && isGlobalTypegpuMetadata(inner.left)) { + return true; + } + + // globalThis.__TYPEGPU_META__ = globalThis.__TYPEGPU_META__ ?? ... return ( - t.isAssignmentExpression(inner, { operator: '??=' }) && isGlobalTypegpuMetadata(inner.left) + t.isAssignmentExpression(inner, { operator: '=' }) && + isGlobalTypegpuMetadata(inner.left) && + t.isLogicalExpression(unwrapParentheses(inner.right), { operator: '??' }) ); } diff --git a/packages/unplugin-typegpu/test/embedded-metadata.test.ts b/packages/unplugin-typegpu/test/embedded-metadata.test.ts index c0b2036e67..8156a8e5fe 100644 --- a/packages/unplugin-typegpu/test/embedded-metadata.test.ts +++ b/packages/unplugin-typegpu/test/embedded-metadata.test.ts @@ -163,6 +163,26 @@ function dualTest( } describe('getEmbeddedTypegpuMetadata', () => { + describe('reads metadata from a lowered nullish assignment', () => { + const code = `\ + const fn = ($ => (globalThis.__TYPEGPU_META__ = (globalThis.__TYPEGPU_META__ ?? new WeakMap())).set( + $.f = () => { 'use gpu'; }, + { + v: 2, + name: 'fn', + ast: { params: [], body: [0, []] }, + externals: {} + } + ) && $.f)({}); + + console.log(fn); + `; + + dualTest(code, (metadata) => { + expect(metadata[0]).toBeDefined(); + }); + }); + test.each([1, 2])('reuses cached v%s metadata', (version) => { const ast = parser.parse( ` @@ -175,6 +195,8 @@ describe('getEmbeddedTypegpuMetadata', () => { externals: {} } ) && $.f)({}); + + console.log(fn); `, { sourceType: 'module' }, ); From d4f0d6636a643ecbe17736c4e5eb86f68af2587d Mon Sep 17 00:00:00 2001 From: Szymon Szulc Date: Tue, 22 Sep 2026 13:04:42 +0200 Subject: [PATCH 09/18] update docs --- .../src/core/embeddedMetadata.ts | 35 ++++--------------- 1 file changed, 7 insertions(+), 28 deletions(-) diff --git a/packages/unplugin-typegpu/src/core/embeddedMetadata.ts b/packages/unplugin-typegpu/src/core/embeddedMetadata.ts index 7794842889..defa821ccd 100644 --- a/packages/unplugin-typegpu/src/core/embeddedMetadata.ts +++ b/packages/unplugin-typegpu/src/core/embeddedMetadata.ts @@ -193,36 +193,15 @@ function parseExternalsPath( } /** - * Returns metadata embedded by unplugin-typegpu for this exact function: + * Returns metadata embedded by unplugin-typegpu for this exact function. * - * @note metadata v1 support is limited. Only the version and name are parsed. - * - * Consider: - * ```ts - * const f = ($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = () => { - * 'use gpu'; - * }, { - * v: 2, - * name: "f", - * ast: { - * params: [], - * body: [0, []] - * }, - * externals: {} - * }) && $.f)({}); - * ``` - * - * Function will return: - * { - * v: 2, - * name: "f", - * ast: { - * params: [], - * body: [0, []] - * }, - * externals: {} - * } + * The returned object contains: + * - `v`: the metadata version + * - `name`: the function name, which may be undefined + * - `function`: an object containing the parsed AST, its source node path, and a map + * from external names to their original property paths * + * @note Metadata v1 support is limited. Only the version and name are parsed. */ export function getEmbeddedTypegpuMetadata( path: NodePath, From 4fcb009457ebcba476d56e3c448430a968c81d86 Mon Sep 17 00:00:00 2001 From: Szymon Szulc Date: Tue, 22 Sep 2026 15:45:17 +0200 Subject: [PATCH 10/18] typo --- .../unplugin-typegpu/src/core/embeddedMetadata.ts | 14 +++++++------- .../test/embedded-metadata.test.ts | 4 ++-- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/packages/unplugin-typegpu/src/core/embeddedMetadata.ts b/packages/unplugin-typegpu/src/core/embeddedMetadata.ts index defa821ccd..df478e9df0 100644 --- a/packages/unplugin-typegpu/src/core/embeddedMetadata.ts +++ b/packages/unplugin-typegpu/src/core/embeddedMetadata.ts @@ -4,7 +4,7 @@ import type { TranspilationResult } from 'tinyest-for-wgsl'; import type { MetadatableFunction } from './common.ts'; const METADATA_PARSING_ERROR_MESSAGE = - 'unplugin-typegpu: Error when parsing metadata: required fields are missing or could not be evaluated.'; + '[unplugin-typegpu] Error when parsing metadata: required fields are missing or could not be evaluated.'; type FunctionAst = Pick; type FunctionExternals = Map>; @@ -229,14 +229,14 @@ export function getEmbeddedTypegpuMetadata( return undefined; } - const unwrappedMetadataPatah = unwrapParenthesesPath(metadataPath); - if (!unwrappedMetadataPatah.isObjectExpression()) { + const unwrappedMetadataPath = unwrapParenthesesPath(metadataPath); + if (!unwrappedMetadataPath.isObjectExpression()) { return undefined; } // get the metadata properties - const versionPath = objectPropertyPath(unwrappedMetadataPatah, 'v'); - const namePath = objectPropertyPath(unwrappedMetadataPatah, 'name'); + const versionPath = objectPropertyPath(unwrappedMetadataPath, 'v'); + const namePath = objectPropertyPath(unwrappedMetadataPath, 'name'); if (versionPath === undefined || namePath === undefined) { throw new Error(METADATA_PARSING_ERROR_MESSAGE); @@ -262,8 +262,8 @@ export function getEmbeddedTypegpuMetadata( return embeddedTypegpuMetadata; } - const astPath = objectPropertyPath(unwrappedMetadataPatah, 'ast'); - const externalsPath = objectPropertyPath(unwrappedMetadataPatah, 'externals'); + const astPath = objectPropertyPath(unwrappedMetadataPath, 'ast'); + const externalsPath = objectPropertyPath(unwrappedMetadataPath, 'externals'); if (astPath === undefined || externalsPath === undefined) { throw new Error(METADATA_PARSING_ERROR_MESSAGE); diff --git a/packages/unplugin-typegpu/test/embedded-metadata.test.ts b/packages/unplugin-typegpu/test/embedded-metadata.test.ts index 8156a8e5fe..8b422b9c3c 100644 --- a/packages/unplugin-typegpu/test/embedded-metadata.test.ts +++ b/packages/unplugin-typegpu/test/embedded-metadata.test.ts @@ -235,7 +235,7 @@ describe('getEmbeddedTypegpuMetadata', () => { expect(() => babelTransform(code, {}, [createBabelMetadataCollector([])]), ).toThrowErrorMatchingInlineSnapshot( - `[Error: unknown file: unplugin-typegpu: Error when parsing metadata: required fields are missing or could not be evaluated.]`, + `[Error: unknown file: [unplugin-typegpu] Error when parsing metadata: required fields are missing or could not be evaluated.]`, ); }); @@ -243,7 +243,7 @@ describe('getEmbeddedTypegpuMetadata', () => { await expect( rollupTransform(code, undefined, [createRollupMetadataCollector([])]), ).rejects.toThrowErrorMatchingInlineSnapshot( - `[Error: unplugin-typegpu: Error when parsing metadata: required fields are missing or could not be evaluated.]`, + `[Error: [unplugin-typegpu] Error when parsing metadata: required fields are missing or could not be evaluated.]`, ); }); }); From 2671548a6c8994b7a0c909ca073bae671a9fd339 Mon Sep 17 00:00:00 2001 From: Szymon Szulc Date: Tue, 22 Sep 2026 16:30:23 +0200 Subject: [PATCH 11/18] polish --- packages/unplugin-typegpu/test/embedded-metadata.test.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/unplugin-typegpu/test/embedded-metadata.test.ts b/packages/unplugin-typegpu/test/embedded-metadata.test.ts index 8b422b9c3c..eb98f494f6 100644 --- a/packages/unplugin-typegpu/test/embedded-metadata.test.ts +++ b/packages/unplugin-typegpu/test/embedded-metadata.test.ts @@ -95,7 +95,7 @@ function createRollupMetadataCollector(metadata: EmbeddedTypegpuMetadata[]): Plu } /** - * Traverses the program AST to extract the expected ASTs from functions containing 'use gpu' directive. + * Traverses the program AST to extract the `TranspilationResult` from functions containing 'use gpu' directive. * * @note For simplicity, shelled functions are omitted. */ From 5bc5fbd04c442295214443cc6c0b99e027417005 Mon Sep 17 00:00:00 2001 From: Szymon Szulc Date: Tue, 22 Sep 2026 16:56:50 +0200 Subject: [PATCH 12/18] one more test --- .../test/embedded-metadata.test.ts | 29 +++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/packages/unplugin-typegpu/test/embedded-metadata.test.ts b/packages/unplugin-typegpu/test/embedded-metadata.test.ts index eb98f494f6..3aa398a584 100644 --- a/packages/unplugin-typegpu/test/embedded-metadata.test.ts +++ b/packages/unplugin-typegpu/test/embedded-metadata.test.ts @@ -698,4 +698,33 @@ describe('getEmbeddedTypegpuMetadata', () => { ); }); }); + + describe('parses metadata assigned to a shelled function without a use gpu directive', () => { + const code = `\ + import { tgpu } from 'typegpu'; + import * as d from 'typegpu/data'; + + const addOneMore = tgpu.fn([d.f32, d.f32], d.f32)((a, b) => { + return a + b + c; + }); + + console.log(addOneMore); + `; + + dualTest(code, (metadata) => { + const embedded = metadata[0]; + expect(embedded?.v).toBe(2); + // shelled functions are named by their shell at runtime, so the metadata name stays undefined + expect(embedded?.name).toBeUndefined(); + expect(embedded?.function?.ast.params).toStrictEqual([ + { type: 'i', name: 'a' }, + { type: 'i', name: 'b' }, + ]); + expect(embedded?.function?.ast.body).toStrictEqual([ + 0, + [[10, [1, [1, 'a', '+', 'b'], '+', 'c']]], + ]); + expect([...(embedded?.function?.externals.keys() ?? [])]).toStrictEqual(['c']); + }); + }); }); From 20f961b16b4024223e042527f6edc6fca9fa85df Mon Sep 17 00:00:00 2001 From: Szymon Szulc Date: Tue, 22 Sep 2026 23:38:59 +0200 Subject: [PATCH 13/18] review changes --- .../unplugin-typegpu/test/embedded-metadata.test.ts | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/packages/unplugin-typegpu/test/embedded-metadata.test.ts b/packages/unplugin-typegpu/test/embedded-metadata.test.ts index 3aa398a584..4e5bc1d445 100644 --- a/packages/unplugin-typegpu/test/embedded-metadata.test.ts +++ b/packages/unplugin-typegpu/test/embedded-metadata.test.ts @@ -163,7 +163,7 @@ function dualTest( } describe('getEmbeddedTypegpuMetadata', () => { - describe('reads metadata from a lowered nullish assignment', () => { + describe('lowered nullish assignment', () => { const code = `\ const fn = ($ => (globalThis.__TYPEGPU_META__ = (globalThis.__TYPEGPU_META__ ?? new WeakMap())).set( $.f = () => { 'use gpu'; }, @@ -178,8 +178,14 @@ describe('getEmbeddedTypegpuMetadata', () => { console.log(fn); `; - dualTest(code, (metadata) => { - expect(metadata[0]).toBeDefined(); + describe('parses metadata', () => { + dualTest(code, (metadata) => { + expect(metadata[0]).toBeDefined(); + }); + }); + + test('does not double-wrap', () => { + expect(babelTransform(code, {})).not.toContain('globalThis.__TYPEGPU_META__ ??='); }); }); From 52e298c6e70e5b7ab069fdadacbb8dd9e821ba4f Mon Sep 17 00:00:00 2001 From: Szymon Szulc Date: Thu, 1 Oct 2026 13:25:37 +0200 Subject: [PATCH 14/18] cyclic dep --- packages/unplugin-typegpu/src/core/version.ts | 5 +++++ 1 file changed, 5 insertions(+) create mode 100644 packages/unplugin-typegpu/src/core/version.ts diff --git a/packages/unplugin-typegpu/src/core/version.ts b/packages/unplugin-typegpu/src/core/version.ts new file mode 100644 index 0000000000..c5773de4d0 --- /dev/null +++ b/packages/unplugin-typegpu/src/core/version.ts @@ -0,0 +1,5 @@ +/** + * Each breaking change to the metadata format requires a bump to this number. + * It's used at runtime by `typegpu` to determine how to interpret a function's metadata. + */ +export const METADATA_FORMAT_VERSION = 2; From 391741f6f171ccfabfdbba4ee0f77942367d46e5 Mon Sep 17 00:00:00 2001 From: Szymon Szulc Date: Thu, 1 Oct 2026 13:41:26 +0200 Subject: [PATCH 15/18] better js doc --- packages/unplugin-typegpu/src/core/embeddedMetadata.ts | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/packages/unplugin-typegpu/src/core/embeddedMetadata.ts b/packages/unplugin-typegpu/src/core/embeddedMetadata.ts index 95669ac516..95eae929dd 100644 --- a/packages/unplugin-typegpu/src/core/embeddedMetadata.ts +++ b/packages/unplugin-typegpu/src/core/embeddedMetadata.ts @@ -208,7 +208,8 @@ function parseExternalsPath( * - `function`: an object containing the parsed AST, its source node path, and a map * from external names to their original property paths * - * @note Metadata v1 support is limited. Only the version and name are parsed. + * @note Support for metadata stamped with a version different from `METADATA_FORMAT_VERSION` is limited. + * Only the version and name are parsed. */ export function getEmbeddedTypegpuMetadata( path: NodePath, From 639dbbca9f639f6b90388be2a184cb28d344f69e Mon Sep 17 00:00:00 2001 From: Szymon Szulc Date: Mon, 5 Oct 2026 12:39:46 +0200 Subject: [PATCH 16/18] review fixes --- .../test/embedded-metadata.test.ts | 96 ++++++++++++++++--- 1 file changed, 82 insertions(+), 14 deletions(-) diff --git a/packages/unplugin-typegpu/test/embedded-metadata.test.ts b/packages/unplugin-typegpu/test/embedded-metadata.test.ts index e47f8fe23f..db63d4ae05 100644 --- a/packages/unplugin-typegpu/test/embedded-metadata.test.ts +++ b/packages/unplugin-typegpu/test/embedded-metadata.test.ts @@ -143,28 +143,22 @@ function extractExternalGetters(metadata: EmbeddedTypegpuMetadata[]) { function dualTest( code: string, - check: ( - metadata: EmbeddedTypegpuMetadata[], - expected: TranspilationResult[], - output: string | null | undefined, - ) => void, + check: (metadata: EmbeddedTypegpuMetadata[], expected: TranspilationResult[]) => void, ) { test('[BABEL]', () => { const expected = extractTranspilationResultFromSource(code); const metadata: EmbeddedTypegpuMetadata[] = []; - const output = babelTransform(code, {}, [createBabelMetadataCollector(metadata)]); - check(metadata, expected, output); + void babelTransform(code, {}, [createBabelMetadataCollector(metadata)]); + check(metadata, expected); }); test('[ROLLUP]', async () => { const expected = extractTranspilationResultFromSource(code); const metadata: EmbeddedTypegpuMetadata[] = []; - const output = await rollupTransform(code, undefined, [ - createRollupMetadataCollector(metadata), - ]); - check(metadata, expected, output); + void (await rollupTransform(code, undefined, [createRollupMetadataCollector(metadata)])); + check(metadata, expected); }); } @@ -642,6 +636,53 @@ describe('getEmbeddedTypegpuMetadata', () => { }); }); + describe('parses switch statements', () => { + const code = `\ + const fn = (a, b) => { + 'use gpu'; + switch (a) { + case 1: + case ext: + break; + case 2: + default: + case 3: + return a; + } + }; + + console.log(fn); + `; + + dualTest(code, (metadata, expected) => { + expect(metadata.map((m) => m.function?.ast.body)).toStrictEqual(expected.map((e) => e.body)); + }); + }); + + describe('points astPath to the ast object of the metadata', () => { + const code = `\ + const fn = (a) => { + 'use gpu'; + return a + 1; + }; + + console.log(fn); + `; + + dualTest(code, (metadata, expected) => { + const astPath = metadata[0]?.function?.astPath; + + const propertyNode = astPath?.parentPath?.node; + + expect(astPath?.isObjectExpression()).toBe(true); + expect(propertyNode).toMatchObject({ type: 'ObjectProperty', key: { name: 'ast' } }); + expect(astPath?.evaluate().value).toStrictEqual({ + params: expected[0]?.params, + body: expected[0]?.body, + }); + }); + }); + describe('parses externals', () => { const code = `\ const noExternals = () => { @@ -716,19 +757,27 @@ describe('getEmbeddedTypegpuMetadata', () => { }); describe('recognizes embedded TypeGPU metadata rewritten by other tools', () => { - const fn = `()=>{"use gpu";let a=1;let b=2;return a+b}`; + // A body that the plugin would fail to transpile (multiple declarations in one statement), + // so re-transforming it would throw. + const fn = `()=>{"use gpu";let a=1,b=2;return a+b}`; const meta = `{v:2,name:"fn",ast:{params:[],body:[0,[]]},externals:{}}`; const variants: Record = { 'original (??=)': `const fn=/*#__PURE__*/($=>(globalThis.__TYPEGPU_META__??=new WeakMap()).set($.f=${fn},${meta})&&$.f)({});`, 'esbuild es2020 (?? + =)': `const fn=(e=>(globalThis.__TYPEGPU_META__??(globalThis.__TYPEGPU_META__=new WeakMap)).set(e.f=${fn},${meta})&&e.f)({});`, 'esbuild es2019 (!= null ?:)': `const fn=(e=>{var a;return((a=globalThis.__TYPEGPU_META__)!=null?a:globalThis.__TYPEGPU_META__=new WeakMap).set(e.f=${fn},${meta})&&e.f})({});`, - 'babel preset-env (!== null && !== void 0 ?:)': `var fn=function($,_g){return((_g=globalThis.__TYPEGPU_META__)!==null&&_g!==void 0?_g:globalThis.__TYPEGPU_META__=new WeakMap()).set($.f=function(){"use gpu";let a=1;let b=2;return a+b},${meta})&&$.f}({});`, + 'babel preset-env (!== null && !== void 0 ?:)': `var fn=function($,_g){return((_g=globalThis.__TYPEGPU_META__)!==null&&_g!==void 0?_g:globalThis.__TYPEGPU_META__=new WeakMap()).set($.f=function(){"use gpu";let a=1,b=2;return a+b},${meta})&&$.f}({});`, 'computed member access': `const fn=(e=>(globalThis["__TYPEGPU_META__"]??=new WeakMap).set((e.f=(${fn})),${meta})&&e.f)({});`, }; describe.each(Object.entries(variants))('%s', (_label, code) => { - dualTest(`${code}\nconsole.log(fn);`, (metadata, _, output) => { + test('[BABEL]', () => { + const metadata: EmbeddedTypegpuMetadata[] = []; + + const output = babelTransform(`${code}\nconsole.log(fn);`, undefined, [ + createBabelMetadataCollector(metadata), + ]); + expect(metadata[0]).toMatchObject({ v: 2, name: 'fn', @@ -740,6 +789,25 @@ describe('recognizes embedded TypeGPU metadata rewritten by other tools', () => code.match(/__TYPEGPU_META__/g)?.length, ); }); + + test('[ROLLUP]', async () => { + const metadata: EmbeddedTypegpuMetadata[] = []; + + const output = await rollupTransform(`${code}\nconsole.log(fn);`, undefined, [ + createRollupMetadataCollector(metadata), + ]); + + expect(metadata[0]).toMatchObject({ + v: 2, + name: 'fn', + function: { + ast: { params: [], body: [0, []] }, + }, + }); + expect(output.match(/__TYPEGPU_META__/g)?.length).toBe( + code.match(/__TYPEGPU_META__/g)?.length, + ); + }); }); test('does not treat unrelated .set calls as embedded metadata', () => { From 5208e45102d4ae4f9de434fd4a3ed0018fedd97b Mon Sep 17 00:00:00 2001 From: Szymon Szulc Date: Thu, 29 Oct 2026 14:08:35 +0100 Subject: [PATCH 17/18] fix unwrapping metadata --- .../src/core/embeddedMetadata.ts | 18 ++++++++++++++++-- .../test/embedded-metadata.test.ts | 19 +++++++++++++++++++ 2 files changed, 35 insertions(+), 2 deletions(-) diff --git a/packages/unplugin-typegpu/src/core/embeddedMetadata.ts b/packages/unplugin-typegpu/src/core/embeddedMetadata.ts index 95eae929dd..7f100b4842 100644 --- a/packages/unplugin-typegpu/src/core/embeddedMetadata.ts +++ b/packages/unplugin-typegpu/src/core/embeddedMetadata.ts @@ -129,6 +129,20 @@ function skipTransparentWrappersUp(path: NodePath): NodePath { return current; } +/** + * Walks down from the given path through parentheses and type-only wrappers, + * returning the innermost wrapped expression. + */ +function skipTransparentWrappersDown(path: NodePath): NodePath { + let current = path; + + while (isTransparentWrapper(current.node)) { + current = current.get('expression') as NodePath; + } + + return current; +} + /** * Returns the value path of the first matching property. If there is no match, returns `undefined`. */ @@ -242,12 +256,12 @@ export function getEmbeddedTypegpuMetadata( } // we check for the metadata object - const metadataPath = callPath.get('arguments.1') as NodePath; + const metadataPath = callPath.get('arguments.1'); if (metadataPath === undefined) { return undefined; } - const unwrappedMetadataPath = skipTransparentWrappersUp(metadataPath); + const unwrappedMetadataPath = skipTransparentWrappersDown(metadataPath as NodePath); if (!unwrappedMetadataPath.isObjectExpression()) { return undefined; } diff --git a/packages/unplugin-typegpu/test/embedded-metadata.test.ts b/packages/unplugin-typegpu/test/embedded-metadata.test.ts index db63d4ae05..cc331755b8 100644 --- a/packages/unplugin-typegpu/test/embedded-metadata.test.ts +++ b/packages/unplugin-typegpu/test/embedded-metadata.test.ts @@ -768,6 +768,7 @@ describe('recognizes embedded TypeGPU metadata rewritten by other tools', () => 'esbuild es2019 (!= null ?:)': `const fn=(e=>{var a;return((a=globalThis.__TYPEGPU_META__)!=null?a:globalThis.__TYPEGPU_META__=new WeakMap).set(e.f=${fn},${meta})&&e.f})({});`, 'babel preset-env (!== null && !== void 0 ?:)': `var fn=function($,_g){return((_g=globalThis.__TYPEGPU_META__)!==null&&_g!==void 0?_g:globalThis.__TYPEGPU_META__=new WeakMap()).set($.f=function(){"use gpu";let a=1,b=2;return a+b},${meta})&&$.f}({});`, 'computed member access': `const fn=(e=>(globalThis["__TYPEGPU_META__"]??=new WeakMap).set((e.f=(${fn})),${meta})&&e.f)({});`, + 'parenthesized metadata': `const fn=(e=>(globalThis.__TYPEGPU_META__??=new WeakMap).set(e.f=${fn},(${meta}))&&e.f)({});`, }; describe.each(Object.entries(variants))('%s', (_label, code) => { @@ -810,6 +811,24 @@ describe('recognizes embedded TypeGPU metadata rewritten by other tools', () => }); }); + test('type-wrapped metadata', () => { + const metadata: EmbeddedTypegpuMetadata[] = []; + const code = `const fn=(e=>(globalThis.__TYPEGPU_META__??=new WeakMap).set(e.f=${fn},(${meta} as any)!)&&e.f)({});`; + + const output = babelTransform(`${code}\nconsole.log(fn);`, undefined, [ + createBabelMetadataCollector(metadata), + ]); + + expect(metadata[0]).toMatchObject({ + v: 2, + name: 'fn', + function: { + ast: { params: [], body: [0, []] }, + }, + }); + expect(output?.match(/__TYPEGPU_META__/g)?.length).toBe(1); + }); + test('does not treat unrelated .set calls as embedded metadata', () => { const output = babelTransform( `const m = new WeakMap(); const fn = ($ => m.set($.f = () => { 'use gpu'; }, ${meta}) && $.f)({});`, From 99777d949485f8fe7f96652f58d76e3b67bb6008 Mon Sep 17 00:00:00 2001 From: Szymon Szulc Date: Fri, 30 Oct 2026 15:41:25 +0100 Subject: [PATCH 18/18] lowering arrow functions --- .../src/core/embeddedMetadata.ts | 2 +- .../test/embedded-metadata.test.ts | 19 +++++++++++++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/packages/unplugin-typegpu/src/core/embeddedMetadata.ts b/packages/unplugin-typegpu/src/core/embeddedMetadata.ts index 7f100b4842..5651bf1858 100644 --- a/packages/unplugin-typegpu/src/core/embeddedMetadata.ts +++ b/packages/unplugin-typegpu/src/core/embeddedMetadata.ts @@ -203,7 +203,7 @@ function parseExternalsPath( ? key.name : undefined; - if (name === undefined || !propertyPath.get('value').isArrowFunctionExpression()) { + if (name === undefined) { return undefined; } diff --git a/packages/unplugin-typegpu/test/embedded-metadata.test.ts b/packages/unplugin-typegpu/test/embedded-metadata.test.ts index cc331755b8..ae42df4c44 100644 --- a/packages/unplugin-typegpu/test/embedded-metadata.test.ts +++ b/packages/unplugin-typegpu/test/embedded-metadata.test.ts @@ -756,6 +756,25 @@ describe('getEmbeddedTypegpuMetadata', () => { }); }); +describe('recognizes external getters downleveled by transform-arrow-functions plugin', () => { + const code = `\ + export const fn = () => { + 'use gpu'; + return ext; + }; + `; + + const lowered = babelTransform(code, undefined, ['transform-arrow-functions']) as string; + + dualTest(lowered, (metadata, expected) => { + expect(extractExternalGetters(metadata)).toStrictEqual( + expected.map((e) => + [...e.externalNames].map(([key, value]) => [key, `{\n return ${value};\n}`]), + ), + ); + }); +}); + describe('recognizes embedded TypeGPU metadata rewritten by other tools', () => { // A body that the plugin would fail to transpile (multiple declarations in one statement), // so re-transforming it would throw.