impr(unplugin-typegpu): Support parsing all metadata fields

packages/unplugin-typegpu/package.json

@@ -345,9 +345,11 @@
     "unplugin": "^3.0.0"
   },
   "devDependencies": {
+    "@babel/generator": "^7.29.0",
     "@babel/standalone": "^7.28.6",
     "@rollup/plugin-virtual": "^3.0.2",
     "@typegpu/tgpu-dev-cli": "workspace:*",
+    "@types/babel__generator": "^7.27.0",
     "@types/babel__standalone": "^7.1.9",
     "@types/babel__traverse": "^7.20.7",
     "@types/bun": "^1.3.11",

packages/unplugin-typegpu/src/babel.ts

@@ -2,8 +2,8 @@ import type { NodePath, TraverseOptions } from '@babel/traverse';
 import defu from 'defu';
 import type { Externals, TranspilationResult } from 'tinyest-for-wgsl';
 import * as t from '@babel/types';
+import { METADATA_FORMAT_VERSION } from './core/version.ts';
 import {
-  METADATA_FORMAT_VERSION,
   type PluginState,
   checkOpts,
   defaultOptions,

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

@@ -6,12 +6,6 @@ import { transpileBabelFn, type TranspilationResult } from 'tinyest-for-wgsl';
 import { getEmbeddedTypegpuMetadata } from './embeddedMetadata.ts';
 import { obfuscate } from './obfuscate.ts';
 
-/**
- * Each breaking change to the metadata format requires a bump to this number.
- * It's used at runtime by `typegpu` to determine how to interpret a function's metadata.
- */
-export const METADATA_FORMAT_VERSION = 2;
-
 export interface Options {
   /** @default [/\.m?[jt]sx?$/] */
   include?: FilterPattern;

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

@@ -1,12 +1,24 @@
 import type { NodePath } from '@babel/traverse';
 import * as t from '@babel/types';
-import { type MetadatableFunction } from './common.ts';
+import type { TranspilationResult } from 'tinyest-for-wgsl';
+import { METADATA_FORMAT_VERSION } from './version.ts';
+import type { MetadatableFunction } from './common.ts';
 
-export interface EmbeddedTypegpuMetadata {
+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 EmbeddedTypegpuMetadata = {
   v: number;
   name: string | undefined;
-  // TODO: parse AST and externals
-}
+  function?: {
+    ast: FunctionAst;
+    astPath: NodePath<t.ObjectExpression>;
+    externals: FunctionExternals;
+  };
+};
 
 const embeddedTypegpuMetadataCache = new WeakMap<
   NodePath<MetadatableFunction>,
@@ -118,46 +130,100 @@ function skipTransparentWrappersUp(path: NodePath): NodePath {
 }
 
 /**
- * Returns the node of the property with the given name in the object, if it exists.
+ * Walks down from the given path through parentheses and type-only wrappers,
+ * returning the innermost wrapped expression.
+ */
+function skipTransparentWrappersDown(path: NodePath): NodePath {
+  let current = path;
+
+  while (isTransparentWrapper(current.node)) {
+    current = current.get('expression') as NodePath;
+  }
+
+  return current;
+}
+
+/**
+ * Returns the value path of the first matching property. If there is no match, returns `undefined`.
  */
-function objectPropertyValue(object: t.ObjectExpression, expectedName: string): t.Node | undefined {
-  for (const property of object.properties) {
-    if (!t.isObjectProperty(property) || property.computed) {
+function objectPropertyPath(
+  objectPath: NodePath<t.ObjectExpression>,
+  name: string,
+): NodePath<t.Expression> | undefined {
+  for (const property of objectPath.get('properties')) {
+    if (!property.isObjectProperty()) {
       continue;
     }
 
-    const name = t.isIdentifier(property.key)
-      ? property.key.name
-      : t.isStringLiteral(property.key)
-        ? property.key.value
-        : undefined;
+    const { key, computed } = property.node;
 
-    if (name === expectedName) {
-      return property.value;
+    const matches =
+      (!computed && t.isIdentifier(key, { name })) || t.isStringLiteral(key, { value: name });
+
+    if (!matches) {
+      continue;
+    }
+
+    const value = property.get('value');
+    if (value.isExpression()) {
+      return value;
     }
   }
 
   return undefined;
 }
+/**
+ * Returns the parsed AST
+ */
+function parseAstPath(astPath: NodePath<t.ObjectExpression>): FunctionAst | undefined {
+  const evaluated = astPath.evaluate();
+
+  return !evaluated.confident ? undefined : (evaluated.value as FunctionAst);
+}
+
+/**
+ * Externals are represented as an object with string keys and external getters.
+ * This function parses such objects and returns a map keyed by string, with values that point to the corresponding object property node paths.
+ */
+function parseExternalsPath(
+  externalsPath: NodePath<t.ObjectExpression>,
+): FunctionExternals | undefined {
+  const externals: FunctionExternals = new Map();
+
+  for (const propertyPath of externalsPath.get('properties')) {
+    if (!propertyPath.isObjectProperty()) {
+      return undefined;
+    }
+
+    const { computed, key } = propertyPath.node;
+
+    const name = t.isStringLiteral(key)
+      ? key.value
+      : !computed && t.isIdentifier(key)
+        ? key.name
+        : undefined;
+
+    if (name === undefined) {
+      return undefined;
+    }
+
+    externals.set(name, propertyPath);
+  }
+
+  return externals;
+}
 
 /**
  * Returns metadata embedded by unplugin-typegpu for this exact function.
- * Externals are required as part of the emitted shape, but are intentionally not evaluated.
  *
- * Consider:
- * ```ts
- * const f = ($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = () => {
- *   'use gpu';
- * }, {
- *   v: 2,
- *   name: "f",
- *   ast: {
- *     params: [],
- *     body: [0, []]
- *   },
- *   externals: {}
- * }) && $.f)({});
- * ```
+ * The returned object contains:
+ * - `v`: the metadata version
+ * - `name`: the function name, which may be undefined
+ * - `function`: an object containing the parsed AST, its source node path, and a map
+ *   from external names to their original property paths
+ *
+ * @note Support for metadata stamped with a version different from `METADATA_FORMAT_VERSION` is limited.
+ *       Only the version and name are parsed.
  */
 export function getEmbeddedTypegpuMetadata(
   path: NodePath<MetadatableFunction>,
@@ -190,31 +256,69 @@ export function getEmbeddedTypegpuMetadata(
   }
 
   // we check for the metadata object
-  const metadataNode = callPath.node.arguments[1];
-  if (!t.isObjectExpression(metadataNode)) {
+  const metadataPath = callPath.get('arguments.1');
+  if (metadataPath === undefined) {
+    return undefined;
+  }
+
+  const unwrappedMetadataPath = skipTransparentWrappersDown(metadataPath as NodePath);
+  if (!unwrappedMetadataPath.isObjectExpression()) {
     return undefined;
   }
 
   // get the metadata properties
-  const versionNode = objectPropertyValue(metadataNode, 'v');
-  const nameNode = objectPropertyValue(metadataNode, 'name');
-  const astNode = objectPropertyValue(metadataNode, 'ast');
-  const externalsNode = objectPropertyValue(metadataNode, 'externals');
+  const versionPath = objectPropertyPath(unwrappedMetadataPath, 'v');
+  const namePath = objectPropertyPath(unwrappedMetadataPath, 'name');
 
-  if (
-    !t.isNumericLiteral(versionNode) ||
-    nameNode === undefined ||
-    astNode === undefined ||
-    externalsNode === undefined
-  ) {
-    return undefined;
+  if (versionPath === undefined || namePath === undefined) {
+    throw new Error(METADATA_PARSING_ERROR_MESSAGE);
+  }
+
+  const versionResult = versionPath.evaluate();
+  const nameResult = namePath.evaluate();
+
+  if (!versionResult.confident || !nameResult.confident) {
+    throw new Error(METADATA_PARSING_ERROR_MESSAGE);
+  }
+
+  const version = versionResult.value as number;
+  const name = nameResult.value as string | undefined;
+
+  if (version !== METADATA_FORMAT_VERSION) {
+    const embeddedTypegpuMetadata = {
+      v: version,
+      name,
+    };
+    embeddedTypegpuMetadataCache.set(path, embeddedTypegpuMetadata);
+    return embeddedTypegpuMetadata;
   }
 
-  const name = t.isStringLiteral(nameNode) ? nameNode.value : undefined;
+  const astPath = objectPropertyPath(unwrappedMetadataPath, 'ast');
+  const externalsPath = objectPropertyPath(unwrappedMetadataPath, 'externals');
+
+  if (astPath === undefined || externalsPath === undefined) {
+    throw new Error(METADATA_PARSING_ERROR_MESSAGE);
+  }
+
+  if (!astPath.isObjectExpression() || !externalsPath.isObjectExpression()) {
+    throw new Error(METADATA_PARSING_ERROR_MESSAGE);
+  }
+
+  const ast = parseAstPath(astPath);
+  const externals = parseExternalsPath(externalsPath);
+
+  if (ast === undefined || externals === undefined) {
+    throw new Error(METADATA_PARSING_ERROR_MESSAGE);
+  }
 
   const embeddedTypegpuMetadata = {
-    v: versionNode.value,
+    v: version,
     name,
+    function: {
+      ast,
+      astPath,
+      externals,
+    },
   };
   embeddedTypegpuMetadataCache.set(path, embeddedTypegpuMetadata);
 

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

@@ -6,13 +6,13 @@ import _traverse, { type NodePath } from '@babel/traverse';
 import type { Externals, TranspilationResult } from 'tinyest-for-wgsl';
 import * as parser from '@babel/parser';
 import * as t from '@babel/types';
+import { METADATA_FORMAT_VERSION } from './version.ts';
 import {
   defaultOptions,
   earlyPruneRegex,
   initPluginState,
   functionVisitor,
   getBlockScope,
-  METADATA_FORMAT_VERSION,
   checkOpts,
 } from './common.ts';
 

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

@@ -0,0 +1,5 @@
+/**
+ * Each breaking change to the metadata format requires a bump to this number.
+ * It's used at runtime by `typegpu` to determine how to interpret a function's metadata.
+ */
+export const METADATA_FORMAT_VERSION = 2;