three/examples/jsm/tsl/shadows/TileShadowNode.js

JavaScript 3D library

import {
	Vector3,
	Object3D,
	ShadowBaseNode,
	Plane,
	Line3,
	DepthTexture,
	Compatibility,
	LessCompare,
	GreaterEqualCompare,
	LinearFilter,
	NearestFilter,
	PCFShadowMap,
	Vector2,
	RedFormat,
	ArrayCamera,
	VSMShadowMap,
	RendererUtils,
	Quaternion,
	UnsignedIntType
} from 'three/webgpu';

import { Fn, int, lightShadowMatrix, mix, normalWorld, reference, renderGroup, shadow, shadowPositionWorld, texture, vec2, vec3, vec4, NodeUpdateType } from 'three/tsl';

const { resetRendererAndSceneState, restoreRendererAndSceneState } = RendererUtils;
let _rendererState;

const _cameraLayers = [];
const _vec3Temp1 = /*@__PURE__*/ new Vector3();
const _vec3Temp2 = /*@__PURE__*/ new Vector3();
const _vec3Temp3 = /*@__PURE__*/ new Vector3();
const _quatTemp1 = /*@__PURE__*/ new Quaternion();

class LwLight extends Object3D {

	constructor() {

		super();
		this.target = new Object3D();

	}

}

/**
 * A class that extends `ShadowBaseNode` to implement tiled shadow mapping.
 * This allows splitting a shadow map into multiple tiles, each with its own light and camera,
 * to improve shadow quality and performance for large scenes.
 *
 * **Note:** This class does not support `VSMShadowMap` at the moment.
 *
 * @class
 * @augments ShadowBaseNode
 * @three_import import { TileShadowNode } from 'three/addons/tsl/shadows/TileShadowNode.js';
 */
class TileShadowNode extends ShadowBaseNode {

	/**
	 * Creates an instance of `TileShadowNode`.
	 *
	 * @param {Light} light - The original light source used for shadow mapping.
	 * @param {Object} [options={}] - Configuration options for the tiled shadow node.
	 * @param {number} [options.tilesX=2] - The number of tiles along the X-axis.
	 * @param {number} [options.tilesY=2] - The number of tiles along the Y-axis.
	 * @param {Object} [options.resolution] - The resolution of the shadow map.
	 * @param {number} [options.depthType=UnsignedIntType] - The type of the tile depth texture. `UnsignedShortType` halves the memory when the depth range allows it.
	 * @param {boolean} [options.debug=false] - Whether to enable debug mode.
	 */
	constructor( light, options = {} ) {

		super( light );

		// Default configuration with sensible defaults
		this.config = {
			tilesX: options.tilesX || 2,
			tilesY: options.tilesY || 2,
			resolution: options.resolution || light.shadow.mapSize,
			depthType: options.depthType || UnsignedIntType,
			debug: options.debug !== undefined ? options.debug : false
		};

		this.debug = this.config.debug;

		this.originalLight = light;
		this.lightPlane = new Plane( new Vector3( 0, 1, 0 ), 0 );
		this.line = new Line3();

		this.initialLightDirection = new Vector3();
		this.updateLightDirection();

		this._cameraFrameId = new WeakMap();

		this.shadowSize = {
			top: light.shadow.camera.top,
			bottom: light.shadow.camera.bottom,
			left: light.shadow.camera.left,
			right: light.shadow.camera.right,
		};

		this.lights = [];
		this._shadowNodes = [];

		this.tiles = this.generateTiles( this.config.tilesX, this.config.tilesY );

	}

	/**
	 * Generates the tiles for the shadow map based on the specified number of tiles along the X and Y axes.
	 *
	 * @param {number} tilesX - The number of tiles along the X-axis.
	 * @param {number} tilesY - The number of tiles along the Y-axis.
	 * @returns {Array<Object>} An array of tile objects, each containing the tile's bounds and index.
	 */
	generateTiles( tilesX, tilesY ) {

		const tiles = [];
		const tileWidth = 1 / tilesX;
		const tileHeight = 1 / tilesY;

		for ( let y = 0; y < tilesY; y ++ ) {

			for ( let x = 0; x < tilesX; x ++ ) {

				tiles.push( {
					x: [ x * tileWidth, ( x + 1 ) * tileWidth ],
					y: [ ( tilesY - 1 - y ) * tileHeight, ( tilesY - y ) * tileHeight ], // Start from top row
					index: y * tilesX + x
				} );

			}

		}

		return tiles;

	}

	/**
	 * Updates the initial light direction based on the light's target position.
	 */
	updateLightDirection() {

		this.initialLightDirection.subVectors(
			this.originalLight.target.getWorldPosition( new Vector3() ),
			this.originalLight.getWorldPosition( new Vector3() )
		).normalize();

	}

	/**
	 * Initializes the tiled shadow node by creating lights, cameras, and shadow maps for each tile.
	 *
	 * @param {Builder} builder - The builder used to create render targets and other resources.
	 */
	init( builder ) {

		const light = this.originalLight;
		const parent = light.parent;

		const width = this.shadowSize.right - this.shadowSize.left;
		const height = this.shadowSize.top - this.shadowSize.bottom;

		const tileCount = this.tiles.length;
		const shadowWidth = this.config.resolution.width;
		const shadowHeight = this.config.resolution.height;

		// Clear existing lights/nodes if re-initializing
		this.disposeLightsAndNodes();

		const depthTexture = new DepthTexture( shadowWidth, shadowHeight, this.config.depthType, undefined, undefined, undefined, undefined, undefined, undefined, undefined, tileCount );
		depthTexture.compareFunction = builder.renderer.reversedDepthBuffer ? GreaterEqualCompare : LessCompare;
		depthTexture.name = 'ShadowDepthArrayTexture';
		// The color attachment is only sampled for transmitted shadows.
		const count = builder.renderer.shadowMap.transmitted === true ? 1 : 0;
		const shadowMap = builder.createRenderTarget( shadowWidth, shadowHeight, { count, format: RedFormat, depth: tileCount, useArrayDepthTexture: true } );
		shadowMap.depthTexture = depthTexture;
		if ( count > 0 ) shadowMap.texture.name = 'ShadowTexture';
		this.shadowMap = shadowMap;
		const cameras = [];


		// Create lights, one for each tile
		for ( let i = 0; i < tileCount; i ++ ) {

			const lwLight = new LwLight();
			lwLight.castShadow = true;
			const lShadow = light.shadow.clone();
			lShadow.filterNode = light.shadow.filterNode;
			const tile = this.tiles[ i ];
			lShadow.camera.left = this.shadowSize.left + width * tile.x[ 0 ];
			lShadow.camera.right = this.shadowSize.left + width * tile.x[ 1 ];
			lShadow.camera.top = this.shadowSize.bottom + height * tile.y[ 1 ];
			lShadow.camera.bottom = this.shadowSize.bottom + height * tile.y[ 0 ];
			lShadow.bias = light.shadow.bias;
			lShadow.camera.near = light.shadow.camera.near;
			lShadow.camera.far = light.shadow.camera.far;
			lShadow.camera.userData.tileIndex = i;
			lwLight.shadow = lShadow;

			if ( parent ) {

				parent.add( lwLight );
				parent.add( lwLight.target );

			} else {

				console.warn( 'TileShadowNode: Original light has no parent during init. Tile lights not added to scene graph directly.' );

			}

			this.syncLightTransformation( lwLight, light );

			this.lights.push( lwLight );
			lShadow.camera.updateMatrixWorld();

			cameras.push( lShadow.camera );
			const shadowNode = shadow( lwLight, lShadow );
			shadowNode.depthLayer = i;
			shadowNode.updateBeforeType = NodeUpdateType.NONE;

			shadowNode.setupRenderTarget = () => {

				return { shadowMap, depthTexture };

			};

			this._shadowNodes.push( shadowNode );

		}

		const cameraArray = new ArrayCamera( cameras );
		this.cameraArray = cameraArray;

		// setup() reads depth through this camera.
		light.shadow.camera.coordinateSystem = builder.renderer.coordinateSystem;
		light.shadow.camera._reversedDepth = builder.renderer.reversedDepthBuffer;

	}

	/**
	 * Updates the light transformations and shadow cameras for each tile.
	 */
	update() {

		const light = this.originalLight;

		const shadowCam = light.shadow.camera;
		const lsMin = new Vector2( shadowCam.left, shadowCam.bottom );
		const lsMax = new Vector2( shadowCam.right, shadowCam.top );
		const fullWidth = lsMax.x - lsMin.x;
		const fullHeight = lsMax.y - lsMin.y;

		for ( let i = 0; i < this.lights.length; i ++ ) {

			const lwLight = this.lights[ i ];
			const tile = this.tiles[ i ];
			this.syncLightTransformation( lwLight, light );
			const lShadow = lwLight.shadow;
			const tileLeft = lsMin.x + tile.x[ 0 ] * fullWidth;
			const tileRight = lsMin.x + tile.x[ 1 ] * fullWidth;
			const tileBottom = lsMin.y + tile.y[ 0 ] * fullHeight;
			const tileTop = lsMin.y + tile.y[ 1 ] * fullHeight;
			lShadow.camera.left = tileLeft;
			lShadow.camera.right = tileRight;
			lShadow.camera.bottom = tileBottom;
			lShadow.camera.top = tileTop;
			lShadow.camera.near = light.shadow.camera.near;
			lShadow.camera.far = light.shadow.camera.far;
			lShadow.camera.updateProjectionMatrix();
			lShadow.camera.updateWorldMatrix( true, false );
			lShadow.camera.updateMatrixWorld( true );
			this._shadowNodes[ i ].shadow.needsUpdate = true;

		}

		// The original light does not render a shadow map, but setup() uses its shadow matrix.

		shadowCam.updateProjectionMatrix();
		light.shadow.updateMatrices( light );

	}

	/**
     * Updates the shadow map rendering.
     * @param {NodeFrame} frame - A reference to the current node frame.
     */
	updateShadow( frame ) {

		const { shadowMap, light } = this;
		const { renderer, scene, camera } = frame;
		const shadowType = renderer.shadowMap.type;
		const currentRenderObjectFunction = renderer.getRenderObjectFunction();

		_rendererState = resetRendererAndSceneState( renderer, scene, _rendererState );
		scene.overrideMaterial = this.getShadowMaterial();
		renderer.setRenderTarget( this.shadowMap );

		const cameraArrayLayers = this.cameraArray.layers.mask;
		this.cameraArray.layers.mask = 0;

		for ( let index = 0; index < this.lights.length; index ++ ) {

			const light = this.lights[ index ];
			const shadow = light.shadow;

			const _shadowCameraLayer = shadow.camera.layers.mask;
			_cameraLayers.push( _shadowCameraLayer );

			if ( ( shadow.camera.layers.mask & 0xFFFFFFFE ) === 0 ) {

				shadow.camera.layers.mask = camera.layers.mask;

			}

			shadow.updateMatrices( light );

			// The tiles render through the array camera, so it has to see every tile's layers.
			this.cameraArray.layers.mask |= shadow.camera.layers.mask;

			renderer.setRenderObjectFunction( this.getShadowRenderObjectFunction( renderer, shadow ) );
			this.shadowMap.setSize( shadow.mapSize.width, shadow.mapSize.height, shadowMap.depth );

		}

		renderer.render( scene, this.cameraArray );
		renderer.setRenderObjectFunction( currentRenderObjectFunction );

		if ( light.isPointLight !== true && shadowType === VSMShadowMap ) {

			console.warn( 'THREE.TileShadowNode: VSM shadow map is not supported yet.' );
			// this.vsmPass( renderer );

		}

		restoreRendererAndSceneState( renderer, scene, _rendererState );

		this.cameraArray.layers.mask = cameraArrayLayers;

		for ( let index = 0; index < this.lights.length; index ++ ) {

			const light = this.lights[ index ];
			const shadow = light.shadow;

			shadow.camera.layers.mask = _cameraLayers[ index ];

		}

		_cameraLayers.length = 0;

	}

	/**
	 * The implementation performs the update of the shadow map if necessary.
	 *
	 * @param {NodeFrame} frame - A reference to the current node frame.
	 */
	updateBefore( frame ) {

		const shadow = this.originalLight.shadow;

		let needsUpdate = shadow.needsUpdate || shadow.autoUpdate;

		if ( needsUpdate ) {

			if ( this._cameraFrameId[ frame.camera ] === frame.frameId ) {

				needsUpdate = false;

			}

			this._cameraFrameId[ frame.camera ] = frame.frameId;

		}

		if ( needsUpdate ) {

			this.update();
			this.updateShadow( frame );

			shadow.needsUpdate = false;

		}

	}

	/**
	 * Synchronizes the transformation of a tile light with the source light.
	 *
	 * @param {LwLight} lwLight - The tile light to synchronize.
	 * @param {Light} sourceLight - The source light to copy transformations from.
	 */
	syncLightTransformation( lwLight, sourceLight ) {

		const sourceWorldPos = sourceLight.getWorldPosition( _vec3Temp1 );
		const targetWorldPos = sourceLight.target.getWorldPosition( _vec3Temp2 );
		const forward = _vec3Temp3.subVectors( targetWorldPos, sourceWorldPos );
		const targetDistance = forward.length();
		forward.normalize();
		lwLight.position.copy( sourceWorldPos );
		lwLight.target.position.copy( sourceWorldPos ).add( forward.multiplyScalar( targetDistance ) );
		lwLight.quaternion.copy( sourceLight.getWorldQuaternion( _quatTemp1 ) );
		lwLight.scale.copy( sourceLight.scale );
		lwLight.updateMatrix();
		lwLight.updateMatrixWorld( true );
		lwLight.target.updateMatrix();
		lwLight.target.updateMatrixWorld( true );

	}

	/**
	 * Sets up the shadow node for rendering.
	 *
	 * @param {Builder} builder - The builder used to set up the shadow node.
	 * @returns {Node} A node representing the shadow value.
	 */
	setup( builder ) {

		if ( this.lights.length === 0 ) {

			this.init( builder );

		}

		// Hardware-filtered depth comparisons for PCF, as in ShadowNode. The shadow type can change between
		// builds, so this is checked every time; disposing the shadow map recreates it with the new filter.

		const { renderer } = builder;
		const filter = renderer.shadowMap.type === PCFShadowMap && renderer.hasCompatibility( Compatibility.TEXTURE_COMPARE ) ? LinearFilter : NearestFilter;
		const depthTexture = this.shadowMap.depthTexture;

		if ( depthTexture.magFilter !== filter ) {

			depthTexture.minFilter = filter;
			depthTexture.magFilter = filter;

			this.shadowMap.dispose();

		}

		return Fn( ( builder ) => {

			this.setupShadowPosition( builder );

			// The fragment's position in the full shadow camera gives its tile and the coordinate within it.

			const { tilesX, tilesY } = this.config;
			const tileShadowNode = this._shadowNodes[ 0 ];
			const tileShadow = this.lights[ 0 ].shadow;
			const shadow = this.originalLight.shadow;
			const normalBias = reference( 'normalBias', 'float', shadow ).setGroup( renderGroup );

			const shadowPosition = lightShadowMatrix( this.originalLight ).mul( vec4( shadowPositionWorld.add( normalWorld.mul( normalBias ) ), 1 ) );
			const shadowCoord = tileShadowNode.setupShadowCoord( builder, shadowPosition );

			const grid = shadowCoord.xy.mul( vec2( tilesX, tilesY ) ); // tiles start from the top row
			const tile = grid.floor().clamp( vec2( 0 ), vec2( tilesX - 1, tilesY - 1 ) );
			const depthLayer = int( tile.y.mul( tilesX ).add( tile.x ) );

			const tileCoord = vec3( grid.sub( tile ), shadowCoord.z );

			const filterFn = tileShadow.filterNode || tileShadowNode.getShadowFilterFn( builder.renderer.shadowMap.type );
			const shadowValue = tileShadowNode.setupShadowFilter( builder, {
				filterFn,
				depthTexture,
				shadowCoord: tileCoord,
				shadow: tileShadow,
				depthLayer
			} );

			const shadowIntensity = reference( 'intensity', 'float', shadow ).setGroup( renderGroup );

			let shadowOutput;

			if ( builder.renderer.shadowMap.transmitted === true ) {

				const shadowColor = texture( this.shadowMap.texture, tileCoord ).depth( depthLayer );

				shadowOutput = mix( 1, shadowValue.rgb.mix( shadowColor, 1 ), shadowIntensity.mul( shadowColor.a ) ).toVar( 'shadowValue' );

			} else {

				shadowOutput = mix( 1, shadowValue, shadowIntensity ).toVar( 'shadowValue' );

			}

			// As in ShadowNode.setup(), which the tiles no longer go through.

			if ( builder.material.receivedShadowNode ) {

				shadowOutput = builder.material.receivedShadowNode( shadowOutput );

			}

			return shadowOutput;

		} )();

	}

	/**
     * Helper method to remove lights and associated nodes/targets.
     * Used internally during dispose and potential re-initialization.
     */
	disposeLightsAndNodes() {

		for ( const light of this.lights ) {

			const parent = light.parent;
			if ( parent ) {

				parent.remove( light.target );
				parent.remove( light );

			}

		}

		this.lights = [];
		this._shadowNodes = [];

		if ( this.shadowMap ) {

			this.shadowMap.dispose(); // Disposes render target and textures
			this.shadowMap = null;

		}

	}


	dispose() {

		// Dispose lights, nodes, and shadow map
		this.disposeLightsAndNodes();
		this.disposeShadowMaterial();
		super.dispose();

	}

}

export { TileShadowNode };