playcanvas/build/playcanvas/src/scene/gsplat/gsplat-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 { UploadStream } from '../../platform/graphics/upload-stream.js';
import { SortWorker } from './gsplat-sort-worker.js';

class GSplatSorter extends EventHandler {
		destroy() {
				this.worker.terminate();
				this.worker = null;
				this.uploadStream.destroy();
				this.uploadStream = null;
		}
		init(target, numSplats, centers, chunks) {
				this.target = target;
				this.centers = centers.slice();
				const orderBuffer = new Uint32Array(numSplats);
				for(let i = 0; i < numSplats; ++i){
						orderBuffer[i] = i;
				}
				this.orderData = new ArrayBuffer(numSplats * 4);
				const obj = {
						order: orderBuffer.buffer,
						centers: centers.buffer,
						chunks: chunks?.buffer
				};
				const transfer = [
						orderBuffer.buffer,
						centers.buffer
				].concat(chunks ? [
						chunks.buffer
				] : []);
				this.worker.postMessage(obj, transfer);
		}
		applyPendingSorted() {
				if (this.pendingSorted) {
						const { count, data } = this.pendingSorted;
						this.pendingSorted = null;
						this.uploadStream.upload(data, this.target);
						return count;
				}
				return -1;
		}
		setMapping(mapping) {
				if (mapping) {
						const centers = new Float32Array(mapping.length * 3);
						for(let i = 0; i < mapping.length; ++i){
								const src = mapping[i] * 3;
								const dst = i * 3;
								centers[dst + 0] = this.centers[src + 0];
								centers[dst + 1] = this.centers[src + 1];
								centers[dst + 2] = this.centers[src + 2];
						}
						this.worker.postMessage({
								centers: centers.buffer,
								mapping: mapping.buffer
						}, [
								centers.buffer,
								mapping.buffer
						]);
				} else {
						const centers = this.centers.slice();
						this.worker.postMessage({
								centers: centers.buffer,
								mapping: null
						}, [
								centers.buffer
						]);
				}
		}
		setCamera(pos, dir) {
				this.worker.postMessage({
						cameraPosition: {
								x: pos.x,
								y: pos.y,
								z: pos.z
						},
						cameraDirection: {
								x: dir.x,
								y: dir.y,
								z: dir.z
						}
				});
		}
		constructor(device, scene){
				super(), this.pendingSorted = null;
				this.scene = scene ?? null;
				this.uploadStream = new UploadStream(device);
				const messageHandler = (message)=>{
						const msgData = message.data ?? message;
						if (this.scene && msgData.sortTime !== undefined) {
								this.scene.fire('gsplat:sorted', msgData.sortTime);
						}
						const newOrder = msgData.order;
						const oldOrder = this.orderData;
						this.worker.postMessage({
								order: oldOrder
						}, [
								oldOrder
						]);
						this.orderData = newOrder;
						this.pendingSorted = {
								count: msgData.count,
								data: new Uint32Array(newOrder)
						};
						this.fire('updated');
				};
				const workerSource = `(${SortWorker.toString()})()`;
				if (platform.environment === 'node') {
						this.worker = new Worker(workerSource, {
								eval: true
						});
						this.worker.on('message', messageHandler);
				} else {
						this.worker = new Worker(URL.createObjectURL(new Blob([
								workerSource
						], {
								type: "application/javascript"
						})));
						this.worker.addEventListener('message', messageHandler);
				}
		}
}

export { GSplatSorter };