playcanvas/build/playcanvas/src/scene/gsplat/gsplat-resource-base.js

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

import { BoundingBox } from '../../core/shape/bounding-box.js';
import { Mesh } from '../mesh.js';
import { ShaderMaterial } from '../materials/shader-material.js';
import { WorkBufferRenderInfo } from '../gsplat-unified/gsplat-work-buffer.js';
import { GSplatStreams } from './gsplat-streams.js';
import { GSplatResourceCleanup } from './gsplat-resource-cleanup.js';

let id = 0;
const tempMap = new Map();
class GSplatResourceBase {
		destroy() {
				if (this.refCount > 0) {
						GSplatResourceCleanup.queueDestroy(this.device, this);
						return;
				}
				this._actualDestroy();
		}
		_actualDestroy() {
				this.streams.destroy();
				this.mesh?.destroy();
				this.workBufferRenderInfos.forEach((info)=>info.destroy());
				this.workBufferRenderInfos.clear();
		}
		incRefCount() {
				this._refCount++;
		}
		decRefCount() {
				this._refCount--;
		}
		get refCount() {
				return this._refCount;
		}
		ensureMesh() {
				if (!this.mesh) {
						this.mesh = GSplatResourceBase.createMesh(this.device);
						this.mesh.aabb.copy(this.aabb);
				}
				this._meshRefCount++;
		}
		releaseMesh() {
				this._meshRefCount--;
				if (this._meshRefCount < 1) {
						this.mesh = null;
				}
		}
		getWorkBufferRenderInfo(colorOnly, workBufferModifier, formatHash, formatDeclarations, workBufferFormat) {
				this.configureMaterialDefines(tempMap);
				tempMap.set('GSPLAT_LOD', '');
				if (colorOnly) tempMap.set('GSPLAT_COLOR_ONLY', '');
				let definesKey = '';
				for (const [k, v] of tempMap){
						if (definesKey) definesKey += ';';
						definesKey += `${k}=${v}`;
				}
				const key = `${formatHash};${workBufferFormat.hash};${workBufferModifier?.hash ?? 0};${definesKey}`;
				let info = this.workBufferRenderInfos.get(key);
				if (!info) {
						const material = new ShaderMaterial();
						this.configureMaterial(material, workBufferModifier, formatDeclarations);
						const chunks = this.device.isWebGPU ? material.shaderChunks.wgsl : material.shaderChunks.glsl;
						const outputStreams = colorOnly ? [
								workBufferFormat.getStream('dataColor')
						] : [
								...workBufferFormat.streams,
								...workBufferFormat.extraStreams
						];
						let outputCode = workBufferFormat.getOutputDeclarations(outputStreams);
						if (colorOnly && workBufferFormat.extraStreams.length > 0) {
								outputCode += `\n${workBufferFormat.getOutputStubs(workBufferFormat.extraStreams)}`;
						}
						chunks.set('gsplatWorkBufferOutputVS', outputCode);
						const writeCode = workBufferFormat.getWriteCode();
						if (writeCode) {
								chunks.set('gsplatWriteVS', writeCode);
						}
						tempMap.forEach((v, k)=>material.setDefine(k, v));
						info = new WorkBufferRenderInfo(this.device, key, material, colorOnly, workBufferFormat);
						this.workBufferRenderInfos.set(key, info);
				}
				tempMap.clear();
				return info;
		}
		static createMesh(device) {
				const splatInstanceSize = GSplatResourceBase.instanceSize;
				const meshPositions = new Float32Array(12 * splatInstanceSize);
				const meshIndices = new Uint32Array(6 * splatInstanceSize);
				for(let i = 0; i < splatInstanceSize; ++i){
						meshPositions.set([
								-1,
								-1,
								i,
								1,
								-1,
								i,
								1,
								1,
								i,
								-1,
								1,
								i
						], i * 12);
						const b = i * 4;
						meshIndices.set([
								0 + b,
								1 + b,
								2 + b,
								0 + b,
								2 + b,
								3 + b
						], i * 6);
				}
				const mesh = new Mesh(device);
				mesh.setPositions(meshPositions, 3);
				mesh.setIndices(meshIndices);
				mesh.update();
				return mesh;
		}
		static get instanceSize() {
				return 128;
		}
		get numSplats() {
				return this.gsplatData.numSplats;
		}
		get format() {
				return this._format;
		}
		getTexture(name) {
				return this.streams.getTexture(name) ?? null;
		}
		get textureDimensions() {
				return this.streams.textureDimensions;
		}
		configureMaterial(material, workBufferModifier, formatDeclarations) {
				this.configureMaterialDefines(material.defines);
				this.streams.syncWithFormat(this.format);
				const chunks = this.device.isWebGPU ? material.shaderChunks.wgsl : material.shaderChunks.glsl;
				chunks.set('gsplatDeclarationsVS', formatDeclarations);
				chunks.set('gsplatReadVS', this.format.getReadCode());
				if (workBufferModifier?.code) {
						chunks.set('gsplatModifyVS', workBufferModifier.code);
				}
				for (const [name, texture] of this.streams.textures){
						material.setParameter(name, texture);
				}
				for (const [name, value] of this.parameters){
						material.setParameter(name, value);
				}
				if (this.textureDimensions.x > 0) {
						material.setParameter('splatTextureSize', this.textureDimensions.x);
				}
		}
		configureMaterialDefines(defines) {}
		instantiate() {}
		constructor(device, gsplatData){
				this.centersVersion = 0;
				this.mesh = null;
				this.id = id++;
				this.workBufferRenderInfos = new Map();
				this._format = null;
				this.parameters = new Map();
				this._refCount = 0;
				this._meshRefCount = 0;
				this.device = device;
				this.gsplatData = gsplatData;
				this.streams = new GSplatStreams(device);
				this.centers = gsplatData.getCenters();
				this.aabb = new BoundingBox();
				gsplatData.calcAabb(this.aabb);
		}
}

export { GSplatResourceBase };