TSL: getNodeType() is unmemoized recursion, so NodeBuilder.build() cost grows with the number of paths through the node graph

[potrepka]

Description

Node.getNodeType( builder ) is resolved by recursing into child nodes on every query, and nothing caches the result for the duration of a build. VarNode.getNodeType returns this.node.getNodeType( builder ), so .toVar() does not cut the recursion either. OperatorNode, MathNode, ConvertNode, SplitNode, JoinNode, ConditionalNode, ContextNode, StackNode and ShaderCallNodeInternal all recurse the same way. Type queries are issued per node per build stage, so the cost of NodeBuilder.build() is proportional to the number of paths through the node graph rather than the number of nodes, and it becomes exponential as soon as a shared subexpression feeds more than one parent.

In our app a generated graph of ~12,000 nodes (a 256-step accumulation chain with shared subexpressions) never returns from NodeBuilder.build(): over 168 million getDataFromNode calls were counted before the tab was killed.

The same applies to layout functions: FnNode.setLayout() stores the layout but never sets shaderNode.nodeType, so ShaderCallNodeInternal.getNodeType re-derives the declared return type through getOutputNode() and the stack machinery on every call, even though the type is written in the layout. Only the Fn( jsFunc, layoutOrType ) constructor argument sets it.

Suggested fix: cache the resolved type per node on the builder for the duration of the build, in the per-node data store getDataFromNode already provides, and have the recursing getNodeType implementations go through it. Only a settled type should be stored (null, undefined and 'void' can still change once a node has been set up). With such a cache on the classes above, the reproduction below goes from not returning to tens of milliseconds, and the generated WGSL is byte-identical before and after (checked on the reproduction and on 65 real materials of up to 340 KB of WGSL each). FnNode.setLayout() could also set this.shaderNode.nodeType = layout.type when given a full layout, so a layout function's calls resolve their type without entering the body.

For reference, r185 already replaced the nodes.includes() scan in NodeBuilder.addNode() and the indexOf() filter in StackNode.build() with set lookups, which removed two O(n²) costs on the same graphs. The type resolution above is what remains in r185.1 (VarNode.js:156, TSLCore.js:464).

Reproduction steps

1. Build the material below, where each step references the previous variable twice, so every step doubles the number of paths through the graph.
2. Time WGSLNodeBuilder.build() (or renderer.compileAsync()) for increasing steps.

Measured in Bun with new WGSLNodeBuilder( mesh, renderer ).build():

steps	r183.2	with a per-build type cache
8	17 ms	8 ms
12	110 ms	9 ms
16	1,569 ms	10 ms
18	6,217 ms	10 ms
64	does not return	18 ms
256	does not return	58 ms

Code

import { float, Fn, uv, vec3 } from 'three/tsl';
import { MeshBasicNodeMaterial } from 'three/webgpu';

const steps = 18;

const material = new MeshBasicNodeMaterial();

material.colorNode = Fn( () => {

	const base = uv().x.mul( 3 ).toVar();
	let sum = float( 0 ).toVar();

	for ( let i = 0; i < steps; i ++ ) {

		sum = sum.add( base.mul( i ).sin() ).mul( sum.add( 1 ).xxx.y ).toVar();

	}

	return vec3( sum, sum.mul( 0.5 ), base );

} )();

Live example

No response

Screenshots

No response

Version

r183.2 (still present in r185.1)

Device

Desktop

Browser

Chrome

OS

macOS

[cmhhelgeson]

Are you able to contribute the code from your testing branch in a PR?