rayzee 9.0.0 → 9.1.0

package/src/Stages/PathTracerStage.js

@@ -24,10 +24,15 @@ import { SceneProcessor } from '../Processor/SceneProcessor.js';
 import { LightSerializer } from '../Processor/LightSerializer';
 
 // Constants
-import { ENGINE_DEFAULTS as DEFAULT_STATE } from '../EngineDefaults.js';
+import { ENGINE_DEFAULTS as DEFAULT_STATE, TRIANGLE_DATA_LAYOUT } from '../EngineDefaults.js';
+import { TRI_GEO_ROWS, TRI_SHADE_ROWS } from '../TSL/Common.js';
 import { getAssetConfig } from '../AssetConfig.js';
+import { cachedObjectURL, DOWNLOAD_POLICY } from '../Storage/DownloadCache.js';
 import { createLogger, fmt } from '../utils/Logger.js';
 
+// Triangles converted per staging pass: 256K is 12 + 8 MB of staging, kept for the stage's life.
+const TRIANGLE_UPLOAD_SLICE = 1 << 18;
+
 const log = createLogger( 'pathtracer' );
 
 // How long after the last move the view still counts as moving. Matches CameraOptimizer's own
@@ -114,26 +119,22 @@ export class PathTracerStage extends RenderStage {
 		// Initialize material data manager
 		this.materialData = new MaterialDataManager( this.sdfs );
 		this.materialData.callbacks.onReset = () => this.reset();
-		// Triangle data carries the per-triangle `side` flag (NORMAL_C.w). The authoritative
-		// CPU copy is `_triangleRecords` when the scene is chunked and triangleStorageAttr.array
-		// otherwise (not sdfs.triangleData, which isn't populated on the PathTracerApp build
-		// path). The patch mutates it in place, then the GPU copy is refreshed.
+		// Material edits patch the per-triangle flags lane in the CPU records (`_triangleRecords`
+		// when chunked, `_triangleFlat` otherwise — not sdfs.triangleData, which the PathTracerApp
+		// build path never fills). The flags live in the shade rows, so only those go up again.
 		this.materialData.callbacks.getTriangleData = () => ( {
-			array: this.triangleStorageAttr?.array,
+			array: this._triangleFlat,
 			records: this._triangleRecords,
 			count: this.triangleCount,
 		} );
 		this.materialData.callbacks.onTriangleDataChanged = () => {
 
-			if ( ! this.triangleStorageAttr ) return;
-			// A chunked attribute owns no CPU array for three.js to re-upload.
-			if ( this._triangleRecords ) uploadStorageChunks( this.renderer, this.triangleStorageAttr, this._triangleRecords.chunks );
-			else this.triangleStorageAttr.needsUpdate = true;
+			if ( this.triangleShadeAttr ) this._uploadTriangles( 0, this._triangleRecordCount, { geo: false } );
 
 		};
 
 		// Initialize environment manager
-		this.environment = new EnvironmentManager( this.scene, this.uniforms );
+		this.environment = new EnvironmentManager( this.scene, this.uniforms, this.renderer );
 		this.environment.callbacks.onReset = () => this.reset();
 		this.environment.callbacks.getSceneTextureNodes = () => this.shaderBuilder.getSceneTextureNodes();
 
@@ -183,9 +184,13 @@ export class PathTracerStage extends RenderStage {
 	 */
 	_initDataTextures() {
 
-		// Triangle data (storage buffer for WebGPU)
-		this.triangleStorageAttr = null;
+		// Triangle data: two GPU buffers (TRI_GEO_ROWS / TRI_SHADE_ROWS rows a triangle), one
+		// { geo, shade } node pair, and the CPU records they were written from.
+		this.triangleGeoAttr = null;
+		this.triangleShadeAttr = null;
 		this._triangleRecords = null;
+		this._triangleFlat = null;
+		this._triangleRecordCount = 0;
 		this.triangleStorageNode = null;
 		this.triangleCount = 0;
 
@@ -238,6 +243,8 @@ export class PathTracerStage extends RenderStage {
 		// Per-triangle bit-trail map (root→leaf Light BVH path, 1 float per triangleIndex); packed
 		// after the emissive entries so the bounce-hit MIS path can re-walk the descent pdf.
 		this._bitTrailMapCache = null;
+		this._bitTrailOffset = 0;
+		this._bitTrailLength = 0;
 
 		// Per-mesh visibility is packed into the TLAS BLAS-pointer leaf's slot [2]
 		// (see TLASBuilder.flatten + BVHTraversal.js). The InstanceTable holds the
@@ -384,7 +391,22 @@ export class PathTracerStage extends RenderStage {
 
 		const { stbnScalarAtlas, stbnVec2Atlas } = getAssetConfig();
 
-		const scalarLoad = loader.loadAsync( stbnScalarAtlas ).then( ( tex ) => {
+		const load = async ( url ) => {
+
+			const source = await cachedObjectURL( url, { policy: DOWNLOAD_POLICY.REVALIDATE } );
+			try {
+
+				return await loader.loadAsync( source.url );
+
+			} finally {
+
+				source.release();
+
+			}
+
+		};
+
+		const scalarLoad = load( stbnScalarAtlas ).then( ( tex ) => {
 
 			this.stbnScalarTexture = configure( tex );
 			stbnScalarTextureNode.value = tex;
@@ -392,7 +414,7 @@ export class PathTracerStage extends RenderStage {
 
 		} );
 
-		const vec2Load = loader.loadAsync( stbnVec2Atlas ).then( ( tex ) => {
+		const vec2Load = load( stbnVec2Atlas ).then( ( tex ) => {
 
 			this.stbnVec2Texture = configure( tex );
 			stbnVec2TextureNode.value = tex;
@@ -543,36 +565,6 @@ export class PathTracerStage extends RenderStage {
 		this.spotLightsData = mockMaterial.uniforms.spotLights.value;
 		this.areaLightsData = mockMaterial.uniforms.areaLights.value;
 
-		// Add sun as directional light if procedural sky is active
-		if ( this.hasSun.value ) {
-
-			const scaledSunIntensity = this.environment.envParams.skySunIntensity * 950.0;
-
-			const sunLight = {
-				intensity: scaledSunIntensity,
-				color: { r: 1.0, g: 1.0, b: 1.0 },
-				userData: {
-					angle: this.sunAngularSize.value
-				},
-				updateMatrixWorld: () => {},
-				getWorldPosition: ( target ) => {
-
-					const sunDir = this.sunDirection.value;
-					return target.set( sunDir.x, sunDir.y, sunDir.z ).multiplyScalar( 1e10 );
-
-				}
-			};
-
-			this.lightSerializer.addDirectionalLight( sunLight );
-			this.lightSerializer.preprocessLights();
-			this.lightSerializer.updateShaderUniforms( mockMaterial );
-
-			this.directionalLightsData = mockMaterial.uniforms.directionalLights.value;
-
-			log.debug( `sun added as directional light · intensity ${scaledSunIntensity.toFixed( 2 )}` );
-
-		}
-
 		// Update TSL uniform buffer nodes from raw Float32Array data
 		this._updateLightBufferNodes();
 
@@ -729,54 +721,212 @@ export class PathTracerStage extends RenderStage {
 	// ===== TEXTURE SETTERS =====
 
 	/**
-	 * Sets the triangle data from raw Float32Array via storage buffer.
-	 * On first call, creates the storage buffer and node.
-	 * On subsequent calls, creates a new attribute with the correct size
-	 * and updates the storage node's value to preserve shader graph references.
-	 * @param {Uint32Array} triangleData - Packed triangle records (uvec4 lanes)
+	 * Sets the triangle data. The records stay whole on the CPU and are split into the geo and
+	 * shade GPU buffers as they go up. Later calls swap the nodes' attributes, which keeps the
+	 * compiled shader graph.
+	 * @param {Uint32Array|import('../Processor/ChunkedRecords.js').ChunkedRecords} triangleData
 	 * @param {number} triangleCount - Number of triangles
 	 */
 	setTriangleData( triangleData, triangleCount ) {
 
 		if ( ! triangleData ) return;
 
-		// Past the ~2 GB array cap the records arrive as several chunks; they still become one
-		// GPU buffer, written at their running byte offsets, so no binding or shader changes.
 		const chunked = triangleData.chunks && triangleData.chunks.length > 1 ? triangleData : null;
 		const flat = chunked ? null : ( triangleData.chunks ? triangleData.chunks[ 0 ] : triangleData );
-		const lanes = chunked ? chunked.recordCount * chunked.lanesPerRecord : flat.length;
-		const vec4Count = lanes / 4;
-
-		const makeAttr = () => chunked
-			? gpuOnlyStorageAttribute( vec4Count, 4, Uint32Array )
-			: new StorageInstancedBufferAttribute( flat, 4 );
+		const records = chunked ? chunked.recordCount : Math.ceil( flat.length / TRIANGLE_DATA_LAYOUT.FLOATS_PER_TRIANGLE );
+		const pre = chunked?._gpuUpload?.triangles ? chunked._gpuUpload : null;
+		const { geo, shade } = pre ?? this._allocateTriangleBuffers( records );
 
 		if ( this.triangleStorageNode ) {
 
-			// Create new attribute with correct size (old one is GC'd, backend WeakMap cleans up GPU buffer)
-			this.triangleStorageAttr = makeAttr();
-
-			// Update storage node references (preserves compiled shader graph)
-			this.triangleStorageNode.value = this.triangleStorageAttr;
-			this.triangleStorageNode.bufferCount = vec4Count;
+			this.triangleStorageNode.geo.value = geo;
+			this.triangleStorageNode.geo.bufferCount = geo.count;
+			this.triangleStorageNode.shade.value = shade;
+			this.triangleStorageNode.shade.bufferCount = shade.count;
 
 		} else {
 
-			// First time: create storage buffer and node
-			this.triangleStorageAttr = makeAttr();
-			this.triangleStorageNode = storage( this.triangleStorageAttr, 'uvec4', vec4Count ).toReadOnly();
+			this.triangleStorageNode = {
+				geo: storage( geo, 'uvec4', geo.count ).toReadOnly(),
+				shade: storage( shade, 'uvec4', shade.count ).toReadOnly(),
+			};
 
 		}
 
+		this.triangleGeoAttr = geo;
+		this.triangleShadeAttr = shade;
 		this._triangleRecords = chunked;
-		if ( chunked ) uploadStorageChunks( this.renderer, this.triangleStorageAttr, chunked.chunks );
-
+		this._triangleFlat = flat;
+		this._triangleRecordCount = records;
 		this.triangleCount = triangleCount;
+		this._uploadTriangles( 0, records, { skipChunks: pre?.uploaded } );
 
 		log.debug( `${fmt.n( this.triangleCount )} triangles (storage buffer)` );
 
 	}
 
+	/** GPU-only geo and shade buffers for `records` triangles, created now. @private */
+	_allocateTriangleBuffers( records ) {
+
+		const n = Math.max( 1, records );
+		this._assertFitsGPU( n * TRI_GEO_ROWS * 16 );
+		const geo = gpuOnlyStorageAttribute( n * TRI_GEO_ROWS, 4, Uint32Array );
+		const shade = gpuOnlyStorageAttribute( n * TRI_SHADE_ROWS, 4, Uint32Array );
+		this.renderer.backend.createStorageAttribute( geo );
+		this.renderer.backend.createStorageAttribute( shade );
+		return { geo, shade };
+
+	}
+
+	/**
+	 * Writes triangles [start, start + count) of the CPU records into the geo and shade buffers:
+	 * the first TRI_GEO_ROWS rows of each record to one, the rest to the other.
+	 * @param {number} start
+	 * @param {number} count
+	 * @param {{ geo?: boolean, shade?: boolean, skipChunks?: Set<number>, source?: Uint32Array, sourceStart?: number }} [options]
+	 *   `source` writes those records instead of reading the stage's own (a chunk being handed over).
+	 * @private
+	 */
+	_uploadTriangles( start, count, { geo = true, shade = true, skipChunks = null, source = null, sourceStart = 0, geoAttr = this.triangleGeoAttr, shadeAttr = this.triangleShadeAttr } = {} ) {
+
+		if ( count <= 0 || ! geoAttr ) return;
+
+		const backend = this.renderer.backend;
+		const geoBuffer = backend.get( geoAttr ).buffer;
+		const shadeBuffer = backend.get( shadeAttr ).buffer;
+		const LANES = TRIANGLE_DATA_LAYOUT.FLOATS_PER_TRIANGLE;
+		const GEO = TRI_GEO_ROWS * 4, SHADE = TRI_SHADE_ROWS * 4;
+		const records = this._triangleRecords;
+		const perChunk = source ? count : records ? records.recordsPerChunk : this._triangleRecordCount;
+		// One staging pair for every upload: a fresh one per chunk was churn that failed to
+		// allocate on a scene already at the edge of the address space.
+		const slice = TRIANGLE_UPLOAD_SLICE;
+		this._triangleStaging ??= { geo: new Uint32Array( slice * GEO ), shade: new Uint32Array( slice * SHADE ) };
+		const geoStage = this._triangleStaging.geo;
+		const shadeStage = this._triangleStaging.shade;
+
+		for ( let at = start, end = start + count; at < end; ) {
+
+			const k = source ? 0 : Math.floor( at / perChunk );
+			const chunkStart = source ? start : k * perChunk;
+			const chunkEnd = Math.min( end, chunkStart + perChunk );
+			if ( skipChunks?.has( k ) ) {
+
+				at = chunkEnd;
+				continue;
+
+			}
+
+			const src = source ?? ( records ? records.chunks[ k ] : this._triangleFlat );
+			if ( ! src ) throw new Error( `triangle chunk ${k} is on disk and was never uploaded` );
+
+			const n = Math.min( slice, chunkEnd - at );
+			let s = ( source ? at - start + sourceStart : at - chunkStart ) * LANES;
+			for ( let i = 0, g = 0, h = 0; i < n; i ++, s += LANES ) {
+
+				if ( geo ) for ( let l = 0; l < GEO; l ++ ) geoStage[ g ++ ] = src[ s + l ];
+				if ( shade ) for ( let l = GEO; l < LANES; l ++ ) shadeStage[ h ++ ] = src[ s + l ];
+
+			}
+
+			if ( geo ) backend.device.queue.writeBuffer( geoBuffer, at * GEO * 4, geoStage, 0, n * GEO );
+			if ( shade ) backend.device.queue.writeBuffer( shadeBuffer, at * SHADE * 4, shadeStage, 0, n * SHADE );
+			at += n;
+
+		}
+
+	}
+
+	/**
+	 * Allocates the GPU buffer for chunked `records` now and returns a function that writes one
+	 * chunk into it, so a build can hand chunks over as it finishes them and let them go (the
+	 * memory spill). setTriangleData / setBVHData then adopt that buffer rather than upload again.
+	 * @param {import('../Processor/ChunkedRecords.js').ChunkedRecords} records
+	 * @returns {function(number): void} uploads chunk k from `records.chunks[ k ]`
+	 */
+	createChunkUploader( records ) {
+
+		this._assertFitsGPU( records.recordCount * records.lanesPerRecord * records.LaneType.BYTES_PER_ELEMENT );
+		const attr = gpuOnlyStorageAttribute( records.recordCount * records.lanesPerRecord / 4, 4, records.LaneType );
+		const backend = this.renderer.backend;
+		backend.createStorageAttribute( attr );
+		const buffer = backend.get( attr ).buffer;
+		const stride = records.recordsPerChunk * records.lanesPerRecord * records.LaneType.BYTES_PER_ELEMENT;
+		const uploaded = new Set();
+		records._gpuUpload = { attr, uploaded };
+
+		return k => {
+
+			const chunk = records.chunks[ k ];
+			backend.device.queue.writeBuffer( buffer, k * stride, chunk, 0, chunk.length );
+			uploaded.add( k );
+
+		};
+
+	}
+
+	/**
+	 * createChunkUploader for triangle records: both GPU buffers are allocated now and each
+	 * chunk is split into them as it is handed over.
+	 * @param {import('../Processor/ChunkedRecords.js').ChunkedRecords} records
+	 * @returns {function(number): void}
+	 */
+	createTriangleChunkUploader( records ) {
+
+		const { geo, shade } = this._allocateTriangleBuffers( records.recordCount );
+		const uploaded = new Set();
+		records._gpuUpload = { triangles: true, geo, shade, uploaded };
+		const perChunk = records.recordsPerChunk;
+
+		return k => {
+
+			const first = k * perChunk;
+			const count = Math.min( perChunk, records.recordCount - first );
+			this._uploadTriangles( first, count, { source: records.chunks[ k ], geoAttr: geo, shadeAttr: shade } );
+			uploaded.add( k );
+
+		};
+
+	}
+
+	/**
+	 * One storage buffer holds all of a scene's triangles, and one its BVH. Past the device's
+	 * limit WebGPU hands back an invalid buffer and every write to it fails quietly.
+	 * @private
+	 */
+	_assertFitsGPU( bytes ) {
+
+		const limits = this.renderer?.backend?.device?.limits;
+		const max = Math.min( limits?.maxBufferSize ?? Infinity, limits?.maxStorageBufferBindingSize ?? Infinity );
+		if ( bytes > max ) throw new Error( `the scene needs a ${fmt.mb( bytes )} GPU buffer, past the ${fmt.mb( max )} this GPU allows for one — load fewer triangles` );
+
+	}
+
+	/** Uploads chunked records, or — for a buffer a build already filled — whatever it has not. @private */
+	_uploadChunked( attr, records ) {
+
+		const pre = records._gpuUpload;
+		if ( ! pre || pre.attr !== attr ) {
+
+			uploadStorageChunks( this.renderer, attr, records.chunks );
+			return;
+
+		}
+
+		const buffer = this.renderer.backend.get( attr ).buffer;
+		const stride = records.recordsPerChunk * records.lanesPerRecord * records.LaneType.BYTES_PER_ELEMENT;
+		for ( let k = 0; k < records.chunks.length; k ++ ) {
+
+			if ( pre.uploaded.has( k ) ) continue;
+			const chunk = records.chunks[ k ];
+			if ( ! chunk ) throw new Error( `chunk ${k} is on disk and was never uploaded` );
+			this.renderer.backend.device.queue.writeBuffer( buffer, k * stride, chunk, 0, chunk.length );
+			pre.uploaded.add( k );
+
+		}
+
+	}
+
 	/**
 	 * Sets the BVH data from raw Float32Array via storage buffer.
 	 * @param {Float32Array} bvhImageData - Raw BVH data from DataTexture.image.data
@@ -791,9 +941,10 @@ export class PathTracerStage extends RenderStage {
 		const flat = chunked ? null : ( bvhImageData.chunks ? bvhImageData.chunks[ 0 ] : bvhImageData );
 		const lanes = chunked ? chunked.recordCount * chunked.lanesPerRecord : flat.length;
 		const vec4Count = lanes / 4;
+		this._assertFitsGPU( lanes * 4 );
 
 		const makeAttr = () => chunked
-			? gpuOnlyStorageAttribute( vec4Count, 4, Float32Array )
+			? chunked._gpuUpload?.attr ?? gpuOnlyStorageAttribute( vec4Count, 4, Float32Array )
 			: new StorageInstancedBufferAttribute( flat, 4 );
 
 		if ( this.bvhStorageNode ) {
@@ -809,8 +960,9 @@ export class PathTracerStage extends RenderStage {
 
 		}
 
+		this.bvhStorageAttr.foldedLeaves = bvhImageData.foldedLeaves === true;
 		this._bvhRecords = chunked;
-		if ( chunked ) uploadStorageChunks( this.renderer, this.bvhStorageAttr, chunked.chunks );
+		if ( chunked ) this._uploadChunked( this.bvhStorageAttr, chunked );
 
 		this.bvhNodeCount = Math.floor( vec4Count / BVH_VEC4_PER_NODE );
 		log.debug( `${fmt.n( this.bvhNodeCount )} BVH nodes (storage buffer)` );
@@ -1019,17 +1171,12 @@ export class PathTracerStage extends RenderStage {
 
 	}
 
-	/** Update triangle positions in the existing GPU buffer (full). */
+	/** Re-uploads every triangle from the CPU records, adopting `triangleData` as them when given. */
 	updateTriangleData( triangleData ) {
 
-		if ( this._triangleRecords ) {
-
-			uploadStorageChunks( this.renderer, this.triangleStorageAttr, this._triangleRecords.chunks );
-			return;
-
-		}
-
-		this._updateStorageBuffer( this.triangleStorageAttr, triangleData?.chunks ? triangleData.chunks[ 0 ] : triangleData );
+		if ( triangleData?.chunks?.length > 1 ) this._triangleRecords = triangleData;
+		else if ( triangleData ) this._triangleFlat = triangleData.chunks ? triangleData.chunks[ 0 ] : triangleData;
+		this._uploadTriangles( 0, this._triangleRecordCount );
 
 	}
 
@@ -1056,29 +1203,14 @@ export class PathTracerStage extends RenderStage {
 	 */
 	updateBufferRanges( triRanges, bvhRanges ) {
 
-		if ( this.triangleStorageAttr && triRanges.length > 0 ) {
-
-			if ( this._triangleRecords ) {
-
-				// Chunked: three.js cannot re-upload an attribute with no CPU array, so write
-				// the dirty lanes straight into the GPU buffer.
-				for ( const r of triRanges ) {
-
-					uploadStorageChunkRange( this.renderer, this.triangleStorageAttr, this._triangleRecords, r.offset, r.count );
+		if ( this.triangleGeoAttr && triRanges.length > 0 ) {
 
-				}
-
-			} else {
+			// Ranges are in lanes of the CPU records; the GPU copy is written a triangle at a time.
+			const lanes = TRIANGLE_DATA_LAYOUT.FLOATS_PER_TRIANGLE;
+			for ( const r of triRanges ) {
 
-				this.triangleStorageAttr.clearUpdateRanges();
-
-				for ( const r of triRanges ) {
-
-					this.triangleStorageAttr.addUpdateRange( r.offset, r.count );
-
-				}
-
-				this.triangleStorageAttr.version ++;
+				const first = Math.floor( r.offset / lanes );
+				this._uploadTriangles( first, Math.ceil( ( r.offset + r.count ) / lanes ) - first );
 
 			}
 
@@ -1438,7 +1570,12 @@ export class PathTracerStage extends RenderStage {
 		const LBVH_STRIDE = 4; // vec4s per LBVH node — must match LightBVHSampling.js
 		const lbvh = this._lbvhDataCache;
 		const emis = this._emissiveDataCache;
-		const trail = this._bitTrailMapCache;
+		// The bit-trail map (4 B a triangle) lives only in the light buffer: without a new one, the
+		// current one is carried over from there rather than kept as a second copy.
+		const held = this.lightStorageAttr?.array;
+		const trail = this._bitTrailMapCache
+			?? ( held && this._bitTrailLength ? held.slice( this._bitTrailOffset, this._bitTrailOffset + this._bitTrailLength ) : null );
+		this._bitTrailMapCache = null;
 		const lbvhLen = lbvh ? lbvh.length : 0;
 		const emisLen = emis ? emis.length : 0;
 		// Bit-trail map packs 4 trails per vec4 → pad to a vec4 boundary.
@@ -1479,6 +1616,8 @@ export class PathTracerStage extends RenderStage {
 		// Offset (in vec4 elements) where the bit-trail map starts (lbvhLen + emisLen are float
 		// counts, both multiples of 4, so this divides cleanly).
 		this.reverseMapVec4Offset.value = ( lbvhLen + emisLen ) / 4;
+		this._bitTrailOffset = lbvhLen + emisLen;
+		this._bitTrailLength = trail ? trail.length : 0;
 
 	}
 
@@ -1548,7 +1687,11 @@ export class PathTracerStage extends RenderStage {
 			log.debug( 'material rebuild started' );
 
 			await this.sdfs.rebuildMaterials( scene );
-			this.updateSceneUniforms();
+			// Materials and textures only. `sdfs` is not the processor that built the scene
+			// (PathTracerApp._sdf is), so re-uploading everything from it put its empty emissive
+			// data and instance table in place of the scene's: emitters stopped being sampled.
+			this.materialData.setMaterialData( this.sdfs.materialData, this.sdfs.materials?.map( m => m.sources ) );
+			this.materialData.loadTexturesFromSdfs();
 			this.shaderBuilder.updateSceneTextures( this );
 			this.updateLights();
 			this.reset();
@@ -1607,8 +1750,12 @@ export class PathTracerStage extends RenderStage {
 		this.placeholderTexture?.dispose();
 
 		// Clear data references
-		this.triangleStorageAttr = null;
+		this.triangleGeoAttr = null;
+		this.triangleShadeAttr = null;
+		this._triangleStaging = null;
 		this._triangleRecords = null;
+		this._triangleFlat = null;
+		this._triangleRecordCount = 0;
 		this.triangleStorageNode = null;
 		this.bvhStorageAttr = null;
 		this._bvhRecords = null;