playcanvas/build/playcanvas.prf/src/scene/gsplat-unified/gsplat-unified-sorter.js

Open-source WebGL/WebGPU 3D engine for the web

import { EventHandler } from '../../core/event-handler.js';
import { platform } from '../../core/platform.js';
import { UnifiedSortWorker } from './gsplat-unified-sort-worker.js';
import { GSplatSortBinWeights } from './gsplat-sort-bin-weights.js';

const _neededIds = new Set();
class GSplatUnifiedSorter extends EventHandler {
		onSorted(message) {
				if (this._destroyed) {
						return;
				}
				const msgData = message.data ?? message;
				if (this.scene && msgData.sortTime !== undefined) {
						this.scene.fire('gsplat:sorted', msgData.sortTime);
				}
				const orderData = new Uint32Array(msgData.order);
				this.jobsInFlight--;
				if (this.pendingSorted) {
						this.releaseOrderData(this.pendingSorted.orderData);
				}
				this.pendingSorted = {
						count: msgData.count,
						version: msgData.version,
						orderData: orderData
				};
		}
		applyPendingSorted() {
				if (this.pendingSorted) {
						const { count, version, orderData } = this.pendingSorted;
						this.pendingSorted = null;
						this.fire('sorted', count, version, orderData);
						this.releaseOrderData(orderData);
				}
		}
		releaseOrderData(orderData) {
				if (orderData.length === this.bufferLength) {
						this.availableOrderData.push(orderData);
				}
		}
		destroy() {
				this._destroyed = true;
				this.pendingSorted = null;
				this.worker.terminate();
				this.worker = null;
		}
		setCenters(id, centers) {
				if (centers) {
						if (!this.centersSet.has(id)) {
								this.centersSet.add(id);
								const centersBuffer = centers.buffer.slice();
								this.worker.postMessage({
										command: 'addCenters',
										id: id,
										centers: centersBuffer
								}, [
										centersBuffer
								]);
						}
				} else {
						if (this.centersSet.has(id)) {
								this.centersSet.delete(id);
								this.worker.postMessage({
										command: 'removeCenters',
										id: id
								});
						}
				}
		}
		updateCentersForSplats(splats) {
				for (const splat of splats){
						const id = splat.resource.id;
						_neededIds.add(id);
						if (!this.centersSet.has(id)) {
								this.setCenters(id, splat.resource.centers);
						}
				}
				for (const id of this.centersSet){
						if (!_neededIds.has(id)) {
								this.setCenters(id, null);
						}
				}
				_neededIds.clear();
		}
		setSortParameters(payload) {
				this.hasNewVersion = true;
				const { textureSize } = payload;
				const newLength = textureSize * textureSize;
				if (newLength !== this.bufferLength) {
						this.bufferLength = newLength;
						this.availableOrderData.length = 0;
				}
				this.worker.postMessage(payload);
		}
		setSortParams(params, radialSorting) {
				if (this.hasNewVersion || this.jobsInFlight === 0) {
						let orderData = this.availableOrderData.pop();
						if (!orderData) {
								orderData = new Uint32Array(this.bufferLength);
						}
						this.jobsInFlight++;
						this.hasNewVersion = false;
						this.worker.postMessage({
								command: 'sort',
								sortParams: params,
								radialSorting: radialSorting,
								order: orderData.buffer
						}, [
								orderData.buffer
						]);
				}
		}
		constructor(scene){
				super(), this.bufferLength = 0, this.availableOrderData = [], this.jobsInFlight = 0, this.hasNewVersion = false, this.pendingSorted = null, this.centersSet = new Set(), this._destroyed = false, this.scene = null;
				this.scene = scene ?? null;
				const workerSource = `
						const GSplatSortBinWeights = ${GSplatSortBinWeights.toString()};
						(${UnifiedSortWorker.toString()})()
				`;
				if (platform.environment === 'node') {
						this.worker = new Worker(workerSource, {
								eval: true
						});
						this.worker.on('message', this.onSorted.bind(this));
				} else {
						this.worker = new Worker(URL.createObjectURL(new Blob([
								workerSource
						], {
								type: "application/javascript"
						})));
						this.worker.addEventListener('message', this.onSorted.bind(this));
				}
		}
}

export { GSplatUnifiedSorter };