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

PlayCanvas WebGL game engine

import { math } from '../../core/math/math.js';
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 _splatsNeedingColorUpdate = [];
const _cameraDeltas = {
		rotationDelta: 0,
		translationDelta: 0
};
const tempOctreesTicked = new Set();
const _lodColorsRaw = [
		[
				1,
				0,
				0
		],
		[
				0,
				1,
				0
		],
		[
				0,
				0,
				1
		],
		[
				1,
				1,
				0
		],
		[
				1,
				0,
				1
		],
		[
				0,
				1,
				1
		],
		[
				1,
				0.5,
				0
		],
		[
				0.5,
				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),
		new Color(0, 1, 1),
		new Color(1, 0.5, 0),
		new Color(0.5, 0, 1)
];
let _randomColorRaw = null;
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.lastLodCameraPos = new Vec3(Infinity, Infinity, Infinity);
				this.lastLodCameraFwd = new Vec3(Infinity, Infinity, Infinity);
				this.lastSortCameraPos = new Vec3(Infinity, Infinity, Infinity);
				this.lastSortCameraFwd = new Vec3(Infinity, Infinity, Infinity);
				this.sortNeeded = true;
				this.lastColorUpdateCameraPos = new Vec3(Infinity, Infinity, Infinity);
				this.lastColorUpdateCameraFwd = new Vec3(Infinity, Infinity, Infinity);
				this.layerPlacements = [];
				this.layerPlacementsDirty = false;
				this.octreeInstances = new Map();
				this.octreeInstancesToDestroy = [];
				this.hasNewOctreeInstances = false;
				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();
		}
		destroy() {
				this._destroyed = true;
				for (const [, worldState] of this.worldStates){
						for (const splat of worldState.splats){
								splat.resource.decRefCount();
						}
						worldState.destroy();
				}
				this.worldStates.clear();
				for (const [, instance] of this.octreeInstances){
						instance.destroy();
				}
				this.octreeInstances.clear();
				for (const instance of this.octreeInstancesToDestroy){
						instance.destroy();
				}
				this.octreeInstancesToDestroy.length = 0;
				this.workBuffer.destroy();
				this.renderer.destroy();
				this.sorter.destroy();
		}
		get material() {
				return this.renderer.material;
		}
		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(this.device, p.resource.octree, p));
										this.hasNewOctreeInstances = true;
								}
								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 = [];
						const { colorUpdateAngle, colorUpdateDistance } = this.scene.gsplat;
						for (const p of this.layerPlacements){
								const splatInfo = new GSplatInfo(this.device, p.resource, p);
								splatInfo.resetColorAccumulators(colorUpdateAngle, colorUpdateDistance);
								splats.push(splatInfo);
						}
						for (const [, inst] of this.octreeInstances){
								inst.activePlacements.forEach((p)=>{
										if (p.resource) {
												const splatInfo = new GSplatInfo(this.device, p.resource, p);
												splatInfo.resetColorAccumulators(colorUpdateAngle, colorUpdateDistance);
												splats.push(splatInfo);
										}
								});
						}
						this.sorter.updateCentersForSplats(splats);
						const newState = new GSplatWorldState(this.device, this.lastWorldStateVersion, splats);
						for (const splat of newState.splats){
								splat.resource.incRefCount();
						}
						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;
						this.sortNeeded = true;
				}
		}
		onSorted(count, version, orderData) {
				for(let v = this.sortedVersion; v < version; v++){
						const oldState = this.worldStates.get(v);
						if (oldState) {
								for (const splat of oldState.splats){
										splat.resource.decRefCount();
								}
								this.worldStates.delete(v);
								oldState.destroy();
						}
				}
				this.sortedVersion = version;
				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);
								}
								this.workBuffer.render(worldState.splats, this.cameraNode, this.getDebugColors());
								const { colorUpdateAngle, colorUpdateDistance } = this.scene.gsplat;
								worldState.splats.forEach((splat)=>{
										splat.update();
										splat.resetColorAccumulators(colorUpdateAngle, colorUpdateDistance);
								});
								this.updateColorCameraTracking();
								if (worldState.pendingReleases && worldState.pendingReleases.length) {
										const cooldownTicks = this.scene.gsplat.cooldownTicks;
										for (const [octree, fileIndex] of worldState.pendingReleases){
												octree.decRefCount(fileIndex, cooldownTicks);
										}
										worldState.pendingReleases.length = 0;
								}
								this.renderer.update(count, textureSize);
						}
						this.workBuffer.setOrderData(orderData);
						this.renderer.frameUpdate(this.scene.gsplat);
				}
		}
		testCameraMovedForLod() {
				const distanceThreshold = this.scene.gsplat.lodUpdateDistance;
				const currentCameraPos = this.cameraNode.getPosition();
				const cameraMoved = this.lastLodCameraPos.distance(currentCameraPos) > distanceThreshold;
				if (cameraMoved) {
						return true;
				}
				let cameraRotated = false;
				const lodUpdateAngleDeg = this.scene.gsplat.lodUpdateAngle;
				if (lodUpdateAngleDeg > 0) {
						if (Number.isFinite(this.lastLodCameraFwd.x)) {
								const currentCameraFwd = this.cameraNode.forward;
								const dot = Math.min(1, Math.max(-1, this.lastLodCameraFwd.dot(currentCameraFwd)));
								const angle = Math.acos(dot);
								const rotThreshold = lodUpdateAngleDeg * math.DEG_TO_RAD;
								cameraRotated = angle > rotThreshold;
						} else {
								cameraRotated = true;
						}
				}
				return cameraMoved || cameraRotated;
		}
		testCameraMovedForSort() {
				const epsilon = 0.001;
				if (this.scene.gsplat.radialSorting) {
						const currentCameraPos = this.cameraNode.getPosition();
						const distance = this.lastSortCameraPos.distance(currentCameraPos);
						return distance > epsilon;
				}
				if (Number.isFinite(this.lastSortCameraFwd.x)) {
						const currentCameraFwd = this.cameraNode.forward;
						const dot = Math.min(1, Math.max(-1, this.lastSortCameraFwd.dot(currentCameraFwd)));
						const angle = Math.acos(dot);
						return angle > epsilon;
				}
				return true;
		}
		updateColorCameraTracking() {
				this.lastColorUpdateCameraPos.copy(this.cameraNode.getPosition());
				this.lastColorUpdateCameraFwd.copy(this.cameraNode.forward);
		}
		getDebugColors() {
				if (this.scene.gsplat.colorizeColorUpdate) {
						_randomColorRaw ??= [];
						const r = Math.random();
						const g = Math.random();
						const b = Math.random();
						for(let i = 0; i < _lodColorsRaw.length; i++){
								_randomColorRaw[i] ??= [
										0,
										0,
										0
								];
								_randomColorRaw[i][0] = r;
								_randomColorRaw[i][1] = g;
								_randomColorRaw[i][2] = b;
						}
						return _randomColorRaw;
				} else if (this.scene.gsplat.colorizeLod) {
						return _lodColorsRaw;
				}
				return undefined;
		}
		calculateColorCameraDeltas() {
				_cameraDeltas.rotationDelta = 0;
				_cameraDeltas.translationDelta = 0;
				if (isFinite(this.lastColorUpdateCameraPos.x)) {
						const currentCameraFwd = this.cameraNode.forward;
						const dot = Math.min(1, Math.max(-1, this.lastColorUpdateCameraFwd.dot(currentCameraFwd)));
						_cameraDeltas.rotationDelta = Math.acos(dot) * math.RAD_TO_DEG;
						const currentCameraPos = this.cameraNode.getPosition();
						_cameraDeltas.translationDelta = this.lastColorUpdateCameraPos.distance(currentCameraPos);
				}
				return _cameraDeltas;
		}
		fireFrameReadyEvent() {
				const ready = this.sortedVersion === this.lastWorldStateVersion;
				let loadingCount = 0;
				for (const [, inst] of this.octreeInstances){
						loadingCount += inst.pendingLoadCount;
				}
				this.director.eventHandler.fire('frame:ready', this.cameraNode.camera, this.renderer.layer, ready, loadingCount);
		}
		update() {
				this.sorter.applyPendingSorted();
				this.renderer.updateViewport(this.cameraNode);
				let fullUpdate = false;
				this.framesTillFullUpdate--;
				if (this.framesTillFullUpdate <= 0) {
						this.framesTillFullUpdate = 10;
						if (this.sorter.jobsInFlight < 3) {
								fullUpdate = true;
						}
				}
				const hasNewInstances = this.hasNewOctreeInstances && this.sorter.jobsInFlight < 3;
				if (hasNewInstances) this.hasNewOctreeInstances = false;
				let anyInstanceNeedsLodUpdate = false;
				let anyOctreeMoved = false;
				let cameraMovedOrRotatedForLod = 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;
						}
						cameraMovedOrRotatedForLod = this.testCameraMovedForLod();
				}
				if (this.testCameraMovedForSort()) {
						this.sortNeeded = true;
				}
				if (this.scene.gsplat.dirty) {
						this.layerPlacementsDirty = true;
						this.renderer.updateOverdrawMode(this.scene.gsplat);
				}
				if (cameraMovedOrRotatedForLod || anyOctreeMoved || this.scene.gsplat.dirty || anyInstanceNeedsLodUpdate || hasNewInstances) {
						for (const [, inst] of this.octreeInstances){
								inst.updateMoved();
						}
						this.lastLodCameraPos.copy(this.cameraNode.getPosition());
						this.lastLodCameraFwd.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);
						}
						if (this.sortNeeded) {
								this.sort(lastState);
								this.sortNeeded = false;
								this.lastSortCameraPos.copy(this.cameraNode.getPosition());
								this.lastSortCameraFwd.copy(this.cameraNode.forward);
						}
				}
				const sortedState = this.worldStates.get(this.sortedVersion);
				if (sortedState) {
						const { colorUpdateAngle, colorUpdateDistance, colorUpdateDistanceLodScale, colorUpdateAngleLodScale } = this.scene.gsplat;
						const { rotationDelta, translationDelta } = this.calculateColorCameraDeltas();
						sortedState.splats.forEach((splat)=>{
								if (splat.update()) {
										_updatedSplats.push(splat);
										splat.resetColorAccumulators(colorUpdateAngle, colorUpdateDistance);
										this.sortNeeded = true;
								} else if (splat.hasSphericalHarmonics) {
										splat.colorAccumulatedRotation += rotationDelta;
										splat.colorAccumulatedTranslation += translationDelta;
										const lodIndex = splat.lodIndex ?? 0;
										const distThreshold = colorUpdateDistance * Math.pow(colorUpdateDistanceLodScale, lodIndex);
										const angleThreshold = colorUpdateAngle * Math.pow(colorUpdateAngleLodScale, lodIndex);
										if (splat.colorAccumulatedRotation >= angleThreshold || splat.colorAccumulatedTranslation >= distThreshold) {
												_splatsNeedingColorUpdate.push(splat);
												splat.resetColorAccumulators(angleThreshold, distThreshold);
										}
								}
						});
						if (_updatedSplats.length > 0) {
								this.workBuffer.render(_updatedSplats, this.cameraNode, this.getDebugColors());
								_updatedSplats.length = 0;
						}
						if (_splatsNeedingColorUpdate.length > 0) {
								this.workBuffer.renderColor(_splatsNeedingColorUpdate, this.cameraNode, this.getDebugColors());
								_splatsNeedingColorUpdate.length = 0;
						}
						this.updateColorCameraTracking();
				}
				if (this.octreeInstances.size) {
						const cooldownTicks = this.scene.gsplat.cooldownTicks;
						for (const [, inst] of this.octreeInstances){
								const octree = inst.octree;
								if (!tempOctreesTicked.has(octree)) {
										tempOctreesTicked.add(octree);
										octree.updateCooldownTick(cooldownTicks);
								}
						}
						tempOctreesTicked.clear();
				}
				this.fireFrameReadyEvent();
				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);
		}
		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 };