fix(@typegpu/three): Proper handling of varying nodes (#2125)

Author: cieplypolar

apps/typegpu-docs/src/examples/threejs/varyings/index.html

@@ -0,0 +1 @@
+<canvas></canvas>

apps/typegpu-docs/src/examples/threejs/varyings/index.ts

@@ -0,0 +1,99 @@
+import * as THREE from 'three/webgpu';
+import * as TSL from 'three/tsl';
+import * as t3 from '@typegpu/three';
+import * as d from 'typegpu/data';
+import * as std from 'typegpu/std';
+
+const canvas = document.querySelector('canvas') as HTMLCanvasElement;
+const renderer = new THREE.WebGPURenderer({ canvas, antialias: true });
+renderer.setSize(canvas.clientWidth, canvas.clientHeight, false);
+renderer.setPixelRatio(window.devicePixelRatio);
+await renderer.init();
+
+const scene = new THREE.Scene();
+scene.background = new THREE.Color(0xe8abbf);
+
+const camera = new THREE.PerspectiveCamera(
+  45,
+  canvas.clientWidth / canvas.clientHeight,
+  0.1,
+  100,
+);
+camera.position.set(0, 7, 7);
+camera.lookAt(0, 0, 0);
+
+const vNormal = TSL.varying(TSL.vec3(), 'vNormal');
+const vNormalAccessor = t3.fromTSL(vNormal, d.vec3f);
+const posAccessor = t3.fromTSL(TSL.positionLocal, d.vec3f);
+
+const updateNormal = (newNormal: d.v3f) => {
+  'use gpu';
+  vNormalAccessor.$.x = newNormal.x;
+  vNormalAccessor.$.y = newNormal.y;
+  vNormalAccessor.$.z = newNormal.z;
+};
+
+const positionNode = t3.toTSL(() => {
+  'use gpu';
+  const frequency = d.f32(3.0);
+  const amplitude = 0.5;
+  const wave = std.sin(posAccessor.$.x * frequency + t3.time.$);
+
+  posAccessor.$.y += wave * amplitude;
+
+  const derivative = std.cos(posAccessor.$.x * frequency + t3.time.$) *
+    amplitude * frequency;
+
+  const newNormalLocal = d.vec3f(-derivative, 1.0, 0);
+
+  updateNormal(newNormalLocal);
+
+  return d.vec3f(posAccessor.$);
+});
+
+const transformedNormalAccessor = t3.fromTSL(
+  TSL.transformNormalToView(vNormal),
+  d.vec3f,
+);
+
+const normalNode = t3.toTSL(() => {
+  'use gpu';
+  return std.normalize(transformedNormalAccessor.$);
+});
+
+const material = new THREE.MeshStandardNodeMaterial({
+  color: 0x9e0d3b,
+  roughness: 0.1,
+  metalness: 0.8,
+  side: THREE.DoubleSide,
+});
+
+material.positionNode = positionNode;
+material.normalNode = normalNode;
+
+const geometry = new THREE.PlaneGeometry(4, 4, 100, 100);
+geometry.rotateX(-Math.PI / 2);
+
+const mesh = new THREE.Mesh(geometry, material);
+scene.add(mesh);
+
+const dirLight = new THREE.DirectionalLight(0xffffff, 3);
+dirLight.position.set(-7, 10, 0);
+scene.add(dirLight);
+scene.add(new THREE.AmbientLight(0x444444));
+
+renderer.setAnimationLoop(() => {
+  renderer.render(scene, camera);
+});
+
+const resizeObserver = new ResizeObserver(() => {
+  camera.aspect = canvas.clientWidth / canvas.clientHeight;
+  camera.updateProjectionMatrix();
+  renderer.setSize(canvas.clientWidth, canvas.clientHeight, false);
+});
+resizeObserver.observe(canvas);
+
+export function onCleanup() {
+  resizeObserver.disconnect();
+  renderer.dispose();
+}

apps/typegpu-docs/src/examples/threejs/varyings/meta.json

@@ -0,0 +1,5 @@
+{
+  "title": "Three.js - Varyings",
+  "category": "threejs",
+  "tags": ["three.js"]
+}

apps/typegpu-docs/src/examples/threejs/varyings/thumbnail.png

@@ -17,16 +17,25 @@ interface TgpuFnNodeData extends THREE.NodeData {
   };
 }
 
-class StageData {
+abstract class StageData {
+  declare readonly type: 'analyze' | 'generate';
   readonly stage: 'vertex' | 'fragment' | 'compute' | null;
-  readonly names: WeakMap<object, string>;
   readonly namespace: Namespace;
-  codeGeneratedThusFar: string;
 
   constructor(stage: 'vertex' | 'fragment' | 'compute' | null) {
     this.stage = stage;
-    this.names = new WeakMap();
     this.namespace = tgpu['~unstable'].namespace();
+  }
+}
+
+class GenerateStageData extends StageData {
+  readonly names: WeakMap<object, string>;
+  readonly type = 'generate';
+  codeGeneratedThusFar: string;
+
+  constructor(stage: 'vertex' | 'fragment' | 'compute' | null) {
+    super(stage);
+    this.names = new WeakMap();
     this.codeGeneratedThusFar = '';
 
     this.namespace.on('name', (event) => {
@@ -37,18 +46,43 @@ class StageData {
   }
 }
 
+class AnalyzeStageData extends StageData {
+  readonly type = 'analyze';
+}
+
 class BuilderData {
-  stageDataMap: Map<'vertex' | 'fragment' | 'compute' | null, StageData>;
+  generateStageDataMap: Map<
+    'vertex' | 'fragment' | 'compute' | null,
+    GenerateStageData
+  >;
+  analyzeStageDataMap: Map<
+    'vertex' | 'fragment' | 'compute' | null,
+    AnalyzeStageData
+  >;
 
   constructor() {
-    this.stageDataMap = new Map();
+    this.generateStageDataMap = new Map();
+    this.analyzeStageDataMap = new Map();
+  }
+
+  getGenerateStageData(
+    stage: 'vertex' | 'fragment' | 'compute' | null,
+  ): GenerateStageData {
+    let stageData = this.generateStageDataMap.get(stage);
+    if (!stageData) {
+      stageData = new GenerateStageData(stage);
+      this.generateStageDataMap.set(stage, stageData);
+    }
+    return stageData;
   }
 
-  getStageData(stage: 'vertex' | 'fragment' | 'compute' | null): StageData {
-    let stageData = this.stageDataMap.get(stage);
+  getAnalyzeStageData(
+    stage: 'vertex' | 'fragment' | 'compute' | null,
+  ): AnalyzeStageData {
+    let stageData = this.analyzeStageDataMap.get(stage);
     if (!stageData) {
-      stageData = new StageData(stage);
-      this.stageDataMap.set(stage, stageData);
+      stageData = new AnalyzeStageData(stage);
+      this.analyzeStageDataMap.set(stage, stageData);
     }
     return stageData;
   }
@@ -108,7 +142,7 @@ class TgpuFnNode<T> extends THREE.Node {
       builderDataMap.set(builder, builderData);
     }
 
-    const stageData = builderData.getStageData(builder.shaderStage);
+    const stageData = builderData.getGenerateStageData(builder.shaderStage);
 
     if (!nodeData.custom) {
       if (currentlyGeneratingFnNodeCtx !== undefined) {
@@ -163,6 +197,42 @@ class TgpuFnNode<T> extends THREE.Node {
     return nodeData.custom.nodeFunction;
   }
 
+  #analyzeFunction(builder: THREE.NodeBuilder) {
+    let builderData = builderDataMap.get(builder);
+
+    if (!builderData) {
+      builderData = new BuilderData();
+      builderDataMap.set(builder, builderData);
+    }
+
+    const stageData = builderData.getAnalyzeStageData(builder.shaderStage);
+
+    const ctx: TgpuFnNodeContext = {
+      builder,
+      stageData,
+      dependencies: [],
+    };
+    currentlyGeneratingFnNodeCtx = ctx;
+    try {
+      tgpu.resolve({
+        names: stageData.namespace,
+        template: '___ID___ fnName',
+        externals: { fnName: this.#impl },
+      });
+    } finally {
+      currentlyGeneratingFnNodeCtx = undefined;
+    }
+  }
+
+  /**
+   * Replicating Three.js `analyze` traversal.
+   * Setting `needsInterpolation` flag to true in varying nodes
+   */
+  analyze(builder: THREE.NodeBuilder, output?: THREE.Node | null) {
+    super.analyze(builder, output);
+    this.#analyzeFunction(builder); // making sure it will find all TSL accessors
+  }
+
   generate(
     builder: THREE.NodeBuilder,
     output: string | null | undefined,
@@ -171,15 +241,16 @@ class TgpuFnNode<T> extends THREE.Node {
 
     const nodeData = builder.getDataFromNode(this) as TgpuFnNodeData;
     const builderData = builderDataMap.get(builder) as BuilderData;
-    const stageData = builderData.getStageData(builder.shaderStage);
+    const stageData = builderData.getGenerateStageData(builder.shaderStage);
 
     // Building dependencies
-    for (const dep of nodeData.custom.dependencies) {
+    const uniqueDeps = [...new Set(nodeData.custom.dependencies)];
+    for (const dep of uniqueDeps) {
       dep.node.build(builder);
     }
     nodeData.custom.priorCode.build(builder);
 
-    for (const dep of nodeData.custom.dependencies) {
+    for (const dep of uniqueDeps) {
       if (!dep.var) {
         continue;
       }
@@ -190,10 +261,9 @@ class TgpuFnNode<T> extends THREE.Node {
       builder.addLineFlowCode(`${varName} = ${varValue};\n`, this);
     }
 
-    if (output === 'property') {
-      return nodeData.custom.functionId;
-    }
-    return `${nodeData.custom.functionId}()`;
+    return output === 'property'
+      ? nodeData.custom.functionId
+      : `${nodeData.custom.functionId}()`;
   }
 }
 
@@ -206,7 +276,7 @@ export function toTSL(
 export class TSLAccessor<T extends d.AnyWgslData, TNode extends THREE.Node> {
   readonly #dataType: T;
 
-  readonly var: TgpuVar<'private', T> | undefined;
+  #var: TgpuVar<'private', T> | undefined;
   readonly node: THREE.TSL.NodeObject<TNode>;
 
   constructor(
@@ -216,31 +286,54 @@ export class TSLAccessor<T extends d.AnyWgslData, TNode extends THREE.Node> {
     this.node = node;
     this.#dataType = dataType;
 
-    // node.isTextureNode - temporary workaround for textures
     if (
-      // @ts-expect-error: The properties exist on the node
-      (!node.isStorageBufferNode && !node.isUniformNode) || node.isTextureNode
+      // @ts-expect-error: they are assigned at runtime
+      (!node.isStorageBufferNode && !node.isUniformNode) ||
+      // @ts-expect-error: it is assigned at runtime
+      node.isTextureNode
     ) {
-      this.var = tgpu.privateVar(dataType);
+      this.#var = tgpu.privateVar(dataType);
     }
   }
 
+  get var(): TgpuVar<'private', T> | undefined {
+    return this.#var;
+  }
+
   get $(): d.InferGPU<T> {
     const ctx = currentlyGeneratingFnNodeCtx;
 
     if (!ctx) {
       throw new Error('Can only access TSL nodes on the GPU.');
     }
 
+    if (ctx.stageData.type === 'analyze') {
+      this.node.traverse((node: THREE.Node) => {
+        node.analyze(ctx.builder);
+      });
+      // dummy return, only for types to match
+      return tgpu['~unstable'].rawCodeSnippet('', this.#dataType, 'runtime').$;
+    }
+
     // biome-ignore lint/suspicious/noExplicitAny: smh
     ctx.dependencies.push(this as any);
 
+    const builtNode = this.node.build(ctx.builder) as string;
+
+    // @ts-expect-error: it is assigned at runtime
+    const trueVaryingNode = this.node.isVaryingNode &&
+      builtNode.includes('varyings.');
+
+    if (trueVaryingNode) {
+      this.#var = undefined; // cannot be checked earlier, ThreeJS is lazy
+    }
+
     if (this.var) {
       return this.var.$;
     }
 
     return tgpu['~unstable'].rawCodeSnippet(
-      this.node.build(ctx.builder) as string,
+      builtNode,
       this.#dataType,
     ).$;
   }
@@ -295,8 +388,14 @@ export const fromTSL = tgpu['~unstable'].comptime<
   if (!sharedBuilder) {
     sharedBuilder = new WGSLNodeBuilder();
   }
-  const nodeType = node.getNodeType(sharedBuilder);
-
+  let nodeType: string | null = null;
+  try { // sometimes it needs information (overrideNodes) from compilation context which is not present
+    nodeType = node.getNodeType(sharedBuilder);
+  } catch (e) {
+    console.warn(
+      `fromTSL: failed to infer node type via getNodeType; skipping type comparison.`,
+    );
+  }
   if (nodeType) {
     const wgslTypeFromTSL = sharedBuilder.getType(nodeType);
     if (wgslTypeFromTSL !== wgslTypeFromTgpu) {