BatchedMesh: Add setCustomSort function, hybrid radix sort implementation to examples (#27213)

* Add radix sort to batched mesh

* Fix options generation

* Add custom sort function, separate radix sort function to examples

* Adjust custom sort application

* revert count and position in demo

* Update demo

* Add docs

* Comments

docs/api/en/objects/BatchedMesh.html

@@ -98,6 +98,12 @@
 			method whenever this instance is no longer used in your app.
 		</p>
 
+		<h3>[method:this setCustomSort]( [param:Function sortFunction] )</h3>
+		<p>
+			Takes a sort a function that is run before render. The function takes a list of items to sort and a camera. The objects
+			in the list include a "z" field to perform a depth-ordered sort with.
+		</p>
+
 		<h3>
 			[method:Matrix4 getMatrixAt]( [param:Integer index], [param:Matrix4 matrix] )
 		</h3>

examples/jsm/utils/SortUtils.js

@@ -0,0 +1,162 @@
+// Hybrid radix sort from
+// - https://gist.github.com/sciecode/93ed864dd77c5c8803c6a86698d68dab
+// - https://github.com/mrdoob/three.js/pull/27202#issuecomment-1817640271
+const POWER = 3;
+const BIT_MAX = 32;
+const BIN_BITS = 1 << POWER;
+const BIN_SIZE = 1 << BIN_BITS;
+const ITERATIONS = BIT_MAX / BIN_BITS;
+
+const bins = new Array( ITERATIONS );
+const caches = new Array( ITERATIONS );
+const bins_buffer = new ArrayBuffer( 2 * ITERATIONS * BIN_SIZE * 4 );
+
+let c = 0;
+for ( let i = 0; i < ITERATIONS; i ++ ) {
+
+	bins[ i ] = new Uint32Array( bins_buffer, c, BIN_SIZE );
+	c += BIN_SIZE * 4;
+	caches[ i ] = new Uint32Array( bins_buffer, c, BIN_SIZE );
+	c += BIN_SIZE * 4;
+
+}
+
+const defaultGet = ( el ) => el;
+
+export const radixSort = ( arr, opt ) => {
+
+	const len = arr.length;
+
+	const options = opt || {};
+	const aux = options.aux || new arr.constructor( len );
+	const get = options.get || defaultGet;
+
+	const data = [ arr, aux ];
+
+	let compare, accumulate, recurse;
+
+	if ( options.reversed ) {
+
+		compare = ( a, b ) => a < b;
+		accumulate = ( bin ) => {
+
+			for ( let j = BIN_SIZE - 2; j >= 0; j -- )
+				bin[ j ] += bin[ j + 1 ];
+
+		};
+
+		recurse = ( cache, depth, start ) => {
+
+			let prev = 0;
+			for ( let j = BIN_SIZE - 1; j >= 0; j -- ) {
+
+				const cur = cache[ j ], diff = cur - prev;
+				if ( diff != 0 ) {
+
+					if ( diff > 32 )
+						radixSortBlock( depth + 1, start + prev, diff );
+					else
+						insertionSortBlock( depth + 1, start + prev, diff );
+					prev = cur;
+
+				}
+
+			}
+
+		};
+
+	} else {
+
+		compare = ( a, b ) => a > b;
+		accumulate = ( bin ) => {
+
+			for ( let j = 1; j < BIN_SIZE; j ++ )
+				bin[ j ] += bin[ j - 1 ];
+
+		};
+
+		recurse = ( cache, depth, start ) => {
+
+			let prev = 0;
+			for ( let j = 0; j < BIN_SIZE; j ++ ) {
+
+				const cur = cache[ j ], diff = cur - prev;
+				if ( diff != 0 ) {
+
+					if ( diff > 32 )
+						radixSortBlock( depth + 1, start + prev, diff );
+					else
+						insertionSortBlock( depth + 1, start + prev, diff );
+					prev = cur;
+
+				}
+
+			}
+
+		};
+
+	}
+
+	const insertionSortBlock = ( depth, start, len ) => {
+
+		const a = data[ depth & 1 ];
+		const b = data[ ( depth + 1 ) & 1 ];
+
+		for ( let j = start + 1; j < start + len; j ++ ) {
+
+			const p = a[ j ], t = get( p );
+			let i = j;
+			while ( i > 0 ) {
+
+				if ( compare( get( a[ i - 1 ] ), t ) )
+					a[ i ] = a[ -- i ];
+				else
+					break;
+
+			}
+
+			a[ i ] = p;
+
+		}
+
+		if ( ( depth & 1 ) == 1 ) {
+
+			for ( let i = start; i < start + len; i ++ )
+				b[ i ] = a[ i ];
+
+		}
+
+	};
+
+	const radixSortBlock = ( depth, start, len ) => {
+
+		const a = data[ depth & 1 ];
+		const b = data[ ( depth + 1 ) & 1 ];
+
+		const shift = ( 3 - depth ) << POWER;
+		const end = start + len;
+
+		const bin = bins[ depth ];
+		const cache = caches[ depth ];
+
+		bin.fill( 0 );
+
+		for ( let j = start; j < end; j ++ )
+			bin[ ( get( a[ j ] ) >> shift ) & ( BIN_SIZE - 1 ) ] ++;
+
+		accumulate( bin );
+
+		cache.set( bin );
+
+		for ( let j = end - 1; j >= start; j -- )
+			b[ start + -- bin[ ( get( a[ j ] ) >> shift ) & ( BIN_SIZE - 1 ) ] ] = a[ j ];
+
+		if ( depth == ITERATIONS - 1 ) return;
+
+		recurse( cache, depth, start );
+
+	};
+
+	radixSortBlock( 0, 0, len );
+
+};

examples/webgl_mesh_batch.html

@@ -35,6 +35,7 @@
 		import { GUI } from 'three/addons/libs/lil-gui.module.min.js';
 
 		import { OrbitControls } from 'three/addons/controls/OrbitControls.js';
+		import { radixSort } from 'three/addons/utils/SortUtils.js';
 
 		let stats, gui, guiStatsEl;
 		let camera, controls, scene, renderer;
@@ -66,6 +67,7 @@
 			sortObjects: true,
 			perObjectFrustumCulled: true,
 			opacity: 1,
+			useCustomSort: true,
 		};
 
 		init();
@@ -117,7 +119,6 @@
 			if ( ! material ) {
 
 				material = new THREE.MeshNormalMaterial();
-				material.opacity = 0.1;
 
 			}
 
@@ -268,6 +269,7 @@
 			} );
 			gui.add( api, 'sortObjects' );
 			gui.add( api, 'perObjectFrustumCulled' );
+			gui.add( api, 'useCustomSort' );
 
 			guiStatsEl = document.createElement( 'li' );
 			guiStatsEl.classList.add( 'gui-stats' );
@@ -280,6 +282,31 @@
 
 		//
 
+		function sortFunction( list, camera ) {
+
+			// initialize options
+			this._options = this._options || {
+				get: el => el.z,
+				aux: new Array( list.length ),
+			};
+
+			const options = this._options;
+			options.reversed = this.material.transparent;
+
+			// convert depth to unsigned 32 bit range
+			const den = camera.far;
+			for ( let i = 0, l = list.length; i < l; i ++ ) {
+
+				const el = list[ i ];
+				el.z = ( 1 << 30 ) * ( el.z / den );
+
+			}
+
+			// perform a fast-sort using the hybrid radix sort function
+			radixSort( list, options );
+
+		}
+
 		function onWindowResize() {
 
 			const width = window.innerWidth;
@@ -347,6 +374,7 @@
 
 				mesh.sortObjects = api.sortObjects;
 				mesh.perObjectFrustumCulled = api.perObjectFrustumCulled;
+				mesh.setCustomSort( api.useCustomSort ? sortFunction : null );
 
 			}
 

src/objects/BatchedMesh.js

@@ -129,6 +129,7 @@ class BatchedMesh extends Mesh {
 		this.sortObjects = true;
 		this.boundingBox = null;
 		this.boundingSphere = null;
+		this.customSort = null;
 
 		this._drawRanges = [];
 		this._reservedRanges = [];
@@ -300,6 +301,13 @@ class BatchedMesh extends Mesh {
 
 	}
 
+	setCustomSort( func ) {
+
+		this.customSort = func;
+		return this;
+
+	}
+
 	computeBoundingBox() {
 
 		if ( this.boundingBox === null ) {
@@ -963,7 +971,7 @@ class BatchedMesh extends Mesh {
 					if ( ! culled ) {
 
 						// get the distance from camera used for sorting
-						const z = _vector.distanceToSquared( _sphere.center );
+						const z = _vector.distanceTo( _sphere.center );
 						_renderList.push( drawRanges[ i ], z );
 
 					}
@@ -974,7 +982,16 @@ class BatchedMesh extends Mesh {
 
 			// Sort the draw ranges and prep for rendering
 			const list = _renderList.list;
-			list.sort( material.transparent ? sortTransparent : sortOpaque );
+			const customSort = this.customSort;
+			if ( customSort === null ) {
+
+				list.sort( material.transparent ? sortTransparent : sortOpaque );
+
+			} else {
+
+				customSort.call( this, list, camera );
+
+			}
 
 			for ( let i = 0, l = list.length; i < l; i ++ ) {