TSL: Add bitcast functions and transpiler support (#31781)

* add bitcast to wgsl/tsl encoders

* alias GLSL functions

* glsl aliasing, change outputType based on inputType

* fix bracket spacing ah-gain

* add js-doc

---------

Co-authored-by: sunag <[email protected]>

Author: cmhhelgeson

examples/jsm/transpiler/WGSLEncoder.js

@@ -80,6 +80,10 @@ const wgslLib = {
 	'textureLod': 'textureSampleLevel',
 	'texelFetch': 'textureLoad',
 	'textureGrad': 'textureSampleGrad',
+	'floatBitsToInt': 'bitcast<i32>',
+	'floatBitsToUint': 'bitcast<u32>',
+	'intBitsToFloat': 'bitcast<f32>',
+	'uintBitsToFloat': 'bitcast<f32>',
 };
 
 class WGSLEncoder {
@@ -170,6 +174,26 @@ class WGSLEncoder {
 
 				code = `${ modFnName }( ${ snippets.join( ', ' ) } )`;
 
+			} else if ( fnName.startsWith( 'bitcast' ) ) {
+
+				const params = node.params.map( p => this.emitExpression( p ) ).join( ',' );
+				const types = node.params.map( p => p.getType() );
+
+				if ( /.*vec[234]/.test( types[ 0 ] ) ) {
+
+					const conversionType = fnName.substring( 8, fnName.length - 1 );
+					const vectorType = types[ 0 ].substring( - 1 );
+
+					code = `bitcast<${ vectorType }<${ conversionType }>>`;
+
+				} else {
+
+					code = fnName;
+
+				}
+
+				code += `( ${ params } )`;
+
 			} else if ( fnName.startsWith( 'texture' ) ) {
 
 				// Handle texture functions separately due to sampler handling

src/Three.TSL.js

@@ -179,6 +179,8 @@ export const faceDirection = TSL.faceDirection;
 export const faceForward = TSL.faceForward;
 export const faceforward = TSL.faceforward;
 export const float = TSL.float;
+export const floatBitsToInt = TSL.floatBitsToInt;
+export const floatBitsToUint = TSL.floatBitsToUint;
 export const floor = TSL.floor;
 export const fog = TSL.fog;
 export const fract = TSL.fract;
@@ -221,6 +223,7 @@ export const instancedBufferAttribute = TSL.instancedBufferAttribute;
 export const instancedDynamicBufferAttribute = TSL.instancedDynamicBufferAttribute;
 export const instancedMesh = TSL.instancedMesh;
 export const int = TSL.int;
+export const intBitsToFloat = TSL.intBitsToFloat;
 export const inverse = TSL.inverse;
 export const inverseSqrt = TSL.inverseSqrt;
 export const inversesqrt = TSL.inversesqrt;
@@ -562,6 +565,7 @@ export const triplanarTexture = TSL.triplanarTexture;
 export const triplanarTextures = TSL.triplanarTextures;
 export const trunc = TSL.trunc;
 export const uint = TSL.uint;
+export const uintBitsToFloat = TSL.uintBitsToFloat;
 export const uniform = TSL.uniform;
 export const uniformArray = TSL.uniformArray;
 export const uniformCubeTexture = TSL.uniformCubeTexture;

src/nodes/math/BitcastNode.js

@@ -19,8 +19,9 @@ class BitcastNode extends TempNode {
 	 *
 	 * @param {Node} valueNode - The value to convert.
 	 * @param {string} conversionType - The type to convert to.
+	 * @param {?string} [inputType = null] - The expected input data type of the bitcast operation.
 	 */
-	constructor( valueNode, conversionType ) {
+	constructor( valueNode, conversionType, inputType = null ) {
 
 		super();
 
@@ -35,10 +36,19 @@ class BitcastNode extends TempNode {
 		 * The type the value will be converted to.
 		 *
 		 * @type {string}
-		 * @default null
 		 */
 		this.conversionType = conversionType;
 
+
+		/**
+		 * The expected input data type of the bitcast operation.
+		 *
+		 *
+		 * @type {string}
+		 * @default null
+		 */
+		this.inputType = inputType;
+
 		/**
 		 * This flag can be used for type testing.
 		 *
@@ -50,12 +60,17 @@ class BitcastNode extends TempNode {
 
 	}
 
-	/**
-	 * The node's type is defined by the conversion type.
-	 *
-	 * @return {string} The node type.
-	 */
-	getNodeType() {
+	getNodeType( builder ) {
+
+		// GLSL aliasing
+		if ( this.inputType !== null ) {
+
+			const valueType = this.valueNode.getNodeType( builder );
+			const valueLength = builder.getTypeLength( valueType );
+
+			return builder.getTypeFromLength( valueLength, this.conversionType );
+
+		}
 
 		return this.conversionType;
 
@@ -65,9 +80,22 @@ class BitcastNode extends TempNode {
 	generate( builder ) {
 
 		const type = this.getNodeType( builder );
-		const inputType = this.valueNode.getNodeType( builder );
+		let inputType = '';
+
+		if ( this.inputType !== null ) {
+
+			const valueType = this.valueNode.getNodeType( builder );
+			const valueTypeLength = builder.getTypeLength( valueType );
 
-		return `${builder.getBitcastMethod( type, inputType )}( ${ this.valueNode.build( builder, inputType ) } )`;
+			inputType = valueTypeLength === 1 ? this.inputType : builder.changeComponentType( valueType, this.inputType );
+
+		} else {
+
+			inputType = this.valueNode.getNodeType( builder );
+
+		}
+
+		return `${ builder.getBitcastMethod( type, inputType ) }( ${ this.valueNode.build( builder, inputType ) } )`;
 
 
 	}
@@ -86,3 +114,43 @@ export default BitcastNode;
  * @returns {Node}
  */
 export const bitcast = /*@__PURE__*/ nodeProxyIntent( BitcastNode ).setParameterLength( 2 );
+
+/**
+ * Bitcasts a float or a vector of floats to a corresponding integer type with the same element size.
+ *
+ * @tsl
+ * @function
+ * @param {Node<float>} value - The float or vector of floats to bitcast.
+ * @returns {BitcastNode}
+ */
+export const floatBitsToInt = ( value ) => new BitcastNode( value, 'int', 'float' );
+
+/**
+ * Bitcasts a float or a vector of floats to a corresponding unsigned integer type with the same element size.
+ *
+ * @tsl
+ * @function
+ * @param {Node<float>} value - The float or vector of floats to bitcast.
+ * @returns {BitcastNode}
+ */
+export const floatBitsToUint = ( value ) => new BitcastNode( value, 'uint', 'float' );
+
+/**
+ * Bitcasts an integer or a vector of integers to a corresponding float type with the same element size.
+ *
+ * @tsl
+ * @function
+ * @param {Node<int>} value - The integer or vector of integers to bitcast.
+ * @returns {BitcastNode}
+ */
+export const intBitsToFloat = ( value ) => new BitcastNode( value, 'float', 'int' );
+
+/**
+ * Bitcast an unsigned integer or a vector of unsigned integers to a corresponding float type with the same element size.
+ *
+ * @tsl
+ * @function
+ * @param {Node<uint>} value - The unsigned integer or vector of unsigned integers to bitcast.
+ * @returns {BitcastNode}
+ */
+export const uintBitsToFloat = ( value ) => new BitcastNode( value, 'float', 'uint' );