TSL: Add switch/case. (#30935)

* TSL: Add switch/case.

* StackNode: Support Case() sequence.

* StackNode: Refactor `Case()`.

Author: Mugen87

src/Three.TSL.js

@@ -13,6 +13,7 @@ export const F_Schlick = TSL.F_Schlick;
 export const Fn = TSL.Fn;
 export const INFINITY = TSL.INFINITY;
 export const If = TSL.If;
+export const Switch = TSL.Switch;
 export const Loop = TSL.Loop;
 export const NodeShaderStage = TSL.NodeShaderStage;
 export const NodeType = TSL.NodeType;

src/nodes/core/StackNode.js

@@ -1,6 +1,6 @@
 import Node from './Node.js';
 import { select } from '../math/ConditionalNode.js';
-import { ShaderNode, nodeProxy, getCurrentStack, setCurrentStack } from '../tsl/TSLBase.js';
+import { ShaderNode, nodeProxy, getCurrentStack, setCurrentStack, nodeObject } from '../tsl/TSLBase.js';
 
 /**
  * Stack is a helper for Nodes that need to produce stack-based code instead of continuous flow.
@@ -57,6 +57,16 @@ class StackNode extends Node {
 		 */
 		this._currentCond = null;
 
+		/**
+		 * The expression node. Only
+		 * relevant for Switch/Case.
+		 *
+		 * @private
+		 * @type {Node}
+		 * @default null
+		 */
+		this._expressionNode = null;
+
 		/**
 		 * This flag can be used for type testing.
 		 *
@@ -143,6 +153,98 @@ class StackNode extends Node {
 
 	}
 
+	/**
+	 * Represents a `switch` statement in TSL.
+	 *
+	 * @param {any} expression - Represents the expression.
+	 * @param {Function} method - TSL code which is executed if the condition evaluates to `true`.
+	 * @return {StackNode} A reference to this stack node.
+	 */
+	Switch( expression ) {
+
+		this._expressionNode = nodeObject( expression );
+
+		return this;
+
+	}
+
+	/**
+	 * Represents a `case` statement in TSL. The TSL version accepts an arbitrary numbers of values.
+	 * The last parameter must be the callback method that should be executed in the `true` case.
+	 *
+	 * @param {...any} params - The values of the `Case()` statement as well as the callback method.
+	 * @return {StackNode} A reference to this stack node.
+	 */
+	Case( ...params ) {
+
+		const caseNodes = [];
+
+		// extract case nodes from the parameter list
+
+		if ( params.length >= 2 ) {
+
+			for ( let i = 0; i < params.length - 1; i ++ ) {
+
+				caseNodes.push( this._expressionNode.equal( nodeObject( params[ i ] ) ) );
+
+			}
+
+		} else {
+
+			throw new Error( 'TSL: Invalid parameter length. Case() requires at least two parameters.' );
+
+		}
+
+		// extract method
+
+		const method = params[ params.length - 1 ];
+		const methodNode = new ShaderNode( method );
+
+		// chain multiple cases when using Case( 1, 2, 3, () => {} )
+
+		let caseNode = caseNodes[ 0 ];
+
+		for ( let i = 1; i < caseNodes.length; i ++ ) {
+
+			caseNode = caseNode.or( caseNodes[ i ] );
+
+		}
+
+		// build condition
+
+		const condNode = select( caseNode, methodNode );
+
+		if ( this._currentCond === null ) {
+
+			this._currentCond = condNode;
+
+			return this.add( this._currentCond );
+
+		} else {
+
+			this._currentCond.elseNode = condNode;
+			this._currentCond = condNode;
+
+			return this;
+
+		}
+
+	}
+
+	/**
+	 * Represents the default code block of a Switch/Case statement.
+	 *
+	 * @param {Function} method - TSL code which is executed in the `else` case.
+	 * @return {StackNode} A reference to this stack node.
+	 */
+	Default( method ) {
+
+		this.Else( method );
+
+		return this;
+
+	}
+
 	build( builder, ...params ) {
 
 		const previousStack = getCurrentStack();

src/nodes/tsl/TSLCore.js

@@ -705,6 +705,7 @@ export const setCurrentStack = ( stack ) => {
 export const getCurrentStack = () => currentStack;
 
 export const If = ( ...params ) => currentStack.If( ...params );
+export const Switch = ( ...params ) => currentStack.Switch( ...params );
 
 export function append( node ) {