WGSLNodeBuilder: Added support for .wrap* using biquadraticTexture (#29828)

* added support for `wrap*` using `biquadraticTexture`

* using `getDataFromNode` for now

* cleanup

* cleanup

Author: sunag

src/renderers/webgpu/nodes/WGSLNodeBuilder.js

@@ -13,11 +13,17 @@ import { getFormat } from '../utils/WebGPUTextureUtils.js';
 import WGSLNodeParser from './WGSLNodeParser.js';
 import { GPUBufferBindingType, GPUStorageTextureAccess } from '../utils/WebGPUConstants.js';
 
-import { NoColorSpace, FloatType } from '../../../constants.js';
+import { NoColorSpace, FloatType, RepeatWrapping, ClampToEdgeWrapping, MirroredRepeatWrapping } from '../../../constants.js';
 
 // GPUShaderStage is not defined in browsers not supporting WebGPU
 const GPUShaderStage = self.GPUShaderStage;
 
+const wrapNames = {
+	[ RepeatWrapping ]: 'repeat',
+	[ ClampToEdgeWrapping ]: 'clamp',
+	[ MirroredRepeatWrapping ]: 'mirror'
+};
+
 const gpuShaderStageLib = {
 	'vertex': GPUShaderStage ? GPUShaderStage.VERTEX : 1,
 	'fragment': GPUShaderStage ? GPUShaderStage.FRAGMENT : 2,
@@ -61,6 +67,8 @@ const wgslTypeLib = {
 	mat4: 'mat4x4<f32>'
 };
 
+const wgslCodeCache = {};
+
 const wgslPolyfill = {
 	tsl_xor: new CodeNode( 'fn tsl_xor( a : bool, b : bool ) -> bool { return ( a || b ) && !( a && b ); }' ),
 	mod_float: new CodeNode( 'fn tsl_mod_float( x : f32, y : f32 ) -> f32 { return x - y * floor( x / y ); }' ),
@@ -71,6 +79,9 @@ const wgslPolyfill = {
 	equals_bvec2: new CodeNode( 'fn tsl_equals_bvec2( a : vec2f, b : vec2f ) -> vec2<bool> { return vec2<bool>( a.x == b.x, a.y == b.y ); }' ),
 	equals_bvec3: new CodeNode( 'fn tsl_equals_bvec3( a : vec3f, b : vec3f ) -> vec3<bool> { return vec3<bool>( a.x == b.x, a.y == b.y, a.z == b.z ); }' ),
 	equals_bvec4: new CodeNode( 'fn tsl_equals_bvec4( a : vec4f, b : vec4f ) -> vec4<bool> { return vec4<bool>( a.x == b.x, a.y == b.y, a.z == b.z, a.w == b.w ); }' ),
+	repeatWrapping_float: new CodeNode( 'fn tsl_repeatWrapping_float( coord: f32 ) -> f32 { return fract( coord ); }' ),
+	mirrorWrapping_float: new CodeNode( 'fn tsl_mirrorWrapping_float( coord: f32 ) -> f32 { let mirrored = fract( coord * 0.5 ) * 2.0; return 1.0 - abs( 1.0 - mirrored ); }' ),
+	clampWrapping_float: new CodeNode( 'fn tsl_clampWrapping_float( coord: f32 ) -> f32 { return clamp( coord, 0.0, 1.0 ); }' ),
 	repeatWrapping: new CodeNode( /* wgsl */`
 fn tsl_repeatWrapping( uv : vec2<f32>, dimension : vec2<u32> ) -> vec2<u32> {
 
@@ -81,9 +92,8 @@ fn tsl_repeatWrapping( uv : vec2<f32>, dimension : vec2<u32> ) -> vec2<u32> {
 }
 ` ),
 	biquadraticTexture: new CodeNode( /* wgsl */`
-fn tsl_biquadraticTexture( map : texture_2d<f32>, coord : vec2f, level : i32 ) -> vec4f {
+fn tsl_biquadraticTexture( map : texture_2d<f32>, coord : vec2f, iRes : vec2u, level : i32 ) -> vec4f {
 
-	let iRes = vec2i( textureDimensions( map, level ) );
 	let res = vec2f( iRes );
 
 	let uvScaled = coord * res;
@@ -95,10 +105,10 @@ fn tsl_biquadraticTexture( map : texture_2d<f32>, coord : vec2f, level : i32 ) -
 	let iuv = floor( uv );
 	let f = fract( uv );
 
-	let rg1 = textureLoad( map, vec2i( iuv + vec2( 0.5, 0.5 ) ) % iRes, level );
-	let rg2 = textureLoad( map, vec2i( iuv + vec2( 1.5, 0.5 ) ) % iRes, level );
-	let rg3 = textureLoad( map, vec2i( iuv + vec2( 0.5, 1.5 ) ) % iRes, level );
-	let rg4 = textureLoad( map, vec2i( iuv + vec2( 1.5, 1.5 ) ) % iRes, level );
+	let rg1 = textureLoad( map, vec2u( iuv + vec2( 0.5, 0.5 ) ) % iRes, level );
+	let rg2 = textureLoad( map, vec2u( iuv + vec2( 1.5, 0.5 ) ) % iRes, level );
+	let rg3 = textureLoad( map, vec2u( iuv + vec2( 0.5, 1.5 ) ) % iRes, level );
+	let rg4 = textureLoad( map, vec2u( iuv + vec2( 1.5, 1.5 ) ) % iRes, level );
 
 	return mix( mix( rg1, rg2, f.x ), mix( rg3, rg4, f.x ), f.y );
 
@@ -229,11 +239,94 @@ class WGSLNodeBuilder extends NodeBuilder {
 
 	}
 
+	generateWrapFunction( texture ) {
+
+		const functionName = `tsl_coord_${ wrapNames[ texture.wrapS ] }S_${ wrapNames[ texture.wrapT ] }T`;
+
+		let nodeCode = wgslCodeCache[ functionName ];
+
+		if ( nodeCode === undefined ) {
+
+			let code = `fn ${ functionName }( coord : vec2f ) -> vec2f {\n\n\treturn vec2f(\n`;
+
+			const addWrapSnippet = ( wrap, axis ) => {
+
+				if ( wrap === RepeatWrapping ) {
+
+					this._include( 'repeatWrapping_float' );
+
+					code += `\t\ttsl_repeatWrapping_float( coord.${ axis } )`;
+
+				} else if ( wrap === ClampToEdgeWrapping ) {
+
+					this._include( 'clampWrapping_float' );
+
+					code += `\t\ttsl_clampWrapping_float( coord.${ axis } )`;
+
+				} else if ( wrap === MirroredRepeatWrapping ) {
+
+					this._include( 'mirrorWrapping_float' );
+
+					code += `\t\ttsl_mirrorWrapping_float( coord.${ axis } )`;
+
+				} else {
+
+					code += `\t\tcoord.${ axis }`;
+
+					console.warn( `WebGPURenderer: Unsupported texture wrap type "${ wrap }" for vertex shader.` );
+
+				}
+
+			};
+
+			addWrapSnippet( texture.wrapS, 'x' );
+
+			code += ',\n';
+
+			addWrapSnippet( texture.wrapT, 'y' );
+
+			code += '\n\t);\n\n}\n';
+
+			wgslCodeCache[ functionName ] = nodeCode = new CodeNode( code );
+
+		}
+
+		nodeCode.build( this );
+
+		return functionName;
+
+	}
+
+	generateTextureDimension( texture, textureProperty, levelSnippet ) {
+
+		const textureData = this.getDataFromNode( texture, this.shaderStage, this.globalCache );
+
+		if ( textureData.dimensionsSnippet === undefined ) textureData.dimensionsSnippet = {};
+
+		let propertyName = textureData.dimensionsSnippet[ levelSnippet ];
+
+		if ( textureData.dimensionsSnippet[ levelSnippet ] === undefined ) {
+
+			propertyName = `textureDimension_${ texture.id }_${ levelSnippet }`;
+
+			this.addLineFlowCode( `let ${ propertyName } = textureDimensions( ${ textureProperty }, i32( ${ levelSnippet } ) );` );
+
+			textureData.dimensionsSnippet[ levelSnippet ] = propertyName;
+
+		}
+
+		return propertyName;
+
+	}
+
 	generateFilteredTexture( texture, textureProperty, uvSnippet, levelSnippet = '0' ) {
 
 		this._include( 'biquadraticTexture' );
 
-		return `tsl_biquadraticTexture( ${ textureProperty }, ${ uvSnippet }, i32( ${ levelSnippet } ) )`;
+		const wrapFunction = this.generateWrapFunction( texture );
+		const textureDimension = this.generateTextureDimension( texture, textureProperty, levelSnippet );
+
+		return `tsl_biquadraticTexture( ${ textureProperty }, ${ wrapFunction }( ${ uvSnippet } ), ${ textureDimension }, i32( ${ levelSnippet } ) )`;
 
 	}