diff --git a/apps/typegpu-docs/astro.config.mjs b/apps/typegpu-docs/astro.config.mjs index 34d9785799..e93fcc9665 100644 --- a/apps/typegpu-docs/astro.config.mjs +++ b/apps/typegpu-docs/astro.config.mjs @@ -215,6 +215,11 @@ export default defineConfig({ label: 'Timing Your Pipelines', slug: 'advanced/timestamp-queries', }, + { + label: 'Minifying & Obfuscating Shaders', + slug: 'advanced/minifying-shaders', + badge: { text: 'new' }, + }, DEV && { label: 'Naming Convention', slug: 'advanced/naming-convention', diff --git a/apps/typegpu-docs/src/content/docs/advanced/minifying-shaders.mdx b/apps/typegpu-docs/src/content/docs/advanced/minifying-shaders.mdx new file mode 100644 index 0000000000..262d2d6998 --- /dev/null +++ b/apps/typegpu-docs/src/content/docs/advanced/minifying-shaders.mdx @@ -0,0 +1,33 @@ +--- +title: Minifying & Obfuscating Shaders +description: How to setup TypeGPU so that resulting shaders will be smaller or obfuscated. +--- + +:::caution[Experimental] +This feature is under development and is yet to reach stability. +::: + +Regular JS minification & obfuscation does not alter the shader code generated by TypeGPU. +This is why we provide specific options for transforming shaders. + +## Plugin obfuscation + +:::note +To reduce both the shader size and readability, it is advised to disable the plugin auto-naming by setting `{ autoNamingEnabled: false }`. +This way, only resources given name via `.$name()` will be named in the resulting shader. +::: + +`unplugin-typegpu` collects function metadata, which is later used for WGSL code generation. +With the `{ EXPERIMENTAL_obfuscate: true }` option, all saved identifiers will be obfuscated: +- In the AST, all parameters and variables will have their names changed to `a`, `b`, `c`, ... +- Externals (the captured scope used by the function) will also have their respective identifiers changed. + +:::caution +Enabling this will obfuscate not only the resulting code, but also error messages appearing during resolution. + +It is not advised to minify shaders during development. +::: + +## Runtime minification + +Coming soon. diff --git a/packages/tinyest-for-wgsl/src/parsers.ts b/packages/tinyest-for-wgsl/src/parsers.ts index 69299ac36d..9e17fef3c4 100644 --- a/packages/tinyest-for-wgsl/src/parsers.ts +++ b/packages/tinyest-for-wgsl/src/parsers.ts @@ -301,7 +301,7 @@ function transpile(ctx: Context, node: JsNode): tinyest.AnyNode { // add it to externals and swap the AST node for an identifier. const externalChain = tryFindExternalChain(ctx, node); if (externalChain) { - ctx.externalNames.add(externalChain); + ctx.externalNames.set(externalChain, externalChain); return externalChain; } } @@ -412,7 +412,7 @@ export function transpileFn(rootNode: JsNode): TranspilationResult { const { params, body } = extractFunctionParts(rootNode); const ctx: Context = { - externalNames: new Set(), + externalNames: new Map(), ignoreExternalDepth: 0, visitedNodes: new Set(), stack: [ @@ -445,7 +445,7 @@ export function transpileFn(rootNode: JsNode): TranspilationResult { export function transpileNode(node: JsNode): tinyest.AnyNode { const ctx: Context = { - externalNames: new Set(), + externalNames: new Map(), ignoreExternalDepth: 0, visitedNodes: new Set(), stack: [ diff --git a/packages/tinyest-for-wgsl/src/types.ts b/packages/tinyest-for-wgsl/src/types.ts index f5a5f5ed4c..5f27600786 100644 --- a/packages/tinyest-for-wgsl/src/types.ts +++ b/packages/tinyest-for-wgsl/src/types.ts @@ -7,7 +7,7 @@ export type Scope = { declaredNames: string[]; }; -export type Externals = Set; +export type Externals = Map; export type Context = { /** Holds a set of all identifiers that were used in code, but were not declared in code. */ diff --git a/packages/tinyest-for-wgsl/tests/helpers.ts b/packages/tinyest-for-wgsl/tests/helpers.ts new file mode 100644 index 0000000000..7bd671e13c --- /dev/null +++ b/packages/tinyest-for-wgsl/tests/helpers.ts @@ -0,0 +1,14 @@ +import babel from '@babel/parser'; +import type { Node } from '@babel/types'; +import * as acorn from 'acorn'; + +export const parseRollup = (code: string) => acorn.parse(code, { ecmaVersion: 'latest' }); +export const parseBabel = (code: string) => + babel.parse(code, { sourceType: 'module', plugins: ['typescript'] }).program.body[0] as Node; + +export function dualTest(test: (p: (code: string) => Node | acorn.AnyNode) => void) { + return () => { + test(parseBabel); + test(parseRollup); + }; +} diff --git a/packages/tinyest-for-wgsl/tests/parsers.test.ts b/packages/tinyest-for-wgsl/tests/parsers.test.ts index d46ceb01f7..7ef040a636 100644 --- a/packages/tinyest-for-wgsl/tests/parsers.test.ts +++ b/packages/tinyest-for-wgsl/tests/parsers.test.ts @@ -1,21 +1,36 @@ -import babel from '@babel/parser'; 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'; -const parseRollup = (code: string) => acorn.parse(code, { ecmaVersion: 'latest' }); -const parseBabel = (code: string) => - babel.parse(code, { sourceType: 'module', plugins: ['typescript'] }).program.body[0] as Node; +describe('transpileFn', () => { + it( + 'handles weird identifiers', + dualTest((p) => { + const { params, body, externalNames } = transpileFn( + p(`() => { + const a = undefined; + const b = Infinity; + const c = NaN; + }`), + ); -function dualTest(test: (p: (code: string) => Node | acorn.AnyNode) => void) { - return () => { - test(parseBabel); - test(parseRollup); - }; -} + expect(params).toStrictEqual([]); + expect(JSON.stringify(body)).toMatchInlineSnapshot( + `"[0,[[13,"a","undefined"],[13,"b","Infinity"],[13,"c","NaN"]]]"`, + ); + // These are identifiers, so they should be in externals. + expect(externalNames).toMatchInlineSnapshot(` + Map { + "undefined" => "undefined", + "Infinity" => "Infinity", + "NaN" => "NaN", + } + `); + }), + ); -describe('transpileFn', () => { it( 'fails when the input is not a function', dualTest((p) => { @@ -30,7 +45,7 @@ describe('transpileFn', () => { expect(params).toStrictEqual([]); expect(JSON.stringify(body)).toMatchInlineSnapshot(`"[0,[]]"`); - expect(externalNames).toMatchInlineSnapshot(`Set {}`); + expect(externalNames).toMatchInlineSnapshot(`Map {}`); }), ); @@ -41,7 +56,7 @@ describe('transpileFn', () => { expect(params).toStrictEqual([]); expect(JSON.stringify(body)).toMatchInlineSnapshot(`"[0,[]]"`); - expect(externalNames).toMatchInlineSnapshot(`Set {}`); + expect(externalNames).toMatchInlineSnapshot(`Map {}`); }), ); @@ -58,8 +73,8 @@ describe('transpileFn', () => { `"[0,[[10,[1,[1,"a","+","b"],"-","c"]]]]"`, ); expect(externalNames).toMatchInlineSnapshot(` - Set { - "c", + Map { + "c" => "c", } `); }), @@ -81,8 +96,8 @@ describe('transpileFn', () => { ); // Only 'c' is external, as 'a' is declared in the same scope. expect(externalNames).toMatchInlineSnapshot(` - Set { - "c", + Map { + "c" => "c", } `); }), @@ -106,8 +121,8 @@ describe('transpileFn', () => { ); // Only 'c' is external, as 'a' is declared in the outer scope. expect(externalNames).toMatchInlineSnapshot(` - Set { - "c", + Map { + "c" => "c", } `); }), @@ -122,8 +137,8 @@ describe('transpileFn', () => { expect(JSON.stringify(body)).toMatchInlineSnapshot(`"[0,[[10,"external.outside.prop"]]]"`); // Only 'external' is external. expect(externalNames).toMatchInlineSnapshot(` - Set { - "external.outside.prop", + Map { + "external.outside.prop" => "external.outside.prop", } `); }), @@ -154,7 +169,7 @@ describe('transpileFn', () => { }, ]); - expect(externalNames).toMatchInlineSnapshot(`Set {}`); + expect(externalNames).toMatchInlineSnapshot(`Map {}`); }), ); @@ -200,7 +215,7 @@ describe('transpileFn', () => { }, ]); - expect(externalNames).toMatchInlineSnapshot(`Set {}`); + expect(externalNames).toMatchInlineSnapshot(`Map {}`); }), ); @@ -225,8 +240,8 @@ describe('transpileFn', () => { ); expect(externalNames).toMatchInlineSnapshot(` - Set { - "a", + Map { + "a" => "a", } `); }), @@ -247,8 +262,8 @@ describe('transpileFn', () => { ); expect(externalNames).toMatchInlineSnapshot(` - Set { - "a", + Map { + "a" => "a", } `); }), @@ -281,19 +296,19 @@ describe('transpileFn', () => { ); expect(externalNames).toMatchInlineSnapshot(` - Set { - "ext.p", - "ext.q.a", - "ext.q.b", - "ext.r.a", - "ext.r", - "ext.s", - "ext.s.a", - "ext.t.fn", - "ext.t.comp", - "ext.t", - "ext.u", - "ext", + Map { + "ext.p" => "ext.p", + "ext.q.a" => "ext.q.a", + "ext.q.b" => "ext.q.b", + "ext.r.a" => "ext.r.a", + "ext.r" => "ext.r", + "ext.s" => "ext.s", + "ext.s.a" => "ext.s.a", + "ext.t.fn" => "ext.t.fn", + "ext.t.comp" => "ext.t.comp", + "ext.t" => "ext.t", + "ext.u" => "ext.u", + "ext" => "ext", } `); @@ -314,8 +329,8 @@ describe('transpileFn', () => { ); expect(externalNames).toMatchInlineSnapshot(` - Set { - "ext", + Map { + "ext" => "ext", } `); }), @@ -334,10 +349,10 @@ describe('transpileFn', () => { ); expect(externalNames).toMatchInlineSnapshot(` - Set { - "ext.value", - "ext.config.multiplier", - "ext.config.zero", + Map { + "ext.value" => "ext.value", + "ext.config.multiplier" => "ext.config.multiplier", + "ext.config.zero" => "ext.config.zero", } `); @@ -369,8 +384,8 @@ describe('transpileFn', () => { const { externalNames } = transpileFn(fn); expect(externalNames).toMatchInlineSnapshot(` - Set { - "this.#v", + Map { + "this.#v" => "this.#v", } `); }), diff --git a/packages/typegpu/src/core/function/extractArgs.ts b/packages/typegpu/src/core/function/extractArgs.ts index 9ff97abd98..6e7050fd99 100644 --- a/packages/typegpu/src/core/function/extractArgs.ts +++ b/packages/typegpu/src/core/function/extractArgs.ts @@ -1,3 +1,5 @@ +import { blankSpaces, lineBreaks } from '../whitespaces.ts'; + interface FunctionArgsInfo { args: ArgInfo[]; ret: ReturnInfo | undefined; @@ -251,22 +253,6 @@ class ParsableString { } } -const lineBreaks = new Set([ - '\u000A', // line feed - '\u000B', // vertical tab - '\u000C', // form feed - '\u000D', // carriage return - '\u0085', // next line - '\u2028', // line separator - '\u2029', // paragraph separator -]); -const blankSpaces = new Set([ - ...lineBreaks, - '\u0020', // space - '\u0009', // horizontal tab - '\u200E', // left-to-right mark - '\u200F', // right-to-left mark -]); const closingParenthesis = new Set([')']); const identifierEndSymbols = new Set([':', ',', ')']); const typeEndSymbols = new Set([',', ')']); diff --git a/packages/typegpu/src/core/resolve/tgpuResolve.ts b/packages/typegpu/src/core/resolve/tgpuResolve.ts index 56b1878104..61fac40ffa 100644 --- a/packages/typegpu/src/core/resolve/tgpuResolve.ts +++ b/packages/typegpu/src/core/resolve/tgpuResolve.ts @@ -26,6 +26,12 @@ export interface TgpuResolveOptions { * @default 'strict' */ names?: 'strict' | 'random' | Namespace | undefined; + /** + * When set to true, the resulting shaders will be stripped from all unnecessary whitespace. + * + * @default false + */ + unstable_minify?: boolean; /** * A function to configure the resolution context. */ @@ -186,6 +192,7 @@ function resolveFromTemplate(options: TgpuExtendedResolveOptions): ResolutionRes externals, unstable_shaderGenerator: shaderGenerator, names = 'strict', + unstable_minify, config, enableExtensions, } = options; @@ -209,12 +216,15 @@ function resolveFromTemplate(options: TgpuExtendedResolveOptions): ResolutionRes toString: () => '', }; + const maybeRoot = tryFindRoot(Object.values(externals)); + return resolveImpl(resolutionObj, { namespace: typeof names === 'string' ? namespace({ names }) : names, + minify: unstable_minify ?? maybeRoot?.minify, enableExtensions, shaderGenerator, config, - root: tryFindRoot(Object.values(externals)), + root: maybeRoot, }); } @@ -225,6 +235,7 @@ function resolveFromArray( const { unstable_shaderGenerator: shaderGenerator, names = 'strict', + unstable_minify, config, enableExtensions, } = options ?? {}; @@ -247,12 +258,15 @@ function resolveFromArray( toString: () => '', }; + const maybeRoot = tryFindRoot(items); + return resolveImpl(resolutionObj, { namespace: typeof names === 'string' ? namespace({ names }) : names, + minify: unstable_minify ?? maybeRoot?.minify, enableExtensions, shaderGenerator, config, - root: tryFindRoot(items), + root: maybeRoot, }); } diff --git a/packages/typegpu/src/core/root/init.ts b/packages/typegpu/src/core/root/init.ts index a3544d67e8..02c81d9721 100644 --- a/packages/typegpu/src/core/root/init.ts +++ b/packages/typegpu/src/core/root/init.ts @@ -286,6 +286,7 @@ class TgpuRootImpl extends WithBindingImpl implements TgpuRoot, ExperimentalTgpu readonly device: GPUDevice; readonly nameRegistrySetting: 'random' | 'strict'; + readonly minify: boolean; readonly shaderGenerator: ShaderGenerator | undefined; #unwrappedBindGroupLayouts = new WeakMemo((key: TgpuBindGroupLayout) => key.unwrap(this)); @@ -299,6 +300,7 @@ class TgpuRootImpl extends WithBindingImpl implements TgpuRoot, ExperimentalTgpu constructor( device: GPUDevice, nameRegistrySetting: 'random' | 'strict', + minify: boolean, ownDevice: boolean, logOptions: LogGeneratorOptions, shaderGenerator?: ShaderGenerator, @@ -307,6 +309,7 @@ class TgpuRootImpl extends WithBindingImpl implements TgpuRoot, ExperimentalTgpu this.device = device; this.nameRegistrySetting = nameRegistrySetting; + this.minify = minify; this.#ownDevice = ownDevice; this.shaderGenerator = shaderGenerator; @@ -672,6 +675,15 @@ export type InitOptions = { device?: (GPUDeviceDescriptor & { optionalFeatures?: Iterable }) | undefined; /** @default 'strict' */ unstable_names?: 'random' | 'strict' | undefined; + /** + * If set to true, rooted resolves will be stripped of all excessive spaces. + * Rooted resolves include pipelines and all their transitive dependencies, + * but not, for example, `tgpu.resolve([MyStruct])`, + * as there may be multiple roots with different settings in one program. + * + * @default false + */ + unstable_minify?: boolean; /** * A custom shader code generator, used when resolving TypeGPU functions. * If not provided, the default WGSL generator will be used. @@ -680,6 +692,7 @@ export type InitOptions = { unstable_logOptions?: LogGeneratorOptions; }; +// TODO: merge these types /** * Options passed into {@link initFromDevice}. */ @@ -687,6 +700,15 @@ export type InitFromDeviceOptions = { device: GPUDevice; /** @default 'strict' */ unstable_names?: 'random' | 'strict' | undefined; + /** + * If set to true, rooted resolves will be stripped of all excessive spaces. + * Rooted resolves include pipelines and all their transitive dependencies, + * but not, for example, `tgpu.resolve([MyStruct])`, + * as there may be multiple roots with different settings in one program. + * + * @default false + */ + unstable_minify?: boolean; /** * A custom shader code generator, used when resolving TypeGPU functions. * If not provided, the default WGSL generator will be used. @@ -718,6 +740,7 @@ export async function init(options?: InitOptions): Promise { adapter: adapterOpt, device: deviceOpt, unstable_names: names = 'strict', + unstable_minify: minify = false, unstable_logOptions: logOptions, unstable_shaderGenerator: shaderGenerator, } = options ?? {}; @@ -755,7 +778,7 @@ export async function init(options?: InitOptions): Promise { requiredFeatures: availableFeatures, }); - return new TgpuRootImpl(device, names, true, logOptions ?? {}, shaderGenerator); + return new TgpuRootImpl(device, names, minify, true, logOptions ?? {}, shaderGenerator); } /** @@ -771,9 +794,10 @@ export function initFromDevice(options: InitFromDeviceOptions): TgpuRoot { const { device, unstable_names: names = 'strict', + unstable_minify: minify = false, unstable_logOptions: logOptions, unstable_shaderGenerator: shaderGenerator, } = options ?? {}; - return new TgpuRootImpl(device, names, false, logOptions ?? {}, shaderGenerator); + return new TgpuRootImpl(device, names, minify, false, logOptions ?? {}, shaderGenerator); } diff --git a/packages/typegpu/src/core/root/rootTypes.ts b/packages/typegpu/src/core/root/rootTypes.ts index e06c2da95b..f6cc792477 100644 --- a/packages/typegpu/src/core/root/rootTypes.ts +++ b/packages/typegpu/src/core/root/rootTypes.ts @@ -968,6 +968,7 @@ export interface TgpuRoot extends Unwrapper, WithBinding { export interface ExperimentalTgpuRoot extends Omit, Withable_Deprecated { readonly nameRegistrySetting: 'strict' | 'random'; + readonly minify: boolean; readonly shaderGenerator?: ShaderGenerator | undefined; /** @deprecated Use `root.createTexture` instead. */ diff --git a/packages/typegpu/src/core/whitespaces.ts b/packages/typegpu/src/core/whitespaces.ts new file mode 100644 index 0000000000..d37f18fde0 --- /dev/null +++ b/packages/typegpu/src/core/whitespaces.ts @@ -0,0 +1,17 @@ +export const lineBreaks = new Set([ + '\u000A', // line feed + '\u000B', // vertical tab + '\u000C', // form feed + '\u000D', // carriage return + '\u0085', // next line + '\u2028', // line separator + '\u2029', // paragraph separator +]); + +export const blankSpaces = new Set([ + ...lineBreaks, + '\u0020', // space + '\u0009', // horizontal tab + '\u200E', // left-to-right mark + '\u200F', // right-to-left mark +]); diff --git a/packages/typegpu/src/internal.ts b/packages/typegpu/src/internal.ts index b1e0fe93f7..7e59b0af7f 100644 --- a/packages/typegpu/src/internal.ts +++ b/packages/typegpu/src/internal.ts @@ -7,6 +7,7 @@ export { UnknownData } from './data/dataTypes.ts'; export { getName } from './shared/meta.ts'; export { WgslGenerator } from './tgsl/wgslGenerator.ts'; export { snip } from './data/snippet.ts'; +export { stringifyNode } from './shared/tseynit.ts'; // types export type { ResolutionCtx, FunctionArgument, TgpuShaderStage } from './types.ts'; diff --git a/packages/typegpu/src/minify.ts b/packages/typegpu/src/minify.ts new file mode 100644 index 0000000000..6c6f171e20 --- /dev/null +++ b/packages/typegpu/src/minify.ts @@ -0,0 +1,89 @@ +import { blankSpaces, lineBreaks } from './core/whitespaces.ts'; + +/** + * Regex for splitting code into tokens. + * We don't separate every WGSL token, for example `main(){` already has no spaces, no need to split it. + * Split if whitespace is encountered, or if either of [:,] is in lookahead. + */ +const splitRegex = new RegExp(`[${[...blankSpaces].join('|')}]+|(?=[:,])`, 'ug'); + +/** + * Regex for detecting tokens that require whitespace separators. + * Exact match is not required, for example: `fn` should be separated from `main()`. + */ +const separatorNeededRegex = /[\p{XID_Continue}]+/u; + +function stripWGSLComments(code: string): string { + let result = ''; + let copiedUpTo = 0; + let offset = 0; + + while (offset < code.length) { + if (code.startsWith('//', offset)) { + result += `${code.slice(copiedUpTo, offset)} `; + offset += 2; + + while (offset < code.length && !lineBreaks.has(code.charAt(offset))) { + offset += 1; + } + + copiedUpTo = offset; + continue; + } + + if (code.startsWith('/*', offset)) { + result += `${code.slice(copiedUpTo, offset)} `; + let depth = 1; + offset += 2; + + while (offset < code.length && depth > 0) { + if (code.startsWith('/*', offset)) { + depth += 1; + offset += 2; + } else if (code.startsWith('*/', offset)) { + depth -= 1; + offset += 2; + } else { + offset += 1; + } + } + + if (depth > 0) { + throw new SyntaxError(`Unterminated block comment found during minification.`); + } + + copiedUpTo = offset; + continue; + } + + offset += 1; + } + + return result + code.slice(copiedUpTo); +} + +/** + * This function accepts a code string, and returns equivalent code + * with unnecessary whitespaces and comments removed. + */ +export function minify(code: string): string { + // Remove comments. + const codeWithoutComments = stripWGSLComments(code); + + // Split into tokens. + const tokens = codeWithoutComments.split(splitRegex); + + // Join and separate if necessary. + let result = ''; + for (let i = 0; i < tokens.length; i++) { + const current = tokens[i] as string; + const next = tokens[i + 1] ?? ''; + + result += current; + if (current.match(separatorNeededRegex) && next.match(separatorNeededRegex)) { + result += ' '; + } + } + + return result; +} diff --git a/packages/typegpu/src/resolutionCtx.ts b/packages/typegpu/src/resolutionCtx.ts index e3889112ea..b315b1d532 100644 --- a/packages/typegpu/src/resolutionCtx.ts +++ b/packages/typegpu/src/resolutionCtx.ts @@ -66,6 +66,7 @@ import type { IOData } from './core/function/fnTypes.ts'; import { AutoStruct } from './data/autoStruct.ts'; import { EntryInputRouter } from './core/function/entryInputRouter.ts'; import { validateIdentifier, sanitizePrimer, bannedTokens } from './nameUtils.ts'; +import { minify } from './minify.ts'; /** * Inserted into bind group entry definitions that belong @@ -84,6 +85,10 @@ export type ResolutionCtxImplOptions = { readonly config?: ((cfg: Configurable) => Configurable) | undefined; readonly root?: ExperimentalTgpuRoot | undefined; readonly namespace: Namespace; + /** + * @default false + */ + readonly minify?: boolean | undefined; }; class ItemStateStackImpl implements ItemStateStack { @@ -202,17 +207,14 @@ class ItemStateStackImpl implements ItemStateStack { return access(); } - const external = layer.externalMap[id]; - if (isNamable(external) && getName(external) === undefined) { - setName(external, id.replaceAll('.', '_')); - } - - if (external !== undefined && external !== null) { + if (id in layer.externalMap) { + const external = layer.externalMap[id]; + if (isNamable(external) && getName(external) === undefined) { + setName(external, id.replaceAll('.', '_')); + } return coerceToSnippet(external); } - // Since functions cannot access resources from the calling scope, we - // return early here. return undefined; } @@ -1209,6 +1211,10 @@ export function resolve(item: Wgsl, options: ResolutionCtxImplOptions): Resoluti ), })); + if (options.minify) { + code = minify(code); + } + return { code, declarations, diff --git a/packages/typegpu/tests/minification.test.ts b/packages/typegpu/tests/minification.test.ts new file mode 100644 index 0000000000..51564428c6 --- /dev/null +++ b/packages/typegpu/tests/minification.test.ts @@ -0,0 +1,244 @@ +import { describe, expect } from 'vitest'; +import { tgpu, d } from 'typegpu'; +import { it } from 'typegpu-testing-utility'; + +describe('minification', () => { + const inner = () => { + 'use gpu'; + return 1; + }; + + const outer = () => { + 'use gpu'; + return inner(); + }; + + const computeFn = tgpu.computeFn({ workgroupSize: [1, 1, 1] })(() => { + 'use gpu'; + outer(); + }); + + it('does not minify if not set to', async () => { + const root = await tgpu.init(); + const pipeline = root.createComputePipeline({ compute: computeFn }); + + const code = tgpu.resolve([pipeline]); + + expect(code).toMatchInlineSnapshot(` + "fn inner() -> i32 { + return 1; + } + + fn outer() -> i32 { + return inner(); + } + + @compute @workgroup_size(1, 1, 1) fn computeFn() { + outer(); + }" + `); + expect(code).toContain(' '); + }); + + it('minifies in resolve', async () => { + const code = tgpu.resolve([inner], { unstable_minify: true }); + + expect(code).toMatchInlineSnapshot(`"fn inner()->i32{return 1;}"`); + expect(code).not.toContain(' '); + }); + + it('minifies in resolveWithContext', async () => { + const code = tgpu.resolveWithContext([inner], { unstable_minify: true }).code; + + expect(code).toMatchInlineSnapshot(`"fn inner()->i32{return 1;}"`); + expect(code).not.toContain(' '); + }); + + it('minifies in resolve with template', async () => { + const code = tgpu.resolve({ + template: 'fn main() { inner(); }', + externals: { inner }, + unstable_minify: true, + }); + + expect(code).toMatchInlineSnapshot(`"fn inner()->i32{return 1; }fn main(){inner();}"`); + expect(code).not.toContain(' '); + }); + + it('minifies raw wgsl implemented functions', async () => { + const rawFn = tgpu.fn([d.u32], d.u32)`(a) => { + return a + 1; + }`; + + const code = tgpu.resolve([rawFn], { unstable_minify: true }); + + expect(code).toMatchInlineSnapshot(`"fn rawFn(a:u32)->u32{return a+1;}"`); + expect(code).not.toContain(' '); + }); + + it('minifies raw code snippets', async () => { + const rawCodeSnippet = tgpu['~unstable'].rawCodeSnippet('1u + 2u', d.u32, 'constant', false); + const fn = () => { + 'use gpu'; + const a = rawCodeSnippet.$; + return a; + }; + + const code = tgpu.resolve([fn], { unstable_minify: true }); + + expect(code).toMatchInlineSnapshot(`"fn fn_1()->u32{const a=1u+2u; return a;}"`); + expect(code).not.toContain(' '); + }); + + it('reduces spaces if items are separated by , or :', async () => { + const helper = (a: number, b: number, c: number) => { + 'use gpu'; + return a + b + c; + }; + + const fn = () => { + 'use gpu'; + return helper(1, 2, 3); + }; + + const code = tgpu.resolve([fn], { unstable_minify: true }); + + expect(code).toMatchInlineSnapshot( + `"fn helper(a:i32,b:i32,c:i32)->i32{return ((a+b)+c);}fn fn_1()->i32{return helper(1i,2i,3i);}"`, + ); + expect(code).not.toContain(' '); + }); + + // it('handles |', () => { + // const rawFn = tgpu.fn([d.u32], d.u32)`(a) => { + // return a | a; + // }`; + + // expect(tgpu.resolve([rawFn], { unstable_minify: true })).toMatchInlineSnapshot( + // `"fn rawFn(a:u32)->u32{return a a;}"`, + // ); + // }); + + // it('handles - -1', () => { + // const rawFn = tgpu.fn([d.u32], d.u32)`(a) => { + // return a - -1; + // }`; + + // expect(tgpu.resolve([rawFn], { unstable_minify: true })).toMatchInlineSnapshot( + // `"fn rawFn(a:u32)->u32{return a--1;}"`, + // ); + // }); + + it('removes line comments', () => { + const rawFn = tgpu.fn([d.u32], d.u32)`(a) => { + // a comment + return a + 1; // my comment /* + // // other comment + } // end of file`; + + expect(tgpu.resolve([rawFn], { unstable_minify: true })).toMatchInlineSnapshot( + `"fn rawFn(a:u32)->u32{return a+1;}"`, + ); + }); + + it('removes block comments', () => { + const rawFn = tgpu.fn([d.u32], d.u32)`(a) => { + /* a comment */return a + 1;/* my comment */ + /* other + comment */ + }`; + + expect(tgpu.resolve([rawFn], { unstable_minify: true })).toMatchInlineSnapshot( + `"fn rawFn(a:u32)->u32{return a+1;}"`, + ); + }); + + it('removes Unicode line comments', () => { + const rawFn = tgpu.fn([d.u32], d.u32)`(a) => { + // 🙂 + return a; + }`; + + expect(tgpu.resolve([rawFn], { unstable_minify: true })).toMatchInlineSnapshot( + `"fn rawFn(a:u32)->u32{return a;}"`, + ); + }); + + it('removes nested block comments and ignores line-comment delimiters', () => { + const rawFn = tgpu.fn([d.u32], d.u32)`(a) => { + /* text + text // text /* text + /////**/ */*/ + return a; + }/* text */`; + + expect(tgpu.resolve([rawFn], { unstable_minify: true })).toMatchInlineSnapshot( + `"fn rawFn(a:u32)->u32{return a;}"`, + ); + }); + + it('keeps comments from joining adjacent tokens', () => { + const rawFn = tgpu.fn([d.u32], d.u32)`(a) => { + return/*text*/a; + }`; + + expect(tgpu.resolve([rawFn], { unstable_minify: true })).toMatchInlineSnapshot( + `"fn rawFn(a:u32)->u32{return a;}"`, + ); + }); + + it('rejects unterminated block comments', () => { + const rawFn = tgpu.fn([d.u32], d.u32)`(a) => { + return a; + } /* text`; + + expect(() => + tgpu.resolve([rawFn], { unstable_minify: true }), + ).toThrowErrorMatchingInlineSnapshot( + `[SyntaxError: Unterminated block comment found during minification.]`, + ); + }); + + it('minifies transitive dependencies in resolve', async () => { + const code = tgpu.resolve([outer], { unstable_minify: true }); + + expect(code).toMatchInlineSnapshot( + `"fn inner()->i32{return 1;}fn outer()->i32{return inner();}"`, + ); + expect(code).not.toContain(' '); + }); + + it('minifies in resolve if root is set to minify', async () => { + const root = await tgpu.init({ unstable_minify: true }); + const pipeline = root.createComputePipeline({ compute: computeFn }); + + const code = tgpu.resolve([pipeline]); + + expect(code).toMatchInlineSnapshot( + `"fn inner()->i32{return 1;}fn outer()->i32{return inner();}@compute @workgroup_size(1,1,1) fn computeFn(){outer();}"`, + ); + expect(code).not.toContain(' '); + }); + + it('does not minify in resolve with minify disabled even if root is set to minify', async () => { + const root = await tgpu.init({ unstable_minify: true }); + const pipeline = root.createComputePipeline({ compute: computeFn }); + + const code = tgpu.resolve([pipeline], { unstable_minify: false }); + + expect(code).toMatchInlineSnapshot(` + "fn inner() -> i32 { + return 1; + } + + fn outer() -> i32 { + return inner(); + } + + @compute @workgroup_size(1, 1, 1) fn computeFn() { + outer(); + }" + `); + expect(code).toContain(' '); + }); +}); diff --git a/packages/typegpu/tests/mutabilityTracking.test.ts b/packages/typegpu/tests/mutabilityTracking.test.ts index 513dad90cd..d31be32e98 100644 --- a/packages/typegpu/tests/mutabilityTracking.test.ts +++ b/packages/typegpu/tests/mutabilityTracking.test.ts @@ -350,15 +350,15 @@ describe('mutability tracking', () => { const resolved = tgpu.resolve([fn]); expect(resolved).toMatchInlineSnapshot(` - "fn item(arg: vec4u) -> u32 { - let a = arg; - { - var a_1 = arg; - a_1.x = 2u; - } - return a.x; - }" - `); + "fn item(arg: vec4u) -> u32 { + let a = arg; + { + var a_1 = arg; + a_1.x = 2u; + } + return a.x; + }" + `); expect(resolved).toContain('let a = arg'); expect(resolved).toContain('var a_1 = arg'); }); diff --git a/packages/typegpu/tests/std/boolean/not.test.ts b/packages/typegpu/tests/std/boolean/not.test.ts index 5a7fb2e57a..80fa45395a 100644 --- a/packages/typegpu/tests/std/boolean/not.test.ts +++ b/packages/typegpu/tests/std/boolean/not.test.ts @@ -60,10 +60,10 @@ describe('not', () => { return not(v); }); expect(tgpu.resolve([testFn])).toMatchInlineSnapshot(` - "fn testFn(v: i32) -> bool { - return !bool(v); - }" - `); + "fn testFn(v: i32) -> bool { + return !bool(v); + }" + `); }); it('generates correct WGSL on a boolean vector runtime-known argument', () => { @@ -178,10 +178,10 @@ describe('not', () => { return not(v); }); expect(tgpu.resolve([testFn])).toMatchInlineSnapshot(` - "fn testFn(v: i32) -> bool { - return !bool(v); - }" - `); + "fn testFn(v: i32) -> bool { + return !bool(v); + }" + `); }); it('generates correct WGSL on a boolean vector runtime-known argument', () => { diff --git a/packages/typegpu/tests/tgsl/wgslGenerator.test.ts b/packages/typegpu/tests/tgsl/wgslGenerator.test.ts index 1f70b19408..a86330b019 100644 --- a/packages/typegpu/tests/tgsl/wgslGenerator.test.ts +++ b/packages/typegpu/tests/tgsl/wgslGenerator.test.ts @@ -1771,10 +1771,10 @@ describe('wgslGenerator', () => { }); expect(tgpu.resolve([testFn])).toMatchInlineSnapshot(` - "fn testFn(b: bool) -> bool { - return !b; - }" - `); + "fn testFn(b: bool) -> bool { + return !b; + }" + `); }); it('handles unary operator `!` on numeric runtime-known operand', () => { @@ -1786,10 +1786,10 @@ describe('wgslGenerator', () => { }); expect(tgpu.resolve([testFn])).toMatchInlineSnapshot(` - "fn testFn(n: i32) -> bool { - return !bool(n); - }" - `); + "fn testFn(n: i32) -> bool { + return !bool(n); + }" + `); }); it('handles unary operator `!` on non-primitive values', ({ root }) => { diff --git a/packages/typegpu/tests/tgslFn.test.ts b/packages/typegpu/tests/tgslFn.test.ts index 2f71e3de9f..cc70250b96 100644 --- a/packages/typegpu/tests/tgslFn.test.ts +++ b/packages/typegpu/tests/tgslFn.test.ts @@ -1,7 +1,7 @@ import { attest } from '@ark/attest'; import { describe, expect } from 'vitest'; import { builtin } from 'typegpu/data'; -import { tgpu, d, type TgpuFn, type TgpuSlot } from 'typegpu'; +import { tgpu, d, type TgpuFn, type TgpuSlot, std } from 'typegpu'; import { it } from 'typegpu-testing-utility'; describe('TGSL tgpu.fn function', () => { @@ -1098,6 +1098,21 @@ describe('tgsl fn when using plugin', () => { `); }); + it('does not accidentally shadow std', () => { + const fn = () => { + 'use gpu'; + const sin = 1; + const a = std.sin(sin); + }; + + expect(tgpu.resolve([fn])).toMatchInlineSnapshot(` + "fn fn_1() { + const sin_1 = 1; + let a = sin(f32(sin_1)); + }" + `); + }); + it('throws a readable error when assigning to a value defined outside of scope', () => { let a = 0; const f = () => { diff --git a/packages/unplugin-typegpu/src/babel.ts b/packages/unplugin-typegpu/src/babel.ts index a4495214c1..eabfd52a60 100644 --- a/packages/unplugin-typegpu/src/babel.ts +++ b/packages/unplugin-typegpu/src/babel.ts @@ -6,6 +6,7 @@ import { METADATA_FORMAT_VERSION, type MetadatableFunction, type PluginState, + checkOpts, defaultOptions, functionVisitor, getBlockScope, @@ -19,8 +20,8 @@ function i(identifier: string): t.Identifier { function externalsToNode(externals: Externals): t.Expression { return t.objectExpression( - Array.from(externals, (key) => { - const chain = key.split('.'); + Array.from(externals, ([key, value]) => { + const chain = value.split('.'); if (!chain[0]) { throw new Error('Internal error, expected chain to not be empty'); } @@ -169,7 +170,7 @@ export default function TypeGPUPlugin() { return { name: 'typegpu', pre(this: PluginState) { - this.opts = defu(this.opts, defaultOptions); + this.opts = checkOpts(defu(this.opts, defaultOptions)); initPluginState(this, { warn: (message) => console.warn(message), assignMetadata, diff --git a/packages/unplugin-typegpu/src/bun.ts b/packages/unplugin-typegpu/src/bun.ts index 1f2e8bb558..1bb4f864af 100644 --- a/packages/unplugin-typegpu/src/bun.ts +++ b/packages/unplugin-typegpu/src/bun.ts @@ -1,10 +1,10 @@ import defu from 'defu'; -import { defaultOptions, earlyPruneRegex, type Options } from './core/common.ts'; +import { checkOpts, defaultOptions, earlyPruneRegex, type Options } from './core/common.ts'; import { unpluginFactory } from './core/factory.ts'; import type { UnpluginBuildContext, UnpluginContext } from 'unplugin'; export default (rawOptions?: Options): Bun.BunPlugin => { - const options = defu(rawOptions, defaultOptions); + const options = checkOpts(defu(rawOptions, defaultOptions)); const include = options.include; if (!(include instanceof RegExp)) { throw new Error( diff --git a/packages/unplugin-typegpu/src/core/common.ts b/packages/unplugin-typegpu/src/core/common.ts index 827e5ce376..19c03f8f89 100644 --- a/packages/unplugin-typegpu/src/core/common.ts +++ b/packages/unplugin-typegpu/src/core/common.ts @@ -3,6 +3,7 @@ import type { NodePath, TraverseOptions } from '@babel/traverse'; import type { FilterPattern } from 'unplugin'; import MagicString from 'magic-string'; import { transpileFn } from 'tinyest-for-wgsl'; +import { obfuscate } from './obfuscate.ts'; /** * Each breaking change to the metadata format requires a bump to this number. @@ -15,7 +16,7 @@ export interface Options { include?: FilterPattern; /** @default undefined */ - exclude?: FilterPattern; + exclude?: FilterPattern | undefined; /** @default undefined */ enforce?: 'post' | 'pre' | undefined; @@ -24,7 +25,15 @@ export interface Options { forceTgpuAlias?: string | undefined; /** @default true */ - autoNamingEnabled?: boolean | undefined; + autoNamingEnabled?: boolean; + + /** + * Obfuscate the generated AST. + * This results in obfuscation of the generated WGSL, and in smaller bundle sizes. + * + * @default false + */ + EXPERIMENTAL_obfuscate?: boolean; /** * Skipping files that don't contain "typegpu", "tgpu" or "use gpu". @@ -36,6 +45,15 @@ export interface Options { earlyPruning?: boolean | undefined; } +export function checkOpts(opts: T): T { + if (opts.EXPERIMENTAL_obfuscate && opts.autoNamingEnabled) { + throw new Error( + `Options 'EXPERIMENTAL_obfuscate' and 'autoNamingEnabled' cannot be enabled at the same time.`, + ); + } + return opts; +} + export type MetadatableFunction = | t.FunctionDeclaration | t.FunctionExpression @@ -106,7 +124,7 @@ export interface PluginState extends TransformMethods { * In Babel, options are assigned to the property `opts` on the plugin state. * We use this pattern everywhere for consistency. */ - opts: Options; + opts: Required; inUseGpuScope: boolean; } @@ -137,7 +155,8 @@ export const defaultOptions = { include: /\.m?[jt]sx?(?:\?.*)?$/, autoNamingEnabled: true, earlyPruning: true, -}; + EXPERIMENTAL_obfuscate: false, +} satisfies Partial; /** * Returns the block scope of a function declaration, if one exists. @@ -471,6 +490,17 @@ function functionOnExit( path.skip(); } +function transpile( + rootNode: Parameters[0], + obf: boolean, +): ReturnType { + const result = transpileFn(rootNode); + if (obf) { + return obfuscate(result); + } + return result; +} + export const functionVisitor: TraverseOptions = { ImportDeclaration(path, state) { gatherTgpuAliases(state, path.node); @@ -529,7 +559,10 @@ export const functionVisitor: TraverseOptions = { ArrowFunctionExpression: { enter(path, state) { if (containsUseGpuDirective(path.node)) { - fnNodeToTranspiledMap.set(path.node, transpileFn(path.node)); + fnNodeToTranspiledMap.set( + path.node, + transpile(path.node, this.opts.EXPERIMENTAL_obfuscate), + ); if (state.inUseGpuScope) { throw new Error(`Nesting 'use gpu' functions is not allowed`); } @@ -542,7 +575,10 @@ export const functionVisitor: TraverseOptions = { FunctionExpression: { enter(path, state) { if (containsUseGpuDirective(path.node)) { - fnNodeToTranspiledMap.set(path.node, transpileFn(path.node)); + fnNodeToTranspiledMap.set( + path.node, + transpile(path.node, this.opts.EXPERIMENTAL_obfuscate), + ); if (state.inUseGpuScope) { throw new Error(`Nesting 'use gpu' functions is not allowed`); } @@ -555,7 +591,10 @@ export const functionVisitor: TraverseOptions = { FunctionDeclaration: { enter(path, state) { if (containsUseGpuDirective(path.node)) { - fnNodeToTranspiledMap.set(path.node, transpileFn(path.node)); + fnNodeToTranspiledMap.set( + path.node, + transpile(path.node, this.opts.EXPERIMENTAL_obfuscate), + ); if (state.inUseGpuScope) { throw new Error(`Nesting 'use gpu' functions is not allowed`); } @@ -584,7 +623,7 @@ export const functionVisitor: TraverseOptions = { t.ArrowFunctionExpression | t.FunctionDeclaration | t.FunctionExpression >, getFunctionName(path.get('arguments.0')), - transpileFn(implementation), + transpile(implementation, this.opts.EXPERIMENTAL_obfuscate), ); } } diff --git a/packages/unplugin-typegpu/src/core/factory.ts b/packages/unplugin-typegpu/src/core/factory.ts index 4668894469..69e7ebb29d 100644 --- a/packages/unplugin-typegpu/src/core/factory.ts +++ b/packages/unplugin-typegpu/src/core/factory.ts @@ -13,6 +13,7 @@ import { functionVisitor, getBlockScope, METADATA_FORMAT_VERSION, + checkOpts, } from './common.ts'; import type { Options, UnpluginPluginState, MetadatableFunction, NodeLocation } from './common.ts'; @@ -33,7 +34,7 @@ function embedJSON(jsValue: unknown) { } function externalsToString(externals: Externals): string { - const entries = Array.from(externals, (key) => `"${key}":() => ${key}`); + const entries = Array.from(externals, ([key, value]) => `"${key}":() => ${value}`); return `{${entries.join(',')}}`; } @@ -151,7 +152,7 @@ const NodeUtils = { }; export const unpluginFactory = ((rawOptions, _meta) => { - const options = defu(rawOptions, defaultOptions); + const options = checkOpts(defu(rawOptions, defaultOptions)); return { name: 'unplugin-typegpu' as const, diff --git a/packages/unplugin-typegpu/src/core/obfuscate.ts b/packages/unplugin-typegpu/src/core/obfuscate.ts new file mode 100644 index 0000000000..7a99648352 --- /dev/null +++ b/packages/unplugin-typegpu/src/core/obfuscate.ts @@ -0,0 +1,215 @@ +import type { transpileFn } from 'tinyest-for-wgsl'; +import * as tinyest from 'tinyest'; +const { NodeTypeCatalog: NODE } = tinyest; + +/** + * Generates all strings consisting of lowercase letters of the given length. + */ +function* fixedLengthNameGenerator(length: number): Generator { + if (length === 0) { + yield ''; + return; + } + + for (let i = 97 /* ASCII a */; i <= 122 /* ASCII z */; i++) { + for (const name of fixedLengthNameGenerator(length - 1)) { + yield `${String.fromCharCode(i)}${name}`; + } + } +} + +/** + * Generates all strings consisting of lowercase letters. + */ +function* nameGenerator(): Generator { + for (let i = 1; ; i++) { + for (const name of fixedLengthNameGenerator(i)) { + yield name; + } + } +} + +class Obfuscator { + #nameMap: Map = new Map(); + #nameGenerator: Generator = nameGenerator(); + + #generateFreshName(): string { + return this.#nameGenerator.next().value; + } + + /** + * If `name` wasn't obfuscated before, give it a new obfuscated name. + * Then, returns the obfuscated version of `name`. + */ + obfuscate(name: string): string { + let obfuscatedName = this.#nameMap.get(name); + if (!obfuscatedName) { + obfuscatedName = this.#generateFreshName(); + this.#nameMap.set(name, obfuscatedName); + } + + return obfuscatedName; + } +} + +class Context { + obfuscator: Obfuscator; + + constructor() { + this.obfuscator = new Obfuscator(); + } +} + +export function obfuscate(fn: ReturnType): ReturnType { + const ctx = new Context(); + + const params = fn.params.map((param) => { + if (param.type === 'i') { + return { ...param, name: ctx.obfuscator.obfuscate(param.name) }; + } + // We cannot obfuscate destructured names, because WGSL generation relies on these names (e.g. `$instanceIndex`). + return { + ...param, + props: param.props.map((prop) => ({ ...prop, alias: ctx.obfuscator.obfuscate(prop.alias) })), + }; + }); + + const body = obf(ctx, fn.body); + + const externalNames = new Map(); + fn.externalNames.forEach((value, key) => externalNames.set(ctx.obfuscator.obfuscate(key), value)); + + return { params, body, externalNames }; +} + +// Nodes like 'continue' and 'break' are still listed +// instead of just falling back to node copy when a node is missing, +// so that types will warn us when a new node is added. +const visitors = { + block(ctx: Context, node: tinyest.Block) { + return [NODE.block, node[1].map((node) => obf(ctx, node))]; + }, + binaryExpr(ctx: Context, node: tinyest.BinaryExpression) { + return [NODE.binaryExpr, obf(ctx, node[1]), node[2], obf(ctx, node[3])]; + }, + assignmentExpr(ctx: Context, node: tinyest.AssignmentExpression) { + return [NODE.assignmentExpr, obf(ctx, node[1]), node[2], obf(ctx, node[3])]; + }, + logicalExpr(ctx: Context, node: tinyest.LogicalExpression) { + return [NODE.logicalExpr, obf(ctx, node[1]), node[2], obf(ctx, node[3])]; + }, + unaryExpr(ctx: Context, node: tinyest.UnaryExpression) { + return [NODE.unaryExpr, node[1], obf(ctx, node[2])]; + }, + numericLiteral(_ctx: Context, node: tinyest.Num) { + return [NODE.numericLiteral, node[1]]; + }, + call(ctx: Context, node: tinyest.Call) { + return [NODE.call, obf(ctx, node[1]), node[2].map((node) => obf(ctx, node))]; + }, + memberAccess(ctx: Context, node: tinyest.MemberAccess) { + return [NODE.memberAccess, obf(ctx, node[1]), /* intentionally omitted */ node[2]]; + }, + indexAccess(ctx: Context, node: tinyest.IndexAccess) { + return [NODE.indexAccess, obf(ctx, node[1]), obf(ctx, node[2])]; + }, + return(ctx: Context, node: tinyest.Return) { + return node.length === 1 ? [NODE.return] : [NODE.return, obf(ctx, node[1])]; + }, + if(ctx: Context, node: tinyest.If) { + return node.length === 3 + ? [NODE.if, obf(ctx, node[1]), obf(ctx, node[2])] + : [NODE.if, obf(ctx, node[1]), obf(ctx, node[2]), obf(ctx, node[3])]; + }, + let(ctx: Context, node: tinyest.Let) { + return node.length === 2 + ? [NODE.let, obf(ctx, node[1])] + : [NODE.let, obf(ctx, node[1]), obf(ctx, node[2])]; + }, + const(ctx: Context, node: tinyest.Const) { + return node.length === 2 + ? [NODE.const, obf(ctx, node[1])] + : [NODE.const, obf(ctx, node[1]), obf(ctx, node[2])]; + }, + for(ctx: Context, node: tinyest.For) { + return [NODE.for, obf(ctx, node[1]), obf(ctx, node[2]), obf(ctx, node[3]), obf(ctx, node[4])]; + }, + while(ctx: Context, node: tinyest.While) { + return [NODE.while, obf(ctx, node[1]), obf(ctx, node[2])]; + }, + continue(_ctx: Context, _node: tinyest.Continue) { + return [NODE.continue]; + }, + break(_ctx: Context, _node: tinyest.Break) { + return [NODE.break]; + }, + forOf(ctx: Context, node: tinyest.ForOf) { + return [NODE.forOf, obf(ctx, node[1]), obf(ctx, node[2]), obf(ctx, node[3])]; + }, + arrayExpr(ctx: Context, node: tinyest.ArrayExpression) { + return [NODE.arrayExpr, node[1].map((node) => obf(ctx, node))]; + }, + preUpdate(ctx: Context, node: tinyest.PreUpdate) { + return [NODE.preUpdate, node[1], obf(ctx, node[2])]; + }, + postUpdate(ctx: Context, node: tinyest.PostUpdate) { + return [NODE.postUpdate, node[1], obf(ctx, node[2])]; + }, + stringLiteral(_ctx: Context, node: tinyest.Str) { + return [NODE.stringLiteral, node[1]]; + }, + objectExpr(ctx: Context, node: tinyest.ObjectExpression) { + return [ + NODE.objectExpr, + Object.fromEntries( + Object.entries(node[1]).map(([key, value]) => [ + /* intentionally omitted */ key, + obf(ctx, value), + ]), + ), + ]; + }, + conditionalExpr(ctx: Context, node: tinyest.ConditionalExpression) { + return [NODE.conditionalExpr, obf(ctx, node[1]), obf(ctx, node[2]), obf(ctx, node[3])]; + }, +} as const satisfies { + [N in keyof typeof NODE]: ( + ctx: Context, + node: Extract, + ) => tinyest.AnyNode; +}; + +const nodeIdToName = new Map(Object.entries(NODE).map(([key, value]) => [value, key])) as Map< + number, + keyof typeof NODE +>; + +/** + * Traverses the AST and generates a new one that is obfuscated. + * Copies old AST when identifiers cannot appear in a subtree, + * e.g. in a member access property, or for operator nodes ('=', '<', ...). + */ +function obf(ctx: Context, node: T): T { + if (node === null) { + return node; + } + + if (typeof node === 'string') { + // If we got here, then this identifier should be obfuscated. + return ctx.obfuscator.obfuscate(node) as T; + } + + if (typeof node === 'boolean') { + return node; + } + + const nodeName: keyof typeof visitors | undefined = nodeIdToName.get(node[0]); + if (nodeName === undefined) { + throw new Error(`Internal error, no name for node type ${node[0]}.`); + } + const visitor = visitors[nodeName] as unknown as ((ctx: Context, node: T) => T) | undefined; + if (!visitor) { + throw new Error(`Internal error, no visitor for node '${nodeName}'.`); + } + return visitor(ctx, node); +} diff --git a/packages/unplugin-typegpu/test/obfuscation.test.ts b/packages/unplugin-typegpu/test/obfuscation.test.ts new file mode 100644 index 0000000000..e90cb595b1 --- /dev/null +++ b/packages/unplugin-typegpu/test/obfuscation.test.ts @@ -0,0 +1,594 @@ +import { type ArrowFunctionExpression } from '@babel/types'; +import { transpileFn } from 'tinyest-for-wgsl'; +import { describe, expect, it, test } from 'vitest'; +import { obfuscate } from '../src/core/obfuscate.ts'; +import babelParser from '@babel/parser'; +import { stringifyNode } from 'typegpu/~internal'; +import { babelTransform, rollupTransform } from './transform.ts'; +import { bunPlugin, rollupPlugin } from '../src/index.ts'; +import { defaultOptions } from '../src/core/common.ts'; + +describe('plugin obfuscation', () => { + describe('assigns obfuscated metadata', () => { + const code = `\ + import { tgpu } from 'typegpu'; + + const external = { n: 1 } + + export const fn = (argument) => { + 'use gpu'; + const variable = 3; + return external.n + argument + variable; + };`; + + test('[BABEL]', () => { + expect(babelTransform(code, { EXPERIMENTAL_obfuscate: true })).toMatchInlineSnapshot(` + "import { tgpu } from 'typegpu'; + const external = { + n: 1 + }; + export const fn = /*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = argument => { + const variable = 3; + return __tsover_add(__tsover_add(external.n, argument), variable); + }, { + v: 2, + name: "fn", + ast: { + params: [{ + type: "i", + name: "a" + }], + body: [0, [[13, "b", [5, "3"]], [10, [1, [1, "c", "+", "a"], "+", "b"]]]] + }, + externals: { + "c": () => external.n + } + }) && $.f)({});" + `); + }); + + test('[ROLLUP]', async () => { + expect(await rollupTransform(code, { EXPERIMENTAL_obfuscate: true })).toMatchInlineSnapshot(` + "import 'typegpu'; + + const external = { n: 1 }; + + const fn = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = ((argument) => { + + const variable = 3; + return __tsover_add(__tsover_add(external.n, argument), variable); + }), { + v: 2, + name: "fn", + ast: {"params":[{"type":"i","name":"a"}],"body":[0,[[13,"b",[5,"3"]],[10,[1,[1,"c","+","a"],"+","b"]]]]}, + externals: {"c":() => external.n} + }) && $.f)({})); + + export { fn }; + " + `); + }); + }); + + describe('weird identifiers', () => { + const code = ` + import { tgpu } from 'typegpu'; + + export const fn = () => { + 'use gpu'; + const a = undefined; + const b = Infinity; + const c = NaN; + }`; + + test('[BABEL]', () => { + expect(babelTransform(code, { EXPERIMENTAL_obfuscate: true })).toMatchInlineSnapshot(` + "import { tgpu } from 'typegpu'; + export const fn = /*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = () => { + const a = undefined; + const b = Infinity; + const c = NaN; + }, { + v: 2, + name: "fn", + ast: { + params: [], + body: [0, [[13, "a", "b"], [13, "c", "d"], [13, "e", "f"]]] + }, + externals: { + "b": () => undefined, + "d": () => Infinity, + "f": () => NaN + } + }) && $.f)({});" + `); + }); + + test('[ROLLUP]', async () => { + expect(await rollupTransform(code, { EXPERIMENTAL_obfuscate: true })).toMatchInlineSnapshot(` + "import 'typegpu'; + + const fn = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (() => { + }), { + v: 2, + name: "fn", + ast: {"params":[],"body":[0,[[13,"a","b"],[13,"c","d"],[13,"e","f"]]]}, + externals: {"b":() => undefined,"d":() => Infinity,"f":() => NaN} + }) && $.f)({})); + + export { fn }; + " + `); + }); + }); + + describe('conflicting options', () => { + test('[BABEL]', () => { + expect(() => + babelTransform('', { EXPERIMENTAL_obfuscate: true, autoNamingEnabled: true }), + ).toThrowErrorMatchingInlineSnapshot( + `[Error: unknown file: Options 'EXPERIMENTAL_obfuscate' and 'autoNamingEnabled' cannot be enabled at the same time.]`, + ); + }); + + test('[ROLLUP]', async () => { + expect(() => + rollupPlugin({ ...defaultOptions, EXPERIMENTAL_obfuscate: true, autoNamingEnabled: true }), + ).toThrowErrorMatchingInlineSnapshot( + `[Error: Options 'EXPERIMENTAL_obfuscate' and 'autoNamingEnabled' cannot be enabled at the same time.]`, + ); + }); + + test('[BUN]', async () => { + expect(() => + bunPlugin({ autoNamingEnabled: true, EXPERIMENTAL_obfuscate: true }), + ).toThrowErrorMatchingInlineSnapshot( + `[Error: Options 'EXPERIMENTAL_obfuscate' and 'autoNamingEnabled' cannot be enabled at the same time.]`, + ); + }); + }); +}); + +// Here, we only test tinyest -> tinyest transformation. +// We could write tinyest by hand, but this is more readable. +function parse(code: string): ArrowFunctionExpression { + const parsed = babelParser.parse(code, { sourceType: 'module', plugins: ['typescript'] }); + const maybeExpressionStatement = parsed.program.body[0]; + if (maybeExpressionStatement?.type !== 'ExpressionStatement') { + throw new Error( + `Invalid parse usage. Expected an expression statement (got ${maybeExpressionStatement?.type}).`, + ); + } + const maybeFunction = maybeExpressionStatement.expression; + if (maybeFunction?.type !== 'ArrowFunctionExpression') { + throw new Error( + `Invalid parse usage. Expected an arrow function expression (got ${maybeFunction?.type}).`, + ); + } + return maybeFunction; +} + +describe('obfuscate', () => { + it('obfuscates used variables', () => { + const code = `() => { const variable = 1; const other = 2; const sensitiveName = 3; }`; + const transpiled = transpileFn(parse(code)); + + const { params, body, externalNames } = obfuscate(transpiled); + + expect(params).toMatchInlineSnapshot(`[]`); + expect(stringifyNode(body)).toMatchInlineSnapshot(` + "{ + const a = 1; + const b = 2; + const c = 3; + }" + `); + expect(externalNames).toMatchInlineSnapshot(`Map {}`); + }); + + it('remembers obfuscated names', () => { + const code = `() => { const variable = 1; return variable; }`; + const transpiled = transpileFn(parse(code)); + + const { params, body, externalNames } = obfuscate(transpiled); + + expect(params).toMatchInlineSnapshot(`[]`); + expect(stringifyNode(body)).toMatchInlineSnapshot(` + "{ + const a = 1; + return a; + }" + `); + expect(externalNames).toMatchInlineSnapshot(`Map {}`); + }); + + it('remembers obfuscated names in computed access', () => { + const code = `() => { const variable = 1; const array = [1, 2]; return array[variable]; }`; + const transpiled = transpileFn(parse(code)); + + const { params, body, externalNames } = obfuscate(transpiled); + + expect(params).toMatchInlineSnapshot(`[]`); + expect(stringifyNode(body)).toMatchInlineSnapshot(` + "{ + const a = 1; + const b = [1, 2]; + return b[a]; + }" + `); + expect(externalNames).toMatchInlineSnapshot(`Map {}`); + }); + + it('remembers obfuscated names in for loops', () => { + const code = `() => { for (let i = 0; i< 10; i++) { return i; } }`; + const transpiled = transpileFn(parse(code)); + + const { params, body, externalNames } = obfuscate(transpiled); + + expect(params).toMatchInlineSnapshot(`[]`); + expect(stringifyNode(body)).toMatchInlineSnapshot(` + "{ + for (let a = 0; a < 10; a++) { + return a; + } + }" + `); + expect(externalNames).toMatchInlineSnapshot(`Map {}`); + }); + + it('handles weird identifiers', () => { + const code = `() => { + const a = undefined; + const b = Infinity; + const c = NaN; + }`; + const transpiled = transpileFn(parse(code)); + + const { params, body, externalNames } = obfuscate(transpiled); + + expect(params).toStrictEqual([]); + expect(stringifyNode(body)).toMatchInlineSnapshot(` + "{ + const a = b; + const c = d; + const e = f; + }" + `); + // These are identifiers, so they should be in externals. + expect(externalNames).toMatchInlineSnapshot(` + Map { + "b" => "undefined", + "d" => "Infinity", + "f" => "NaN", + } + `); + }); + + it('obfuscates parameters', () => { + const code = `(param1, param2) => { return param2 + param1; }`; + const transpiled = transpileFn(parse(code)); + + const { params, body, externalNames } = obfuscate(transpiled); + + expect(params).toMatchInlineSnapshot(` + [ + { + "name": "a", + "type": "i", + }, + { + "name": "b", + "type": "i", + }, + ] + `); + expect(stringifyNode(body)).toMatchInlineSnapshot(` + "{ + return b + a; + }" + `); + expect(externalNames).toMatchInlineSnapshot(`Map {}`); + }); + + it('obfuscates destructured parameters', () => { + const code = `(param, { prop }) => { return param + prop; }`; + const transpiled = transpileFn(parse(code)); + + const { params, body, externalNames } = obfuscate(transpiled); + + expect(params).toMatchInlineSnapshot(` + [ + { + "name": "a", + "type": "i", + }, + { + "props": [ + { + "alias": "b", + "name": "prop", + }, + ], + "type": "d", + }, + ] + `); + expect(stringifyNode(body)).toMatchInlineSnapshot(` + "{ + return a + b; + }" + `); + expect(externalNames).toMatchInlineSnapshot(`Map {}`); + }); + + it('obfuscates destructured parameters with aliases', () => { + const code = `(param, { prop, other: alias }) => { return param + prop + alias; }`; + const transpiled = transpileFn(parse(code)); + + const { params, body, externalNames } = obfuscate(transpiled); + + expect(params).toMatchInlineSnapshot(` + [ + { + "name": "a", + "type": "i", + }, + { + "props": [ + { + "alias": "b", + "name": "prop", + }, + { + "alias": "c", + "name": "other", + }, + ], + "type": "d", + }, + ] + `); + expect(stringifyNode(body)).toMatchInlineSnapshot(` + "{ + return (a + b) + c; + }" + `); + expect(externalNames).toMatchInlineSnapshot(`Map {}`); + }); + + it('does not obfuscate struct props', () => { + const code = `(param) => { let struct; return param.prop + struct.field; }`; + const transpiled = transpileFn(parse(code)); + + const { params, body, externalNames } = obfuscate(transpiled); + + expect(params).toMatchInlineSnapshot(` + [ + { + "name": "a", + "type": "i", + }, + ] + `); + expect(stringifyNode(body)).toMatchInlineSnapshot(` + "{ + let b; + return a.prop + b.field; + }" + `); + expect(externalNames).toMatchInlineSnapshot(`Map {}`); + }); + + it('does not obfuscate struct keys', () => { + const code = `(param) => { let struct = { field: 1 }; return struct.field; }`; + const transpiled = transpileFn(parse(code)); + + const { params, body, externalNames } = obfuscate(transpiled); + + expect(params).toMatchInlineSnapshot(` + [ + { + "name": "a", + "type": "i", + }, + ] + `); + expect(stringifyNode(body)).toMatchInlineSnapshot(` + "{ + let b = { field: 1 }; + return b.field; + }" + `); + expect(externalNames).toMatchInlineSnapshot(`Map {}`); + }); + + it("obfuscates 'this'", () => { + const code = `() => { return this.prop1.prop2; }`; + const transpiled = transpileFn(parse(code)); + + const { params, body, externalNames } = obfuscate(transpiled); + + expect(params).toMatchInlineSnapshot(`[]`); + expect(stringifyNode(body)).toMatchInlineSnapshot(` + "{ + return a; + }" + `); + expect(externalNames).toMatchInlineSnapshot(` + Map { + "a" => "this.prop1.prop2", + } + `); + }); + + it('obfuscates externals', () => { + const code = `() => { + const var1 = ext.value; + const var2 = ext.config.multiplier; + const var3 = ext.config.zero; + const var4 = ext.config.multiplier; + }`; + const transpiled = transpileFn(parse(code)); + + const { params, body, externalNames } = obfuscate(transpiled); + + expect(params).toMatchInlineSnapshot(`[]`); + expect(stringifyNode(body)).toMatchInlineSnapshot(` + "{ + const a = b; + const c = d; + const e = f; + const g = d; + }" + `); + expect(externalNames).toMatchInlineSnapshot(` + Map { + "b" => "ext.value", + "d" => "ext.config.multiplier", + "f" => "ext.config.zero", + } + `); + }); + + it('obfuscates complex externals', () => { + const code = `() => { + const h = ext.t.fn().prop; + const i = ext.t.comp['computed'].prop; + const j = ext.t.$.prop; + const k = (ext).prop; + }`; + const transpiled = transpileFn(parse(code)); + + const { params, body, externalNames } = obfuscate(transpiled); + + expect(params).toMatchInlineSnapshot(`[]`); + expect(stringifyNode(body)).toMatchInlineSnapshot(` + "{ + const a = b().prop; + const c = d["computed"].prop; + const e = f.$.prop; + const g = h; + }" + `); + expect(externalNames).toMatchInlineSnapshot(` + Map { + "b" => "ext.t.fn", + "d" => "ext.t.comp", + "f" => "ext.t", + "h" => "ext.prop", + } + `); + }); + + it('correctly handles variable shadowing', () => { + const code = `() => { + const variable = 1; + { + const variable = 2; + if (false) { + return variable; + } + } + return variable; + }`; + const transpiled = transpileFn(parse(code)); + + const { params, body, externalNames } = obfuscate(transpiled); + + expect(params).toMatchInlineSnapshot(`[]`); + expect(stringifyNode(body)).toMatchInlineSnapshot(` + "{ + const a = 1; + { + const a = 2; + if (false) { + return a; + } + } + return a; + }" + `); + expect(externalNames).toMatchInlineSnapshot(`Map {}`); + }); + + it('correctly handles parameter shadowing', () => { + const code = `(parameter) => { + { + const parameter = 2; + if (false) { + return parameter; + } + } + return parameter; + }`; + const transpiled = transpileFn(parse(code)); + + const { params, body, externalNames } = obfuscate(transpiled); + + expect(params).toMatchInlineSnapshot(` + [ + { + "name": "a", + "type": "i", + }, + ] + `); + expect(stringifyNode(body)).toMatchInlineSnapshot(` + "{ + { + const a = 2; + if (false) { + return a; + } + } + return a; + }" + `); + expect(externalNames).toMatchInlineSnapshot(`Map {}`); + }); + + it('correctly handles external shadowing', () => { + const code = `() => { + const variable = external; + { + const external = 1; + return external; + } + return external; + }`; + const transpiled = transpileFn(parse(code)); + + const { params, body, externalNames } = obfuscate(transpiled); + + expect(params).toMatchInlineSnapshot(`[]`); + expect(stringifyNode(body)).toMatchInlineSnapshot(` + "{ + const a = b; + { + const b = 1; + return b; + } + return b; + }" + `); + expect(externalNames).toMatchInlineSnapshot(` + Map { + "b" => "external", + } + `); + }); + + it('supports more than 26 names', () => { + const code = `() => { ${Array.from({ length: 100 }, (_, i) => `let v${i};`).join('\n')} }`; + const transpiled = transpileFn(parse(code)); + + const { params, body, externalNames } = obfuscate(transpiled); + + expect(params).toMatchInlineSnapshot(`[]`); + const stringifiedBody = stringifyNode(body); + expect(stringifiedBody).toContain('z'); + expect(stringifiedBody).toContain('aa'); + expect(stringifiedBody).toContain('ab'); + expect(externalNames).toMatchInlineSnapshot(`Map {}`); + }); +});