playcanvas/build/playcanvas/src/scene/gsplat-unified/gsplat-octree-instance.js

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

import { math } from '../../core/math/math.js';
import { Mat4 } from '../../core/math/mat4.js';
import { Vec2 } from '../../core/math/vec2.js';
import { Vec3 } from '../../core/math/vec3.js';
import { BoundingBox } from '../../core/shape/bounding-box.js';
import { Color } from '../../core/math/color.js';
import { GSplatPlacement } from './gsplat-placement.js';
import { GsplatAllocId } from './gsplat-alloc-id.js';

const _invWorldMat = new Mat4();
const _localCameraPos = new Vec3();
const _localCameraFwd = new Vec3();
const _dirToNode = new Vec3();
const _tempCompletedUrls = [];
new BoundingBox();
const REF_TAN_HALF_FOV = Math.tan(22.5 * math.DEG_TO_RAD);
[
		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 NodeInfo {
		resetLod() {
				this.currentLod = -1;
				this.optimalLod = -1;
		}
		constructor(){
				this.currentLod = -1;
				this.optimalLod = -1;
				this.worldDistance = 0;
				this.colorAccumulatedTranslation = 0;
				this.inst = null;
				this.lods = null;
				this.allocId = GsplatAllocId.get();
		}
}
class GSplatOctreeInstance {
		get pendingLoadCount() {
				let count = this.pending.size + this.prefetchPending.size;
				if (this.octree.environmentUrl && !this.environmentPlacement) {
						count++;
				}
				return count;
		}
		destroy(skipRefCounting = false) {
				if (!skipRefCounting && this.octree && !this.octree.destroyed) {
						const filesToDecRef = this.getFileDecrements();
						for (const fileIndex of filesToDecRef){
								this.octree.decRefCount(fileIndex, 0);
						}
						for (const fileIndex of this.pending){
								if (!this.filePlacements[fileIndex]) {
										this.octree.unloadResource(fileIndex);
								}
						}
						for (const fileIndex of this.prefetchPending){
								if (!this.filePlacements[fileIndex]) {
										this.octree.unloadResource(fileIndex);
								}
						}
						if (this.environmentPlacement) {
								this.octree.decEnvironmentRefCount();
						}
				}
				this.pending.clear();
				this.pendingDecrements.clear();
				this.filePlacements.length = 0;
				if (this.environmentPlacement) {
						this.activePlacements.delete(this.environmentPlacement);
						this.environmentPlacement = null;
				}
				this._deviceLostEvent?.off();
				this._deviceLostEvent = null;
		}
		_onDeviceLost() {
				for(let i = 0; i < this.filePlacements.length; i++){
						if (this.filePlacements[i]) {
								this.octree.decRefCount(i, 0);
						}
				}
				this.filePlacements.fill(null);
				this.activePlacements.clear();
				this.pending.clear();
				this.pendingDecrements.clear();
				this.removedCandidates.clear();
				this.prefetchPending.clear();
				this.pendingVisibleAdds.clear();
				for (const nodeInfo of this.nodeInfos){
						nodeInfo.resetLod();
				}
				if (this.environmentPlacement) {
						this.activePlacements.delete(this.environmentPlacement);
						this.environmentPlacement = null;
						this.octree.unloadEnvironmentResource();
				}
				this.dirtyModifiedPlacements = true;
				this.dirtyPlacementSetChanged = true;
				this.needsLodUpdate = true;
		}
		getFileDecrements() {
				const toRelease = [];
				for(let i = 0; i < this.filePlacements.length; i++){
						if (this.filePlacements[i]) {
								toRelease.push(i);
						}
				}
				return toRelease;
		}
		selectDesiredLodIndex(node, optimalLodIndex, maxLod, lodUnderfillLimit) {
				if (lodUnderfillLimit > 0) {
						const allowedMaxCoarseLod = Math.min(maxLod, optimalLodIndex + lodUnderfillLimit);
						for(let lod = optimalLodIndex; lod <= allowedMaxCoarseLod; lod++){
								const fi = node.lods[lod].fileIndex;
								if (fi !== -1 && this.octree.getFileResource(fi)) {
										return lod;
								}
						}
						for(let lod = allowedMaxCoarseLod; lod >= optimalLodIndex; lod--){
								const fi = node.lods[lod].fileIndex;
								if (fi !== -1) {
										return lod;
								}
						}
				}
				return optimalLodIndex;
		}
		prefetchNextLod(node, desiredLodIndex, optimalLodIndex) {
				if (desiredLodIndex === -1 || optimalLodIndex === -1) return;
				if (desiredLodIndex === optimalLodIndex) {
						const fi = node.lods[optimalLodIndex].fileIndex;
						if (fi !== -1) {
								this.octree.ensureFileResource(fi);
								if (!this.octree.getFileResource(fi)) {
										this.prefetchPending.add(fi);
								}
						}
						return;
				}
				const targetLod = Math.max(optimalLodIndex, desiredLodIndex - 1);
				for(let lod = targetLod; lod >= optimalLodIndex; lod--){
						const fi = node.lods[lod].fileIndex;
						if (fi !== -1) {
								this.octree.ensureFileResource(fi);
								if (!this.octree.getFileResource(fi)) {
										this.prefetchPending.add(fi);
								}
								break;
						}
				}
		}
		updateLod(cameraNode, params) {
				const maxLod = this.octree.lodLevels - 1;
				const { lodBaseDistance, lodMultiplier } = this.placement;
				const { lodRangeMin, lodRangeMax } = params;
				const rangeMin = Math.max(0, Math.min(lodRangeMin ?? 0, maxLod));
				const rangeMax = Math.max(rangeMin, Math.min(lodRangeMax ?? maxLod, maxLod));
				const uniformScale = this.placement.node.getWorldTransform().getScale().x;
				this.evaluateNodeLods(cameraNode, maxLod, lodBaseDistance, lodMultiplier, rangeMin, rangeMax, params, uniformScale);
				this.applyLodChanges(maxLod, params);
		}
		evaluateNodeLods(cameraNode, maxLod, lodBaseDistance, lodMultiplier, rangeMin, rangeMax, params, uniformScale) {
				const { lodBehindPenalty } = params;
				const camera = cameraNode.camera;
				let tanHalfVFov = Math.tan(camera.fov * 0.5 * math.DEG_TO_RAD);
				if (camera.horizontalFov) {
						tanHalfVFov /= camera.aspectRatio;
				}
				const tanHalfHFov = tanHalfVFov * camera.aspectRatio;
				const fovScale = Math.min(tanHalfVFov, tanHalfHFov) / REF_TAN_HALF_FOV;
				const invLogMult = 1.0 / Math.log(lodMultiplier);
				const worldCameraPosition = cameraNode.getPosition();
				const octreeWorldTransform = this.placement.node.getWorldTransform();
				_invWorldMat.copy(octreeWorldTransform).invert();
				const localCameraPosition = _invWorldMat.transformPoint(worldCameraPosition, _localCameraPos);
				const worldCameraForward = cameraNode.forward;
				const localCameraForward = _invWorldMat.transformVector(worldCameraForward, _localCameraFwd).normalize();
				const nodes = this.octree.nodes;
				const nodeInfos = this.nodeInfos;
				let totalSplats = 0;
				for(let nodeIndex = 0; nodeIndex < nodes.length; nodeIndex++){
						const node = nodes[nodeIndex];
						const nodeInfo = nodeInfos[nodeIndex];
						node.bounds.closestPoint(localCameraPosition, _dirToNode);
						_dirToNode.sub(localCameraPosition);
						const actualDistance = _dirToNode.length();
						let penalizedDistance = actualDistance;
						if (lodBehindPenalty > 1 && actualDistance > 0.01) {
								const dotOverDistance = localCameraForward.dot(_dirToNode) / actualDistance;
								if (dotOverDistance < 0) {
										const t = -dotOverDistance;
										const factor = 1 + t * (lodBehindPenalty - 1);
										penalizedDistance = actualDistance * factor;
								}
						}
						const fovAdjustedDistance = penalizedDistance * fovScale;
						let optimalLodIndex;
						if (fovAdjustedDistance < lodBaseDistance) {
								optimalLodIndex = 0;
						} else {
								const rawLod = 1 + Math.log(fovAdjustedDistance / lodBaseDistance) * invLogMult;
								optimalLodIndex = Math.min(maxLod, rawLod | 0);
						}
						if (optimalLodIndex < rangeMin) optimalLodIndex = rangeMin;
						if (optimalLodIndex > rangeMax) optimalLodIndex = rangeMax;
						nodeInfo.optimalLod = optimalLodIndex;
						nodeInfo.worldDistance = fovAdjustedDistance * uniformScale;
						const lod = nodes[nodeIndex].lods[optimalLodIndex];
						if (lod && lod.count) {
								totalSplats += lod.count;
						}
				}
				return totalSplats;
		}
		evaluateOptimalLods(cameraNode, params, budgetScale = 1) {
				const maxLod = this.octree.lodLevels - 1;
				const { lodBaseDistance, lodMultiplier } = this.placement;
				const { lodRangeMin, lodRangeMax } = params;
				const rangeMin = Math.max(0, Math.min(lodRangeMin ?? 0, maxLod));
				const rangeMax = Math.max(rangeMin, Math.min(lodRangeMax ?? maxLod, maxLod));
				this.rangeMin = rangeMin;
				this.rangeMax = rangeMax;
				const uniformScale = this.placement.node.getWorldTransform().getScale().x;
				const effectiveBase = lodBaseDistance * budgetScale;
				const effectiveMult = Math.max(1.2, lodMultiplier * Math.pow(budgetScale, -0.2));
				return this.evaluateNodeLods(cameraNode, maxLod, effectiveBase, effectiveMult, rangeMin, rangeMax, params, uniformScale);
		}
		applyLodChanges(maxLod, params) {
				const nodes = this.octree.nodes;
				const { lodUnderfillLimit = 0 } = params;
				for(let nodeIndex = 0; nodeIndex < nodes.length; nodeIndex++){
						const node = nodes[nodeIndex];
						const nodeInfo = this.nodeInfos[nodeIndex];
						const optimalLodIndex = nodeInfo.optimalLod;
						const currentLodIndex = nodeInfo.currentLod;
						const desiredLodIndex = this.selectDesiredLodIndex(node, optimalLodIndex, maxLod, lodUnderfillLimit);
						if (desiredLodIndex !== currentLodIndex) {
								const currentFileIndex = currentLodIndex >= 0 ? node.lods[currentLodIndex].fileIndex : -1;
								const desiredFileIndex = desiredLodIndex >= 0 ? node.lods[desiredLodIndex].fileIndex : -1;
								const wasVisible = currentFileIndex !== -1;
								const willBeVisible = desiredFileIndex !== -1;
								const pendingEntry = this.pendingDecrements.get(nodeIndex);
								if (pendingEntry) {
										if (pendingEntry.newFileIndex !== desiredFileIndex) {
												const prevPendingPlacement = this.filePlacements[pendingEntry.newFileIndex];
												if (prevPendingPlacement) {
														this.decrementFileRef(pendingEntry.newFileIndex, nodeIndex);
												}
												if (wasVisible && willBeVisible) {
														this.pendingDecrements.set(nodeIndex, {
																oldFileIndex: pendingEntry.oldFileIndex,
																newFileIndex: desiredFileIndex
														});
												} else {
														this.pendingDecrements.delete(nodeIndex);
												}
										}
								}
								if (!wasVisible && willBeVisible) {
										const prevPendingFi = this.pendingVisibleAdds.get(nodeIndex);
										if (prevPendingFi !== undefined && prevPendingFi !== desiredFileIndex) {
												this.decrementFileRef(prevPendingFi, nodeIndex);
												this.pendingVisibleAdds.delete(nodeIndex);
										}
										this.incrementFileRef(desiredFileIndex, nodeIndex, desiredLodIndex);
										const newPlacement = this.filePlacements[desiredFileIndex];
										if (newPlacement?.resource) {
												nodeInfo.currentLod = desiredLodIndex;
												this.pendingVisibleAdds.delete(nodeIndex);
										} else {
												this.pendingVisibleAdds.set(nodeIndex, desiredFileIndex);
										}
								} else if (wasVisible && !willBeVisible) {
										const pendingEntry2 = this.pendingDecrements.get(nodeIndex);
										if (pendingEntry2) {
												this.decrementFileRef(pendingEntry2.newFileIndex, nodeIndex);
												this.pendingDecrements.delete(nodeIndex);
										}
										this.decrementFileRef(currentFileIndex, nodeIndex);
										nodeInfo.currentLod = -1;
										this.pendingVisibleAdds.delete(nodeIndex);
								} else if (wasVisible && willBeVisible) {
										this.incrementFileRef(desiredFileIndex, nodeIndex, desiredLodIndex);
										const newPlacement = this.filePlacements[desiredFileIndex];
										if (newPlacement?.resource) {
												this.decrementFileRef(currentFileIndex, nodeIndex);
												this.pendingDecrements.delete(nodeIndex);
												nodeInfo.currentLod = desiredLodIndex;
												this.pendingVisibleAdds.delete(nodeIndex);
										} else {
												this.pendingDecrements.set(nodeIndex, {
														oldFileIndex: currentFileIndex,
														newFileIndex: desiredFileIndex
												});
												this.pendingVisibleAdds.delete(nodeIndex);
										}
								}
						}
						this.prefetchNextLod(node, desiredLodIndex, optimalLodIndex);
				}
		}
		incrementFileRef(fileIndex, nodeIndex, lodIndex) {
				if (fileIndex === -1) return;
				let placement = this.filePlacements[fileIndex];
				if (!placement) {
						placement = new GSplatPlacement(null, this.placement.node, lodIndex, null, this.placement);
						this.filePlacements[fileIndex] = placement;
						const removeScheduled = this.removedCandidates.delete(fileIndex);
						if (!removeScheduled) {
								this.octree.incRefCount(fileIndex);
						}
						if (!this.addFilePlacement(fileIndex)) {
								this.octree.ensureFileResource(fileIndex);
								this.pending.add(fileIndex);
						}
				}
				const nodes = this.octree.nodes;
				const node = nodes[nodeIndex];
				const lod = node.lods[lodIndex];
				const interval = new Vec2(lod.offset, lod.offset + lod.count - 1);
				placement.intervals.set(nodeIndex, interval);
				this.dirtyModifiedPlacements = true;
		}
		decrementFileRef(fileIndex, nodeIndex) {
				if (fileIndex === -1) return;
				const placement = this.filePlacements[fileIndex];
				if (!placement) {
						return;
				}
				if (placement) {
						placement.intervals.delete(nodeIndex);
						this.dirtyModifiedPlacements = true;
						if (placement.intervals.size === 0) {
								if (placement.resource) {
										this.activePlacements.delete(placement);
										if (this.activePlacements.size === 0) {
												this.dirtyPlacementSetChanged = true;
										}
								}
								this.removedCandidates.add(fileIndex);
								this.filePlacements[fileIndex] = null;
								this.pending.delete(fileIndex);
						}
				}
		}
		addFilePlacement(fileIndex) {
				const res = this.octree.getFileResource(fileIndex);
				if (res) {
						const placement = this.filePlacements[fileIndex];
						if (placement) {
								placement.resource = res;
								if (this.activePlacements.size === 0) {
										this.dirtyPlacementSetChanged = true;
								}
								this.activePlacements.add(placement);
								this.dirtyModifiedPlacements = true;
								this.removedCandidates.delete(fileIndex);
								return true;
						}
				}
				return false;
		}
		testMoved(threshold) {
				const position = this.placement.node.getPosition();
				const length = position.distance(this.previousPosition);
				if (length > threshold) {
						return true;
				}
				return false;
		}
		updateMoved() {
				this.previousPosition.copy(this.placement.node.getPosition());
		}
		update() {
				if (this.placement.lodDirty) {
						this.placement.lodDirty = false;
						this.needsLodUpdate = true;
				}
				if (this.pending.size) {
						for (const fileIndex of this.pending){
								this.octree.ensureFileResource(fileIndex);
								if (this.addFilePlacement(fileIndex)) {
										_tempCompletedUrls.push(fileIndex);
										for (const [nodeIndex, { oldFileIndex, newFileIndex }] of this.pendingDecrements){
												if (newFileIndex === fileIndex) {
														this.decrementFileRef(oldFileIndex, nodeIndex);
														this.pendingDecrements.delete(nodeIndex);
														let newLodIndex = 0;
														const nodeLods = this.octree.nodes[nodeIndex].lods;
														for(let li = 0; li < nodeLods.length; li++){
																if (nodeLods[li].fileIndex === newFileIndex) {
																		newLodIndex = li;
																		break;
																}
														}
														this.nodeInfos[nodeIndex].currentLod = newLodIndex;
												}
										}
								}
						}
						if (_tempCompletedUrls.length > 0) {
								this.needsLodUpdate = true;
						}
						for (const fileIndex of _tempCompletedUrls){
								this.pending.delete(fileIndex);
						}
						_tempCompletedUrls.length = 0;
				}
				this.pollPrefetchCompletions();
				if (this.octree.environmentUrl && !this.environmentPlacement) {
						this.octree.ensureEnvironmentResource();
						const envResource = this.octree.environmentResource;
						if (envResource) {
								this.environmentPlacement = new GSplatPlacement(envResource, this.placement.node, 0, null, this.placement);
								this.environmentPlacement.aabb.copy(envResource.aabb);
								this.activePlacements.add(this.environmentPlacement);
								this.dirtyModifiedPlacements = true;
								this.dirtyPlacementSetChanged = true;
						}
				}
				const dirty = this.dirtyModifiedPlacements;
				this.dirtyModifiedPlacements = false;
				return dirty;
		}
		consumePlacementSetChanged() {
				const changed = this.dirtyPlacementSetChanged;
				this.dirtyPlacementSetChanged = false;
				return changed;
		}
		debugRender(scene) {}
		consumeNeedsLodUpdate() {
				const v = this.needsLodUpdate;
				this.needsLodUpdate = false;
				return v;
		}
		pollPrefetchCompletions() {
				if (this.prefetchPending.size) {
						for (const fileIndex of this.prefetchPending){
								this.octree.ensureFileResource(fileIndex);
								if (this.octree.getFileResource(fileIndex)) {
										_tempCompletedUrls.push(fileIndex);
								}
						}
						if (_tempCompletedUrls.length > 0) {
								this.needsLodUpdate = true;
						}
						for (const fileIndex of _tempCompletedUrls){
								this.prefetchPending.delete(fileIndex);
						}
						_tempCompletedUrls.length = 0;
				}
		}
		constructor(device, octree, placement){
				this.activePlacements = new Set();
				this.dirtyModifiedPlacements = false;
				this.dirtyPlacementSetChanged = false;
				this.pending = new Set();
				this.pendingDecrements = new Map();
				this.removedCandidates = new Set();
				this.rangeMin = 0;
				this.rangeMax = 0;
				this.previousPosition = new Vec3();
				this.needsLodUpdate = false;
				this.prefetchPending = new Set();
				this.pendingVisibleAdds = new Map();
				this.environmentPlacement = null;
				this._deviceLostEvent = null;
				this.device = device;
				this.octree = octree;
				this.placement = placement;
				this.nodeInfos = new Array(octree.nodes.length);
				for(let i = 0; i < octree.nodes.length; i++){
						const nodeInfo = new NodeInfo();
						nodeInfo.inst = this;
						this.nodeInfos[i] = nodeInfo;
				}
				const numFiles = octree.files.length;
				this.filePlacements = new Array(numFiles).fill(null);
				if (octree.environmentUrl) {
						octree.incEnvironmentRefCount();
						octree.ensureEnvironmentResource();
				}
				this._deviceLostEvent = device.on('devicelost', this._onDeviceLost, this);
		}
}

export { GSplatOctreeInstance, NodeInfo };