Impr(unplugin-typegpu): Support for early pruning in Babel plugin

packages/unplugin-typegpu/src/babel.ts

@@ -9,6 +9,7 @@ import {
   defaultOptions,
   functionVisitor,
   getBlockScope,
+  getEarlyPruneRegex,
   initPluginState,
 } from './core/common.ts';
 import type { FunctionAst, FunctionExternals } from './core/embeddedMetadata.ts';
@@ -199,6 +200,15 @@ export default function TypeGPUPlugin() {
           return;
         }
 
+        const code = path.hub.getCode();
+        if (
+          state.opts.earlyPruning &&
+          code &&
+          getEarlyPruneRegex(state.opts.unstable_obfuscate).every((pattern) => !pattern.test(code))
+        ) {
+          return;
+        }
+
         path.traverse(functionVisitor, state);
       },
     } satisfies TraverseOptions<PluginState>,

packages/unplugin-typegpu/test/auto-naming.test.ts

@@ -283,7 +283,8 @@ describe('[BABEL] auto naming', () => {
       };
     `;
 
-    expect(babelTransform(code, { autoNamingEnabled: true })).toMatchInlineSnapshot(`
+    expect(babelTransform(code, { autoNamingEnabled: true, earlyPruning: false }))
+      .toMatchInlineSnapshot(`
       "const mySchemas = {
         myStruct: /*#__PURE__*/(globalThis.__TYPEGPU_AUTONAME__ ?? (a => a))(d.struct({
           a: d.vec3f

packages/unplugin-typegpu/test/early-pruning.test.ts

@@ -0,0 +1,25 @@
+import { describe, expect, it } from 'vitest';
+import { babelTransform } from './transform.ts';
+
+describe('[BABEL] early pruning', () => {
+  const code = `const buffer = root.createBuffer(d.f32);`;
+
+  it('assigns an auto-name when early pruning is disabled', () => {
+    expect(babelTransform(code, { autoNamingEnabled: true, earlyPruning: false })).toContain(
+      '__TYPEGPU_AUTONAME__',
+    );
+  });
+
+  it('skips files without keywords', () => {
+    expect(babelTransform(code, { autoNamingEnabled: true, earlyPruning: true })).toBe(code);
+  });
+
+  it.each(['typegpu', 'tgpu', "'use gpu'", '"use gpu"'])(
+    'processes files containing %s',
+    (marker) => {
+      expect(
+        babelTransform(`// ${marker}\n${code}`, { autoNamingEnabled: true, earlyPruning: true }),
+      ).toContain('__TYPEGPU_AUTONAME__');
+    },
+  );
+});

packages/unplugin-typegpu/test/obfuscation.test.ts

@@ -6,7 +6,7 @@ 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';
+import * as common from '../src/core/common.ts';
 
 describe('plugin obfuscation', () => {
   describe('assigns obfuscated metadata', () => {
@@ -309,21 +309,46 @@ describe('plugin obfuscation', () => {
     });
   });
 
-  test('does not prune files that only contain embedded metadata when obfuscation is on', async () => {
+  describe('does not prune files that only contain embedded metadata when obfuscation is on', () => {
     const code = `\
       const external = { n: 1 };
       export const fn = () => {
         'use gpu';
         return external.n;
       };`;
 
-    const built = await rollupTransform(code, { unstable_obfuscate: false });
+    test('[BABEL]', () => {
+      const built = babelTransform(code, { unstable_obfuscate: false }) as string;
+      const withoutDirective = built.replace(/^[ \t]*['"]use gpu['"];?\r?\n/gm, '');
+      expect(withoutDirective).not.toMatch(/use gpu|t(ype)?gpu/);
+
+      expect(babelTransform(withoutDirective, { unstable_obfuscate: true })).toMatchInlineSnapshot(`
+          "const external = {
+            n: 1
+          };
+          export const fn = /*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = () => {
+            return external.n;
+          }, {
+            v: 2,
+            name: "fn",
+            ast: {
+              params: [],
+              body: [0, [[10, "a"]]]
+            },
+            externals: {
+              "a": () => external.n
+            }
+          }) && $.f)({});"
+        `);
+    });
 
-    const withoutDirective = built.replace(/^[ \t]*['"]use gpu['"];?\r?\n/gm, '');
-    expect(withoutDirective).not.toMatch(/use gpu|t(ype)?gpu/);
+    test('[ROLLUP]', async () => {
+      const built = await rollupTransform(code, { unstable_obfuscate: false });
+      const withoutDirective = built.replace(/^[ \t]*['"]use gpu['"];?\r?\n/gm, '');
+      expect(withoutDirective).not.toMatch(/use gpu|t(ype)?gpu/);
 
-    expect(await rollupTransform(withoutDirective, { unstable_obfuscate: true }))
-      .toMatchInlineSnapshot(`
+      expect(await rollupTransform(withoutDirective, { unstable_obfuscate: true }))
+        .toMatchInlineSnapshot(`
       "const external = { n: 1 };
             const fn = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (() => {
               return external.n;
@@ -336,22 +361,36 @@ describe('plugin obfuscation', () => {
 
       export { fn };
       "
-    `);
+      `);
+    });
   });
 
-  test('prunes files that only contain embedded metadata when obfuscation is off', async () => {
+  describe('prunes files that only contain embedded metadata when obfuscation is off', () => {
     const code = `\
       export const fn = () => {
         'use gpu';
         return 42;
       };`;
 
-    const built = await rollupTransform(code, { unstable_obfuscate: false });
-    const withoutDirective = built.replace(/^[ \t]*['"]use gpu['"];?\r?\n/gm, '');
+    test('[BABEL]', () => {
+      const built = babelTransform(code, { unstable_obfuscate: false }) as string;
+      const withoutDirective = built.replace(/^[ \t]*['"]use gpu['"];?\r?\n/gm, '');
+      expect(withoutDirective).not.toMatch(/use gpu|t(ype)?gpu/);
+
+      using visitor = vi.spyOn(common, 'functionVisitor', 'get');
+      babelTransform(withoutDirective, { unstable_obfuscate: false });
+      expect(visitor).not.toHaveBeenCalled();
+    });
 
-    using parse = vi.spyOn(babelParser, 'parse');
-    await rollupTransform(withoutDirective, { unstable_obfuscate: false });
-    expect(parse).not.toHaveBeenCalled();
+    test('[ROLLUP]', async () => {
+      const built = await rollupTransform(code, { unstable_obfuscate: false });
+      const withoutDirective = built.replace(/^[ \t]*['"]use gpu['"];?\r?\n/gm, '');
+      expect(withoutDirective).not.toMatch(/use gpu|t(ype)?gpu/);
+
+      using visitor = vi.spyOn(common, 'functionVisitor', 'get');
+      await rollupTransform(withoutDirective, { unstable_obfuscate: false });
+      expect(visitor).not.toHaveBeenCalled();
+    });
   });
 
   describe('conflicting options', () => {
@@ -365,7 +404,11 @@ describe('plugin obfuscation', () => {
 
     test('[ROLLUP]', async () => {
       expect(() =>
-        rollupPlugin({ ...defaultOptions, unstable_obfuscate: true, autoNamingEnabled: true }),
+        rollupPlugin({
+          ...common.defaultOptions,
+          unstable_obfuscate: true,
+          autoNamingEnabled: true,
+        }),
       ).toThrowErrorMatchingInlineSnapshot(
         `[Error: Options 'unstable_obfuscate' and 'autoNamingEnabled' cannot be enabled at the same time.]`,
       );