TSL Transpiler: Add Support for Switch Statements (#31272)

* init branch

* work on switch transpiler

* work

* alternate approach

* remove logs, add default case

* rework switch to use parseBody and getGroupDelta

* fix read bug

* remove comments, add continue and break

* add import and improve code style

* skip break statements in switch cases

* add multiples case conditions support

---------

examples/jsm/transpiler/AST.js

@@ -214,6 +214,26 @@ export class Discard {
 
 }
 
+export class Continue {
+
+	constructor() {
+
+		this.isContinue = true;
+
+	}
+
+}
+
+export class Break {
+
+	constructor() {
+
+		this.isBreak = true;
+
+	}
+
+}
+
 export class Accessor {
 
 	constructor( property ) {
@@ -278,3 +298,37 @@ export class For {
 	}
 
 }
+
+export class Switch {
+
+	constructor( discriminant ) {
+
+		this.body = [];
+
+		this.discriminant = discriminant;
+		this.case = null;
+		this.isSwitch = true;
+
+	}
+
+}
+
+export class SwitchCase {
+
+	constructor( caseCondition ) {
+
+		// Condition for the case body to execute
+		this.caseCondition = caseCondition;
+
+		// Body of the case statement
+		this.body = [];
+
+		// Next case to fall to if current case fails
+		this.nextCase = null;
+
+		this.isDefault = caseCondition === null ? true : false;
+		this.isSwitchCase = true;
+
+	}
+
+}

examples/jsm/transpiler/GLSLDecoder.js

@@ -1,4 +1,4 @@
-import { Program, FunctionDeclaration, For, AccessorElements, Ternary, Varying, DynamicElement, StaticElement, FunctionParameter, Unary, Conditional, VariableDeclaration, Operator, Number, String, FunctionCall, Return, Accessor, Uniform, Discard } from './AST.js';
+import { Program, FunctionDeclaration, Switch, For, AccessorElements, Ternary, Varying, DynamicElement, StaticElement, FunctionParameter, Unary, Conditional, VariableDeclaration, Operator, Number, String, FunctionCall, Return, Accessor, Uniform, Discard, SwitchCase, Continue, Break } from './AST.js';
 
 const unaryOperators = [
 	'+', '-', '~', '!', '++', '--'
@@ -63,7 +63,7 @@ function getFunctionName( str ) {
 function getGroupDelta( str ) {
 
 	if ( str === '(' || str === '[' || str === '{' ) return 1;
-	if ( str === ')' || str === ']' || str === '}' ) return - 1;
+	if ( str === ')' || str === ']' || str === '}' || str === 'case' || str === 'default' ) return - 1;
 
 	return 0;
 
@@ -492,6 +492,14 @@ class GLSLDecoder {
 
 				return new Discard();
 
+			} else if ( firstToken.str === 'continue' ) {
+
+				return new Continue();
+
+			} else if ( firstToken.str === 'break' ) {
+
+				return new Break();
+
 			}
 
 			const secondToken = tokens[ 1 ];
@@ -801,6 +809,110 @@ class GLSLDecoder {
 
 	}
 
+	parseSwitch() {
+
+		const parseSwitchExpression = () => {
+
+			this.readToken(); // Skip 'switch'
+
+			const switchDeterminantTokens = this.readTokensUntil( ')' );
+
+			// Parse expresison between parentheses. Index 1: char after '('. Index -1: char before ')'
+			return this.parseExpressionFromTokens( switchDeterminantTokens.slice( 1, - 1 ) );
+
+
+		};
+
+		const parseSwitchBlock = ( switchStatement ) => {
+
+			// Validate curly braces
+			if ( this.getToken().str === '{' ) {
+
+				this.readToken(); // Skip '{'
+
+			} else {
+
+				throw new Error( 'Expected \'{\' after switch(...) ' );
+
+			}
+
+			if ( this.getToken() && ( this.getToken().str === 'case' || this.getToken().str === 'default' ) ) {
+
+				switchStatement.case = this.parseSwitchCase();
+
+			} else {
+
+				this.parseBlock( switchStatement );
+
+			}
+
+
+		};
+
+		const switchStatement = new Switch( parseSwitchExpression() );
+
+		parseSwitchBlock( switchStatement );
+
+		return switchStatement;
+
+
+	}
+
+	parseSwitchCase() {
+
+		const parseCaseExpression = ( token ) => {
+
+			const caseTypeToken = token ? token : this.readToken(); // Skip 'case' or 'default
+
+			const caseTokens = this.readTokensUntil( ':' );
+
+			// No case condition on default
+			if ( caseTypeToken.str === 'default' ) {
+
+				return null;
+
+			}
+
+			return this.parseExpressionFromTokens( caseTokens.slice( 0, - 1 ) );
+
+		};
+
+		let lastReadToken = null;
+
+		// No '{' so use different approach
+		const parseCaseBlock = ( caseStatement ) => {
+
+			lastReadToken = this.parseBlock( caseStatement );
+
+		};
+
+		// Parse case condition
+		const caseCondition = parseCaseExpression();
+		const switchCase = new SwitchCase( caseCondition );
+
+		// Get case body
+		parseCaseBlock( switchCase );
+
+		let currentCase = switchCase;
+
+		// If block ended with case, then continue chaining cases, otherwise, ended with '}' and no more case blocks to parse
+		while ( lastReadToken.str === 'case' || lastReadToken.str === 'default' ) {
+
+			const previousCase = currentCase;
+
+			// case and default already skipped at block end, so need to pass it in as last read token
+			currentCase = new SwitchCase( parseCaseExpression( lastReadToken ) );
+
+			previousCase.nextCase = currentCase;
+
+			parseCaseBlock( currentCase );
+
+		}
+
+		return switchCase;
+
+	}
+
 	parseIf() {
 
 		const parseIfExpression = () => {
@@ -841,10 +953,13 @@ class GLSLDecoder {
 
 			this.readToken(); // skip 'else'
 
+			// Assign the current if/else statement as the previous within the chain of conditionals
 			const previous = current;
 
+			// If an 'else if' statement, parse the conditional within the if
 			if ( this.getToken().str === 'if' ) {
 
+				// Current conditional now equal to next conditional in the chain
 				current = new Conditional( parseIfExpression() );
 
 			} else {
@@ -853,8 +968,10 @@ class GLSLDecoder {
 
 			}
 
+			// n - 1 conditional's else statement assigned to new if/else statement
 			previous.elseConditional = current;
 
+			// Parse conditional of latest if statement
 			parseIfBlock( current );
 
 		}
@@ -885,9 +1002,10 @@ class GLSLDecoder {
 
 			if ( groupIndex < 0 ) {
 
-				this.readToken(); // skip '}'
+				this.readToken(); // skip '}', ']', 'case', or other block ending tokens'
 
-				break;
+				// Return skipped token
+				return token;
 
 			}
 
@@ -931,6 +1049,10 @@ class GLSLDecoder {
 
 					statement = this.parseFor();
 
+				} else if ( token.str === 'switch' ) {
+
+					statement = this.parseSwitch();
+
 				} else {
 
 					statement = this.parseExpression();

examples/jsm/transpiler/TSLEncoder.js

@@ -66,6 +66,8 @@ class TSLEncoder {
 		this._currentProperties = {};
 		this._lastStatement = null;
 
+		this.block = null;
+
 	}
 
 	addImport( name ) {
@@ -253,6 +255,17 @@ class TSLEncoder {
 
 			code = 'Discard()';
 
+		} else if ( node.isBreak ) {
+
+			this.addImport( 'Break' );
+
+			code = 'Break()';
+
+		} else if ( node.isContinue ) {
+
+			this.addImport( 'Continue' );
+			code = 'Continue()';
+
 		} else if ( node.isAccessorElements ) {
 
 			code = this.emitExpression( node.object );
@@ -293,6 +306,10 @@ class TSLEncoder {
 
 			code = this.emitFor( node );
 
+		} else if ( node.isSwitch ) {
+
+			code = this.emitSwitch( node );
+
 		} else if ( node.isVariableDeclaration ) {
 
 			code = this.emitVariables( node );
@@ -378,6 +395,12 @@ class TSLEncoder {
 
 		for ( const statement of body ) {
 
+			if ( this.block && this.block.isSwitchCase ) {
+
+				if ( statement.isBreak ) continue; // skip break statements in switch cases
+
+			}
+
 			code += this.emitExtraLine( statement );
 			code += this.tab + this.emitExpression( statement );
 
@@ -507,6 +530,69 @@ ${ this.tab }} )`;
 
 	}
 
+
+	emitSwitch( switchNode ) {
+
+		const discriminantString = this.emitExpression( switchNode.discriminant );
+
+		this.tab += '\t';
+
+		let switchString = `Switch( ${ discriminantString } )\n${ this.tab }`;
+
+		let caseNode = switchNode.case;
+
+		const previousBlock = this.block;
+
+		while ( caseNode !== null ) {
+
+			this.block = caseNode;
+
+			let caseBodyString;
+
+			if ( ! caseNode.isDefault ) {
+
+				const caseConditions = [ this.emitExpression( caseNode.caseCondition ) ];
+
+				while ( caseNode.body.length === 0 && caseNode.nextCase !== null && caseNode.nextCase.isDefault !== true ) {
+
+					caseNode = caseNode.nextCase;
+
+					caseConditions.push( this.emitExpression( caseNode.caseCondition ) );
+
+				}
+
+				caseBodyString = this.emitBody( caseNode.body );
+
+				switchString += `.Case( ${ caseConditions.join( ', ' ) }, `;
+
+			} else {
+
+				caseBodyString = this.emitBody( caseNode.body );
+
+				switchString += '.Default( ';
+
+			}
+
+			switchString += `() => {
+
+${ caseBodyString }
+
+${ this.tab }} )`;
+
+			caseNode = caseNode.nextCase;
+
+		}
+
+		this.block = previousBlock;
+
+		this.tab = this.tab.slice( 0, - 1 );
+
+		this.imports.add( 'Switch' );
+
+		return switchString;
+
+	}
+
 	emitFor( node ) {
 
 		const { initialization, condition, afterthought } = node;
@@ -754,7 +840,7 @@ ${ this.tab }} )`;
 
 		if ( statement.isReturn ) return '\n';
 
-		const isExpression = ( st ) => st.isFunctionDeclaration !== true && st.isFor !== true && st.isConditional !== true;
+		const isExpression = ( st ) => st.isFunctionDeclaration !== true && st.isFor !== true && st.isConditional !== true && st.isSwitch !== true;
 		const lastExp = isExpression( last );
 		const currExp = isExpression( statement );