Change the way AST is assigned by plugin (#599)

packages/rollup-plugin/src/index.ts

@@ -137,8 +137,6 @@ export default function typegpu(options?: TypegpuPluginOptions): TypegpuPlugin {
                 // Assuming that every call to `.does` is related to TypeGPU
                 // because shells can be created separately from calls to `tgpu`,
                 // making it hard to detect.
-                // TODO: We can improve this by first checking if $__ast exists on this object
-                // at runtime, before calling it.
                 node.callee.property.name === 'does')
             ) {
               const implementation = node.arguments[0];
@@ -165,16 +163,20 @@ export default function typegpu(options?: TypegpuPluginOptions): TypegpuPlugin {
           expr.implementation,
         );
 
+        // Wrap the implementation in a call to `tgpu.__assignAst` to associate the AST with the implementation.
+        magicString.appendLeft(expr.implementation.start, 'tgpu.__assignAst(');
         magicString.appendRight(
-          expr.varDecl.end,
-          `.$__ast(${embedJSON(argNames)}, ${embedJSON(body)})`,
+          expr.implementation.end,
+          `, ${embedJSON({ argNames, body, externalNames })}`,
         );
 
         if (externalNames.length > 0) {
           magicString.appendRight(
-            expr.varDecl.end,
-            `.$uses({ ${externalNames.join(', ')} })`,
+            expr.implementation.end,
+            `, {${externalNames.join(', ')}})`,
           );
+        } else {
+          magicString.appendRight(expr.implementation.end, ', undefined)');
         }
       }
 

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

@@ -0,0 +1,32 @@
+import type { Block } from 'tinyest';
+
+export type Ast = {
+  argNames: string[];
+  body: Block;
+  externalNames: string[];
+};
+
+export type AstInfo = {
+  ast: Ast;
+  externals?: Record<string, unknown> | undefined;
+};
+
+const functionToAstMap = new WeakMap<
+  (...args: unknown[]) => unknown,
+  AstInfo
+>();
+
+export function getPrebuiltAstFor(
+  fn: (...args: unknown[]) => unknown,
+): AstInfo | undefined {
+  return functionToAstMap.get(fn);
+}
+
+export function assignAst<T extends (...args: unknown[]) => unknown>(
+  fn: T,
+  ast: Ast,
+  externals?: Record<string, unknown> | undefined,
+): T {
+  functionToAstMap.set(fn, { ast, externals });
+  return fn;
+}

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

@@ -1,11 +1,12 @@
 import { MissingLinksError } from '../../errors';
 import type { AnyTgpuData, ResolutionCtx, Resource } from '../../types';
+import { getPrebuiltAstFor } from './astUtils';
 import {
   type ExternalMap,
   applyExternals,
   replaceExternalsInWgsl,
 } from './externals';
-import type { Implementation, TranspilationResult } from './fnTypes';
+import type { Implementation } from './fnTypes';
 
 export interface TgpuFnShellBase<Args extends unknown[], Return> {
   readonly argTypes: Args;
@@ -15,7 +16,6 @@ export interface TgpuFnShellBase<Args extends unknown[], Return> {
 interface FnCore {
   label: string | undefined;
   applyExternals(newExternals: ExternalMap): void;
-  setAst(ast: TranspilationResult): void;
   resolve(ctx: ResolutionCtx, fnAttribute?: string): string;
 }
 
@@ -30,7 +30,6 @@ export function createFnCore(
    * entry fn).
    */
   const externalsToApply: ExternalMap[] = [];
-  let prebuiltAst: TranspilationResult | null = null;
 
   return {
     label: undefined as string | undefined,
@@ -39,12 +38,9 @@ export function createFnCore(
       externalsToApply.push(newExternals);
     },
 
-    setAst(ast: TranspilationResult): void {
-      prebuiltAst = ast;
-    },
-
     resolve(ctx: ResolutionCtx, fnAttribute = ''): string {
       const externalMap: ExternalMap = {};
+
       for (const externals of externalsToApply) {
         applyExternals(externalMap, externals);
       }
@@ -60,7 +56,13 @@ export function createFnCore(
 
         ctx.addDeclaration(`${fnAttribute}fn ${id}${replacedImpl}`);
       } else {
-        const ast = prebuiltAst ?? ctx.transpileFn(String(implementation));
+        // get data generated by the plugin
+        const pluginData = getPrebuiltAstFor(implementation);
+
+        if (pluginData?.externals) {
+          applyExternals(externalMap, pluginData.externals);
+        }
+        const ast = pluginData?.ast ?? ctx.transpileFn(String(implementation));
 
         // Verifying all required externals are present.
         const missingExternals = ast.externalNames.filter(

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

@@ -1,4 +1,3 @@
-import type { Block } from 'tinyest';
 import type { AnyBuiltin } from '../../builtin';
 import type { TgpuNamable } from '../../namable';
 import type { ResolutionCtx, TgpuResolvable } from '../../types';
@@ -35,7 +34,6 @@ export interface TgpuComputeFn extends TgpuResolvable, TgpuNamable {
   readonly shell: TgpuComputeFnShell;
 
   $uses(dependencyMap: Record<string, unknown>): this;
-  $__ast(argNames: string[], body: Block): this;
 }
 
 export interface ComputeFnOptions {
@@ -95,13 +93,6 @@ function createComputeFn(
       return this;
     },
 
-    $__ast(argNames: string[], body: Block): This {
-      // When receiving a pre-built $__ast, we are receiving $uses alongside it, so
-      // we do not need to verify external names.
-      core.setAst({ argNames, body, externalNames: [] });
-      return this;
-    },
-
     $name(newLabel: string): This {
       core.label = newLabel;
       return this;

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

@@ -1,4 +1,3 @@
-import type { Block } from 'tinyest';
 import { inGPUMode } from '../../gpuMode';
 import type { TgpuNamable } from '../../namable';
 import { valueList } from '../../resolutionUtils';
@@ -52,7 +51,6 @@ interface TgpuFnBase<
   readonly shell: TgpuFnShell<Args, Return>;
 
   $uses(dependencyMap: Record<string, unknown>): this;
-  $__ast(argNames: string[], body: Block): this;
   with<T>(slot: TgpuSlot<T>, value: Eventual<T>): TgpuFn<Args, Return>;
 }
 
@@ -115,13 +113,6 @@ function createFn<
       return this;
     },
 
-    $__ast(argNames: string[], body: Block): This {
-      // When receiving a pre-built $__ast, we are receiving $uses alongside it, so
-      // we do not need to verify external names.
-      core.setAst({ argNames, body, externalNames: [] });
-      return this;
-    },
-
     $name(newLabel: string): This {
       core.label = newLabel;
       return this;
@@ -179,11 +170,6 @@ function createBoundFunction<
       return this;
     },
 
-    $__ast(argNames: string[], body: Block): This {
-      innerFn.$__ast(argNames, body);
-      return this;
-    },
-
     $name(newLabel: string): This {
       innerFn.$name(newLabel);
       return this;

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

@@ -1,4 +1,3 @@
-import type { Block } from 'tinyest';
 import type { Vec4f } from '../../data';
 import type { TgpuNamable } from '../../namable';
 import type { ResolutionCtx, TgpuResolvable } from '../../types';
@@ -43,7 +42,6 @@ export interface TgpuFragmentFn<
   readonly shell: TgpuFragmentFnShell<Varying, Output>;
 
   $uses(dependencyMap: Record<string, unknown>): this;
-  $__ast(argNames: string[], body: Block): this;
 }
 
 /**
@@ -99,13 +97,6 @@ function createFragmentFn(
       return this;
     },
 
-    $__ast(argNames: string[], body: Block): This {
-      // When receiving a pre-built $__ast, we are receiving $uses alongside it, so
-      // we do not need to verify external names.
-      core.setAst({ argNames, body, externalNames: [] });
-      return this;
-    },
-
     $name(newLabel: string): This {
       core.label = newLabel;
       return this;

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

@@ -1,4 +1,3 @@
-import type { Block } from 'tinyest';
 import { type TgpuStruct, location, struct } from '../../data';
 import { getCustomLocation, isBuiltin } from '../../data/attributes';
 import type { TgpuNamable } from '../../namable';
@@ -52,7 +51,6 @@ export interface TgpuVertexFn<
   readonly Output: IOLayoutToOutputStruct<Output>;
 
   $uses(dependencyMap: Record<string, unknown>): this;
-  $__ast(argNames: string[], body: Block): this;
 }
 
 /**
@@ -147,13 +145,6 @@ function createVertexFn(
       return this;
     },
 
-    $__ast(argNames: string[], body: Block): This {
-      // When receiving a pre-built $__ast, we are receiving $uses alongside it, so
-      // we do not need to verify external names.
-      core.setAst({ argNames, body, externalNames: [] });
-      return this;
-    },
-
     $name(newLabel: string): This {
       core.label = newLabel;
       Output.$name(`${newLabel}_Output`);

packages/typegpu/src/experimental/index.ts

@@ -2,6 +2,7 @@
  * @module typegpu/experimental
  */
 
+import { assignAst } from '../core/function/astUtils';
 import { computeFn } from '../core/function/tgpuComputeFn';
 import { fn, procedure } from '../core/function/tgpuFn';
 import { fragmentFn } from '../core/function/tgpuFragmentFn';
@@ -48,6 +49,12 @@ export const tgpu = {
   read,
   write,
 };
+
+// Hidden API, used only by tooling (e.g., rollup plugin).
+Object.assign(tgpu, {
+  __assignAst: assignAst,
+});
+
 export default tgpu;
 
 export * from '../errors';