three.js/examples/jsm/transpiler/TSLEncoder.js

import { REVISION } from 'three/webgpu';
import * as TSL from 'three/tsl';

import { VariableDeclaration, Accessor } from './AST.js';

const opLib = {
	'=': 'assign',
	'+': 'add',
	'-': 'sub',
	'*': 'mul',
	'/': 'div',
	'%': 'remainder',
	'<': 'lessThan',
	'>': 'greaterThan',
	'<=': 'lessThanEqual',
	'>=': 'greaterThanEqual',
	'==': 'equal',
	'&&': 'and',
	'||': 'or',
	'^^': 'xor',
	'&': 'bitAnd',
	'|': 'bitOr',
	'^': 'bitXor',
	'<<': 'shiftLeft',
	'>>': 'shiftRight',
	'+=': 'addAssign',
	'-=': 'subAssign',
	'*=': 'mulAssign',
	'/=': 'divAssign',
	'%=': 'remainderAssign',
	'^=': 'bitXorAssign',
	'&=': 'bitAndAssign',
	'|=': 'bitOrAssign',
	'<<=': 'shiftLeftAssign',
	'>>=': 'shiftRightAssign'
};

const unaryLib = {
	'+': '', // positive
	'-': 'negate',
	'~': 'bitNot',
	'!': 'not',
	'++': 'increment', // incrementBefore
	'--': 'decrement' // decrementBefore
};

const isPrimitive = ( value ) => /^(true|false|-?(\d|\.\d))/.test( value );

class TSLEncoder {

	constructor() {

		this.tab = '';
		this.imports = new Set();
		this.global = new Set();
		this.overloadings = new Map();
		this.iife = false;
		this.uniqueNames = false;
		this.reference = false;

		this._currentProperties = {};
		this._lastStatement = null;

	}

	addImport( name ) {

		// import only if it's a node

		name = name.split( '.' )[ 0 ];

		if ( TSL[ name ] !== undefined && this.global.has( name ) === false && this._currentProperties[ name ] === undefined ) {

			this.imports.add( name );

		}

	}

	emitUniform( node ) {

		let code = `const ${ node.name } = `;

		if ( this.reference === true ) {

			this.addImport( 'reference' );

			this.global.add( node.name );

			//code += `reference( '${ node.name }', '${ node.type }', uniforms )`;

			// legacy
			code += `reference( 'value', '${ node.type }', uniforms[ '${ node.name }' ] )`;

		} else {

			this.addImport( 'uniform' );

			this.global.add( node.name );

			code += `uniform( '${ node.type }' )`;

		}

		return code;

	}

	emitExpression( node ) {

		let code;

		if ( node.isAccessor ) {

			this.addImport( node.property );

			code = node.property;

		} else if ( node.isNumber ) {

			if ( node.type === 'int' || node.type === 'uint' ) {

				code = node.type + '( ' + node.value + ' )';

				this.addImport( node.type );

			} else {

				code = node.value;

			}

		} else if ( node.isString ) {

			code = '\'' + node.value + '\'';

		} else if ( node.isOperator ) {

			const opFn = opLib[ node.type ] || node.type;

			const left = this.emitExpression( node.left );
			const right = this.emitExpression( node.right );

			if ( isPrimitive( left ) && isPrimitive( right ) ) {

				return left + ' ' + node.type + ' ' + right;

			}

			if ( isPrimitive( left ) ) {

				code = opFn + '( ' + left + ', ' + right + ' )';

				this.addImport( opFn );

			} else if ( opFn === '.' ) {

				code = left + opFn + right;

			} else {

				code = left + '.' + opFn + '( ' + right + ' )';

			}

		} else if ( node.isFunctionCall ) {

			const params = [];

			for ( const parameter of node.params ) {

				params.push( this.emitExpression( parameter ) );

			}

			this.addImport( node.name );

			const paramsStr = params.length > 0 ? ' ' + params.join( ', ' ) + ' ' : '';

			code = `${ node.name }(${ paramsStr })`;

		} else if ( node.isReturn ) {

			code = 'return';

			if ( node.value ) {

				code += ' ' + this.emitExpression( node.value );

			}

		} else if ( node.isAccessorElements ) {

			code = this.emitExpression( node.object );

			for ( const element of node.elements ) {

				if ( element.isStaticElement ) {

					code += '.' + this.emitExpression( element.value );

				} else if ( element.isDynamicElement ) {

					const value = this.emitExpression( element.value );

					if ( isPrimitive( value ) ) {

						code += `[ ${ value } ]`;

					} else {

						code += `.element( ${ value } )`;

					}

				}

			}

		} else if ( node.isDynamicElement ) {

			code = this.emitExpression( node.value );

		} else if ( node.isStaticElement ) {

			code = this.emitExpression( node.value );

		} else if ( node.isFor ) {

			code = this.emitFor( node );

		} else if ( node.isVariableDeclaration ) {

			code = this.emitVariables( node );

		} else if ( node.isUniform ) {

			code = this.emitUniform( node );

		} else if ( node.isVarying ) {

			code = this.emitVarying( node );

		} else if ( node.isTernary ) {

			code = this.emitTernary( node );

		} else if ( node.isConditional ) {

			code = this.emitConditional( node );

		} else if ( node.isUnary && node.expression.isNumber ) {

			code = node.expression.type + '( ' + node.type + ' ' + node.expression.value + ' )';

			this.addImport( node.expression.type );

		} else if ( node.isUnary ) {

			let type = unaryLib[ node.type ];

			if ( node.after === false && ( node.type === '++' || node.type === '--' ) ) {

				type += 'Before';

			}

			const exp = this.emitExpression( node.expression );

			if ( isPrimitive( exp ) ) {

				this.addImport( type );

				code = type + '( ' + exp + ' )';

			} else {

				code = exp + '.' + type + '()';

			}

		} else {

			console.warn( 'Unknown node type', node );

		}

		if ( ! code ) code = '/* unknown statement */';

		return code;

	}

	emitBody( body ) {

		this.setLastStatement( null );

		let code = '';

		this.tab += '\t';

		for ( const statement of body ) {

			code += this.emitExtraLine( statement );
			code += this.tab + this.emitExpression( statement );

			if ( code.slice( - 1 ) !== '}' ) code += ';';

			code += '\n';

			this.setLastStatement( statement );

		}

		code = code.slice( 0, - 1 ); // remove the last extra line

		this.tab = this.tab.slice( 0, - 1 );

		return code;


	}

	emitTernary( node ) {

		const condStr = this.emitExpression( node.cond );
		const leftStr = this.emitExpression( node.left );
		const rightStr = this.emitExpression( node.right );

		this.addImport( 'select' );

		return `select( ${ condStr }, ${ leftStr }, ${ rightStr } )`;

	}

	emitConditional( node ) {

		const condStr = this.emitExpression( node.cond );
		const bodyStr = this.emitBody( node.body );

		let ifStr = `If( ${ condStr }, () => {

${ bodyStr }

${ this.tab }} )`;

		let current = node;

		while ( current.elseConditional ) {

			const elseBodyStr = this.emitBody( current.elseConditional.body );

			if ( current.elseConditional.cond ) {

				const elseCondStr = this.emitExpression( current.elseConditional.cond );

				ifStr += `.ElseIf( ${ elseCondStr }, () => {

${ elseBodyStr }

${ this.tab }} )`;

			} else {

				ifStr += `.Else( () => {

${ elseBodyStr }

${ this.tab }} )`;

			}

			current = current.elseConditional;


		}

		this.imports.add( 'If' );

		return ifStr;

	}

	emitLoop( node ) {

		const start = this.emitExpression( node.initialization.value );
		const end = this.emitExpression( node.condition.right );

		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 loopStr = `Loop( { start: ${ start }, end: ${ end + nameParam + typeParam + conditionParam + updateParam } }, ( { ${ name } } ) => {\n\n`;

		loopStr += this.emitBody( node.body ) + '\n\n';

		loopStr += this.tab + '} )';

		this.imports.add( 'Loop' );

		return loopStr;

	}

	emitFor( node ) {

		const { initialization, condition, afterthought } = node;

		if ( ( initialization && initialization.isVariableDeclaration && initialization.next === null ) &&
			( condition && condition.left.isAccessor && condition.left.property === initialization.name ) &&
			( afterthought && afterthought.isUnary ) &&
			( initialization.name === afterthought.expression.property )
		) {

			return this.emitLoop( node );

		}

		return this.emitForWhile( node );

	}

	emitForWhile( node ) {

		const initialization = this.emitExpression( node.initialization );
		const condition = this.emitExpression( node.condition );
		const afterthought = this.emitExpression( node.afterthought );

		this.tab += '\t';

		let forStr = '{\n\n' + this.tab + initialization + ';\n\n';
		forStr += `${ this.tab }While( ${ condition }, () => {\n\n`;

		forStr += this.emitBody( node.body ) + '\n\n';

		forStr += this.tab + '\t' + afterthought + ';\n\n';

		forStr += this.tab + '} )\n\n';

		this.tab = this.tab.slice( 0, - 1 );

		forStr += this.tab + '}';

		this.imports.add( 'While' );

		return forStr;

	}

	emitVariables( node, isRoot = true ) {

		const { name, type, value, next } = node;

		const valueStr = value ? this.emitExpression( value ) : '';

		let varStr = isRoot ? 'const ' : '';
		varStr += name;

		if ( value ) {

			if ( value.isFunctionCall && value.name === type ) {

				varStr += ' = ' + valueStr;

			} else {

				varStr += ` = ${ type }( ${ valueStr } )`;

			}

		} else {

			varStr += ` = ${ type }()`;

		}

		if ( node.immutable === false ) {

			varStr += '.toVar()';

		}

		if ( next ) {

			varStr += ', ' + this.emitVariables( next, false );

		}

		this.addImport( type );

		return varStr;

	}

	emitVarying( node ) {

		const { name, type } = node;

		this.addImport( 'varying' );
		this.addImport( type );

		return `const ${ name } = varying( ${ type }(), '${ name }' )`;

	}

	emitOverloadingFunction( nodes ) {

		const { name } = nodes[ 0 ];

		this.addImport( 'overloadingFn' );

		const prefix = this.iife === false ? 'export ' : '';

		return `${ prefix }const ${ name } = /*#__PURE__*/ overloadingFn( [ ${ nodes.map( node => node.name + '_' + nodes.indexOf( node ) ).join( ', ' ) } ] );\n`;

	}

	emitFunction( node ) {

		const { name, type } = node;

		this._currentProperties = { name: node };

		const params = [];
		const inputs = [];
		const mutableParams = [];

		let hasPointer = false;

		for ( const param of node.params ) {

			let str = `{ name: '${ param.name }', type: '${ param.type }'`;

			let name = param.name;

			if ( param.immutable === false && ( param.qualifier !== 'inout' && param.qualifier !== 'out' ) ) {

				name = name + '_immutable';

				mutableParams.push( param );

			}

			if ( param.qualifier ) {

				if ( param.qualifier === 'inout' || param.qualifier === 'out' ) {

					hasPointer = true;

				}

				str += ', qualifier: \'' + param.qualifier + '\'';

			}

			inputs.push( str + ' }' );
			params.push( name );

			this._currentProperties[ name ] = param;

		}

		for ( const param of mutableParams ) {

			node.body.unshift( new VariableDeclaration( param.type, param.name, new Accessor( param.name + '_immutable' ) ) );

		}

		const paramsStr = params.length > 0 ? ' [ ' + params.join( ', ' ) + ' ] ' : '';
		const bodyStr = this.emitBody( node.body );

		let fnName = name;
		let overloadingNodes = null;

		if ( this.overloadings.has( name ) ) {

			const overloadings = this.overloadings.get( name );

			if ( overloadings.length > 1 ) {

				const index = overloadings.indexOf( node );

				fnName += '_' + index;

				if ( index === overloadings.length - 1 ) {

					overloadingNodes = overloadings;

				}

			}

		}

		const prefix = this.iife === false ? 'export ' : '';

		let funcStr = `${ prefix }const ${ fnName } = /*#__PURE__*/ Fn( (${ paramsStr }) => {

${ bodyStr }

${ this.tab }} )`;

		const layoutInput = inputs.length > 0 ? '\n\t\t' + this.tab + inputs.join( ',\n\t\t' + this.tab ) + '\n\t' + this.tab : '';

		if ( node.layout !== false && hasPointer === false ) {

			const uniqueName = this.uniqueNames ? fnName + '_' + Math.random().toString( 36 ).slice( 2 ) : fnName;

			funcStr += `.setLayout( {
${ this.tab }\tname: '${ uniqueName }',
${ this.tab }\ttype: '${ type }',
${ this.tab }\tinputs: [${ layoutInput }]
${ this.tab }} )`;

		}

		funcStr += ';\n';

		this.imports.add( 'Fn' );

		this.global.add( node.name );

		if ( overloadingNodes !== null ) {

			funcStr += '\n' + this.emitOverloadingFunction( overloadingNodes );

		}

		return funcStr;

	}

	setLastStatement( statement ) {

		this._lastStatement = statement;

	}

	emitExtraLine( statement ) {

		const last = this._lastStatement;
		if ( last === null ) return '';

		if ( statement.isReturn ) return '\n';

		const isExpression = ( st ) => st.isFunctionDeclaration !== true && st.isFor !== true && st.isConditional !== true;
		const lastExp = isExpression( last );
		const currExp = isExpression( statement );

		if ( lastExp !== currExp || ( ! lastExp && ! currExp ) ) return '\n';

		return '';

	}

	emit( ast ) {

		let code = '\n';

		if ( this.iife ) this.tab += '\t';

		const overloadings = this.overloadings;

		for ( const statement of ast.body ) {

			if ( statement.isFunctionDeclaration ) {

				if ( overloadings.has( statement.name ) === false ) {

					overloadings.set( statement.name, [] );

				}

				overloadings.get( statement.name ).push( statement );

			}

		}

		for ( const statement of ast.body ) {

			code += this.emitExtraLine( statement );

			if ( statement.isFunctionDeclaration ) {

				code += this.tab + this.emitFunction( statement );

			} else {

				code += this.tab + this.emitExpression( statement ) + ';\n';

			}

			this.setLastStatement( statement );

		}

		const imports = [ ...this.imports ];
		const exports = [ ...this.global ];

		let header = '// Three.js Transpiler r' + REVISION + '\n\n';
		let footer = '';

		if ( this.iife ) {

			header += '( function ( TSL, uniforms ) {\n\n';

			header += imports.length > 0 ? '\tconst { ' + imports.join( ', ' ) + ' } = TSL;\n' : '';
			footer += exports.length > 0 ? '\treturn { ' + exports.join( ', ' ) + ' };\n' : '';

			footer += '\n} );';

		} else {

			header += imports.length > 0 ? 'import { ' + imports.join( ', ' ) + ' } from \'three/tsl\';\n' : '';

		}

		return header + code + footer;

	}

}

export default TSLEncoder;