TSL: Introduce `premultipliedGaussianBlur` (#29604)

examples/jsm/tsl/display/GaussianBlurNode.js

@@ -1,5 +1,5 @@
 import { RenderTarget, Vector2, PostProcessingUtils } from 'three';
-import { TempNode, nodeObject, Fn, float, NodeUpdateType, uv, uniform, convertToTexture, vec2, vec4, QuadMesh, passTexture, mul, NodeMaterial } from 'three/tsl';
+import { TempNode, nodeObject, Fn, If, float, NodeUpdateType, uv, uniform, convertToTexture, vec2, vec4, QuadMesh, passTexture, mul, NodeMaterial } from 'three/tsl';
 
 // WebGPU: The use of a single QuadMesh for both gaussian blur passes results in a single RenderObject with a SampledTexture binding that
 // alternates between source textures and triggers creation of new BindGroups and BindGroupLayouts every frame.
@@ -9,6 +9,32 @@ const _quadMesh2 = /*@__PURE__*/ new QuadMesh();
 
 let _rendererState;
 
+const premult = /*@__PURE__*/ Fn( ( [ color ] ) => {
+
+	return vec4( color.rgb.mul( color.a ), color.a );
+
+} ).setLayout( {
+	name: 'premult',
+	type: 'vec4',
+	inputs: [
+		{ name: 'color', type: 'vec4' }
+	]
+} );
+
+const unpremult = /*@__PURE__*/ Fn( ( [ color ] ) => {
+
+	If( color.a.equal( 0.0 ), () => vec4( 0.0 ) );
+
+	return vec4( color.rgb.div( color.a ), color.a );
+
+} ).setLayout( {
+	name: 'unpremult',
+	type: 'vec4',
+	inputs: [
+		{ name: 'color', type: 'vec4' }
+	]
+} );
+
 class GaussianBlurNode extends TempNode {
 
 	static get type() {
@@ -39,6 +65,22 @@ class GaussianBlurNode extends TempNode {
 
 		this.resolution = new Vector2( 1, 1 );
 
+		this.premultipliedAlpha = false;
+
+	}
+
+	setPremultipliedAlpha( value ) {
+
+		this.premultipliedAlpha = value;
+
+		return this;
+
+	}
+
+	getPremultipliedAlpha() {
+
+		return this.premultipliedAlpha;
+
 	}
 
 	setSize( width, height ) {
@@ -123,7 +165,21 @@ class GaussianBlurNode extends TempNode {
 		const uvNode = textureNode.uvNode || uv();
 		const directionNode = vec2( this.directionNode || 1 );
 
-		const sampleTexture = ( uv ) => textureNode.uv( uv );
+		let sampleTexture, output;
+
+		if ( this.premultipliedAlpha ) {
+
+			// https://lisyarus.github.io/blog/posts/blur-coefficients-generator.html
+
+			sampleTexture = ( uv ) => premult( textureNode.uv( uv ) );
+			output = ( color ) => unpremult( color );
+
+		} else {
+
+			sampleTexture = ( uv ) => textureNode.uv( uv );
+			output = ( color ) => color;
+
+		}
 
 		const blur = Fn( () => {
 
@@ -143,15 +199,15 @@ class GaussianBlurNode extends TempNode {
 
 				const uvOffset = vec2( direction.mul( invSize.mul( x ) ) ).toVar();
 
-				const sample1 = vec4( sampleTexture( uvNode.add( uvOffset ) ) );
-				const sample2 = vec4( sampleTexture( uvNode.sub( uvOffset ) ) );
+				const sample1 = sampleTexture( uvNode.add( uvOffset ) );
+				const sample2 = sampleTexture( uvNode.sub( uvOffset ) );
 
 				diffuseSum.addAssign( sample1.add( sample2 ).mul( w ) );
 				weightSum.addAssign( mul( 2.0, w ) );
 
 			}
 
-			return diffuseSum.div( weightSum );
+			return output( diffuseSum.div( weightSum ) );
 
 		} );
 
@@ -199,3 +255,4 @@ class GaussianBlurNode extends TempNode {
 export default GaussianBlurNode;
 
 export const gaussianBlur = ( node, directionNode, sigma ) => nodeObject( new GaussianBlurNode( convertToTexture( node ), directionNode, sigma ) );
+export const premultipliedGaussianBlur = ( node, directionNode, sigma ) => nodeObject( new GaussianBlurNode( convertToTexture( node ), directionNode, sigma ).setPremultipliedAlpha( true ) );