SceneUtils: Add sortInstancedMesh(mesh, fn) (#24113)

* SceneUtils: Add sortInstancedMesh(mesh, fn)

* Update SceneUtils.js

Co-authored-by: Michael Herzog <[email protected]>

examples/jsm/utils/BufferGeometryUtils.js

@@ -366,6 +366,28 @@ function mergeBufferAttributes( attributes ) {
 
 }
 
+/**
+ * @param {BufferAttribute}
+ * @return {BufferAttribute}
+ */
+export function deepCloneAttribute( attribute ) {
+
+	if ( attribute.isInstancedInterleavedBufferAttribute || attribute.isInterleavedBufferAttribute ) {
+
+		return deinterleaveAttribute( attribute );
+
+	}
+
+	if ( attribute.isInstancedBufferAttribute ) {
+
+		return new InstancedBufferAttribute().copy( attribute );
+
+	}
+
+	return new BufferAttribute().copy( attribute );
+
+}
+
 /**
  * @param {Array<BufferAttribute>} attributes
  * @return {Array<InterleavedBufferAttribute>}

examples/jsm/utils/SceneUtils.js

@@ -1,11 +1,16 @@
 import {
-	Group,
-	Mesh,
 	BufferAttribute,
-	BufferGeometry
+	BufferGeometry,
+	Color,
+	Group,
+	Matrix4,
+	Mesh
 } from 'three';
 
-import { mergeGroups } from './BufferGeometryUtils.js';
+import { mergeGroups, deepCloneAttribute } from './BufferGeometryUtils.js';
+
+const _color = /*@__PURE__*/new Color();
+const _matrix = /*@__PURE__*/new Matrix4();
 
 function createMeshesFromInstancedMesh( instancedMesh ) {
 
@@ -118,8 +123,81 @@ function createMultiMaterialObject( geometry, materials ) {
 
 }
 
+/**
+ * @param {InstancedMesh}
+ * @param {function(int, int):int}
+ */
+function sortInstancedMesh( mesh, compareFn ) {
+
+	// store copy of instanced attributes for lookups
+
+	const instanceMatrixRef = deepCloneAttribute( mesh.instanceMatrix );
+	const instanceColorRef = mesh.instanceColor ? deepCloneAttribute( mesh.instanceColor ) : null;
+
+	const attributeRefs = new Map();
+
+	for ( const name in mesh.geometry.attributes ) {
+
+		const attribute = mesh.geometry.attributes[ name ];
+
+		if ( attribute.isInstancedBufferAttribute ) {
+
+			attributeRefs.set( attribute, deepCloneAttribute( attribute ) );
+
+		}
+
+	}
+
+
+	// compute sort order
+
+	const tokens = [];
+
+	for ( let i = 0; i < mesh.count; i ++ ) tokens.push( i );
+
+	tokens.sort( compareFn );
+
+
+	// apply sort order
+
+	for ( let i = 0; i < tokens.length; i ++ ) {
+
+		const refIndex = tokens[ i ];
+
+		_matrix.fromArray( instanceMatrixRef.array, refIndex * mesh.instanceMatrix.itemSize );
+		_matrix.toArray( mesh.instanceMatrix.array, i * mesh.instanceMatrix.itemSize );
+
+		if ( mesh.instanceColor ) {
+
+			_color.fromArray( instanceColorRef.array, refIndex * mesh.instanceColor.itemSize );
+			_color.toArray( mesh.instanceColor.array, i * mesh.instanceColor.itemSize );
+
+		}
+
+		for ( const name in mesh.geometry.attributes ) {
+
+			const attribute = mesh.geometry.attributes[ name ];
+
+			if ( attribute.isInstancedBufferAttribute ) {
+
+				const attributeRef = attributeRefs.get( attribute );
+
+				attribute.setX( i, attributeRef.getX( refIndex ) );
+				if ( attribute.itemSize > 1 ) attribute.setY( i, attributeRef.getY( refIndex ) );
+				if ( attribute.itemSize > 2 ) attribute.setZ( i, attributeRef.getZ( refIndex ) );
+				if ( attribute.itemSize > 3 ) attribute.setW( i, attributeRef.getW( refIndex ) );
+
+			}
+
+		}
+
+	}
+
+}
+
 export {
 	createMeshesFromInstancedMesh,
 	createMeshesFromMultiMaterialMesh,
-	createMultiMaterialObject
+	createMultiMaterialObject,
+	sortInstancedMesh
 };