feat: Add source mapping option to unplugin

packages/typegpu/src/core/function/fnCore.ts

@@ -1,3 +1,4 @@
+import { stripSourceMap, type Block, type SourceMap } from 'tinyest';
 import { getAttributesString } from '../../data/attributes.ts';
 import { undecorate } from '../../data/dataTypes.ts';
 import { type ResolvedSnippet, snip } from '../../data/snippet.ts';
@@ -221,6 +222,7 @@ export function createFnCore(
           "Missing metadata for tgpu.fn function body (either missing 'use gpu' directive, or misconfigured `unplugin-typegpu`)",
         );
       }
+      const [body, _] = stripSourceMap(ast.body) as [Block, SourceMap];
 
       // If an entrypoint implementation has a second argument, it represents the output schema.
       // We look at the identifier chosen by the user and add it to externals.
@@ -242,7 +244,7 @@ export function createFnCore(
         entryInput,
         params: ast.params,
         returnType,
-        body: ast.body,
+        body,
         externalMap: mergeFunctionExternals(externals),
       });
 

packages/typegpu/src/shared/normalizeMetadata.ts

@@ -1,4 +1,4 @@
-import type { Block, FuncParameter } from 'tinyest';
+import type { Block, FuncParameter, SourceMappedNode } from 'tinyest';
 import { safeStringify } from './stringify.ts';
 
 export interface RawMetadataV1 {
@@ -13,7 +13,7 @@ export interface RawMetadataV1 {
 export interface RawMetadataV2 {
   v: 2;
   name: string;
-  ast: { params: FuncParameter[]; body: Block };
+  ast: { params: FuncParameter[]; body: Block | SourceMappedNode };
   externals: { [key: string]: () => unknown };
 }
 
@@ -26,7 +26,7 @@ export type RawMetadata = RawMetadataV1 | RawMetadataV2;
  * Holds normalized function metadata required for WGSL generation
  */
 export interface Metadata {
-  ast: { params: FuncParameter[]; body: Block };
+  ast: { params: FuncParameter[]; body: Block | SourceMappedNode };
   externals: () => Record<string, unknown>;
 }
 

packages/typegpu/tests/internal/metadata.test.ts

@@ -9,6 +9,7 @@ import type {
   RawMetadataV1,
   RawMetadataV2,
 } from '../../src/shared/normalizeMetadata.ts';
+import type { Block, SourceMap } from 'tinyest';
 
 // functions in this file intentionally have no 'use gpu' directive
 describe('meta', () => {
@@ -311,4 +312,30 @@ describe('meta', () => {
       }"
     `);
   });
+
+  it('correctly strips source-mapped functions', () => {
+    const fn = () => {};
+    const ast: Block = [
+      NODE.block,
+      [[NODE.const, [NODE.identifier, 'a'], [NODE.numericLiteral, '1']]],
+    ];
+    const sourceMap: SourceMap = new Map();
+    sourceMap.get = () => [1, 1];
+    const meta: RawMetadataV2 = {
+      v: 2,
+      name: 'fn',
+      externals: {},
+      ast: {
+        params: [],
+        body: tinyest.embedSourceMap(ast, sourceMap),
+      },
+    };
+    assignMetadata(fn, meta);
+
+    expect(tgpu.resolve([fn])).toMatchInlineSnapshot(`
+      "fn fn_1() {
+        const a = 1;
+      }"
+    `);
+  });
 });

packages/typegpu/tests/internal/tseynit.test.ts

@@ -9,7 +9,7 @@ function getBodyAst(fn: () => void) {
   if (!meta) {
     throw new Error('Expected ast to be defined');
   }
-  return meta.ast.body;
+  return tinyest.stripSourceMap(meta.ast.body)[0];
 }
 
 describe('ast to JS transformation', () => {

packages/unplugin-typegpu/package.json

@@ -335,6 +335,7 @@
     "@babel/traverse": "^7.29.0",
     "@babel/types": "catalog:",
     "@jridgewell/remapping": "^2.3.5",
+    "@jridgewell/trace-mapping": "^0.3.31",
     "ast-kit": "^2.2.0",
     "defu": "^6.1.4",
     "magic-string": "^0.30.21",
@@ -354,7 +355,8 @@
     "@types/picomatch": "^4.0.1",
     "rollup": "catalog:build",
     "tsdown": "catalog:build",
-    "typescript": "catalog:types"
+    "typescript": "catalog:types",
+    "webpack": "^5.104.1"
   },
   "peerDependencies": {
     "typegpu": "workspace:^"

packages/unplugin-typegpu/src/babel.ts

@@ -1,15 +1,17 @@
 import type { NodePath, TraverseOptions } from '@babel/traverse';
 import defu from 'defu';
-import type { Externals, TranspilationResult } from 'tinyest-for-wgsl';
+import type { Externals } from 'tinyest-for-wgsl';
 import * as t from '@babel/types';
 import {
   METADATA_FORMAT_VERSION,
   type PluginState,
+  type PluginTranspilationResult,
   checkOpts,
   defaultOptions,
   functionVisitor,
   getBlockScope,
   initPluginState,
+  nodePosition,
 } from './core/common.ts';
 import { createFilterForId } from './core/filter.ts';
 
@@ -42,7 +44,7 @@ function assignMetadata(
   this: PluginState,
   path: NodePath<t.FunctionDeclaration | t.ArrowFunctionExpression | t.FunctionExpression>,
   name: string | undefined,
-  ast: TranspilationResult,
+  ast: PluginTranspilationResult,
 ): void {
   const metadata = t.objectExpression([
     t.objectProperty(i('v'), t.numericLiteral(METADATA_FORMAT_VERSION)),
@@ -166,6 +168,7 @@ export default function TypeGPUPlugin() {
     name: 'typegpu',
     pre(this: PluginState) {
       this.opts = checkOpts(defu(this.opts, defaultOptions));
+      this.originalPositionFor = nodePosition;
       initPluginState(this, {
         warn: (message) => console.warn(message),
         assignMetadata,

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

@@ -5,6 +5,7 @@ import MagicString from 'magic-string';
 import { transpileBabelFn, type TranspilationResult } from 'tinyest-for-wgsl';
 import { getEmbeddedTypegpuMetadata } from './embeddedMetadata.ts';
 import { obfuscate } from './obfuscate.ts';
+import { embedSourceMap, type Block, type SourceMappedNode } from 'tinyest';
 
 /**
  * Each breaking change to the metadata format requires a bump to this number.
@@ -36,6 +37,13 @@ export interface Options {
    */
   unstable_obfuscate?: boolean;
 
+  /**
+   * Keep the source map info so that resolve errors link to an appropriate source.
+   *
+   * @default false
+   */
+  unstable_sourceMaps?: boolean;
+
   /**
    * Skipping files that don't contain "typegpu", "tgpu" or "use gpu".
    * In case this early pruning hinders transformation, you
@@ -52,6 +60,11 @@ export function checkOpts<T extends Options>(opts: T): T {
       `Options 'unstable_obfuscate' and 'autoNamingEnabled' cannot be enabled at the same time.`,
     );
   }
+  if (opts.unstable_obfuscate && opts.unstable_sourceMaps) {
+    throw new Error(
+      `Options 'unstable_obfuscate' and 'unstable_sourceMaps' cannot be enabled at the same time.`,
+    );
+  }
   return opts;
 }
 
@@ -60,6 +73,10 @@ export type MetadatableFunction =
   | t.FunctionExpression
   | t.ArrowFunctionExpression;
 
+export type PluginTranspilationResult = Omit<TranspilationResult, 'body'> & {
+  body: Block | SourceMappedNode;
+};
+
 export interface TransformMethods {
   warn(message: string): void;
 
@@ -96,7 +113,7 @@ export interface TransformMethods {
     this: PluginState,
     path: NodePath<MetadatableFunction>,
     name: string | undefined,
-    ast: TranspilationResult,
+    ast: PluginTranspilationResult,
   ): void;
 
   wrapInAutoName(this: PluginState, path: NodePath<t.Expression>, name: string): void;
@@ -126,8 +143,21 @@ export interface PluginState extends TransformMethods {
   opts: Required<Options>;
 
   inUseGpuScope: boolean;
+
+  /**
+   * Used for source mapping.
+   * Babel keeps the correct combined source map on the node,
+   * but for rollup we need to calculate it ourself.
+   */
+  originalPositionFor: NodePositionProvider;
 }
 
+/**
+ * Resolves the position of a node in the file the user actually authored,
+ * or `undefined` if it cannot be mapped back.
+ */
+export type NodePositionProvider = (node: t.Node) => [line: number, column: number] | undefined;
+
 export interface NodeLocation {
   start?: number | null;
   end?: number | null;
@@ -155,6 +185,7 @@ export const defaultOptions = {
   autoNamingEnabled: true,
   earlyPruning: true,
   unstable_obfuscate: false,
+  unstable_sourceMaps: false,
 } satisfies Partial<Options>;
 
 /**
@@ -456,6 +487,13 @@ const operators = {
   '%=': '__tsover_mod',
 };
 
+export const nodePosition: NodePositionProvider = (node) => {
+  if (!node.loc) {
+    return undefined;
+  }
+  return [node.loc.start.line, node.loc.start.column];
+};
+
 function containsUseGpuDirective(
   node: t.FunctionDeclaration | t.FunctionExpression | t.ArrowFunctionExpression,
 ): boolean {
@@ -466,7 +504,7 @@ function containsUseGpuDirective(
 
 const fnNodeToTranspiledMap = new WeakMap<
   t.FunctionDeclaration | t.FunctionExpression | t.ArrowFunctionExpression,
-  TranspilationResult
+  PluginTranspilationResult
 >();
 
 function functionOnExit(
@@ -490,11 +528,18 @@ function functionOnExit(
 }
 
 function transpile(
+  ctx: PluginState,
   rootNode: Parameters<typeof transpileBabelFn>[0],
-  obf: boolean,
-): TranspilationResult {
+): PluginTranspilationResult {
+  if (ctx.opts.unstable_sourceMaps) {
+    const result = transpileBabelFn(rootNode, {
+      verboseNodes: true,
+      sourceMap: ctx.originalPositionFor,
+    });
+    return { ...result, body: embedSourceMap(result.body, result.sourceMap) };
+  }
   const result = transpileBabelFn(rootNode);
-  if (obf) {
+  if (ctx.opts.unstable_obfuscate) {
     return obfuscate(result);
   }
   return result;
@@ -558,7 +603,7 @@ export const functionVisitor: TraverseOptions<PluginState> = {
   ArrowFunctionExpression: {
     enter(path, state) {
       if (containsUseGpuDirective(path.node) && getEmbeddedTypegpuMetadata(path) === undefined) {
-        fnNodeToTranspiledMap.set(path.node, transpile(path.node, this.opts.unstable_obfuscate));
+        fnNodeToTranspiledMap.set(path.node, transpile(this, path.node));
         if (state.inUseGpuScope) {
           throw new Error(`Nesting 'use gpu' functions is not allowed`);
         }
@@ -571,7 +616,7 @@ export const functionVisitor: TraverseOptions<PluginState> = {
   FunctionExpression: {
     enter(path, state) {
       if (containsUseGpuDirective(path.node) && getEmbeddedTypegpuMetadata(path) === undefined) {
-        fnNodeToTranspiledMap.set(path.node, transpile(path.node, this.opts.unstable_obfuscate));
+        fnNodeToTranspiledMap.set(path.node, transpile(this, path.node));
         if (state.inUseGpuScope) {
           throw new Error(`Nesting 'use gpu' functions is not allowed`);
         }
@@ -584,7 +629,7 @@ export const functionVisitor: TraverseOptions<PluginState> = {
   FunctionDeclaration: {
     enter(path, state) {
       if (containsUseGpuDirective(path.node) && getEmbeddedTypegpuMetadata(path) === undefined) {
-        fnNodeToTranspiledMap.set(path.node, transpile(path.node, this.opts.unstable_obfuscate));
+        fnNodeToTranspiledMap.set(path.node, transpile(this, path.node));
         if (state.inUseGpuScope) {
           throw new Error(`Nesting 'use gpu' functions is not allowed`);
         }
@@ -613,7 +658,7 @@ export const functionVisitor: TraverseOptions<PluginState> = {
               t.ArrowFunctionExpression | t.FunctionDeclaration | t.FunctionExpression
             >,
             getFunctionName(path.get('arguments.0')),
-            transpile(implementation, this.opts.unstable_obfuscate),
+            transpile(this, implementation),
           );
         }
       }