fix(unplugin-typegpu): source-map misalignment when removing directives (& @ark/attest)

.gitignore

@@ -9,6 +9,7 @@ coverage
 # development
 .devcontainer
 .vscode
+.attest
 
 # production
 dist

package.json

@@ -6,12 +6,13 @@
   "license": "MIT",
   "scripts": {
     "dev": "DEV=true pnpm --parallel -r dev:watch",
-    "dev:test": "vitest",
+    "dev:test": "ATTEST_skipTypes=1 vitest",
     "check": "biome check --write .",
     "coverage": "vitest --coverage",
-    "test": "pnpm run '/^test:.*/'",
+    "test": "pnpm run test:types && pnpm run test:spec-and-attest",
     "test:types": "pnpm run -r --parallel test:types",
-    "test:spec": "vitest run",
+    "test:spec": "ATTEST_skipTypes=1 vitest run",
+    "test:spec-and-attest": "vitest run",
     "nightly-build": "pnpm --filter typegpu prepublishOnly --skip-publish-tag-check"
   },
   "engines": {

packages/typegpu/package.json

@@ -65,14 +65,16 @@
   },
   "homepage": "https://typegpu.com",
   "devDependencies": {
+    "@ark/attest": "^0.45.10",
     "@typegpu/tgpu-dev-cli": "workspace:*",
     "@types/node": "^22.13.14",
     "@webgpu/types": "catalog:",
+    "arktype": "catalog:",
     "jiti": "^2.4.2",
-    "unplugin-typegpu": "workspace:*",
+    "tgpu-jit": "workspace:*",
     "tsup": "catalog:",
     "typescript": "catalog:",
-    "tgpu-jit": "workspace:*",
+    "unplugin-typegpu": "workspace:*",
     "wesl": "0.6.0-rc1",
     "wgpu-matrix": "^3.4.0"
   },

packages/typegpu/setupVitest.ts

@@ -0,0 +1,6 @@
+import { setup } from '@ark/attest';
+
+export default () =>
+  setup({
+    formatCmd: 'pnpm check',
+  });

packages/typegpu/tests/array.test.ts

@@ -1,3 +1,4 @@
+import { attest } from '@ark/attest';
 import { BufferReader, BufferWriter } from 'typed-binary';
 import { describe, expect, it } from 'vitest';
 import { readData, writeData } from '../src/data/dataIO.ts';
@@ -9,6 +10,11 @@ import type { Infer } from '../src/shared/repr.ts';
 import { parse, parseResolved } from './utils/parseResolved.ts';
 
 describe('array', () => {
+  it('produces a visually pleasant type', () => {
+    const TestArray = d.arrayOf(d.vec3u, 3);
+    attest(TestArray).type.toString.snap('WgslArray<Vec3u>');
+  });
+
   it('takes element alignment into account when measuring', () => {
     const TestArray = d.arrayOf(d.vec3u, 3);
     expect(d.sizeOf(TestArray)).toEqual(48);

packages/typegpu/vitest.config.mts

@@ -1,13 +1,16 @@
 import { createJiti } from 'jiti';
-import type TypeGPUPlugin from 'unplugin-typegpu/rollup';
+import type TypeGPUPlugin from 'unplugin-typegpu/vite';
 import { defineConfig } from 'vitest/config';
 
 const jiti = createJiti(import.meta.url);
 const typegpu = await jiti.import<typeof TypeGPUPlugin>(
-  'unplugin-typegpu/rollup',
+  'unplugin-typegpu/vite',
   { default: true },
 );
 
 export default defineConfig({
-  plugins: [typegpu({ include: [/.*\.test\.ts/], forceTgpuAlias: 'tgpu' })],
+  plugins: [typegpu({ forceTgpuAlias: 'tgpu' })],
+  test: {
+    globalSetup: ['setupVitest.ts'],
+  },
 });

packages/unplugin-typegpu/package.json

@@ -105,11 +105,13 @@
     "prepublishOnly": "tgpu-dev-cli prepack"
   },
   "dependencies": {
-    "@babel/standalone": "^7.26.6",
+    "@babel/standalone": "^7.27.0",
     "estree-walker": "^3.0.3",
-    "magic-string": "^0.30.11",
+    "magic-string-ast": "^0.9.1",
+    "picomatch": "^4.0.2",
+    "defu": "^6.1.4",
     "tinyest-for-wgsl": "workspace:~0.1.0-alpha.9",
-    "unplugin": "^2.2.0"
+    "unplugin": "^2.3.1"
   },
   "peerDependencies": {
     "typegpu": "^0.5.4"
@@ -122,6 +124,7 @@
     "@types/babel__standalone": "^7.1.9",
     "@types/babel__template": "^7.4.4",
     "@types/babel__traverse": "^7.20.7",
+    "@types/picomatch": "^4.0.0",
     "acorn": "^8.14.1",
     "rollup": "~4.37.0",
     "tsup": "catalog:",

packages/unplugin-typegpu/src/babel.ts

@@ -6,13 +6,14 @@ import { transpileFn } from 'tinyest-for-wgsl';
 import {
   type Context,
   type KernelDirective,
-  type TypegpuPluginOptions,
+  type Options,
+  codeFilterRegexes,
   embedJSON,
   gatherTgpuAliases,
   isShellImplementationCall,
   kernelDirectives,
-  shouldSkipFile,
 } from './common.ts';
+import { createFilterForId } from './filter.ts';
 
 // NOTE: @babel/standalone does expose internal packages, as specified in the docs, but the
 // typing for @babel/standalone does not expose them.
@@ -185,14 +186,23 @@ export default function () {
   return {
     visitor: {
       Program(path, state) {
-        // @ts-ignore
-        const code: string | undefined = state.file?.code;
-        // @ts-ignore
-        const options: TypegpuPluginOptions | undefined = state.opts;
-        // @ts-ignore
-        const id: string | undefined = state.filename;
-
-        if (shouldSkipFile(options, id, code)) {
+        // biome-ignore lint/suspicious/noExplicitAny: <oh babel babel...>
+        const code: string | undefined = (state as any).file?.code;
+        // biome-ignore lint/suspicious/noExplicitAny: <oh babel babel...>
+        const options: Options | undefined = (state as any).opts;
+        // biome-ignore lint/suspicious/noExplicitAny: <oh babel babel...>
+        const id: string | undefined = (state as any).filename;
+
+        const filter = createFilterForId(options);
+        if (id && filter && !filter?.(id)) {
+          return;
+        }
+
+        if (
+          !options?.forceTgpuAlias &&
+          code &&
+          !codeFilterRegexes.some((reg) => reg.test(code))
+        ) {
           return;
         }
 

packages/unplugin-typegpu/src/common.ts

@@ -1,5 +1,6 @@
 import type * as babel from '@babel/types';
 import type * as acorn from 'acorn';
+import type { FilterPattern } from 'unplugin';
 
 export type Context = {
   /**
@@ -10,11 +11,17 @@ export type Context = {
   fileId?: string | undefined;
 };
 
-export interface TypegpuPluginOptions {
-  include?: 'all' | RegExp[];
+export interface Options {
+  include?: FilterPattern;
+  exclude?: FilterPattern;
+  enforce?: 'post' | 'pre' | undefined;
   forceTgpuAlias?: string;
 }
 
+export const defaultOptions = {
+  include: [/\.m?[jt]sx?$/],
+};
+
 export function embedJSON(jsValue: unknown) {
   return JSON.stringify(jsValue)
     .replace(/\u2028/g, '\\u2028')
@@ -53,30 +60,16 @@ const typegpuImportRegex = /import.*from\s*['"]typegpu.*['"]/;
 const typegpuDynamicImportRegex = /import\s*\(\s*['"]\s*typegpu.*['"]/;
 const typegpuRequireRegex = /require\s*\(\s*['"]\s*typegpu.*['"]\s*\)/;
 
-export function shouldSkipFile(
-  options: TypegpuPluginOptions | undefined,
-  id: string | undefined,
-  code: string | undefined,
-) {
-  if (code && !options?.include) {
-    if (
-      !typegpuImportRegex.test(code) &&
-      !typegpuRequireRegex.test(code) &&
-      !typegpuDynamicImportRegex.test(code)
-    ) {
-      // No imports to `typegpu` or its sub modules, exiting early.
-      return true;
-    }
-  } else if (
-    options?.include &&
-    options.include !== 'all' &&
-    (!id || !options.include.some((pattern) => pattern.test(id)))
-  ) {
-    return true;
-  }
-
-  return false;
-}
+/**
+ * Regexes used to efficiently determine if a file is
+ * meant to be processed by our plugin. We assume every file
+ * that should be processed imports `typegpu` in some way.
+ */
+export const codeFilterRegexes = [
+  typegpuImportRegex,
+  typegpuDynamicImportRegex,
+  typegpuRequireRegex,
+];
 
 export function gatherTgpuAliases(
   node: acorn.ImportDeclaration | babel.ImportDeclaration,

packages/unplugin-typegpu/src/filter.ts

@@ -0,0 +1,159 @@
+// Copied from https://github.com/unjs/unplugin/blob/f514bf8e2d751b48b3a5acb1077eb1292d9711b3/src/utils/filter.ts#L2
+// Used only in the babel version of the plugin, since others can rely on native unplugin functionality.
+
+import { resolve } from 'node:path';
+import picomatch from 'picomatch';
+import type {
+  Arrayable,
+  Nullable,
+  StringFilter,
+  StringOrRegExp,
+} from 'unplugin';
+
+function toArray<T>(array?: Nullable<Arrayable<T>>): Array<T> {
+  const result = array || [];
+  if (Array.isArray(result)) return result;
+  return [result];
+}
+
+const BACKSLASH_REGEX = /\\/g;
+function normalize(path: string): string {
+  return path.replace(BACKSLASH_REGEX, '/');
+}
+
+const ABSOLUTE_PATH_REGEX = /^(?:\/|(?:[A-Z]:)?[/\\|])/i;
+function isAbsolute(path: string): boolean {
+  return ABSOLUTE_PATH_REGEX.test(path);
+}
+
+export type PluginFilter = (input: string) => boolean;
+export type TransformHookFilter = (id: string, code: string) => boolean;
+
+interface NormalizedStringFilter {
+  include?: StringOrRegExp[] | undefined;
+  exclude?: StringOrRegExp[] | undefined;
+}
+
+function getMatcherString(glob: string, cwd: string) {
+  if (glob.startsWith('**') || isAbsolute(glob)) {
+    return normalize(glob);
+  }
+
+  const resolved = resolve(cwd, glob);
+  return normalize(resolved);
+}
+
+function patternToIdFilter(pattern: StringOrRegExp): PluginFilter {
+  if (pattern instanceof RegExp) {
+    return (id: string) => {
+      const normalizedId = normalize(id);
+      const result = pattern.test(normalizedId);
+      pattern.lastIndex = 0;
+      return result;
+    };
+  }
+  const cwd = process.cwd();
+  const glob = getMatcherString(pattern, cwd);
+  const matcher = picomatch(glob, { dot: true });
+  return (id: string) => {
+    const normalizedId = normalize(id);
+    return matcher(normalizedId);
+  };
+}
+
+function patternToCodeFilter(pattern: StringOrRegExp): PluginFilter {
+  if (pattern instanceof RegExp) {
+    return (code: string) => {
+      const result = pattern.test(code);
+      pattern.lastIndex = 0;
+      return result;
+    };
+  }
+  return (code: string) => code.includes(pattern);
+}
+
+function createFilter(
+  exclude: PluginFilter[] | undefined,
+  include: PluginFilter[] | undefined,
+): PluginFilter | undefined {
+  if (!exclude && !include) {
+    return;
+  }
+
+  return (input) => {
+    if (exclude?.some((filter) => filter(input))) {
+      return false;
+    }
+    if (include?.some((filter) => filter(input))) {
+      return true;
+    }
+    return !(include && include.length > 0);
+  };
+}
+
+function normalizeFilter(filter: StringFilter): NormalizedStringFilter {
+  if (typeof filter === 'string' || filter instanceof RegExp) {
+    return {
+      include: [filter],
+    };
+  }
+  if (Array.isArray(filter)) {
+    return {
+      include: filter,
+    };
+  }
+  return {
+    exclude: filter.exclude ? toArray(filter.exclude) : undefined,
+    include: filter.include ? toArray(filter.include) : undefined,
+  };
+}
+
+function createIdFilter(
+  filter: StringFilter | undefined,
+): PluginFilter | undefined {
+  if (!filter) return;
+  const { exclude, include } = normalizeFilter(filter);
+  const excludeFilter = exclude?.map(patternToIdFilter);
+  const includeFilter = include?.map(patternToIdFilter);
+  return createFilter(excludeFilter, includeFilter);
+}
+
+function createCodeFilter(
+  filter: StringFilter | undefined,
+): PluginFilter | undefined {
+  if (!filter) return;
+  const { exclude, include } = normalizeFilter(filter);
+  const excludeFilter = exclude?.map(patternToCodeFilter);
+  const includeFilter = include?.map(patternToCodeFilter);
+  return createFilter(excludeFilter, includeFilter);
+}
+
+export function createFilterForId(
+  filter: StringFilter | undefined,
+): PluginFilter | undefined {
+  const filterFunction = createIdFilter(filter);
+  return filterFunction ? (id) => !!filterFunction(id) : undefined;
+}
+
+function createFilterForTransform(
+  idFilter: StringFilter | undefined,
+  codeFilter: StringFilter | undefined,
+): TransformHookFilter | undefined {
+  if (!idFilter && !codeFilter) return;
+  const idFilterFunction = createIdFilter(idFilter);
+  const codeFilterFunction = createCodeFilter(codeFilter);
+  return (id, code) => {
+    let fallback = true;
+    if (idFilterFunction) {
+      fallback &&= idFilterFunction(id);
+    }
+    if (!fallback) {
+      return false;
+    }
+
+    if (codeFilterFunction) {
+      fallback &&= codeFilterFunction(code);
+    }
+    return fallback;
+  };
+}