feat: Obfuscate variables in plugin

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',

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 result in changes to 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.

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: [

packages/tinyest-for-wgsl/src/types.ts

@@ -7,7 +7,7 @@ export type Scope = {
   declaredNames: string[];
 };
 
-export type Externals = Set<string>;
+export type Externals = Map<string, string>;
 
 export type Context = {
   /** Holds a set of all identifiers that were used in code, but were not declared in code. */

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);
+  };
+}

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",
         }
       `);
     }),

packages/typegpu/src/internal.ts

@@ -5,6 +5,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';

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;
       }
 

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');
   });