fix(@typegpu/three): Allow TgpuFnNode to transitively nest

packages/typegpu-three/package.json

@@ -33,17 +33,21 @@
   },
   "scripts": {
     "build": "unbuild",
-    "test:types": "pnpm tsc --p ./tsconfig.json --noEmit",
+    "test": "vitest run",
+    "test:types": "pnpm tsc --p ./tsconfig.json --noEmit && pnpm tsc --p ./tsconfig.test.json --noEmit",
     "prepublishOnly": "tgpu-dev-cli prepack"
   },
   "devDependencies": {
     "@typegpu/tgpu-dev-cli": "workspace:*",
     "@types/three": "catalog:types",
     "@webgpu/types": "catalog:types",
+    "jiti": "catalog:build",
     "typegpu": "workspace:*",
+    "typegpu-testing-utility": "workspace:*",
     "typescript": "catalog:types",
     "unbuild": "catalog:build",
-    "unplugin-typegpu": "workspace:*"
+    "unplugin-typegpu": "workspace:*",
+    "vitest": "catalog:test"
   },
   "peerDependencies": {
     "three": ">0.126.0",

packages/typegpu-three/src/typegpu-node.ts

@@ -3,8 +3,7 @@ import type NodeVarying from 'three/src/nodes/core/NodeVarying.js';
 import type VaryingNode from 'three/src/nodes/core/VaryingNode.js';
 import * as THREE from 'three/webgpu';
 import * as TSL from 'three/tsl';
-import { tgpu, type Namespace, type TgpuVar } from 'typegpu';
-import * as d from 'typegpu/data';
+import { tgpu, d, type Namespace, type TgpuVar, type ResolvedDeclaration } from 'typegpu';
 import WGSLNodeBuilder from 'three/src/renderers/webgpu/nodes/WGSLNodeBuilder.js';
 
 /**
@@ -31,14 +30,18 @@ abstract class StageData {
 }
 
 class GenerateStageData extends StageData {
-  readonly names: WeakMap<object, string>;
   readonly type = 'generate';
-  codeGeneratedThusFar: string;
+  /**
+   * Keeping track of all declarations resolved by a
+   * specific builder, this helps to find the declaration
+   * of a function that might have been previously used
+   * transitively, but then is passed directly into toTSL
+   */
+  existingDeclarations: ResolvedDeclaration[];
 
   constructor(stage: 'vertex' | 'fragment' | 'compute' | null) {
     super(stage);
-    this.names = new WeakMap();
-    this.codeGeneratedThusFar = '';
+    this.existingDeclarations = [];
   }
 }
 
@@ -115,6 +118,16 @@ interface TgpuFnNodeContext {
 
 let currentlyGeneratingFnNodeCtx: TgpuFnNodeContext | undefined;
 
+function withGeneratingFnNodeCtx<T>(ctx: TgpuFnNodeContext, callback: () => T): T {
+  const previous = currentlyGeneratingFnNodeCtx;
+  currentlyGeneratingFnNodeCtx = ctx;
+  try {
+    return callback();
+  } finally {
+    currentlyGeneratingFnNodeCtx = previous;
+  }
+}
+
 function forceExplicitVoidReturn(codeIn: string) {
   if (codeIn.includes('->')) {
     // Has return type, so we don't need to force it
@@ -161,48 +174,55 @@ class TgpuFnNode<T> extends THREE.Node {
     const stageData = builderData.getGenerateStageData(builder.shaderStage);
 
     if (!nodeData.custom) {
-      if (currentlyGeneratingFnNodeCtx !== undefined) {
-        console.warn('[@typegpu/three] Nested function generation detected');
-      }
-
       const ctx: TgpuFnNodeContext = {
         builder,
         stageData,
         dependencies: [],
       };
-      currentlyGeneratingFnNodeCtx = ctx;
-      let resolved: string;
-      try {
-        resolved = tgpu.resolve({
+
+      const resolved = withGeneratingFnNodeCtx(ctx, () => {
+        const { code, declarations } = tgpu.resolveWithContext([this.#impl], {
+          names: stageData.namespace,
+        });
+
+        // Resolving this.#impl as second time in the same
+        // namespace resolved to only its identifier
+        const functionId = tgpu.resolve({
           names: stageData.namespace,
-          template: '___ID___ fnName',
-          externals: { fnName: this.#impl },
+          template: 'impl',
+          externals: { impl: this.#impl },
           unstable_minify: false, // TODO(#2826): investigate
         });
-      } finally {
-        currentlyGeneratingFnNodeCtx = undefined;
-      }
 
-      const [code = '', functionId] = resolved.split('___ID___').map((s) => s.trim());
-      stageData.codeGeneratedThusFar += code;
-      let lastFnStart = stageData.codeGeneratedThusFar.indexOf(`\nfn ${functionId}`);
-      if (lastFnStart === -1) {
-        // We're starting with the function declaration
-        lastFnStart = 0;
-      }
+        return {
+          code,
+          declarations,
+          functionId,
+        };
+      });
+
+      stageData.existingDeclarations.push(...resolved.declarations);
 
       // Extracting the function code
-      const fnCode = stageData.codeGeneratedThusFar.slice(lastFnStart).trim();
+      const fnDeclaration = stageData.existingDeclarations.find(
+        (decl) => decl.name === resolved.functionId,
+      )?.code;
+
+      if (!fnDeclaration) {
+        throw new Error(
+          `[@typegpu/three] Internal error, function declaration wasn't found in the generated shader code.`,
+        );
+      }
 
       nodeData.custom = {
-        functionId: functionId ?? '',
+        functionId: resolved.functionId,
         nodeFunction: builder.parser.parseFunction(
           // TODO: Upstream a fix to Three.js that accepts functions with no return type
-          forceExplicitVoidReturn(fnCode),
+          forceExplicitVoidReturn(fnDeclaration),
         ),
         // Including code that was resolved before the function as another node
         // that this node depends on
-        priorCode: TSL.code(code),
+        priorCode: TSL.code(resolved.code),
         dependencies: ctx.dependencies,
       };
     }
@@ -225,17 +245,15 @@ class TgpuFnNode<T> extends THREE.Node {
       stageData,
       dependencies: [],
     };
-    currentlyGeneratingFnNodeCtx = ctx;
-    try {
+
+    withGeneratingFnNodeCtx(ctx, () =>
       tgpu.resolve({
         names: stageData.namespace,
-        template: '___ID___ fnName',
-        externals: { fnName: this.#impl },
+        template: 'impl',
+        externals: { impl: this.#impl },
         unstable_minify: false, // TODO(#2826): investigate
-      });
-    } finally {
-      currentlyGeneratingFnNodeCtx = undefined;
-    }
+      }),
+    );
   }
 
   /**

packages/typegpu-three/tests/typegpu-node.test.ts

@@ -0,0 +1,108 @@
+import * as THREE from 'three/webgpu';
+import * as TSL from 'three/tsl';
+import WGSLNodeBuilder from 'three/src/renderers/webgpu/nodes/WGSLNodeBuilder.js';
+import { describe, expect, it } from 'vitest';
+import { tgpu, d } from 'typegpu';
+import { fromTSL, toTSL } from '@typegpu/three';
+
+class ObservableFloatNode extends THREE.Node {
+  analyzeCount = 0;
+  generateCount = 0;
+
+  getNodeType() {
+    return 'float';
+  }
+
+  analyze() {
+    this.analyzeCount += 1;
+  }
+
+  generate() {
+    this.generateCount += 1;
+    return '1.0';
+  }
+}
+
+function observableAccessor() {
+  const node = new ObservableFloatNode();
+  return {
+    node,
+    accessor: fromTSL(TSL.nodeObject(node), d.f32),
+  };
+}
+
+function builderFor(stage: 'analyze' | 'generate') {
+  const builder = new WGSLNodeBuilder();
+  builder.setShaderStage('fragment');
+  builder.setBuildStage(stage);
+  return builder;
+}
+
+describe('TypeGPU node generation context', () => {
+  it.each(['analyze', 'generate'] as const)(
+    'restores the outer context after nested %s traversal',
+    (stage) => {
+      const before = observableAccessor();
+      const inner = observableAccessor();
+      const after = observableAccessor();
+
+      const innerNode = toTSL(() => {
+        'use gpu';
+        return inner.accessor.$;
+      });
+      const innerNodeAccessor = fromTSL(innerNode, d.f32);
+
+      const outerNode = toTSL(() => {
+        'use gpu';
+        return before.accessor.$ + innerNodeAccessor.$ + after.accessor.$;
+      });
+
+      expect(() => outerNode.build(builderFor(stage))).not.toThrow();
+
+      if (stage === 'analyze') {
+        expect(before.node.analyzeCount).toBe(1);
+        expect(inner.node.analyzeCount).toBe(1);
+        expect(after.node.analyzeCount).toBe(1);
+      } else {
+        expect(before.node.generateCount).toBe(1);
+        expect(inner.node.generateCount).toBe(1);
+        expect(after.node.generateCount).toBe(1);
+      }
+    },
+  );
+
+  it('restores the outer context when nested generation throws and is caught', () => {
+    const before = observableAccessor();
+    const after = observableAccessor();
+    const fail = tgpu.comptime(() => {
+      throw new Error('inner failure');
+    });
+
+    class CatchingNode extends THREE.Node {
+      getNodeType() {
+        return 'float';
+      }
+
+      generate(builder: THREE.NodeBuilder) {
+        const throwingInner = toTSL(() => {
+          'use gpu';
+          fail();
+          return d.f32(0);
+        });
+
+        expect(() => throwingInner.getNodeType(builder)).toThrow('inner failure');
+        return '2.0';
+      }
+    }
+
+    const catchingAccessor = fromTSL(TSL.nodeObject(new CatchingNode()), d.f32);
+    const outerNode = toTSL(() => {
+      'use gpu';
+      return before.accessor.$ + catchingAccessor.$ + after.accessor.$;
+    });
+
+    expect(() => outerNode.build(builderFor('generate'))).not.toThrow();
+    expect(before.node.generateCount).toBe(1);
+    expect(after.node.generateCount).toBe(1);
+  });
+});

packages/typegpu-three/tsconfig.test.json

@@ -0,0 +1,4 @@
+{
+  "extends": "./tsconfig.json",
+  "include": ["src/**/*", "tests/**/*", "vitest.config.mts"]
+}

packages/typegpu-three/vitest.config.mts

@@ -0,0 +1,16 @@
+import { createJiti } from 'jiti';
+import type TypeGPUPlugin from 'unplugin-typegpu/vite';
+import { defineConfig } from 'vitest/config';
+import { typegpuBuiltAliases } from 'typegpu-testing-utility/config';
+
+const jiti = createJiti(import.meta.url);
+const typegpu = await jiti.import<typeof TypeGPUPlugin>('unplugin-typegpu/vite', {
+  default: true,
+});
+
+export default defineConfig({
+  plugins: [typegpu({ forceTgpuAlias: 'tgpu', earlyPruning: false })],
+  resolve: {
+    alias: typegpuBuiltAliases(),
+  },
+});

packages/typegpu/src/indexNamedExports.ts

@@ -51,6 +51,7 @@ export { warn } from './tgpuLogger.ts';
 // types
 
 export type { ResolvableObject } from './types.ts';
+export type { ResolvedDeclaration } from './resolutionCtx.ts';
 export type {
   Configurable,
   TgpuGuardedComputePipeline,