diff --git a/packages/unplugin-typegpu/package.json b/packages/unplugin-typegpu/package.json index 4b9c9d342c..9a089f3814 100644 --- a/packages/unplugin-typegpu/package.json +++ b/packages/unplugin-typegpu/package.json @@ -345,9 +345,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/babel.ts b/packages/unplugin-typegpu/src/babel.ts index e6529dd1b4..733cfb41a7 100644 --- a/packages/unplugin-typegpu/src/babel.ts +++ b/packages/unplugin-typegpu/src/babel.ts @@ -2,8 +2,8 @@ import type { NodePath, TraverseOptions } from '@babel/traverse'; import defu from 'defu'; import type { Externals, TranspilationResult } from 'tinyest-for-wgsl'; import * as t from '@babel/types'; +import { METADATA_FORMAT_VERSION } from './core/version.ts'; import { - METADATA_FORMAT_VERSION, type PluginState, checkOpts, defaultOptions, diff --git a/packages/unplugin-typegpu/src/core/common.ts b/packages/unplugin-typegpu/src/core/common.ts index e829e84fed..d6e449cfd9 100644 --- a/packages/unplugin-typegpu/src/core/common.ts +++ b/packages/unplugin-typegpu/src/core/common.ts @@ -6,12 +6,6 @@ import { transpileBabelFn, type TranspilationResult } from 'tinyest-for-wgsl'; import { getEmbeddedTypegpuMetadata } from './embeddedMetadata.ts'; import { obfuscate } from './obfuscate.ts'; -/** - * 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; - export interface Options { /** @default [/\.m?[jt]sx?$/] */ include?: FilterPattern; diff --git a/packages/unplugin-typegpu/src/core/embeddedMetadata.ts b/packages/unplugin-typegpu/src/core/embeddedMetadata.ts index 4108ab30c5..5651bf1858 100644 --- a/packages/unplugin-typegpu/src/core/embeddedMetadata.ts +++ b/packages/unplugin-typegpu/src/core/embeddedMetadata.ts @@ -1,12 +1,24 @@ 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 { METADATA_FORMAT_VERSION } from './version.ts'; +import type { MetadatableFunction } from './common.ts'; -export interface EmbeddedTypegpuMetadata { +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; name: string | undefined; - // TODO: parse AST and externals -} + function?: { + ast: FunctionAst; + astPath: NodePath; + externals: FunctionExternals; + }; +}; const embeddedTypegpuMetadataCache = new WeakMap< NodePath, @@ -118,46 +130,100 @@ function skipTransparentWrappersUp(path: NodePath): NodePath { } /** - * Returns the node of the property with the given name in the object, if it exists. + * 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`. */ -function objectPropertyValue(object: t.ObjectExpression, expectedName: string): t.Node | undefined { - for (const property of object.properties) { - if (!t.isObjectProperty(property) || property.computed) { +function objectPropertyPath( + objectPath: NodePath, + name: string, +): NodePath | undefined { + for (const property of objectPath.get('properties')) { + if (!property.isObjectProperty()) { continue; } - const name = t.isIdentifier(property.key) - ? property.key.name - : t.isStringLiteral(property.key) - ? property.key.value - : undefined; + const { key, computed } = property.node; - if (name === expectedName) { - return property.value; + const matches = + (!computed && t.isIdentifier(key, { name })) || t.isStringLiteral(key, { value: name }); + + if (!matches) { + continue; + } + + const value = property.get('value'); + if (value.isExpression()) { + return value; } } return undefined; } +/** + * Returns the parsed AST + */ +function parseAstPath(astPath: NodePath): FunctionAst | undefined { + const evaluated = astPath.evaluate(); + + 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) { + return undefined; + } + + externals.set(name, propertyPath); + } + + return externals; +} /** * 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 - * const f = ($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = () => { - * 'use gpu'; - * }, { - * v: 2, - * name: "f", - * ast: { - * params: [], - * body: [0, []] - * }, - * externals: {} - * }) && $.f)({}); - * ``` + * 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 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, @@ -190,31 +256,69 @@ export function getEmbeddedTypegpuMetadata( } // we check for the metadata object - const metadataNode = callPath.node.arguments[1]; - if (!t.isObjectExpression(metadataNode)) { + const metadataPath = callPath.get('arguments.1'); + if (metadataPath === undefined) { + return undefined; + } + + const unwrappedMetadataPath = skipTransparentWrappersDown(metadataPath as NodePath); + if (!unwrappedMetadataPath.isObjectExpression()) { return undefined; } // get the metadata properties - const versionNode = objectPropertyValue(metadataNode, 'v'); - const nameNode = objectPropertyValue(metadataNode, 'name'); - const astNode = objectPropertyValue(metadataNode, 'ast'); - const externalsNode = objectPropertyValue(metadataNode, 'externals'); + const versionPath = objectPropertyPath(unwrappedMetadataPath, 'v'); + const namePath = objectPropertyPath(unwrappedMetadataPath, 'name'); - if ( - !t.isNumericLiteral(versionNode) || - nameNode === undefined || - astNode === undefined || - externalsNode === undefined - ) { - return undefined; + if (versionPath === undefined || namePath === undefined) { + throw new Error(METADATA_PARSING_ERROR_MESSAGE); + } + + const versionResult = versionPath.evaluate(); + const nameResult = namePath.evaluate(); + + if (!versionResult.confident || !nameResult.confident) { + throw new Error(METADATA_PARSING_ERROR_MESSAGE); + } + + const version = versionResult.value as number; + const name = nameResult.value as string | undefined; + + if (version !== METADATA_FORMAT_VERSION) { + const embeddedTypegpuMetadata = { + v: version, + name, + }; + embeddedTypegpuMetadataCache.set(path, embeddedTypegpuMetadata); + return embeddedTypegpuMetadata; } - const name = t.isStringLiteral(nameNode) ? nameNode.value : undefined; + const astPath = objectPropertyPath(unwrappedMetadataPath, 'ast'); + const externalsPath = objectPropertyPath(unwrappedMetadataPath, 'externals'); + + if (astPath === undefined || externalsPath === undefined) { + throw new Error(METADATA_PARSING_ERROR_MESSAGE); + } + + if (!astPath.isObjectExpression() || !externalsPath.isObjectExpression()) { + throw new Error(METADATA_PARSING_ERROR_MESSAGE); + } + + const ast = parseAstPath(astPath); + const externals = parseExternalsPath(externalsPath); + + if (ast === undefined || externals === undefined) { + throw new Error(METADATA_PARSING_ERROR_MESSAGE); + } const embeddedTypegpuMetadata = { - v: versionNode.value, + v: version, name, + function: { + ast, + astPath, + externals, + }, }; embeddedTypegpuMetadataCache.set(path, embeddedTypegpuMetadata); diff --git a/packages/unplugin-typegpu/src/core/factory.ts b/packages/unplugin-typegpu/src/core/factory.ts index 8c1a424dfa..5a4445fdb9 100644 --- a/packages/unplugin-typegpu/src/core/factory.ts +++ b/packages/unplugin-typegpu/src/core/factory.ts @@ -6,13 +6,13 @@ import _traverse, { type NodePath } from '@babel/traverse'; import type { Externals, TranspilationResult } from 'tinyest-for-wgsl'; import * as parser from '@babel/parser'; import * as t from '@babel/types'; +import { METADATA_FORMAT_VERSION } from './version.ts'; import { defaultOptions, earlyPruneRegex, initPluginState, functionVisitor, getBlockScope, - METADATA_FORMAT_VERSION, checkOpts, } from './common.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; diff --git a/packages/unplugin-typegpu/test/embedded-metadata.test.ts b/packages/unplugin-typegpu/test/embedded-metadata.test.ts index 1cf8487cf8..ae42df4c44 100644 --- a/packages/unplugin-typegpu/test/embedded-metadata.test.ts +++ b/packages/unplugin-typegpu/test/embedded-metadata.test.ts @@ -1,9 +1,11 @@ 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'; +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 type { MetadatableFunction } from '../src/core/common.ts'; import { type EmbeddedTypegpuMetadata, getEmbeddedTypegpuMetadata, @@ -14,6 +16,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[], @@ -45,7 +72,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, // to catch cases where `unwrapParentheses` or `parentPathSkippingParentheses` are needed + }); traverse(ast, { ArrowFunctionExpression(path) { @@ -64,42 +94,683 @@ function createRollupMetadataCollector(metadata: EmbeddedTypegpuMetadata[]): Plu }; } -describe('collects embedded TypeGPU metadata', () => { - const code = `\ - const fn1 = () => { - 'use gpu'; - }; +/** + * Traverses the program AST to extract the `TranspilationResult` from functions containing 'use gpu' directive. + * + * @note For simplicity, shelled functions are omitted. + */ +function extractTranspilationResultFromSource(code: string) { + const ast = parser.parse(code, { + sourceType: 'module', + plugins: ['typescript'], + }); - const fn2 = () => { - 'use gpu'; - 'worklet'; - }; + const expected: TranspilationResult[] = []; - const fn3 = () => { - 'worklet'; - 'use gpu'; - }; + function collect(path: NodePath) { + const body = path.node.body; - console.log(fn1, fn2, fn3); - `; + if ( + body.type !== 'BlockStatement' || + !body.directives.some((directive) => directive.value.value === 'use gpu') + ) { + return; + } + + const result = transpileBabelFn(path.node); + expected.push(result); + } + + traverse(ast, { + ArrowFunctionExpression: collect, + FunctionExpression: collect, + FunctionDeclaration: collect, + }); + + 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, +) { + test('[BABEL]', () => { + const expected = extractTranspilationResultFromSource(code); - test('babel', () => { const metadata: EmbeddedTypegpuMetadata[] = []; + void babelTransform(code, {}, [createBabelMetadataCollector(metadata)]); + check(metadata, expected); + }); + + test('[ROLLUP]', async () => { + const expected = extractTranspilationResultFromSource(code); + + const metadata: EmbeddedTypegpuMetadata[] = []; + void (await rollupTransform(code, undefined, [createRollupMetadataCollector(metadata)])); + check(metadata, expected); + }); +} - babelTransform(code, undefined, [createBabelMetadataCollector(metadata)]); +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)({}); - expect(JSON.stringify(metadata)).toMatchInlineSnapshot( - `"[{"v":2,"name":"fn1"},{"v":2,"name":"fn2"},{"v":2,"name":"fn3"}]"`, + console.log(fn); + `, + { sourceType: 'module' }, ); + + const functionPath = getFirstUseGpuFunctionPath(ast); + const firstRead = getEmbeddedTypegpuMetadata(functionPath); + expect(firstRead).toBeDefined(); + expect(getEmbeddedTypegpuMetadata(functionPath)).toBe(firstRead); }); - test('rollup', async () => { - const metadata: EmbeddedTypegpuMetadata[] = []; + 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: {} }", + ], + [ + '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 = () => { + '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 = () => { + 'use gpu'; + return 1; + }, { + v: 1, + name: 'fn', + }) && $.f)({}); - await rollupTransform(code, undefined, [createRollupMetadataCollector(metadata)]); + console.log(fn); + `; + + dualTest(code, (metadata) => { + expect(metadata).toStrictEqual([{ v: 1, name: 'fn' }]); + }); + }); + + describe('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, expected) => { + expect(metadata.map((m) => m.function?.ast.params)).toStrictEqual( + expected.map((e) => e.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, expected) => { + expect(metadata.map((m) => m.function?.ast.body)).toStrictEqual(expected.map((e) => e.body)); + }); + }); + + describe('parses blocks', () => { + const code = `\ + const fn = (a) => { + 'use gpu'; + let value = a; + { + value += 1; + { + value += 2; + } + } + }; + + console.log(fn); + `; + + dualTest(code, (metadata, expected) => { + expect(metadata.map((m) => m.function?.ast.body)).toStrictEqual(expected.map((e) => e.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, expected) => { + expect(metadata.map((m) => m.function?.ast.body)).toStrictEqual(expected.map((e) => e.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, expected) => { + expect(metadata.map((m) => m.function?.ast.body)).toStrictEqual(expected.map((e) => e.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, expected) => { + expect(metadata.map((m) => m.function?.ast.body)).toStrictEqual(expected.map((e) => e.body)); + }); + }); + + describe('parses unary expressions', () => { + const code = `\ + const fn = (a) => { + 'use gpu'; + const negation = !a; + }; + + console.log(fn); + `; + + dualTest(code, (metadata, expected) => { + expect(metadata.map((m) => m.function?.ast.body)).toStrictEqual(expected.map((e) => e.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, expected) => { + expect(metadata.map((m) => m.function?.ast.body)).toStrictEqual(expected.map((e) => e.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, expected) => { + expect(metadata.map((m) => m.function?.ast.body)).toStrictEqual(expected.map((e) => e.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, expected) => { + expect(metadata.map((m) => m.function?.ast.body)).toStrictEqual(expected.map((e) => e.body)); + }); + }); - expect(JSON.stringify(metadata)).toMatchInlineSnapshot( - `"[{"v":2,"name":"fn1"},{"v":2,"name":"fn2"},{"v":2,"name":"fn3"}]"`, + 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, expected) => { + expect(metadata.map((m) => m.function?.ast.body)).toStrictEqual(expected.map((e) => e.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, expected) => { + expect(metadata.map((m) => m.function?.ast.body)).toStrictEqual(expected.map((e) => e.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, expected) => { + expect(metadata.map((m) => m.function?.ast.body)).toStrictEqual(expected.map((e) => e.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, expected) => { + expect(metadata.map((m) => m.function?.ast.body)).toStrictEqual(expected.map((e) => e.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, expected) => { + expect(metadata.map((m) => m.function?.ast.body)).toStrictEqual(expected.map((e) => e.body)); + }); + }); + + describe('parses post-update expressions', () => { + const code = `\ + const fn = () => { + 'use gpu'; + let i = 0; + i++; + i--; + }; + + console.log(fn); + `; + + dualTest(code, (metadata, expected) => { + expect(metadata.map((m) => m.function?.ast.body)).toStrictEqual(expected.map((e) => e.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, expected) => { + expect(metadata.map((m) => m.function?.ast.body)).toStrictEqual(expected.map((e) => e.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, expected) => { + expect(metadata.map((m) => m.function?.ast.body)).toStrictEqual(expected.map((e) => e.body)); + }); + }); + + describe('parses conditional expressions', () => { + const code = `\ + const fn = (a, b) => { + 'use gpu'; + return a ? b : 1; + }; + + console.log(fn); + `; + + dualTest(code, (metadata, expected) => { + expect(metadata.map((m) => m.function?.ast.body)).toStrictEqual(expected.map((e) => e.body)); + }); + }); + + 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 = () => { + '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]), + ); + }); + }); + + 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']); + }); + }); +}); + +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}`]), + ), ); }); }); @@ -116,36 +787,67 @@ 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) => { - test('babel', () => { + test('[BABEL]', () => { const metadata: EmbeddedTypegpuMetadata[] = []; const output = babelTransform(`${code}\nconsole.log(fn);`, undefined, [ createBabelMetadataCollector(metadata), ]); - expect(metadata).toStrictEqual([{ v: 2, name: 'fn' }]); + 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('rollup', async () => { + test('[ROLLUP]', async () => { const metadata: EmbeddedTypegpuMetadata[] = []; const output = await rollupTransform(`${code}\nconsole.log(fn);`, undefined, [ createRollupMetadataCollector(metadata), ]); - expect(metadata).toStrictEqual([{ v: 2, name: 'fn' }]); + 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('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)({});`, diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index dc72149561..0ae22d469a 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -166,7 +166,7 @@ importers: devDependencies: '@types/bun': specifier: latest - version: 1.4.3 + version: 1.4.4 apps/infra-benchmarks: devDependencies: @@ -1099,6 +1099,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 @@ -1108,6 +1111,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 @@ -1341,6 +1347,10 @@ packages: resolution: {integrity: sha512-CGOfOJqWjg2qW/Mb6zNsDm+u5vFQ8DxXfbM09z69p5Z6+mE1ikP2jUXw+j42Pf1XTYED2Rni5f95npYeuwMDQA==} engines: {node: '>=6.9.0'} + '@babel/generator@7.29.7': + resolution: {integrity: sha512-DkXD5OJQaAQIdZ1bt3UZdEnHAn9Imd3IVBdX03UFe+ony9Ojw5pzr9YVKGDY1jt+Gcn/FnGkNf8r+Vj5NOJWtQ==} + engines: {node: '>=6.9.0'} + '@babel/generator@7.29.8': resolution: {integrity: sha512-gZbepsdh3WDtgZKWL+vTPh71LSBrm/Y4/QDZBVCcYfmeTEEuoOYwlSy+G1StfJg+/Zy550u/3TATbm7qDbbMtg==} engines: {node: '>=6.9.0'} @@ -4576,6 +4586,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==} @@ -4591,6 +4604,9 @@ packages: '@types/bun@1.4.3': resolution: {integrity: sha512-N67Pqy8JXbAY/j0tNraS3V5mubGvc6rrrNDmDCtZLlFzccbUJPDb0tvgUBC5q4m8ybkP93Jcn1QB+RYQ6PvBhw==} + '@types/bun@1.4.4': + resolution: {integrity: sha512-9wkRDXUXK124pEnAutE6n/l18j+i9GC5IAHK8TmX/2nlySETmky0OU6yL5U8pW1AoKh++AN9auDCr647BfurjQ==} + '@types/chai@5.2.3': resolution: {integrity: sha512-Mw558oeA9fFbv65/y4mHtXDs9bPnFMZAL/jxdPFUpOHHIXX91mcgEHbS5Lahr+pwZFR8A7GQleRWeI6cGFC2UA==} @@ -5285,6 +5301,9 @@ packages: bun-types@1.4.3: resolution: {integrity: sha512-tinRJNGOmUN+tEAsn2y2e942n9JHRHR42iUYzDH75NmhJqVBokGSQkkfSFdNWwsANIdrK6ACH9t01M452YyTwg==} + bun-types@1.4.4: + resolution: {integrity: sha512-yjrJh5oGX7mMPKGd19YXyIjmEuj4G4zNC30bs7PM9Az7NQ3ndF3j6hbeRdvJTRQEhJ/FYQwW0a2moweEi2hMxQ==} + bun@1.3.10: resolution: {integrity: sha512-S/CXaXXIyA4CMjdMkYQ4T2YMqnAn4s0ysD3mlsY4bUiOCqGlv28zck4Wd4H4kpvbekx15S9mUeLQ7Uxd0tYTLA==} cpu: [arm64, x64] @@ -10165,6 +10184,14 @@ snapshots: transitivePeerDependencies: - supports-color + '@babel/generator@7.29.7': + dependencies: + '@babel/parser': 7.29.9 + '@babel/types': 7.29.8 + '@jridgewell/gen-mapping': 0.3.13 + '@jridgewell/trace-mapping': 0.3.31 + jsesc: 3.1.0 + '@babel/generator@7.29.8': dependencies: '@babel/parser': 7.29.9 @@ -12822,6 +12849,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.8 + '@types/babel__generator@7.6.8': dependencies: '@babel/types': 7.29.8 @@ -12848,6 +12879,10 @@ snapshots: dependencies: bun-types: 1.4.3 + '@types/bun@1.4.4': + dependencies: + bun-types: 1.4.4 + '@types/chai@5.2.3': dependencies: '@types/deep-eql': 4.0.2 @@ -13783,6 +13818,11 @@ snapshots: '@types/node': 24.10.0 undici-types: 7.16.0 + bun-types@1.4.4: + dependencies: + '@types/node': 24.10.0 + undici-types: 7.16.0 + bun@1.3.10: optionalDependencies: '@oven/bun-darwin-aarch64': 1.3.10