diff --git a/packages/tinyest-for-wgsl/src/parsers.ts b/packages/tinyest-for-wgsl/src/parsers.ts index d69ef7e622..b0b889b244 100644 --- a/packages/tinyest-for-wgsl/src/parsers.ts +++ b/packages/tinyest-for-wgsl/src/parsers.ts @@ -16,7 +16,7 @@ import { extractFunctionParts } from './functionParts.ts'; const { NodeTypeCatalog: NODE } = tinyest; -function createContext(params: tinyest.FuncParameter[]): Context { +function createContext(params: tinyest.FuncParameter[], opts: TranspilationOptions): Context { return { externalNames: new Map(), ignoreExternalDepth: 0, @@ -30,6 +30,7 @@ function createContext(params: tinyest.FuncParameter[]): Context { ), }, ], + opts, }; } @@ -51,6 +52,9 @@ function createParser(ast: AstKind) { const externalChain = tryFindExternalChain(ctx, node); if (externalChain) { ctx.externalNames.set(externalChain, externalChain); + if (ctx.opts.verboseNodes) { + return [NODE.identifier, externalChain]; + } return externalChain; } } @@ -60,9 +64,9 @@ function createParser(ast: AstKind) { }; return { - transpileFn(rootNode: JsNode): TranspilationResult { + transpileFn(rootNode: JsNode, options: TranspilationOptions): TranspilationResult { const { params, body } = extractFunctionParts(rootNode); - const ctx = createContext(params); + const ctx = createContext(params, options); const tinyestBody = transpile(ctx, body); @@ -81,8 +85,8 @@ function createParser(ast: AstKind) { }; }, - transpileNode(node: JsNode): tinyest.AnyNode { - return transpile(createContext([]), node); + transpileNode(node: JsNode, options: TranspilationOptions): tinyest.AnyNode { + return transpile(createContext([], options), node); }, }; } @@ -100,8 +104,8 @@ export function transpileFn( rootNode: babel.Node, options: TranspilationOptions<'babel'>, ): TranspilationResult; -export function transpileFn(rootNode: JsNode, { ast }: TranspilationOptions): TranspilationResult { - return parsers[ast].transpileFn(rootNode); +export function transpileFn(rootNode: JsNode, opts: TranspilationOptions): TranspilationResult { + return parsers[opts.ast].transpileFn(rootNode, opts); } export function transpileNode( @@ -112,6 +116,6 @@ export function transpileNode( rootNode: babel.Node, options: TranspilationOptions<'babel'>, ): tinyest.AnyNode; -export function transpileNode(rootNode: JsNode, { ast }: TranspilationOptions): tinyest.AnyNode { - return parsers[ast].transpileNode(rootNode); +export function transpileNode(rootNode: JsNode, opts: TranspilationOptions): tinyest.AnyNode { + return parsers[opts.ast].transpileNode(rootNode, opts); } diff --git a/packages/tinyest-for-wgsl/src/transpilers.ts b/packages/tinyest-for-wgsl/src/transpilers.ts index 719684dddc..9b63d133e5 100644 --- a/packages/tinyest-for-wgsl/src/transpilers.ts +++ b/packages/tinyest-for-wgsl/src/transpilers.ts @@ -45,7 +45,10 @@ export const baseTranspilers = { : [NODE.return]; }, - Identifier(_ctx, node) { + Identifier(ctx, node) { + if (ctx.opts.verboseNodes) { + return [NODE.identifier, node.name]; + } return node.name; }, @@ -92,13 +95,9 @@ export const baseTranspilers = { // If the property is not computed, we don't want to register identifiers as external. ctx.ignoreExternalDepth++; - const property = transpile(ctx, node.property) as tinyest.Expression; + const property = transpile(ctx, node.property) as tinyest.Identifier; ctx.ignoreExternalDepth--; - if (typeof property !== 'string') { - throw new Error('Expected identifier as property access key.'); - } - return [NODE.memberAccess, object, property]; }, @@ -147,14 +146,10 @@ export const baseTranspilers = { const decl = node.declarations[0]; ctx.ignoreExternalDepth++; - const id = transpile(ctx, decl.id); + const id = transpile(ctx, decl.id) as tinyest.Identifier; ctx.ignoreExternalDepth--; - if (typeof id !== 'string') { - throw new Error('Invalid variable declaration, expected identifier.'); - } - - ctx.stack[ctx.stack.length - 1]?.declaredNames.push(id); + ctx.stack[ctx.stack.length - 1]?.declaredNames.push(extractId(id)); const init = decl.init ? (transpile(ctx, decl.init) as tinyest.Expression) : undefined; @@ -221,7 +216,7 @@ export const baseTranspilers = { } satisfies Pick, SharedTranspilers>; const acornSpecificTranspilers = { - Literal(_ctx, node) { + Literal(ctx, node) { if (node.regex) { throw new Error('Regular expression literals are not representable in WGSL.'); } @@ -229,6 +224,9 @@ const acornSpecificTranspilers = { return [NODE.nullLiteral]; } if (typeof node.value === 'boolean') { + if (ctx.opts.verboseNodes) { + return [NODE.booleanLiteral, node.value]; + } return node.value; } if (typeof node.value === 'string') { @@ -298,7 +296,10 @@ const babelSpecificTranspilers = { return [NODE.numericLiteral, String(Number(node.value))]; }, - BooleanLiteral(_ctx, node) { + BooleanLiteral(ctx, node) { + if (ctx.opts.verboseNodes) { + return [NODE.booleanLiteral, node.value]; + } return node.value; }, @@ -362,3 +363,10 @@ export const babelTranspilers = { ...(baseTranspilers as Pick, SharedTranspilers>), ...babelSpecificTranspilers, } satisfies Transpilers; + +function extractId(ident: tinyest.Identifier): string { + if (typeof ident === 'string') { + return ident; + } + return ident[1]; +} diff --git a/packages/tinyest-for-wgsl/src/types.ts b/packages/tinyest-for-wgsl/src/types.ts index fc812258c8..837aad21f8 100644 --- a/packages/tinyest-for-wgsl/src/types.ts +++ b/packages/tinyest-for-wgsl/src/types.ts @@ -22,6 +22,7 @@ export type Context = { */ visitedNodes: Set; stack: Scope[]; + opts: TranspilationOptions; }; export type TranspilationResult = { @@ -50,4 +51,11 @@ export type AstKind = 'acorn' | 'babel'; export type TranspilationOptions = { ast: TAst; + /** + * With this option enabled, identifiers and boolean literals will be wrapped + * in dedicated nodes, instead of being transpiled as string/boolean. + * + * @default false + */ + verboseNodes?: boolean; }; diff --git a/packages/tinyest-for-wgsl/tests/helpers.ts b/packages/tinyest-for-wgsl/tests/helpers.ts index 745b1b8977..ac909fb0c3 100644 --- a/packages/tinyest-for-wgsl/tests/helpers.ts +++ b/packages/tinyest-for-wgsl/tests/helpers.ts @@ -1,7 +1,7 @@ import babel from '@babel/parser'; import type { Node } from '@babel/types'; import * as acorn from 'acorn'; -import { transpileFn, type TranspilationResult } from 'tinyest-for-wgsl'; +import { transpileFn, type TranspilationOptions, type TranspilationResult } from 'tinyest-for-wgsl'; export const parseRollup = (code: string) => acorn.parse(code, { ecmaVersion: 'latest' }); export const parseBabel = (code: string) => @@ -10,11 +10,15 @@ export const parseBabel = (code: string) => export function dualTest( test: ( p: (code: string) => TNode, - transpileFn: (node: TNode) => TranspilationResult, + transpileFn: (node: TNode, options?: Partial) => TranspilationResult, ) => void, ) { return () => { - test(parseBabel, (node) => transpileFn(node, { ast: 'babel' })); - test(parseRollup, (node) => transpileFn(node, { ast: 'acorn' })); + test(parseBabel, (node, options) => + transpileFn(node, { ast: 'babel', ...options } as TranspilationOptions<'babel'>), + ); + test(parseRollup, (node, options) => + transpileFn(node, { ast: 'acorn', ...options } as TranspilationOptions<'acorn'>), + ); }; } diff --git a/packages/tinyest-for-wgsl/tests/verboseNodes.test.ts b/packages/tinyest-for-wgsl/tests/verboseNodes.test.ts new file mode 100644 index 0000000000..efce4fb93f --- /dev/null +++ b/packages/tinyest-for-wgsl/tests/verboseNodes.test.ts @@ -0,0 +1,152 @@ +import type { ClassDeclaration, ClassProperty, Expression, Node } from '@babel/types'; +import * as acorn from 'acorn'; +import { describe, expect, it } from 'vitest'; +import { transpileFn } from '../src/parsers.ts'; +import { dualTest, parseBabel } from './helpers.ts'; + +describe('verbose nodes', () => { + it( + 'uses nodes for identifiers', + dualTest((p, transpileFn) => { + const { params, body, externalNames } = transpileFn( + p(`(a, b, c) => { + return a + b + c; + }`), + { verboseNodes: true }, + ); + + expect(params).toMatchInlineSnapshot(` + [ + { + "name": "a", + "type": "i", + }, + { + "name": "b", + "type": "i", + }, + { + "name": "c", + "type": "i", + }, + ] + `); + expect(JSON.stringify(body)).toMatchInlineSnapshot( + `"[0,[[10,[1,[1,[9,"a"],"+",[9,"b"]],"+",[9,"c"]]]]]"`, + ); + expect(externalNames).toMatchInlineSnapshot(`Map {}`); + }), + ); + + it( + 'uses nodes for boolean literals', + dualTest((p, transpileFn) => { + const { params, body, externalNames } = transpileFn( + p(`() => { + return true && false; + }`), + { verboseNodes: true }, + ); + + expect(params).toMatchInlineSnapshot(`[]`); + expect(JSON.stringify(body)).toMatchInlineSnapshot( + `"[0,[[10,[3,[107,true],"&&",[107,false]]]]]"`, + ); + expect(externalNames).toMatchInlineSnapshot(`Map {}`); + }), + ); + + it( + 'uses nodes for const declarations', + dualTest((p, transpileFn) => { + const { params, body, externalNames } = transpileFn( + p(`() => { + const a = 1; + }`), + { verboseNodes: true }, + ); + + expect(params).toMatchInlineSnapshot(`[]`); + expect(JSON.stringify(body)).toMatchInlineSnapshot(`"[0,[[13,[9,"a"],[5,"1"]]]]"`); + expect(externalNames).toMatchInlineSnapshot(`Map {}`); + }), + ); + + it( + 'uses nodes for let declarations', + dualTest((p, transpileFn) => { + const { params, body, externalNames } = transpileFn( + p(`() => { + let a = 1; + }`), + { verboseNodes: true }, + ); + + expect(params).toMatchInlineSnapshot(`[]`); + expect(JSON.stringify(body)).toMatchInlineSnapshot(`"[0,[[12,[9,"a"],[5,"1"]]]]"`); + expect(externalNames).toMatchInlineSnapshot(`Map {}`); + }), + ); + + it( + 'uses nodes for member expressions', + dualTest((p, transpileFn) => { + const { params, body, externalNames } = transpileFn( + p(`() => { + const o = {}; + return o.prop; + }`), + { verboseNodes: true }, + ); + + expect(params).toMatchInlineSnapshot(`[]`); + expect(JSON.stringify(body)).toMatchInlineSnapshot( + `"[0,[[13,[9,"o"],[104,{}]],[10,[7,[9,"o"],[9,"prop"]]]]]"`, + ); + expect(externalNames).toMatchInlineSnapshot(`Map {}`); + }), + ); + + it( + 'uses nodes for externals', + dualTest((p, transpileFn) => { + const { params, body, externalNames } = transpileFn( + p(`() => { + return ext + ext.prop; + }`), + { verboseNodes: true }, + ); + + expect(params).toMatchInlineSnapshot(`[]`); + expect(JSON.stringify(body)).toMatchInlineSnapshot( + `"[0,[[10,[1,[9,"ext"],"+",[9,"ext.prop"]]]]]"`, + ); + expect(externalNames).toMatchInlineSnapshot(` + Map { + "ext" => "ext", + "ext.prop" => "ext.prop", + } + `); + }), + ); + + it( + 'does not use nodes for object expressions', + dualTest((p, transpileFn) => { + const { params, body, externalNames } = transpileFn( + p(`() => { + return { p: ext }; + }`), + { verboseNodes: true }, + ); + + expect(params).toMatchInlineSnapshot(`[]`); + expect(JSON.stringify(body)).toMatchInlineSnapshot(`"[0,[[10,[104,{"p":[9,"ext"]}]]]]"`); + expect(externalNames).toMatchInlineSnapshot(` + Map { + "ext" => "ext", + } + `); + }), + ); +}); diff --git a/packages/tinyest/src/nodes.ts b/packages/tinyest/src/nodes.ts index 942afe4d78..de25e80bfe 100644 --- a/packages/tinyest/src/nodes.ts +++ b/packages/tinyest/src/nodes.ts @@ -13,6 +13,7 @@ export const NodeTypeCatalog = { call: 6, memberAccess: 7, indexAccess: 8, + identifier: 9, // regular return: 10, @@ -33,10 +34,14 @@ export const NodeTypeCatalog = { objectExpr: 104, conditionalExpr: 105, nullLiteral: 106, + booleanLiteral: 107, } as const; export type NodeTypeCatalog = typeof NodeTypeCatalog; +/** Identifier can either be encoded as a node, or as a plain string */ +export type Identifier = string | readonly [type: NodeTypeCatalog['identifier'], id: string]; + /** * Represents a return statement */ @@ -60,15 +65,15 @@ export type Block = readonly [type: NodeTypeCatalog['block'], Statement[]]; * Represents a let statement */ export type Let = - | readonly [type: NodeTypeCatalog['let'], identifier: string] - | readonly [type: NodeTypeCatalog['let'], identifier: string, value: Expression]; + | readonly [type: NodeTypeCatalog['let'], identifier: Identifier] + | readonly [type: NodeTypeCatalog['let'], identifier: Identifier, value: Expression]; /** * Represents a const statement */ export type Const = - | readonly [type: NodeTypeCatalog['const'], identifier: string] - | readonly [type: NodeTypeCatalog['const'], identifier: string, value: Expression]; + | readonly [type: NodeTypeCatalog['const'], identifier: Identifier] + | readonly [type: NodeTypeCatalog['const'], identifier: Identifier, value: Expression]; export type For = readonly [ type: NodeTypeCatalog['for'], @@ -205,7 +210,7 @@ export type ConditionalExpression = readonly [ export type MemberAccess = readonly [ type: NodeTypeCatalog['memberAccess'], object: Expression, - member: string, + member: Identifier, ]; export type IndexAccess = readonly [ @@ -241,11 +246,13 @@ export type Str = readonly [type: NodeTypeCatalog['stringLiteral'], string]; /** null literal */ export type Null = readonly [type: NodeTypeCatalog['nullLiteral']]; -export type Literal = Num | Str | boolean | Null; +/** boolean can either be encoded as a node, or as a plain literal */ +export type Bool = boolean | readonly [type: NodeTypeCatalog['booleanLiteral'], boolean]; + +export type Literal = Num | Str | Bool | Null; -/** Identifiers are just strings, since string literals are rare in WGSL, and identifiers are everywhere. */ export type Expression = - | string + | Identifier | BinaryExpression | AssignmentExpression | LogicalExpression diff --git a/packages/typegpu/src/shared/tseynit.ts b/packages/typegpu/src/shared/tseynit.ts index ef7a8479df..b8e2700049 100644 --- a/packages/typegpu/src/shared/tseynit.ts +++ b/packages/typegpu/src/shared/tseynit.ts @@ -36,14 +36,14 @@ function stringifyStatement(node: tinyest.Statement, ident: string): string { if (node[0] === NODE.let) { if (node[2] !== undefined) { - return `${ident}let ${node[1]} = ${stringifyExpression(node[2], ident)};`; + return `${ident}let ${stringifyExpression(node[1], ident)} = ${stringifyExpression(node[2], ident)};`; } return `${ident}let ${node[1]};`; } if (node[0] === NODE.const) { if (node[2] !== undefined) { - return `${ident}const ${node[1]} = ${stringifyExpression(node[2], ident)};`; + return `${ident}const ${stringifyExpression(node[1], ident)} = ${stringifyExpression(node[2], ident)};`; } return `${ident}const ${node[1]};`; } @@ -90,6 +90,14 @@ function stringifyExpression(node: tinyest.Expression, ident: string): string { return `${node}`; } + if (node[0] === NODE.identifier) { + return node[1]; + } + + if (node[0] === NODE.booleanLiteral) { + return `${node[1]}`; + } + if (node[0] === NODE.numericLiteral) { return node[1]; } @@ -129,9 +137,9 @@ function stringifyExpression(node: tinyest.Expression, ident: string): string { if (node[0] === NODE.memberAccess) { if (Array.isArray(node[1]) && node[1][0] === NODE.numericLiteral) { - return `(${stringifyExpression(node[1], ident)}).${node[2]}`; + return `(${stringifyExpression(node[1], ident)}).${stringifyExpression(node[2], ident)}`; } - return `${wrapIfComplex(node[1], ident)}.${node[2]}`; + return `${wrapIfComplex(node[1], ident)}.${stringifyExpression(node[2], ident)}`; } if (node[0] === NODE.indexAccess) { @@ -172,6 +180,8 @@ function isExpression(node: tinyest.AnyNode): node is tinyest.Expression { if ( typeof node === 'string' || typeof node === 'boolean' || + node[0] === NODE.identifier || + node[0] === NODE.booleanLiteral || node[0] === NODE.numericLiteral || node[0] === NODE.stringLiteral || node[0] === NODE.arrayExpr || diff --git a/packages/typegpu/src/tgsl/wgslGenerator.ts b/packages/typegpu/src/tgsl/wgslGenerator.ts index 33d64ec786..c60a5042eb 100644 --- a/packages/typegpu/src/tgsl/wgslGenerator.ts +++ b/packages/typegpu/src/tgsl/wgslGenerator.ts @@ -434,6 +434,14 @@ export class WgslGenerator implements ShaderGenerator { return snip(expression, bool, /* origin */ 'constant', false); } + if (expression[0] === NODE.identifier) { + return this._identifier(expression[1]); + } + + if (expression[0] === NODE.booleanLiteral) { + return snip(expression[1], bool, /* origin */ 'constant', false); + } + if (expression[0] === NODE.logicalExpr) { const [_, lhs, op, rhs] = expression; const lhsExpr = this._expression(lhs); @@ -674,8 +682,9 @@ export class WgslGenerator implements ShaderGenerator { if (expression[0] === NODE.memberAccess) { // Member Access - const [_, targetNode, property] = expression; + const [_, targetNode, propertyNode] = expression; const target = this._expression(targetNode); + const property = extractId(propertyNode); const accessed = accessProp(target, property); if (!accessed) { @@ -1312,7 +1321,8 @@ Try 'return ${typeStr}(${str});' instead. } protected _letStatement(statement: tinyest.Let): ResolvedStatement { - const [_, rawId, eqNode] = statement; + const [_, rawIdNode, eqNode] = statement; + const rawId = extractId(rawIdNode); if (eqNode === undefined) { throw new Error( @@ -1385,7 +1395,8 @@ Try 'return ${typeStr}(${str});' instead. } protected _constStatement(statement: tinyest.Const): ResolvedStatement { - const [_, rawId, eqNode] = statement; + const [_, rawIdNode, eqNode] = statement; + const rawId = extractId(rawIdNode); if (eqNode === undefined) { throw new Error( @@ -1684,7 +1695,7 @@ ${this.ctx.pre}else ${alternate}`, const shouldUnroll = iterableExpr.value instanceof UnrollableIterable; const iterableSnippet = shouldUnroll ? iterableExpr.value.snippet : iterableExpr; const range = forOfUtils.getRangeSnippets(this.ctx, iterableSnippet, shouldUnroll); - const originalLoopVarName = loopVar[1]; + const originalLoopVarName = extractId(loopVar[1]); const blockified = blockifySingleStatement(body); if (shouldUnroll) { @@ -1934,3 +1945,10 @@ function extractObject(expr: tinyest.Expression): string | undefined { return object; } } + +function extractId(ident: tinyest.Identifier): string { + if (typeof ident === 'string') { + return ident; + } + return ident[1]; +} diff --git a/packages/typegpu/tests/internal/tseynit.test.ts b/packages/typegpu/tests/internal/tseynit.test.ts index 96cccef9b0..3dde362bf3 100644 --- a/packages/typegpu/tests/internal/tseynit.test.ts +++ b/packages/typegpu/tests/internal/tseynit.test.ts @@ -3,6 +3,8 @@ import * as tinyest from 'tinyest'; import { getFunctionMetadata } from '../../src/shared/meta.ts'; import { stringifyNode } from '../../src/shared/tseynit.ts'; import { tgpu, d } from '../../src/index.js'; +import type { BinaryExpression } from 'tinyest'; +import type { LogicalExpression } from 'tinyest'; function getBodyAst(fn: () => void) { const meta = getFunctionMetadata(fn); @@ -355,5 +357,37 @@ describe('ast to JS transformation', () => { }" `); }); + + it('handles boolean node', () => { + const NODE = tinyest.NodeTypeCatalog; + const ast: LogicalExpression = [ + NODE.logicalExpr, + [NODE.booleanLiteral, true], + '||', + [NODE.booleanLiteral, false], + ]; + + expect(stringifyNode(ast)).toMatchInlineSnapshot(`"(true) || (false)"`); + }); + + it('handles identifier node', () => { + const NODE = tinyest.NodeTypeCatalog; + const ast: tinyest.Block = [ + NODE.block, + [ + [NODE.let, [NODE.identifier, 'ident1'], [NODE.identifier, 'other1']], + [NODE.const, [NODE.identifier, 'ident2'], [NODE.identifier, 'other2']], + [NODE.memberAccess, [NODE.identifier, 'ident3'], [NODE.identifier, 'other3']], + ], + ]; + + expect(stringifyNode(ast)).toMatchInlineSnapshot(` + "{ + let ident1 = other1; + const ident2 = other2; + (ident3).other3; + }" + `); + }); }); }); diff --git a/packages/unplugin-typegpu/src/core/obfuscate.ts b/packages/unplugin-typegpu/src/core/obfuscate.ts index d547efef11..5afcb5e0f1 100644 --- a/packages/unplugin-typegpu/src/core/obfuscate.ts +++ b/packages/unplugin-typegpu/src/core/obfuscate.ts @@ -175,6 +175,18 @@ const visitors = { nullLiteral(_: Context, node: tinyest.Null) { return node; }, + booleanLiteral(_: Context, node: tinyest.Bool) { + if (typeof node === 'boolean') { + return node; + } + return [NODE.booleanLiteral, node[1]]; + }, + identifier(ctx: Context, node: tinyest.Identifier) { + if (typeof node === 'string') { + return obf(ctx, node); + } + return [NODE.identifier, obf(ctx, node[1])]; + }, } as const satisfies { [N in keyof typeof NODE]: ( ctx: Context, diff --git a/packages/unplugin-typegpu/test/obfuscation.test.ts b/packages/unplugin-typegpu/test/obfuscation.test.ts index 3105e515e3..500b909d3c 100644 --- a/packages/unplugin-typegpu/test/obfuscation.test.ts +++ b/packages/unplugin-typegpu/test/obfuscation.test.ts @@ -617,4 +617,41 @@ describe('obfuscate', () => { expect(stringifiedBody).toContain('ab'); expect(externalNames).toMatchInlineSnapshot(`Map {}`); }); + + it('supports boolean nodes', () => { + const code = `() => { return true || false; }`; + const transpiled = _transpileFn(parse(code), { ast: 'babel', verboseNodes: true }); + + const { params, body, externalNames } = obfuscate(transpiled); + + expect(params).toMatchInlineSnapshot(`[]`); + expect(stringifyNode(body)).toMatchInlineSnapshot(` + "{ + return (true) || (false); + }" + `); + expect(externalNames).toMatchInlineSnapshot(`Map {}`); + }); + + it('supports identifier nodes', () => { + const code = `(a) => { return a; }`; + const transpiled = _transpileFn(parse(code), { ast: 'babel', verboseNodes: true }); + + const { params, body, externalNames } = obfuscate(transpiled); + + expect(params).toMatchInlineSnapshot(` + [ + { + "name": "a", + "type": "i", + }, + ] + `); + expect(stringifyNode(body)).toMatchInlineSnapshot(` + "{ + return a; + }" + `); + expect(externalNames).toMatchInlineSnapshot(`Map {}`); + }); });