TSL: Forces assignment of a function call if a loop is detected (#31961)

* forces assignment of a function call if a loop is detected

* Update LoopNode.js

* Update LoopNode.js

src/nodes/core/ContextNode.js

@@ -94,9 +94,7 @@ class ContextNode extends Node {
 
 	analyze( builder ) {
 
-		const previousContext = builder.getContext();
-
-		builder.setContext( { ...builder.context, ...this.value } );
+		const previousContext = builder.addContext( this.value );
 
 		this.node.build( builder );
 
@@ -106,9 +104,7 @@ class ContextNode extends Node {
 
 	setup( builder ) {
 
-		const previousContext = builder.getContext();
-
-		builder.setContext( { ...builder.context, ...this.value } );
+		const previousContext = builder.addContext( this.value );
 
 		this.node.build( builder );
 
@@ -118,9 +114,7 @@ class ContextNode extends Node {
 
 	generate( builder, output ) {
 
-		const previousContext = builder.getContext();
-
-		builder.setContext( { ...builder.context, ...this.value } );
+		const previousContext = builder.addContext( this.value );
 
 		const snippet = this.node.build( builder, output );
 

src/nodes/core/NodeBuilder.js

@@ -899,6 +899,22 @@ class NodeBuilder {
 
 	}
 
+	/**
+	 * Adds context data to the builder's current context.
+	 *
+	 * @param {Object} context - The context to add.
+	 * @return {Object} The previous context.
+	 */
+	addContext( context ) {
+
+		const previousContext = this.getContext();
+
+		this.setContext( { ...this.context, ...context } );
+
+		return previousContext;
+
+	}
+
 	/**
 	 * Gets a context used in shader construction that can be shared across different materials.
 	 * This is necessary since the renderer cache can reuse shaders generated in one material and use them in another.

src/nodes/core/StackNode.js

@@ -261,15 +261,9 @@ class StackNode extends Node {
 
 		for ( const childNode of this.getChildren() ) {
 
-			if ( childNode.isVarNode && childNode.intent === true ) {
+			if ( childNode.isVarNode && childNode.isAssign( builder ) !== true ) {
 
-				const properties = builder.getNodeProperties( childNode );
-
-				if ( properties.assign !== true ) {
-
-					continue;
-
-				}
+				continue;
 
 			}
 
@@ -302,15 +296,9 @@ class StackNode extends Node {
 
 		for ( const node of this.nodes ) {
 
-			if ( node.isVarNode && node.intent === true ) {
+			if ( node.isVarNode && node.isAssign( builder ) !== true ) {
 
-				const properties = builder.getNodeProperties( node );
-
-				if ( properties.assign !== true ) {
-
-					continue;
-
-				}
+				continue;
 
 			}
 

src/nodes/core/VarNode.js

@@ -146,14 +146,53 @@ class VarNode extends Node {
 
 	}
 
+	isAssign( builder ) {
+
+		if ( this.intent !== true ) return true;
+
+		//
+
+		const properties = builder.getNodeProperties( this );
+
+		let assign = properties.assign;
+
+		if ( assign !== true ) {
+
+			if ( this.node.isShaderCallNodeInternal && this.node.shaderNode.getLayout() === null ) {
+
+				if ( builder.context.fnCall && builder.context.fnCall.shaderNode ) {
+
+					const nodeType = this.node.getNodeType( builder );
+
+					if ( nodeType !== 'void' ) {
+
+						const shaderNodeData = builder.getDataFromNode( this.node.shaderNode );
+
+						if ( shaderNodeData.hasLoop ) {
+
+							assign = true;
+
+						}
+
+					}
+
+				}
+
+			}
+
+		}
+
+		return assign;
+
+	}
+
 	build( ...params ) {
 
 		if ( this.intent === true ) {
 
 			const builder = params[ 0 ];
-			const properties = builder.getNodeProperties( this );
 
-			if ( properties.assign !== true ) {
+			if ( this.isAssign( builder ) !== true ) {
 
 				return this.node.build( ...params );
 

src/nodes/tsl/TSLCore.js

@@ -488,6 +488,7 @@ class ShaderCallNodeInternal extends Node {
 		//
 
 		const previousSubBuildFn = builder.subBuildFn;
+		const previousContext = builder.addContext( { fnCall: this } );
 
 		builder.subBuildFn = subBuild;
 
@@ -565,6 +566,7 @@ class ShaderCallNodeInternal extends Node {
 		}
 
 		builder.subBuildFn = previousSubBuildFn;
+		builder.setContext( previousContext );
 
 		if ( shaderNode.once ) {
 
@@ -608,6 +610,8 @@ class ShaderCallNodeInternal extends Node {
 		const subBuildOutput = builder.getSubBuildOutput( this );
 		const outputNode = this.getOutputNode( builder );
 
+		const previousContext = builder.addContext( { fnCall: this } );
+
 		if ( buildStage === 'setup' ) {
 
 			const subBuildInitialized = builder.getSubBuildProperty( 'initialized', this );
@@ -655,6 +659,8 @@ class ShaderCallNodeInternal extends Node {
 
 		}
 
+		builder.setContext( previousContext );
+
 		return result;
 
 	}
@@ -788,9 +794,15 @@ class ShaderNodeInternal extends Node {
 
 	}
 
+	getLayout() {
+
+		return this.layout;
+
+	}
+
 	call( rawInputs = null ) {
 
-		return nodeObject( new ShaderCallNodeInternal( this, rawInputs ) );
+		return new ShaderCallNodeInternal( this, rawInputs );
 
 	}
 

src/nodes/utils/LoopNode.js

@@ -1,6 +1,6 @@
 import Node from '../core/Node.js';
 import { expression } from '../code/ExpressionNode.js';
-import { nodeObject, nodeArray, Fn } from '../tsl/TSLBase.js';
+import { nodeArray, Fn } from '../tsl/TSLBase.js';
 import { error } from '../../utils.js';
 
 /**
@@ -139,6 +139,13 @@ class LoopNode extends Node {
 
 		this.getProperties( builder );
 
+		if ( builder.context.fnCall ) {
+
+			const shaderNodeData = builder.getDataFromNode( builder.context.fnCall.shaderNode );
+			shaderNodeData.hasLoop = true;
+
+		}
+
 	}
 
 	generate( builder ) {
@@ -333,7 +340,7 @@ export default LoopNode;
  * @param {...any} params - A list of parameters.
  * @returns {LoopNode}
  */
-export const Loop = ( ...params ) => nodeObject( new LoopNode( nodeArray( params, 'int' ) ) ).toStack();
+export const Loop = ( ...params ) => new LoopNode( nodeArray( params, 'int' ) ).toStack();
 
 /**
  * TSL function for creating a `Continue()` expression.