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/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/resolutionCtx.ts b/packages/typegpu/src/resolutionCtx.ts index e3889112ea..feb285b694 100644 --- a/packages/typegpu/src/resolutionCtx.ts +++ b/packages/typegpu/src/resolutionCtx.ts @@ -202,17 +202,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; } 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 {}`); + }); +});