TSL: Add expression support for loop( { update: ... } ) (mrdoob#30947)

* add expression support for `loop()` update

* Raymarching: update with new approach

* TSLEncoder: update with new approach

* rev

* rev

Author: sunag

examples/jsm/transpiler/TSLEncoder.js

@@ -466,12 +466,34 @@ ${ this.tab }} )`;
 		const name = node.initialization.name;
 		const type = node.initialization.type;
 		const condition = node.condition.type;
-		const update = node.afterthought.type;
 
 		const nameParam = name !== 'i' ? `, name: '${ name }'` : '';
 		const typeParam = type !== 'int' ? `, type: '${ type }'` : '';
 		const conditionParam = condition !== '<' ? `, condition: '${ condition }'` : '';
-		const updateParam = update !== '++' ? `, update: '${ update }'` : '';
+
+		let updateParam = '';
+
+		if ( node.afterthought.isUnary ) {
+
+			if ( node.afterthought.type !== '++' ) {
+
+				updateParam = `, update: '${ node.afterthought.type }'`;
+
+			}
+
+		} else if ( node.afterthought.isOperator ) {
+
+			if ( node.afterthought.right.isAccessor || node.afterthought.right.isNumber ) {
+
+				updateParam = `, update: ${ this.emitExpression( node.afterthought.right ) }`;
+
+			} else {
+
+				updateParam = `, update: ( { i } ) => ${ this.emitExpression( node.afterthought ) }`;
+
+			}
+
+		}
 
 		let loopStr = `Loop( { start: ${ start }, end: ${ end + nameParam + typeParam + conditionParam + updateParam } }, ( { ${ name } } ) => {\n\n`;
 

examples/jsm/tsl/utils/Raymarching.js

@@ -53,13 +53,13 @@ export const RaymarchingBox = ( steps, callback ) => {
 	bounds.assign( vec2( max( bounds.x, 0.0 ), bounds.y ) );
 
 	const inc = vec3( rayDir.abs().reciprocal() ).toVar();
-	const delta = float( min( inc.x, min( inc.y, inc.z ) ) ).toVar( 'rayDelta' ); // used 'rayDelta' name in loop
+	const delta = float( min( inc.x, min( inc.y, inc.z ) ) ).toVar();
 
 	delta.divAssign( float( steps ) );
 
 	const positionRay = vec3( vOrigin.add( bounds.x.mul( rayDir ) ) ).toVar();
 
-	Loop( { type: 'float', start: bounds.x, end: bounds.y, update: '+= rayDelta' }, () => {
+	Loop( { type: 'float', start: bounds.x, end: bounds.y, update: delta }, () => {
 
 		callback( { positionRay } );
 

src/nodes/utils/LoopNode.js

@@ -1,6 +1,6 @@
 import Node from '../core/Node.js';
 import { expression } from '../code/ExpressionNode.js';
-import { nodeObject, nodeArray } from '../tsl/TSLBase.js';
+import { nodeObject, nodeArray, Fn } from '../tsl/TSLBase.js';
 
 /**
  * This module offers a variety of ways to implement loops in TSL. In it's basic form it's:
@@ -101,9 +101,17 @@ class LoopNode extends Node {
 
 		const stack = builder.addStack(); // TODO: cache() it
 
-		properties.returnsNode = this.params[ this.params.length - 1 ]( inputs, stack, builder );
+		properties.returnsNode = this.params[ this.params.length - 1 ]( inputs, builder );
 		properties.stackNode = stack;
 
+		const baseParam = this.params[ 0 ];
+
+		if ( baseParam.isNode !== true && typeof baseParam.update === 'function' ) {
+
+			properties.updateNode = Fn( this.params[ 0 ].update )( inputs );
+
+		}
+
 		builder.removeStack();
 
 		return properties;
@@ -222,30 +230,69 @@ class LoopNode extends Node {
 				const startSnippet = internalParam.start;
 				const endSnippet = internalParam.end;
 
-				let declarationSnippet = '';
-				let conditionalSnippet = '';
-				let updateSnippet = '';
+				let updateSnippet;
+
+				const deltaOperator = () => condition.includes( '<' ) ? '+=' : '-=';
+
+				if ( update !== undefined && update !== null ) {
+
+					switch ( typeof update ) {
+
+						case 'function':
+
+							const flow = builder.flowStagesNode( properties.updateNode, 'void' );
+							const snippet = flow.code.replace( /\t|;/g, '' );
+
+							updateSnippet = snippet;
+
+							break;
+
+						case 'number':
+
+							updateSnippet = name + ' ' + deltaOperator() + ' ' + builder.generateConst( type, update );
+
+							break;
 
-				if ( ! update ) {
+						case 'string':
+
+							updateSnippet = name + ' ' + update;
+
+							break;
+
+						default:
+
+							if ( update.isNode ) {
+
+								updateSnippet = name + ' ' + deltaOperator() + ' ' + update.build( builder );
+
+							} else {
+
+								console.error( 'THREE.TSL: \'Loop( { update: ... } )\' is not a function, string or number.' );
+
+								updateSnippet = 'break /* invalid update */';
+
+							}
+
+					}
+
+				} else {
 
 					if ( type === 'int' || type === 'uint' ) {
 
-						if ( condition.includes( '<' ) ) update = '++';
-						else update = '--';
+						update = condition.includes( '<' ) ? '++' : '--';
 
 					} else {
 
-						if ( condition.includes( '<' ) ) update = '+= 1.';
-						else update = '-= 1.';
+						update = deltaOperator() + ' 1.';
 
 					}
 
-				}
+					updateSnippet = name + ' ' + update;
 
-				declarationSnippet += builder.getVar( type, name ) + ' = ' + startSnippet;
+				}
 
-				conditionalSnippet += name + ' ' + condition + ' ' + endSnippet;
-				updateSnippet += name + ' ' + update;
+				const declarationSnippet = builder.getVar( type, name ) + ' = ' + startSnippet;
+				const conditionalSnippet = name + ' ' + condition + ' ' + endSnippet;
 
 				loopSnippet = `for ( ${ declarationSnippet }; ${ conditionalSnippet }; ${ updateSnippet } )`;