BatchedMesh: add getColorAt and setColorAt (#28255)

* BatchedMesh: add getColorAt and setColorAt

* pr fixes and suggestions

* bring back previous condition to recompile shader

* add colorSpace to DataTexture

* change instance to geometry in docs

docs/api/en/objects/BatchedMesh.html

@@ -129,6 +129,19 @@ <h3>[method:this setCustomSort]( [param:Function sortFunction] )</h3>
 			in the list include a "z" field to perform a depth-ordered sort with.
 		</p>
 
+		<h3>
+			[method:undefined getColorAt]( [param:Integer index], [param:Color color] )
+		</h3>
+		<p>
+			[page:Integer index]: The index of a geometry. Values have to be in the
+			range [0, count].
+		</p>
+		<p>
+			[page:Color color]: This color object will be set to the color of the
+			defined geometry.
+		</p>
+		<p>Get the color of the defined geometry.</p>
+
 		<h3>
 			[method:Matrix4 getMatrixAt]( [param:Integer index], [param:Matrix4 matrix] )
 		</h3>
@@ -151,6 +164,18 @@ <h3>
 		</p>
 		<p>Get whether the given instance is marked as "visible" or not.</p>
 
+		<h3>
+			[method:undefined setColorAt]( [param:Integer index], [param:Color color] )
+		</h3>
+		<p>
+			[page:Integer index]: The index of a geometry. Values have to be in the
+			range [0, count].
+		</p>
+		<p>[page:Color color]: The color of a single geometry.</p>
+		<p>
+			Sets the given color to the defined geometry.
+		</p>
+
 		<h3>
 			[method:this setMatrixAt]( [param:Integer index], [param:Matrix4 matrix] )
 		</h3>

examples/jsm/renderers/common/RenderObject.js

@@ -225,6 +225,7 @@ export default class RenderObject {
 		if ( object.isBatchedMesh ) {
 
 			cacheKey += object._matricesTexture.uuid + ',';
+			cacheKey += object._colorsTexture.uuid + ',';
 
 		}
 

src/core/Object3D.js

@@ -761,6 +761,8 @@ class Object3D extends EventDispatcher {
 
 			object.matricesTexture = this._matricesTexture.toJSON( meta );
 
+			if ( this._colorsTexture !== null ) object.colorsTexture = this._colorsTexture.toJSON( meta );
+
 			if ( this.boundingSphere !== null ) {
 
 				object.boundingSphere = {

src/loaders/ObjectLoader.js

@@ -946,6 +946,7 @@ class ObjectLoader extends Loader {
 				object._geometryCount = data.geometryCount;
 
 				object._matricesTexture = getTexture( data.matricesTexture.uuid );
+				if ( data.colorsTexture !== undefined ) object._colorsTexture = getTexture( data.colorsTexture.uuid );
 
 				break;
 

src/objects/BatchedMesh.js

@@ -5,6 +5,7 @@ import { FloatType } from '../constants.js';
 import { Matrix4 } from '../math/Matrix4.js';
 import { Mesh } from './Mesh.js';
 import { RGBAFormat } from '../constants.js';
+import { ColorManagement } from '../math/ColorManagement.js';
 import { Box3 } from '../math/Box3.js';
 import { Sphere } from '../math/Sphere.js';
 import { Frustum } from '../math/Frustum.js';
@@ -160,6 +161,9 @@ class BatchedMesh extends Mesh {
 
 		this._initMatricesTexture();
 
+		// Local color per geometry by using data texture
+		this._colorsTexture = null;
+
 	}
 
 	_initMatricesTexture() {
@@ -182,6 +186,19 @@ class BatchedMesh extends Mesh {
 
 	}
 
+	_initColorsTexture() {
+
+		let size = Math.sqrt( this._maxGeometryCount );
+		size = Math.ceil( size );
+
+		const colorsArray = new Float32Array( size * size * 4 ); // 4 floats per RGBA pixel
+		const colorsTexture = new DataTexture( colorsArray, size, size, RGBAFormat, FloatType );
+		colorsTexture.colorSpace = ColorManagement.workingColorSpace;
+
+		this._colorsTexture = colorsTexture;
+
+	}
+
 	_initializeGeometry( reference ) {
 
 		const geometry = this.geometry;
@@ -746,6 +763,49 @@ class BatchedMesh extends Mesh {
 
 	}
 
+	setColorAt( geometryId, color ) {
+
+		if ( this._colorsTexture === null ) {
+
+			this._initColorsTexture();
+
+		}
+
+		// @TODO: Map geometryId to index of the arrays because
+		//        optimize() can make geometryId mismatch the index
+
+		const active = this._active;
+		const colorsTexture = this._colorsTexture;
+		const colorsArray = this._colorsTexture.image.data;
+		const geometryCount = this._geometryCount;
+		if ( geometryId >= geometryCount || active[ geometryId ] === false ) {
+
+			return this;
+
+		}
+
+		color.toArray( colorsArray, geometryId * 4 );
+		colorsTexture.needsUpdate = true;
+
+		return this;
+
+	}
+
+	getColorAt( geometryId, color ) {
+
+		const active = this._active;
+		const colorsArray = this._colorsTexture.image.data;
+		const geometryCount = this._geometryCount;
+		if ( geometryId >= geometryCount || active[ geometryId ] === false ) {
+
+			return null;
+
+		}
+
+		return color.fromArray( colorsArray, geometryId * 4 );
+
+	}
+
 	setVisibleAt( geometryId, value ) {
 
 		const visibility = this._visibility;
@@ -886,6 +946,13 @@ class BatchedMesh extends Mesh {
 		this._matricesTexture = source._matricesTexture.clone();
 		this._matricesTexture.image.data = this._matricesTexture.image.slice();
 
+		if ( this._colorsTexture !== null ) {
+
+			this._colorsTexture = source._colorsTexture.clone();
+			this._colorsTexture.image.data = this._colorsTexture.image.slice();
+
+		}
+
 		return this;
 
 	}
@@ -897,6 +964,14 @@ class BatchedMesh extends Mesh {
 
 		this._matricesTexture.dispose();
 		this._matricesTexture = null;
+
+		if ( this._colorsTexture !== null ) {
+
+			this._colorsTexture.dispose();
+			this._colorsTexture = null;
+
+		}
+
 		return this;
 
 	}

src/renderers/WebGLRenderer.js

@@ -1716,6 +1716,7 @@ class WebGLRenderer {
 
 			materialProperties.outputColorSpace = parameters.outputColorSpace;
 			materialProperties.batching = parameters.batching;
+			materialProperties.batchingColor = parameters.batchingColor;
 			materialProperties.instancing = parameters.instancing;
 			materialProperties.instancingColor = parameters.instancingColor;
 			materialProperties.instancingMorph = parameters.instancingMorph;
@@ -1805,6 +1806,14 @@ class WebGLRenderer {
 
 					needsProgramChange = true;
 
+				} else if ( object.isBatchedMesh && materialProperties.batchingColor === true && object.colorTexture === null ) {
+
+					needsProgramChange = true;
+
+				} else if ( object.isBatchedMesh && materialProperties.batchingColor === false && object.colorTexture !== null ) {
+
+					needsProgramChange = true;
+
 				} else if ( object.isInstancedMesh && materialProperties.instancing === false ) {
 
 					needsProgramChange = true;
@@ -1997,6 +2006,13 @@ class WebGLRenderer {
 				p_uniforms.setOptional( _gl, object, 'batchingTexture' );
 				p_uniforms.setValue( _gl, 'batchingTexture', object._matricesTexture, textures );
 
+				p_uniforms.setOptional( _gl, object, 'batchingColorTexture' );
+				if ( object._colorsTexture !== null ) {
+
+					p_uniforms.setValue( _gl, 'batchingColorTexture', object._colorsTexture, textures );
+
+				}
+
 			}
 
 			const morphAttributes = geometry.morphAttributes;

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

@@ -16,4 +16,19 @@ export default /* glsl */`
 
 	}
 #endif
+
+#ifdef USE_BATCHING_COLOR
+
+	uniform sampler2D batchingColorTexture;
+	vec3 getBatchingColor( const in float i ) {
+
+		int size = textureSize( batchingColorTexture, 0 ).x;
+		int j = int( i );
+		int x = j % size;
+		int y = j / size;
+		return texelFetch( batchingColorTexture, ivec2( x, y ), 0 ).rgb;
+
+	}
+
+#endif
 `;

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

@@ -3,7 +3,7 @@ export default /* glsl */`
 
 	varying vec4 vColor;
 
-#elif defined( USE_COLOR ) || defined( USE_INSTANCING_COLOR )
+#elif defined( USE_COLOR ) || defined( USE_INSTANCING_COLOR ) || defined( USE_BATCHING_COLOR )
 
 	varying vec3 vColor;
 

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

@@ -3,7 +3,7 @@ export default /* glsl */`
 
 	vColor = vec4( 1.0 );
 
-#elif defined( USE_COLOR ) || defined( USE_INSTANCING_COLOR )
+#elif defined( USE_COLOR ) || defined( USE_INSTANCING_COLOR ) || defined( USE_BATCHING_COLOR )
 
 	vColor = vec3( 1.0 );
 
@@ -19,5 +19,13 @@ export default /* glsl */`
 
 	vColor.xyz *= instanceColor.xyz;
 
+#endif
+
+#ifdef USE_BATCHING_COLOR
+
+	vec3 batchingColor = getBatchingColor( batchId );
+
+	vColor.xyz *= batchingColor.xyz;
+
 #endif
 `;

src/renderers/webgl/WebGLProgram.js

@@ -522,6 +522,7 @@ function WebGLProgram( renderer, cacheKey, parameters, bindingStates ) {
 
 			parameters.extensionClipCullDistance ? '#define USE_CLIP_DISTANCE' : '',
 			parameters.batching ? '#define USE_BATCHING' : '',
+			parameters.batchingColor ? '#define USE_BATCHING_COLOR' : '',
 			parameters.instancing ? '#define USE_INSTANCING' : '',
 			parameters.instancingColor ? '#define USE_INSTANCING_COLOR' : '',
 			parameters.instancingMorph ? '#define USE_INSTANCING_MORPH' : '',
@@ -773,7 +774,7 @@ function WebGLProgram( renderer, cacheKey, parameters, bindingStates ) {
 			parameters.thicknessMap ? '#define USE_THICKNESSMAP' : '',
 
 			parameters.vertexTangents && parameters.flatShading === false ? '#define USE_TANGENT' : '',
-			parameters.vertexColors || parameters.instancingColor ? '#define USE_COLOR' : '',
+			parameters.vertexColors || parameters.instancingColor || parameters.batchingColor ? '#define USE_COLOR' : '',
 			parameters.vertexAlphas ? '#define USE_COLOR_ALPHA' : '',
 			parameters.vertexUv1s ? '#define USE_UV1' : '',
 			parameters.vertexUv2s ? '#define USE_UV2' : '',

src/renderers/webgl/WebGLPrograms.js

@@ -192,6 +192,7 @@ function WebGLPrograms( renderer, cubemaps, cubeuvmaps, extensions, capabilities
 			precision: precision,
 
 			batching: IS_BATCHEDMESH,
+			batchingColor: IS_BATCHEDMESH && object._colorsTexture !== null,
 			instancing: IS_INSTANCEDMESH,
 			instancingColor: IS_INSTANCEDMESH && object.instanceColor !== null,
 			instancingMorph: IS_INSTANCEDMESH && object.morphTexture !== null,
@@ -510,6 +511,8 @@ function WebGLPrograms( renderer, cubemaps, cubeuvmaps, extensions, capabilities
 			_programLayers.enable( 19 );
 		if ( parameters.dispersion )
 			_programLayers.enable( 20 );
+		if ( parameters.batchingColor )
+			_programLayers.enable( 21 );
 
 		array.push( _programLayers.mask );
 		_programLayers.disableAll();