feat(unplugin-typegpu): Support obfuscation of already built sources

apps/typegpu-docs/src/content/docs/advanced/minifying-shaders.mdx

@@ -53,6 +53,14 @@ fn item(a: f32, b: f32) -> f32 {
 Obfuscation does not affect raw-WGSL implemented functions.
 :::
 
+:::caution[Experimental]
+`unplugin-typegpu` also supports obfuscating sources that have already been built with the plugin.
+
+Process those files again with `{ unstable_obfuscate: true, autoNamingEnabled: false }` to obfuscate their embedded shader metadata.
+
+Names added during the original build are preserved, including those added with `{ autoNamingEnabled: true }`.
+:::
+
 
 ## Removing whitespaces in runtime
 
@@ -86,4 +94,4 @@ fn coneVolume(radius:f32,height:f32)->f32{let baseArea=(3.141592653589793f*pow(r
 Of course, minification and obfuscation can be combined to achieve the following:
 ```wgsl
 fn item(a:f32,b:f32)->f32{let c=(3.141592653589793f*pow(a,2f));let e=((c*b)/3f);return e;}
-```
+```

apps/typegpu-docs/src/content/docs/tooling/unplugin-typegpu.mdx

@@ -182,9 +182,9 @@ interface Options {
   autoNamingEnabled?: boolean;
 
   /**
-   * Skipping files that don't contain 'typegpu', 'tgpu' or 'use gpu'.
-   * In case this early pruning hinders transformation, you
-   * can disable it.
+   * Skipping files that don't contain "typegpu", "tgpu" or "use gpu".
+   * When `unstable_obfuscate` is enabled, files containing `__TYPEGPU_META__` are also processed.
+   * In case this early pruning hinders transformation, you can disable it.
    *
    * @default true
    */

packages/unplugin-typegpu/src/babel.ts

@@ -11,6 +11,7 @@ import {
   getBlockScope,
   initPluginState,
 } from './core/common.ts';
+import type { FunctionAst, FunctionExternals } from './core/embeddedMetadata.ts';
 import { createFilterForId } from './core/filter.ts';
 
 function i(identifier: string): t.Identifier {
@@ -126,6 +127,22 @@ function assignMetadata(
   path.skip();
 }
 
+function replaceWithObfuscated(
+  this: PluginState,
+  astPath: NodePath<t.ObjectExpression>,
+  obfuscatedAst: FunctionAst,
+  embeddedExternals: FunctionExternals,
+  externalKeyToObfuscatedMap: Map<string, string>,
+): void {
+  astPath.replaceWith(t.valueToNode(obfuscatedAst));
+
+  for (const [key, propertyPath] of embeddedExternals) {
+    propertyPath
+      .get('key')
+      .replaceWith(t.stringLiteral(externalKeyToObfuscatedMap.get(key) as string));
+  }
+}
+
 function wrapInAutoName(path: NodePath<t.Expression>, name: string): void {
   // /*#__PURE__*/(globalThis.__TYPEGPU_AUTONAME__ ?? (a => a))(<node>, '<name>')
   const callExpr = t.callExpression(
@@ -169,6 +186,7 @@ export default function TypeGPUPlugin() {
       initPluginState(this, {
         warn: (message) => console.warn(message),
         assignMetadata,
+        replaceWithObfuscated,
         wrapInAutoName,
         replaceWithAssignmentOverload,
         replaceWithBinaryOverload,

packages/unplugin-typegpu/src/bun.ts

@@ -1,10 +1,11 @@
 import defu from 'defu';
-import { checkOpts, defaultOptions, earlyPruneRegex, type Options } from './core/common.ts';
+import { checkOpts, defaultOptions, getEarlyPruneRegex, 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 = checkOpts(defu(rawOptions, defaultOptions));
+  const earlyPruneRegex = getEarlyPruneRegex(options.unstable_obfuscate);
   const include = options.include;
   if (!(include instanceof RegExp)) {
     throw new Error(
@@ -32,7 +33,8 @@ export default (rawOptions?: Options): Bun.BunPlugin => {
         }
 
         const result = rawPlugin.transform.handler.apply(
-          {} as UnpluginBuildContext & UnpluginContext,
+          { warn: (message: string) => console.warn(message) } as UnpluginBuildContext &
+            UnpluginContext,
           [codeIn, args.path],
         );
 

packages/unplugin-typegpu/src/core/common.ts

@@ -3,8 +3,14 @@ import type { NodePath, TraverseOptions } from '@babel/traverse';
 import type { FilterPattern } from 'unplugin';
 import MagicString from 'magic-string';
 import { transpileBabelFn, type TranspilationResult } from 'tinyest-for-wgsl';
+import type {
+  EmbeddedTypegpuMetadata,
+  FunctionAst,
+  FunctionExternals,
+} from './embeddedMetadata.ts';
 import { getEmbeddedTypegpuMetadata } from './embeddedMetadata.ts';
-import { obfuscate } from './obfuscate.ts';
+import { obfuscate, obfuscateAst } from './obfuscate.ts';
+import { METADATA_FORMAT_VERSION } from './version.ts';
 
 export interface Options {
   /** @default [/\.m?[jt]sx?$/] */
@@ -32,8 +38,8 @@ export interface Options {
 
   /**
    * Skipping files that don't contain "typegpu", "tgpu" or "use gpu".
-   * In case this early pruning hinders transformation, you
-   * can disable it.
+   * When `unstable_obfuscate` is enabled, files containing `__TYPEGPU_META__` are also processed.
+   * In case this early pruning hinders transformation, you can disable it.
    *
    * @default true
    */
@@ -93,6 +99,19 @@ export interface TransformMethods {
     ast: TranspilationResult,
   ): void;
 
+  /**
+   * Replaces the embedded AST and the keys of the externals object of an
+   * already transformed function with their obfuscated versions.
+   * Externals' values (getters) are left untouched.
+   */
+  replaceWithObfuscated(
+    this: PluginState,
+    astPath: NodePath<t.ObjectExpression>,
+    obfuscatedAst: FunctionAst,
+    embeddedExternals: FunctionExternals,
+    externalKeyToObfuscatedMap: Map<string, string>,
+  ): void;
+
   wrapInAutoName(this: PluginState, path: NodePath<t.Expression>, name: string): void;
 
   replaceWithAssignmentOverload(path: NodePath<t.AssignmentExpression>, runtimeFn: string): void;
@@ -120,6 +139,9 @@ export interface PluginState extends TransformMethods {
   opts: Required<Options>;
 
   inUseGpuScope: boolean;
+
+  /** Used to warn about unsupported metadata versions at most once per file. */
+  warnedAboutMetadataVersion: boolean;
 }
 
 export interface NodeLocation {
@@ -138,11 +160,18 @@ export function initPluginState(state: PluginState, methods: TransformMethods):
   state.tgpuAliases = new Set<string>(state.opts.forceTgpuAlias ? [state.opts.forceTgpuAlias] : []);
   state.autoNamingEnabled = state.opts.autoNamingEnabled ?? true;
   state.inUseGpuScope = false;
+  state.warnedAboutMetadataVersion = false;
   Object.assign(state, methods);
 }
 
 /** Regular expressions used for early pruning (to avoid unnecessary parsing, which is expensive) */
-export const earlyPruneRegex = [/["']use gpu["']/, /t(ype)?gpu/];
+export const getEarlyPruneRegex = (obfuscate: boolean) => {
+  const baseRegex = [/["']use gpu["']/, /t(ype)?gpu/];
+  if (obfuscate) {
+    baseRegex.push(/__TYPEGPU_META__/);
+  }
+  return baseRegex;
+};
 
 export const defaultOptions = {
   include: /\.m?[jt]sx?(?:\?.*)?$/,
@@ -463,12 +492,39 @@ const fnNodeToTranspiledMap = new WeakMap<
   TranspilationResult
 >();
 
-function functionOnExit(
-  path: NodePath<t.FunctionDeclaration | t.FunctionExpression | t.ArrowFunctionExpression>,
-  state: PluginState,
-) {
+function reobfuscateEmbedded(state: PluginState, metadata: EmbeddedTypegpuMetadata) {
+  if (metadata.v === METADATA_FORMAT_VERSION && metadata.function !== undefined) {
+    const { ast, astPath, externals } = metadata.function;
+    const obfuscated = obfuscateAst(ast, externals.keys());
+    state.replaceWithObfuscated(astPath, obfuscated.ast, externals, obfuscated.externalKeyMap);
+    return;
+  }
+
+  if (state.warnedAboutMetadataVersion) {
+    return;
+  }
+
+  state.warnedAboutMetadataVersion = true;
+  state.warn(
+    `[unplugin-typegpu] Cannot obfuscate already transformed functions in ${state.filename ?? '<unknown>'}: unsupported metadata version ${metadata.v} (expected ${METADATA_FORMAT_VERSION}).`,
+  );
+}
+
+function functionOnExit(path: NodePath<MetadatableFunction>, state: PluginState) {
   const node = path.node;
-  if (!containsUseGpuDirective(node) || getEmbeddedTypegpuMetadata(path) !== undefined) {
+  const hasDirective = containsUseGpuDirective(node);
+  const obfuscate = state.opts.unstable_obfuscate;
+
+  if (!hasDirective && !obfuscate) {
+    return;
+  }
+
+  const metadata = getEmbeddedTypegpuMetadata(path);
+  if (metadata !== undefined && obfuscate) {
+    reobfuscateEmbedded(state, metadata);
+  }
+
+  if (!hasDirective || metadata !== undefined) {
     return;
   }
 

packages/unplugin-typegpu/src/core/embeddedMetadata.ts

@@ -7,17 +7,19 @@ import type { MetadatableFunction } from './common.ts';
 const METADATA_PARSING_ERROR_MESSAGE =
   '[unplugin-typegpu] Error when parsing metadata: required fields are missing or could not be evaluated.';
 
-type FunctionAst = Pick<TranspilationResult, 'params' | 'body'>;
-type FunctionExternals = Map<string, NodePath<t.ObjectProperty>>;
+export type FunctionAst = Pick<TranspilationResult, 'params' | 'body'>;
+export type FunctionExternals = Map<string, NodePath<t.ObjectProperty>>;
+
+type EmbeddedFunctionMetadata = {
+  ast: FunctionAst;
+  astPath: NodePath<t.ObjectExpression>;
+  externals: FunctionExternals;
+};
 
 export type EmbeddedTypegpuMetadata = {
   v: number;
   name: string | undefined;
-  function?: {
-    ast: FunctionAst;
-    astPath: NodePath<t.ObjectExpression>;
-    externals: FunctionExternals;
-  };
+  function?: EmbeddedFunctionMetadata;
 };
 
 const embeddedTypegpuMetadataCache = new WeakMap<

packages/unplugin-typegpu/src/core/factory.ts

@@ -9,14 +9,14 @@ import * as t from '@babel/types';
 import { METADATA_FORMAT_VERSION } from './version.ts';
 import {
   defaultOptions,
-  earlyPruneRegex,
+  getEarlyPruneRegex,
   initPluginState,
   functionVisitor,
   getBlockScope,
   checkOpts,
 } from './common.ts';
-
 import type { Options, UnpluginPluginState, MetadatableFunction, NodeLocation } from './common.ts';
+import type { FunctionAst, FunctionExternals } from './embeddedMetadata.ts';
 
 // I love CommonJS 💔
 let traverse = _traverse;
@@ -97,6 +97,20 @@ function assignMetadata(
   }
 }
 
+function replaceWithObfuscated(
+  this: UnpluginPluginState,
+  astPath: NodePath<t.ObjectExpression>,
+  obfuscatedAst: FunctionAst,
+  embeddedExternals: FunctionExternals,
+  externalKeyToObfuscatedMap: Map<string, string>,
+): void {
+  this.overwrite(astPath.node, embedJSON(obfuscatedAst));
+
+  for (const [key, propertyPath] of embeddedExternals) {
+    this.overwrite(propertyPath.node.key, `"${externalKeyToObfuscatedMap.get(key)}"`);
+  }
+}
+
 function wrapInAutoName(
   this: UnpluginPluginState,
   path: NodePath<t.Expression>,
@@ -152,7 +166,7 @@ export const unpluginFactory = ((rawOptions, _meta) => {
       filter: options.earlyPruning
         ? {
             id: options,
-            code: earlyPruneRegex,
+            code: getEarlyPruneRegex(options.unstable_obfuscate),
           }
         : {
             id: options,
@@ -188,6 +202,7 @@ export const unpluginFactory = ((rawOptions, _meta) => {
         initPluginState(state, {
           warn: (message) => this.warn(message),
           assignMetadata,
+          replaceWithObfuscated,
           wrapInAutoName,
           replaceWithAssignmentOverload,
           replaceWithBinaryOverload,

packages/unplugin-typegpu/src/core/obfuscate.ts

@@ -1,5 +1,7 @@
 import type { TranspilationResult } from 'tinyest-for-wgsl';
 import * as tinyest from 'tinyest';
+import type { FunctionAst } from './embeddedMetadata.ts';
+
 const { NodeTypeCatalog: NODE } = tinyest;
 
 /**
@@ -60,7 +62,12 @@ class Context {
   }
 }
 
-export function obfuscate(fn: TranspilationResult): TranspilationResult {
+type ObfuscationResult = {
+  ast: FunctionAst;
+  externalKeyMap: Map<string, string>;
+};
+
+export function obfuscateAst(fn: FunctionAst, externalKeys: Iterable<string>): ObfuscationResult {
   const ctx = new Context();
 
   const params = fn.params.map((param) => {
@@ -76,10 +83,21 @@ export function obfuscate(fn: TranspilationResult): TranspilationResult {
 
   const body = obf(ctx, fn.body);
 
-  const externalNames = new Map();
-  fn.externalNames.forEach((value, key) => externalNames.set(ctx.obfuscator.obfuscate(key), value));
+  const externalKeyMap = new Map(
+    Array.from(externalKeys, (key) => [key, ctx.obfuscator.obfuscate(key)] as const),
+  );
+
+  return { ast: { params, body }, externalKeyMap };
+}
+
+export function obfuscate(fn: TranspilationResult): TranspilationResult {
+  const { ast, externalKeyMap } = obfuscateAst(fn, fn.externalNames.keys());
+
+  const externalNames = new Map(
+    Array.from(fn.externalNames, ([key, value]) => [externalKeyMap.get(key) as string, value]),
+  );
 
-  return { ...fn, params, body, externalNames };
+  return { ...fn, ...ast, externalNames };
 }
 
 // Nodes like 'continue' and 'break' are still listed

packages/unplugin-typegpu/src/node.ts

@@ -2,7 +2,7 @@ import { registerHooks, type ModuleHooks } from 'node:module';
 import { fileURLToPath } from 'node:url';
 import remapping from '@jridgewell/remapping';
 import type { UnpluginBuildContext, UnpluginContext } from 'unplugin';
-import { earlyPruneRegex, type Options } from './core/common.ts';
+import { getEarlyPruneRegex, type Options } from './core/common.ts';
 import { unpluginFactory } from './core/factory.ts';
 import { createFilterForId } from './core/filter.ts';
 
@@ -52,6 +52,7 @@ export default function install(rawOptions?: Options): ModuleHooks {
     { framework: 'unloader' },
   );
   const filter = createFilterForId(plugin.transform.filter.id);
+  const earlyPruneRegex = getEarlyPruneRegex(rawOptions?.unstable_obfuscate ?? false);
   const transformContext = {
     warn: (message: string) => console.warn(`[unplugin-typegpu] ${message}`),
   } as UnpluginBuildContext & UnpluginContext;