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

PlayCanvas WebGL game engine

import { Mat4 } from '../../core/math/mat4.js';
import { Vec3 } from '../../core/math/vec3.js';
import { GraphNode } from '../graph-node.js';
import { GSplatInfo } from './gsplat-info.js';
import { GSplatUnifiedSorter } from './gsplat-unified-sorter.js';
import { GSplatWorkBuffer } from './gsplat-work-buffer.js';
import { GSplatRenderer } from './gsplat-renderer.js';
import { GSplatOctreeInstance } from './gsplat-octree-instance.js';
import { GSplatOctreeResource } from './gsplat-octree.resource.js';
import { GSplatWorldState } from './gsplat-world-state.js';
import { Color } from '../../core/math/color.js';

const cameraPosition = new Vec3();
const cameraDirection = new Vec3();
const translation = new Vec3();
const invModelMat = new Mat4();
const tempNonOctreePlacements = new Set();
const tempOctreePlacements = new Set();
const _updatedSplats = [];
const tempOctreesTicked = new Set();
const _lodColorsRaw = [
		[
				1,
				0,
				0
		],
		[
				0,
				1,
				0
		],
		[
				0,
				0,
				1
		],
		[
				1,
				1,
				0
		],
		[
				1,
				0,
				1
		]
];
[
		new Color(1, 0, 0),
		new Color(0, 1, 0),
		new Color(0, 0, 1),
		new Color(1, 1, 0),
		new Color(1, 0, 1)
];
class GSplatManager {
		constructor(device, director, layer, cameraNode){
				this.node = new GraphNode('GSplatManager');
				this.worldStates = new Map();
				this.lastWorldStateVersion = 0;
				this.sortedVersion = 0;
				this.framesTillFullUpdate = 0;
				this.lastCameraPos = new Vec3(Infinity, Infinity, Infinity);
				this.lastCameraFwd = new Vec3(Infinity, Infinity, Infinity);
				this.layerPlacements = [];
				this.layerPlacementsDirty = false;
				this.octreeInstances = new Map();
				this.octreeInstancesToDestroy = [];
				this.device = device;
				this.scene = director.scene;
				this.director = director;
				this.cameraNode = cameraNode;
				this.workBuffer = new GSplatWorkBuffer(device);
				this.renderer = new GSplatRenderer(device, this.node, this.cameraNode, layer, this.workBuffer);
				this.sorter = this.createSorter();
				this.cooldownTicks = this.director.assetLoader.cooldownTicks;
		}
		destroy() {
				this.workBuffer.destroy();
				this.renderer.destroy();
				this.sorter.destroy();
		}
		createSorter() {
				const sorter = new GSplatUnifiedSorter();
				sorter.on('sorted', (count, version, orderData)=>{
						this.onSorted(count, version, orderData);
				});
				return sorter;
		}
		reconcile(placements) {
				tempNonOctreePlacements.clear();
				for (const p of placements){
						if (p.resource instanceof GSplatOctreeResource) {
								if (!this.octreeInstances.has(p)) {
										this.octreeInstances.set(p, new GSplatOctreeInstance(p.resource.octree, p, this.director.assetLoader));
								}
								tempOctreePlacements.add(p);
						} else {
								tempNonOctreePlacements.add(p);
						}
				}
				for (const [placement, inst] of this.octreeInstances){
						if (!tempOctreePlacements.has(placement)) {
								this.octreeInstances.delete(placement);
								this.layerPlacementsDirty = true;
								this.octreeInstancesToDestroy.push(inst);
						}
				}
				this.layerPlacementsDirty = this.layerPlacements.length !== tempNonOctreePlacements.size;
				if (!this.layerPlacementsDirty) {
						for(let i = 0; i < this.layerPlacements.length; i++){
								const existing = this.layerPlacements[i];
								if (!tempNonOctreePlacements.has(existing)) {
										this.layerPlacementsDirty = true;
										break;
								}
						}
				}
				this.layerPlacements.length = 0;
				for (const p of tempNonOctreePlacements){
						this.layerPlacements.push(p);
				}
				tempNonOctreePlacements.clear();
				tempOctreePlacements.clear();
		}
		updateWorldState() {
				const worldChanged = this.layerPlacementsDirty || this.worldStates.size === 0;
				if (worldChanged) {
						this.lastWorldStateVersion++;
						const splats = [];
						for (const p of this.layerPlacements){
								const splatInfo = new GSplatInfo(this.device, p.resource, p);
								splats.push(splatInfo);
						}
						for (const [, inst] of this.octreeInstances){
								inst.activePlacements.forEach((p)=>{
										if (p.resource) {
												splats.push(new GSplatInfo(this.device, p.resource, p));
										}
								});
						}
						splats.forEach((splat)=>{
								this.sorter.setCenters(splat.resource.id, splat.resource.centers);
						});
						const newState = new GSplatWorldState(this.device, this.lastWorldStateVersion, splats);
						for (const [, inst] of this.octreeInstances){
								if (inst.removedCandidates && inst.removedCandidates.size) {
										for (const fileIndex of inst.removedCandidates){
												newState.pendingReleases.push([
														inst.octree,
														fileIndex
												]);
										}
										inst.removedCandidates.clear();
								}
						}
						if (this.octreeInstancesToDestroy.length) {
								for (const inst of this.octreeInstancesToDestroy){
										const toRelease = inst.getFileDecrements();
										for (const fileIndex of toRelease){
												newState.pendingReleases.push([
														inst.octree,
														fileIndex
												]);
										}
										inst.destroy();
								}
								this.octreeInstancesToDestroy.length = 0;
						}
						this.worldStates.set(this.lastWorldStateVersion, newState);
						this.layerPlacementsDirty = false;
				}
		}
		onSorted(count, version, orderData) {
				this.sortedVersion = version;
				const oldState = this.worldStates.get(version - 1);
				if (oldState) {
						this.worldStates.delete(version - 1);
						oldState.destroy();
				}
				const worldState = this.worldStates.get(version);
				if (worldState) {
						if (!worldState.sortedBefore) {
								worldState.sortedBefore = true;
								const textureSize = worldState.textureSize;
								if (textureSize !== this.workBuffer.textureSize) {
										this.workBuffer.resize(textureSize);
										this.renderer.setMaxNumSplats(textureSize * textureSize);
								}
								const colorize = this.scene.gsplat.colorizeLod;
								this.workBuffer.render(worldState.splats, this.cameraNode, colorize ? _lodColorsRaw : undefined);
								if (worldState.pendingReleases && worldState.pendingReleases.length) {
										for (const [octree, fileIndex] of worldState.pendingReleases){
												octree.decRefCount(fileIndex, this.cooldownTicks);
										}
										worldState.pendingReleases.length = 0;
								}
								this.renderer.setNumSplats(count);
						}
						this.workBuffer.setOrderData(orderData);
				}
		}
		testCameraMoved() {
				const distanceThreshold = this.scene.gsplat.lodUpdateDistance;
				const currentCameraPos = this.cameraNode.getPosition();
				const cameraMoved = this.lastCameraPos.distance(currentCameraPos) > distanceThreshold;
				if (cameraMoved) {
						return true;
				}
				let cameraRotated = false;
				const lodUpdateAngleDeg = this.scene.gsplat.lodUpdateAngle;
				if (lodUpdateAngleDeg > 0) {
						if (Number.isFinite(this.lastCameraFwd.x)) {
								const currentCameraFwd = this.cameraNode.forward;
								const dot = Math.min(1, Math.max(-1, this.lastCameraFwd.dot(currentCameraFwd)));
								const angle = Math.acos(dot);
								const rotThreshold = lodUpdateAngleDeg * Math.PI / 180;
								cameraRotated = angle > rotThreshold;
						} else {
								cameraRotated = true;
						}
				}
				return cameraMoved || cameraRotated;
		}
		update() {
				let fullUpdate = false;
				this.framesTillFullUpdate--;
				if (this.framesTillFullUpdate <= 0) {
						this.framesTillFullUpdate = 10;
						if (this.sorter.jobsInFlight < 3) {
								fullUpdate = true;
						}
				}
				let anyInstanceNeedsLodUpdate = false;
				let anyOctreeMoved = false;
				let cameraMovedOrRotated = false;
				if (fullUpdate) {
						for (const [, inst] of this.octreeInstances){
								const isDirty = inst.update(this.scene);
								this.layerPlacementsDirty ||= isDirty;
								const instNeeds = inst.consumeNeedsLodUpdate();
								anyInstanceNeedsLodUpdate ||= instNeeds;
						}
						const threshold = this.scene.gsplat.lodUpdateDistance;
						for (const [, inst] of this.octreeInstances){
								const moved = inst.testMoved(threshold);
								anyOctreeMoved ||= moved;
						}
						cameraMovedOrRotated = this.testCameraMoved();
				}
				if (this.scene.gsplat.dirty) {
						this.layerPlacementsDirty = true;
				}
				if (cameraMovedOrRotated || anyOctreeMoved || this.scene.gsplat.dirty || anyInstanceNeedsLodUpdate) {
						for (const [, inst] of this.octreeInstances){
								inst.updateMoved();
						}
						this.lastCameraPos.copy(this.cameraNode.getPosition());
						this.lastCameraFwd.copy(this.cameraNode.forward);
						for (const [, inst] of this.octreeInstances){
								inst.updateLod(this.cameraNode, this.scene.gsplat);
						}
				}
				this.updateWorldState();
				const lastState = this.worldStates.get(this.lastWorldStateVersion);
				if (lastState) {
						if (!lastState.sortParametersSet) {
								lastState.sortParametersSet = true;
								const payload = this.prepareSortParameters(lastState);
								this.sorter.setSortParameters(payload);
						}
						this.sort(lastState);
				}
				const sortedState = this.worldStates.get(this.sortedVersion);
				if (sortedState) {
						sortedState.splats.forEach((splat)=>{
								if (splat.update()) {
										_updatedSplats.push(splat);
								}
						});
						if (_updatedSplats.length > 0) {
								const colorize = this.scene.gsplat.colorizeLod;
								this.workBuffer.render(_updatedSplats, this.cameraNode, colorize ? _lodColorsRaw : undefined);
								_updatedSplats.length = 0;
						}
				}
				if (this.octreeInstances.size) {
						for (const [, inst] of this.octreeInstances){
								const octree = inst.octree;
								if (!tempOctreesTicked.has(octree)) {
										tempOctreesTicked.add(octree);
										octree.updateCooldownTick(this.director.assetLoader);
								}
						}
						tempOctreesTicked.clear();
				}
				const { textureSize } = this.workBuffer;
				return textureSize * textureSize;
		}
		sort(lastState) {
				const cameraNode = this.cameraNode;
				const cameraMat = cameraNode.getWorldTransform();
				cameraMat.getTranslation(cameraPosition);
				cameraMat.getZ(cameraDirection).normalize();
				const sorterRequest = [];
				lastState.splats.forEach((splat)=>{
						const modelMat = splat.node.getWorldTransform();
						invModelMat.copy(modelMat).invert();
						const uniformScale = modelMat.getScale().x;
						const transformedDirection = invModelMat.transformVector(cameraDirection).normalize();
						const transformedPosition = invModelMat.transformPoint(cameraPosition);
						modelMat.getTranslation(translation);
						const offset = translation.sub(cameraPosition).dot(cameraDirection);
						const aabbMin = splat.aabb.getMin();
						const aabbMax = splat.aabb.getMax();
						sorterRequest.push({
								transformedDirection,
								transformedPosition,
								offset,
								scale: uniformScale,
								modelMat: modelMat.data.slice(),
								aabbMin: [
										aabbMin.x,
										aabbMin.y,
										aabbMin.z
								],
								aabbMax: [
										aabbMax.x,
										aabbMax.y,
										aabbMax.z
								]
						});
				});
				this.sorter.setSortParams(sorterRequest, this.scene.gsplat.radialSorting);
				this.renderer.updateViewport(cameraNode);
		}
		prepareSortParameters(worldState) {
				return {
						command: 'intervals',
						textureSize: worldState.textureSize,
						totalUsedPixels: worldState.totalUsedPixels,
						version: worldState.version,
						ids: worldState.splats.map((splat)=>splat.resource.id),
						lineStarts: worldState.splats.map((splat)=>splat.lineStart),
						padding: worldState.splats.map((splat)=>splat.padding),
						intervals: worldState.splats.map((splat)=>splat.intervals)
				};
		}
}

export { GSplatManager };