TSL: Introduce automatic use of .toVar() to assign() (#31459)

* introduce automatic use of `.toVar()` to `assign()`

* fix nodeObj return

* move `builder.getArrayCount()` -> `node.getArrayCount()`

* Update NodeBuilder.js

* add `builder.flowBuildStage()`

* add docs and rename `intention` -> `intent`

* rename `nodeProxyIntention` -> `nodeProxyIntent`

* Update VarNode.js

* add `VarIntent`

* Update NodeMaterial.js

* improve setup check

Author: sunag

src/Three.TSL.js

@@ -519,6 +519,7 @@ export const uvec2 = TSL.uvec2;
 export const uvec3 = TSL.uvec3;
 export const uvec4 = TSL.uvec4;
 export const Var = TSL.Var;
+export const VarIntent = TSL.VarIntent;
 export const varying = TSL.varying;
 export const varyingProperty = TSL.varyingProperty;
 export const vec2 = TSL.vec2;

src/materials/nodes/NodeMaterial.js

@@ -1071,7 +1071,7 @@ class NodeMaterial extends Material {
 
 			output.assign( outputNode );
 
-			outputNode = vec4( fogNode.toStack() );
+			outputNode = vec4( fogNode.toVar() );
 
 		}
 

src/nodes/core/ArrayNode.js

@@ -58,6 +58,18 @@ class ArrayNode extends TempNode {
 
 	}
 
+	/**
+	 * Returns the number of elements in the node array.
+	 *
+	 * @param {NodeBuilder} builder - The current node builder.
+	 * @return {number} The number of elements in the node array.
+	 */
+	getArrayCount( /*builder*/ ) {
+
+		return this.count;
+
+	}
+
 	/**
 	 * Returns the node's type.
 	 *

src/nodes/core/AssignNode.js

@@ -97,6 +97,9 @@ class AssignNode extends TempNode {
 
 		const { targetNode, sourceNode } = this;
 
+		const targetProperties = builder.getNodeProperties( targetNode );
+		targetProperties.assign = true;
+
 		const properties = builder.getNodeProperties( this );
 		properties.sourceNode = sourceNode;
 		properties.targetNode = targetNode.context( { assign: true } );

src/nodes/core/Node.js

@@ -480,6 +480,18 @@ class Node extends EventDispatcher {
 
 	}
 
+	/**
+	 * Returns the number of elements in the node array.
+	 *
+	 * @param {NodeBuilder} builder - The current node builder.
+	 * @return {?number} The number of elements in the node array.
+	 */
+	getArrayCount( /*builder*/ ) {
+
+		return null;
+
+	}
+
 	/**
 	 * Represents the setup stage which is the first step of the build process, see {@link Node#build} method.
 	 * This method is often overwritten in derived modules to prepare the node which is used as the output/result.

src/nodes/core/NodeBuilder.js

@@ -1747,23 +1747,6 @@ class NodeBuilder {
 
 	}
 
-	/**
-	 * Returns the array length.
-	 *
-	 * @param {Node} node - The node.
-	 * @return {?number} The array length.
-	 */
-	getArrayCount( node ) {
-
-		let count = null;
-
-		if ( node.isArrayNode ) count = node.count;
-		else if ( node.isVarNode && node.node.isArrayNode ) count = node.node.count;
-
-		return count;
-
-	}
-
 	/**
 	 * Returns an instance of {@link NodeVar} for the given variable node.
 	 *
@@ -1807,7 +1790,7 @@ class NodeBuilder {
 
 			//
 
-			const count = this.getArrayCount( node );
+			const count = node.getArrayCount( this );
 
 			nodeVar = new NodeVar( name, type, readOnly, count );
 
@@ -2233,6 +2216,28 @@ class NodeBuilder {
 
 	}
 
+	/**
+	 * Executes the node in a specific build stage.
+	 *
+	 * @param {Node} node - The node to execute.
+	 * @param {string} buildStage - The build stage to execute the node in.
+	 * @param {Node|string|null} output - Expected output type. For example 'vec3'.
+	 * @return {Node|string|null} The result of the node build.
+	 */
+	flowBuildStage( node, buildStage, output = null ) {
+
+		const previousBuildStage = this.getBuildStage();
+
+		this.setBuildStage( buildStage );
+
+		const result = node.build( this, output );
+
+		this.setBuildStage( previousBuildStage );
+
+		return result;
+
+	}
+
 	/**
 	 * Runs the node flow through all the steps of creation, 'setup', 'analyze', 'generate'.
 	 *

src/nodes/core/StackNode.js

@@ -245,6 +245,36 @@ class StackNode extends Node {
 
 	}
 
+	setup( builder ) {
+
+		const nodeProperties = builder.getNodeProperties( this );
+
+		let index = 0;
+
+		for ( const childNode of this.getChildren() ) {
+
+			if ( childNode.isVarNode && childNode.intent === true ) {
+
+				const properties = builder.getNodeProperties( childNode );
+
+				if ( properties.assign !== true ) {
+
+					continue;
+
+				}
+
+			}
+
+			nodeProperties[ 'node' + index ++ ] = childNode;
+
+		}
+
+		// return a outputNode if exists or null
+
+		return nodeProperties.outputNode || null;
+
+	}
+
 	build( builder, ...params ) {
 
 		const previousBuildStack = builder.currentStack;
@@ -258,6 +288,18 @@ class StackNode extends Node {
 
 		for ( const node of this.nodes ) {
 
+			if ( node.isVarNode && node.intent === true ) {
+
+				const properties = builder.getNodeProperties( node );
+
+				if ( properties.assign !== true ) {
+
+					continue;
+
+				}
+
+			}
+
 			if ( buildStage === 'setup' ) {
 
 				node.build( builder );

src/nodes/core/VarNode.js

@@ -1,5 +1,5 @@
 import Node from './Node.js';
-import { addMethodChaining, nodeProxy } from '../tsl/TSLCore.js';
+import { addMethodChaining, getCurrentStack, nodeProxy } from '../tsl/TSLCore.js';
 
 /**
  * Class for representing shader variables as nodes. Variables are created from
@@ -81,6 +81,45 @@ class VarNode extends Node {
 		 */
 		this.parents = true;
 
+		/**
+		 * This flag is used to indicate that this node is used for intent.
+		 *
+		 * @type {boolean}
+		 * @default false
+		 */
+		this.intent = false;
+
+	}
+
+	/**
+	 * Sets the intent flag for this node.
+	 *
+	 * This flag is used to indicate that this node is used for intent
+	 * and should not be built directly. Instead, it is used to indicate that
+	 * the node should be treated as a variable intent.
+	 *
+	 * It's useful for assigning variables without needing creating a new variable node.
+	 *
+	 * @param {boolean} value - The value to set for the intent flag.
+	 * @returns {VarNode} This node.
+	 */
+	setIntent( value ) {
+
+		this.intent = value;
+
+		return this;
+
+	}
+
+	/**
+	 * Returns the intent flag of this node.
+	 *
+	 * @return {boolean} The intent flag.
+	 */
+	getIntent() {
+
+		return this.intent;
+
 	}
 
 	getMemberType( builder, name ) {
@@ -101,6 +140,31 @@ class VarNode extends Node {
 
 	}
 
+	getArrayCount( builder ) {
+
+		return this.node.getArrayCount( builder );
+
+	}
+
+	build( ...params ) {
+
+		if ( this.intent === true ) {
+
+			const builder = params[ 0 ];
+			const properties = builder.getNodeProperties( this );
+
+			if ( properties.assign !== true ) {
+
+				return this.node.build( ...params );
+
+			}
+
+		}
+
+		return super.build( ...params );
+
+	}
+
 	generate( builder ) {
 
 		const { node, name, readOnly } = this;
@@ -138,7 +202,7 @@ class VarNode extends Node {
 
 			} else {
 
-				const count = builder.getArrayCount( node );
+				const count = node.getArrayCount( builder );
 
 				declarationPrefix = `const ${ builder.getVar( nodeVar.type, propertyName, count ) }`;
 
@@ -189,10 +253,35 @@ export const Var = ( node, name = null ) => createVar( node, name ).toStack();
  */
 export const Const = ( node, name = null ) => createVar( node, name, true ).toStack();
 
+//
+//
+
+/**
+ * TSL function for creating a var intent node.
+ *
+ * @tsl
+ * @function
+ * @param {Node} node - The node for which a variable should be created.
+ * @param {?string} name - The name of the variable in the shader.
+ * @returns {VarNode}
+ */
+export const VarIntent = ( node ) => {
+
+	if ( getCurrentStack() === null ) {
+
+		return node;
+
+	}
+
+	return createVar( node ).setIntent( true ).toStack();
+
+};
+
 // Method chaining
 
 addMethodChaining( 'toVar', Var );
 addMethodChaining( 'toConst', Const );
+addMethodChaining( 'toVarIntent', VarIntent );
 
 // Deprecated
 

src/nodes/math/ConditionalNode.js

@@ -73,7 +73,7 @@ class ConditionalNode extends Node {
 
 			// fallback setup
 
-			this.setup( builder );
+			builder.flowBuildStage( this, 'setup' );
 
 			return this.getNodeType( builder );