BatchedMesh: Add built-in material support (#27115)

* Add support for multi draw

* Fix multidraw

* Add batching material properties

* Adjust batching define

* Move normal shaders

* Move vertex transform shadres

* Move batching pars

* Move batching matrix fetch

* Switch _batch_id_ to batchId

* Remove custom shader callback

* Remove unused custom uniforms

examples/jsm/objects/BatchedMesh.js

@@ -9,7 +9,7 @@ import {
 	RGBAFormat
 } from 'three';
 
-const ID_ATTR_NAME = '_batch_id_';
+const ID_ATTR_NAME = 'batchId';
 const _identityMatrix = new Matrix4();
 const _zeroScaleMatrix = new Matrix4().set(
 	0, 0, 0, 0,
@@ -18,48 +18,6 @@ const _zeroScaleMatrix = new Matrix4().set(
 	0, 0, 0, 1,
 );
 
-// Custom shaders
-const batchingParsVertex = /* glsl */`
-#ifdef BATCHING
-	attribute float ${ ID_ATTR_NAME };
-	uniform highp sampler2D batchingTexture;
-	mat4 getBatchingMatrix( const in float i ) {
-
-		int size = textureSize( batchingTexture, 0 ).x;
-		int j = int( i ) * 4;
-		int x = j % size;
-		int y = j / size;
-		vec4 v1 = texelFetch( batchingTexture, ivec2( x, y ), 0 );
-		vec4 v2 = texelFetch( batchingTexture, ivec2( x + 1, y ), 0 );
-		vec4 v3 = texelFetch( batchingTexture, ivec2( x + 2, y ), 0 );
-		vec4 v4 = texelFetch( batchingTexture, ivec2( x + 3, y ), 0 );
-		return mat4( v1, v2, v3, v4 );
-
-	}
-#endif
-`;
-
-const batchingbaseVertex = /* glsl */`
-#ifdef BATCHING
-	mat4 batchingMatrix = getBatchingMatrix( ${ ID_ATTR_NAME } );
-#endif
-`;
-
-const batchingnormalVertex = /* glsl */`
-#ifdef BATCHING
-	objectNormal = vec4( batchingMatrix * vec4( objectNormal, 0.0 ) ).xyz;
-	#ifdef USE_TANGENT
-		objectTangent = vec4( batchingMatrix * vec4( objectTangent, 0.0 ) ).xyz;
-	#endif
-#endif
-`;
-
-const batchingVertex = /* glsl */`
-#ifdef BATCHING
-	transformed = ( batchingMatrix * vec4( transformed, 1.0 ) ).xyz;
-#endif
-`;
-
 // @TODO: SkinnedMesh support?
 // @TODO: Future work if needed. Move into the core. Can be optimized more with WEBGL_multi_draw.
 
@@ -126,12 +84,7 @@ class BatchedMesh extends Mesh {
 		// @TODO: Calculate the entire binding box and make frustumCulled true
 		this.frustumCulled = false;
 
-		this._customUniforms = {
-			batchingTexture: { value: null }
-		};
-
 		this._initMatricesTexture();
-		this._initShader();
 
 	}
 
@@ -152,53 +105,6 @@ class BatchedMesh extends Mesh {
 		const matricesTexture = new DataTexture( matricesArray, size, size, RGBAFormat, FloatType );
 
 		this._matricesTexture = matricesTexture;
-		this._customUniforms.batchingTexture.value = this._matricesTexture;
-
-	}
-
-	_initShader() {
-
-		const material = this.material;
-		const currentOnBeforeCompile = material.onBeforeCompile;
-		const customUniforms = this._customUniforms;
-
-		material.onBeforeCompile = function onBeforeCompile( parameters, renderer ) {
-
-			// Is this replacement stable across any materials?
-			parameters.vertexShader = parameters.vertexShader
-				.replace(
-					'#include <skinning_pars_vertex>',
-					'#include <skinning_pars_vertex>\n'
-						+ batchingParsVertex
-				)
-				.replace(
-					'#include <uv_vertex>',
-					'#include <uv_vertex>\n'
-						+ batchingbaseVertex
-				)
-				.replace(
-					'#include <skinnormal_vertex>',
-					'#include <skinnormal_vertex>\n'
-						+ batchingnormalVertex
-				)
-				.replace(
-					'#include <skinning_vertex>',
-					'#include <skinning_vertex>\n'
-						+ batchingVertex
-				);
-
-			for ( const uniformName in customUniforms ) {
-
-				parameters.uniforms[ uniformName ] = customUniforms[ uniformName ];
-
-			}
-
-			currentOnBeforeCompile.call( this, parameters, renderer );
-
-		};
-
-		material.defines = material.defines || {};
-		material.defines.BATCHING = false;
 
 	}
 
@@ -706,13 +612,11 @@ class BatchedMesh extends Mesh {
 
 	}
 
-	onBeforeRender( _renderer, _scene, _camera, _geometry, material/*, _group*/ ) {
-
-		material.defines.BATCHING = true;
+	onBeforeRender( _renderer, _scene, _camera, geometry ) {
 
 		// the indexed version of the multi draw function requires specifying the start
 		// offset in bytes.
-		const index = _geometry.getIndex();
+		const index = geometry.getIndex();
 		const bytesPerElement = index === null ? 1 : index.array.BYTES_PER_ELEMENT;
 
 		const visible = this._visible;
@@ -742,12 +646,6 @@ class BatchedMesh extends Mesh {
 
 	}
 
-	onAfterRender( _renderer, _scene, _camera, _geometry, material/*, _group*/ ) {
-
-		material.defines.BATCHING = false;
-
-	}
-
 }
 
 export { BatchedMesh };

src/renderers/WebGLRenderer.js

@@ -1719,6 +1719,7 @@ class WebGLRenderer {
 			const materialProperties = properties.get( material );
 
 			materialProperties.outputColorSpace = parameters.outputColorSpace;
+			materialProperties.batching = parameters.batching;
 			materialProperties.instancing = parameters.instancing;
 			materialProperties.instancingColor = parameters.instancingColor;
 			materialProperties.skinning = parameters.skinning;
@@ -1799,6 +1800,14 @@ class WebGLRenderer {
 
 					needsProgramChange = true;
 
+				} else if ( object.isBatchedMesh && materialProperties.batching === false ) {
+
+					needsProgramChange = true;
+
+				} else if ( ! object.isBatchedMesh && materialProperties.batching === true ) {
+
+					needsProgramChange = true;
+
 				} else if ( object.isInstancedMesh && materialProperties.instancing === false ) {
 
 					needsProgramChange = true;
@@ -1986,6 +1995,13 @@ class WebGLRenderer {
 
 			}
 
+			if ( object.isBatchedMesh ) {
+
+				p_uniforms.setOptional( _gl, object, 'batchingTexture' );
+				p_uniforms.setValue( _gl, 'batchingTexture', object._matricesTexture, textures );
+
+			}
+
 			const morphAttributes = geometry.morphAttributes;
 
 			if ( morphAttributes.position !== undefined || morphAttributes.normal !== undefined || ( morphAttributes.color !== undefined && capabilities.isWebGL2 === true ) ) {

src/renderers/shaders/ShaderChunk.js

@@ -6,6 +6,7 @@ import alphatest_fragment from './ShaderChunk/alphatest_fragment.glsl.js';
 import alphatest_pars_fragment from './ShaderChunk/alphatest_pars_fragment.glsl.js';
 import aomap_fragment from './ShaderChunk/aomap_fragment.glsl.js';
 import aomap_pars_fragment from './ShaderChunk/aomap_pars_fragment.glsl.js';
+import batching_pars_vertex from './ShaderChunk/batching_pars_vertex.glsl.js';
 import begin_vertex from './ShaderChunk/begin_vertex.glsl.js';
 import beginnormal_vertex from './ShaderChunk/beginnormal_vertex.glsl.js';
 import bsdfs from './ShaderChunk/bsdfs.glsl.js';
@@ -131,6 +132,7 @@ export const ShaderChunk = {
 	alphatest_pars_fragment: alphatest_pars_fragment,
 	aomap_fragment: aomap_fragment,
 	aomap_pars_fragment: aomap_pars_fragment,
+	batching_pars_vertex: batching_pars_vertex,
 	begin_vertex: begin_vertex,
 	beginnormal_vertex: beginnormal_vertex,
 	bsdfs: bsdfs,

src/renderers/shaders/ShaderChunk/batching_pars_vertex.glsl.js

@@ -0,0 +1,19 @@
+export default /* glsl */`
+#ifdef USE_BATCHING
+	attribute float batchId;
+	uniform highp sampler2D batchingTexture;
+	mat4 getBatchingMatrix( const in float i ) {
+
+		int size = textureSize( batchingTexture, 0 ).x;
+		int j = int( i ) * 4;
+		int x = j % size;
+		int y = j / size;
+		vec4 v1 = texelFetch( batchingTexture, ivec2( x, y ), 0 );
+		vec4 v2 = texelFetch( batchingTexture, ivec2( x + 1, y ), 0 );
+		vec4 v3 = texelFetch( batchingTexture, ivec2( x + 2, y ), 0 );
+		vec4 v4 = texelFetch( batchingTexture, ivec2( x + 3, y ), 0 );
+		return mat4( v1, v2, v3, v4 );
+
+	}
+#endif
+`;

src/renderers/shaders/ShaderChunk/defaultnormal_vertex.glsl.js

@@ -1,4 +1,24 @@
 export default /* glsl */`
+
+#ifdef USE_BATCHING
+
+	// this is in lieu of a per-instance normal-matrix
+	// shear transforms in the instance matrix are not supported
+
+	mat4 batchingMatrix = getBatchingMatrix( batchId );
+	mat3 bm = mat3( batchingMatrix );
+	objectNormal /= vec3( dot( bm[ 0 ], bm[ 0 ] ), dot( bm[ 1 ], bm[ 1 ] ), dot( bm[ 2 ], bm[ 2 ] ) );
+	objectNormal = bm * objectNormal;
+
+	#ifdef USE_TANGENT
+
+		objectTangent /= vec3( dot( bm[ 0 ], bm[ 0 ] ), dot( bm[ 1 ], bm[ 1 ] ), dot( bm[ 2 ], bm[ 2 ] ) );
+		objectTangent = bm * objectTangent;
+
+	#endif
+
+#endif
+
 vec3 transformedNormal = objectNormal;
 
 #ifdef USE_INSTANCING

src/renderers/shaders/ShaderChunk/project_vertex.glsl.js

@@ -1,6 +1,12 @@
 export default /* glsl */`
 vec4 mvPosition = vec4( transformed, 1.0 );
 
+#ifdef USE_BATCHING
+
+	mvPosition = batchingMatrix * mvPosition;
+
+#endif
+
 #ifdef USE_INSTANCING
 
 	mvPosition = instanceMatrix * mvPosition;

src/renderers/shaders/ShaderChunk/worldpos_vertex.glsl.js

@@ -3,6 +3,12 @@ export default /* glsl */`
 
 	vec4 worldPosition = vec4( transformed, 1.0 );
 
+	#ifdef USE_BATCHING
+
+		worldPosition = batchingMatrix * worldPosition;
+
+	#endif
+
 	#ifdef USE_INSTANCING
 
 		worldPosition = instanceMatrix * worldPosition;

src/renderers/shaders/ShaderLib/depth.glsl.js

@@ -1,5 +1,6 @@
 export const vertex = /* glsl */`
 #include <common>
+#include <batching_pars_vertex>
 #include <uv_pars_vertex>
 #include <displacementmap_pars_vertex>
 #include <morphtarget_pars_vertex>

src/renderers/shaders/ShaderLib/distanceRGBA.glsl.js

@@ -4,6 +4,7 @@ export const vertex = /* glsl */`
 varying vec3 vWorldPosition;
 
 #include <common>
+#include <batching_pars_vertex>
 #include <uv_pars_vertex>
 #include <displacementmap_pars_vertex>
 #include <morphtarget_pars_vertex>

src/renderers/shaders/ShaderLib/meshbasic.glsl.js

@@ -1,5 +1,6 @@
 export const vertex = /* glsl */`
 #include <common>
+#include <batching_pars_vertex>
 #include <uv_pars_vertex>
 #include <envmap_pars_vertex>
 #include <color_pars_vertex>

src/renderers/shaders/ShaderLib/meshlambert.glsl.js

@@ -4,6 +4,7 @@ export const vertex = /* glsl */`
 varying vec3 vViewPosition;
 
 #include <common>
+#include <batching_pars_vertex>
 #include <uv_pars_vertex>
 #include <displacementmap_pars_vertex>
 #include <envmap_pars_vertex>

src/renderers/shaders/ShaderLib/meshmatcap.glsl.js

@@ -4,6 +4,7 @@ export const vertex = /* glsl */`
 varying vec3 vViewPosition;
 
 #include <common>
+#include <batching_pars_vertex>
 #include <uv_pars_vertex>
 #include <color_pars_vertex>
 #include <displacementmap_pars_vertex>

src/renderers/shaders/ShaderLib/meshnormal.glsl.js

@@ -8,6 +8,7 @@ export const vertex = /* glsl */`
 #endif
 
 #include <common>
+#include <batching_pars_vertex>
 #include <uv_pars_vertex>
 #include <displacementmap_pars_vertex>
 #include <normal_pars_vertex>

src/renderers/shaders/ShaderLib/meshphong.glsl.js

@@ -4,6 +4,7 @@ export const vertex = /* glsl */`
 varying vec3 vViewPosition;
 
 #include <common>
+#include <batching_pars_vertex>
 #include <uv_pars_vertex>
 #include <displacementmap_pars_vertex>
 #include <envmap_pars_vertex>

src/renderers/shaders/ShaderLib/meshphysical.glsl.js

@@ -10,6 +10,7 @@ varying vec3 vViewPosition;
 #endif
 
 #include <common>
+#include <batching_pars_vertex>
 #include <uv_pars_vertex>
 #include <displacementmap_pars_vertex>
 #include <color_pars_vertex>

src/renderers/shaders/ShaderLib/meshtoon.glsl.js

@@ -4,6 +4,7 @@ export const vertex = /* glsl */`
 varying vec3 vViewPosition;
 
 #include <common>
+#include <batching_pars_vertex>
 #include <uv_pars_vertex>
 #include <displacementmap_pars_vertex>
 #include <color_pars_vertex>

src/renderers/shaders/ShaderLib/shadow.glsl.js

@@ -1,5 +1,6 @@
 export const vertex = /* glsl */`
 #include <common>
+#include <batching_pars_vertex>
 #include <fog_pars_vertex>
 #include <morphtarget_pars_vertex>
 #include <skinning_pars_vertex>

src/renderers/webgl/WebGLProgram.js

@@ -503,6 +503,7 @@ function WebGLProgram( renderer, cacheKey, parameters, bindingStates ) {
 
 			customDefines,
 
+			parameters.batching ? '#define USE_BATCHING' : '',
 			parameters.instancing ? '#define USE_INSTANCING' : '',
 			parameters.instancingColor ? '#define USE_INSTANCING_COLOR' : '',
 

src/renderers/webgl/WebGLPrograms.js

@@ -108,6 +108,7 @@ function WebGLPrograms( renderer, cubemaps, cubeuvmaps, extensions, capabilities
 		const currentRenderTarget = renderer.getRenderTarget();
 
 		const IS_INSTANCEDMESH = object.isInstancedMesh === true;
+		const IS_BATCHEDMESH = object.isBatchedMesh === true;
 
 		const HAS_MAP = !! material.map;
 		const HAS_MATCAP = !! material.matcap;
@@ -193,6 +194,7 @@ function WebGLPrograms( renderer, cubemaps, cubeuvmaps, extensions, capabilities
 
 			precision: precision,
 
+			batching: IS_BATCHEDMESH,
 			instancing: IS_INSTANCEDMESH,
 			instancingColor: IS_INSTANCEDMESH && object.instanceColor !== null,
 
@@ -503,6 +505,8 @@ function WebGLPrograms( renderer, cubemaps, cubeuvmaps, extensions, capabilities
 			_programLayers.enable( 17 );
 		if ( parameters.alphaHash )
 			_programLayers.enable( 18 );
+		if ( parameters.batching )
+			_programLayers.enable( 19 );
 
 		array.push( _programLayers.mask );
 		_programLayers.disableAll();