TSL: Function Layout (#26889)

* TSL: Function Interface

* cleanup

* cleanup

* cleanup

* Add support for GLSL

* cleanup

examples/jsm/nodes/code/FunctionNode.js

@@ -81,7 +81,7 @@ class FunctionNode extends CodeNode {
 
 		}
 
-		nodeCode.code = code;
+		nodeCode.code = code + '\n';
 
 		if ( output === 'property' ) {
 

examples/jsm/nodes/core/NodeBuilder.js

@@ -5,6 +5,7 @@ import NodeVar from './NodeVar.js';
 import NodeCode from './NodeCode.js';
 import NodeKeywords from './NodeKeywords.js';
 import NodeCache from './NodeCache.js';
+import PropertyNode from './PropertyNode.js';
 import { createNodeMaterialFromType } from '../materials/NodeMaterial.js';
 import { NodeUpdateType, defaultBuildStages, shaderStages } from './constants.js';
 
@@ -77,14 +78,14 @@ class NodeBuilder {
 		this.flowCode = { vertex: '', fragment: '', compute: [] };
 		this.uniforms = { vertex: [], fragment: [], compute: [], index: 0 };
 		this.structs = { vertex: [], fragment: [], compute: [], index: 0 };
-		this.codes = { vertex: [], fragment: [], compute: [] };
 		this.bindings = { vertex: [], fragment: [], compute: [] };
 		this.bindingsOffset = { vertex: 0, fragment: 0, compute: 0 };
 		this.bindingsArray = null;
 		this.attributes = [];
 		this.bufferAttributes = [];
 		this.varyings = [];
-		this.vars = { vertex: [], fragment: [], compute: [] };
+		this.codes = {};
+		this.vars = {};
 		this.flow = { code: '' };
 		this.chaining = [];
 		this.stack = stack();
@@ -684,7 +685,7 @@ class NodeBuilder {
 
 		if ( nodeVar === undefined ) {
 
-			const vars = this.vars[ shaderStage ];
+			const vars = this.vars[ shaderStage ] || ( this.vars[ shaderStage ] = [] );
 
 			if ( name === null ) name = 'nodeVar' + vars.length;
 
@@ -731,7 +732,7 @@ class NodeBuilder {
 
 		if ( nodeCode === undefined ) {
 
-			const codes = this.codes[ shaderStage ];
+			const codes = this.codes[ shaderStage ] || ( this.codes[ shaderStage ] = [] );
 			const index = codes.length;
 
 			nodeCode = new NodeCode( 'nodeCode' + index, type );
@@ -806,12 +807,67 @@ class NodeBuilder {
 
 	}
 
+	flowShaderNode( shaderNode ) {
+
+		const layout = shaderNode.layout;
+		const inputs = {};
+
+		for ( const input of layout.inputs ) {
+
+			inputs[ input.name ] = new PropertyNode( input.type, input.name, false )
+
+		}
+
+		//
+
+		shaderNode.layout = null;
+
+		const callNode = shaderNode.call( inputs );
+		const flowData = this.flowStagesNode( callNode, layout.type );
+
+		shaderNode.layout = layout;
+
+		return flowData;
+
+	}
+
+	flowStagesNode( node, output = null ) {
+
+		const previousFlow = this.flow;
+		const previousVars = this.vars;
+		const previousBuildStage = this.buildStage;
+
+		const flow = {
+			code: ''
+		};
+
+		this.flow = flow;
+		this.vars = {};
+
+		for ( const buildStage of defaultBuildStages ) {
+
+			this.setBuildStage( buildStage );
+
+			flow.result = node.build( this, output );
+
+		}
+
+		flow.vars = this.getVars( this.shaderStage );
+
+		this.flow = previousFlow;
+		this.vars = previousVars;
+		this.setBuildStage( previousBuildStage );
+
+		return flow;
+
+	}
+
 	flowChildNode( node, output = null ) {
 
 		const previousFlow = this.flow;
 
 		const flow = {
-			code: '',
+			code: ''
 		};
 
 		this.flow = flow;
@@ -876,9 +932,13 @@ class NodeBuilder {
 
 		const vars = this.vars[ shaderStage ];
 
-		for ( const variable of vars ) {
+		if ( vars !== undefined ) {
+
+			for ( const variable of vars ) {
 
-			snippet += `${ this.getVar( variable.type, variable.name ) }; `;
+				snippet += `${ this.getVar( variable.type, variable.name ) }; `;
+
+			}
 
 		}
 
@@ -898,9 +958,13 @@ class NodeBuilder {
 
 		let code = '';
 
-		for ( const nodeCode of codes ) {
+		if ( codes !== undefined ) {
+
+			for ( const nodeCode of codes ) {
 
-			code += nodeCode.code + '\n';
+				code += nodeCode.code + '\n';
+
+			}
 
 		}
 

examples/jsm/nodes/core/PropertyNode.js

@@ -3,11 +3,12 @@ import { nodeImmutable, nodeObject } from '../shadernode/ShaderNode.js';
 
 class PropertyNode extends Node {
 
-	constructor( nodeType, name = null ) {
+	constructor( nodeType, name = null, declare = true ) {
 
 		super( nodeType );
 
 		this.name = name;
+		this.declare = declare;
 
 		this.isPropertyNode = true;
 
@@ -27,6 +28,8 @@ class PropertyNode extends Node {
 
 	generate( builder ) {
 
+		if ( this.declare === false ) return this.name;
+
 		return builder.getPropertyName( builder.getVarFromNode( this, this.name ) );
 
 	}

examples/jsm/nodes/functions/BSDF/BRDF_Sheen.js

@@ -4,7 +4,7 @@ import { sheen, sheenRoughness } from '../../core/PropertyNode.js';
 import { tslFn, float } from '../../shadernode/ShaderNode.js';
 
 // https://github.com/google/filament/blob/master/shaders/src/brdf.fs
-const D_Charlie = ( roughness, dotNH ) => {
+const D_Charlie = tslFn( ( { roughness, dotNH } ) => {
 
 	const alpha = roughness.pow2();
 
@@ -15,15 +15,29 @@ const D_Charlie = ( roughness, dotNH ) => {
 
 	return float( 2.0 ).add( invAlpha ).mul( sin2h.pow( invAlpha.mul( 0.5 ) ) ).div( 2.0 * Math.PI );
 
-};
+} ).setLayout( {
+	name: 'D_Charlie',
+	type: 'float',
+	inputs: [
+		{ name: 'roughness', type: 'float' },
+		{ name: 'dotNH', type: 'float' }
+	]
+} );
 
 // https://github.com/google/filament/blob/master/shaders/src/brdf.fs
-const V_Neubelt = ( dotNV, dotNL ) => {
+const V_Neubelt = tslFn( ( { dotNV, dotNL } ) => {
 
 	// Neubelt and Pettineo 2013, "Crafting a Next-gen Material Pipeline for The Order: 1886"
 	return float( 1.0 ).div( float( 4.0 ).mul( dotNL.add( dotNV ).sub( dotNL.mul( dotNV ) ) ) );
 
-};
+} ).setLayout( {
+	name: 'V_Neubelt',
+	type: 'float',
+	inputs: [
+		{ name: 'dotNV', type: 'float' },
+		{ name: 'dotNL', type: 'float' }
+	]
+} );
 
 const BRDF_Sheen = tslFn( ( { lightDirection } ) => {
 
@@ -33,8 +47,8 @@ const BRDF_Sheen = tslFn( ( { lightDirection } ) => {
 	const dotNV = transformedNormalView.dot( positionViewDirection ).clamp();
 	const dotNH = transformedNormalView.dot( halfDir ).clamp();
 
-	const D = D_Charlie( sheenRoughness, dotNH );
-	const V = V_Neubelt( dotNV, dotNL );
+	const D = D_Charlie( { roughness: sheenRoughness, dotNH } );
+	const V = V_Neubelt( { dotNV, dotNL } );
 
 	return sheen.mul( D ).mul( V );
 

examples/jsm/nodes/functions/BSDF/DFGApprox.js

@@ -1,16 +1,10 @@
-import { transformedNormalView } from '../../accessors/NormalNode.js';
-import { positionViewDirection } from '../../accessors/PositionNode.js';
 import { tslFn, vec2, vec4 } from '../../shadernode/ShaderNode.js';
 
 // Analytical approximation of the DFG LUT, one half of the
 // split-sum approximation used in indirect specular lighting.
 // via 'environmentBRDF' from "Physically Based Shading on Mobile"
 // https://www.unrealengine.com/blog/physically-based-shading-on-mobile
-const DFGApprox = tslFn( ( inputs ) => {
-
-	const { roughness } = inputs;
-
-	const dotNV = inputs.dotNV || transformedNormalView.dot( positionViewDirection ).clamp(); // @ TODO: Move to core dotNV
+const DFGApprox = tslFn( ( { roughness, dotNV } ) => {
 
 	const c0 = vec4( - 1, - 0.0275, - 0.572, 0.022 );
 
@@ -24,6 +18,13 @@ const DFGApprox = tslFn( ( inputs ) => {
 
 	return fab;
 
+} ).setLayout( {
+	name: 'DFGApprox',
+	type: 'vec2',
+	inputs: [
+		{ name: 'roughness', type: 'float' },
+		{ name: 'dotNV', type: 'vec3' }
+	]
 } );
 
 export default DFGApprox;

examples/jsm/nodes/functions/BSDF/D_GGX.js

@@ -3,9 +3,7 @@ import { tslFn } from '../../shadernode/ShaderNode.js';
 // Microfacet Models for Refraction through Rough Surfaces - equation (33)
 // http://graphicrants.blogspot.com/2013/08/specular-brdf-reference.html
 // alpha is "roughness squared" in Disney’s reparameterization
-const D_GGX = tslFn( ( inputs ) => {
-
-	const { alpha, dotNH } = inputs;
+const D_GGX = tslFn( ( { alpha, dotNH } ) => {
 
 	const a2 = alpha.pow2();
 
@@ -13,6 +11,13 @@ const D_GGX = tslFn( ( inputs ) => {
 
 	return a2.div( denom.pow2() ).mul( 1 / Math.PI );
 
+} ).setLayout( {
+	name: 'D_GGX',
+	type: 'float',
+	inputs: [
+		{ name: 'alpha', type: 'float' },
+		{ name: 'dotNH', type: 'float' }
+	]
 } ); // validated
 
 export default D_GGX;

examples/jsm/nodes/functions/BSDF/Schlick_to_F0.js

@@ -8,6 +8,14 @@ const Schlick_to_F0 = tslFn( ( { f, f90, dotVH } ) => {
 
 	return f.sub( vec3( f90 ).mul( x5 ) ).div( x5.oneMinus() );
 
+} ).setLayout( {
+	name: 'Schlick_to_F0',
+	type: 'vec3',
+	inputs: [
+		{ name: 'f', type: 'vec3' },
+		{ name: 'f90', type: 'float' },
+		{ name: 'dotVH', type: 'float' }
+	]
 } );
 
 export default Schlick_to_F0;

examples/jsm/nodes/functions/BSDF/V_GGX_SmithCorrelated.js

@@ -15,6 +15,14 @@ const V_GGX_SmithCorrelated = tslFn( ( inputs ) => {
 
 	return div( 0.5, gv.add( gl ).max( EPSILON ) );
 
+} ).setLayout( {
+	name: 'V_GGX_SmithCorrelated',
+	type: 'float',
+	inputs: [
+		{ name: 'alpha', type: 'float' },
+		{ name: 'dotNL', type: 'float' },
+		{ name: 'dotNV', type: 'float' }
+	]
 } ); // validated
 
 export default V_GGX_SmithCorrelated;

examples/jsm/nodes/functions/PhysicalLightingModel.js

@@ -9,7 +9,7 @@ import LightingModel from '../core/LightingModel.js';
 import { diffuseColor, specularColor, roughness, clearcoat, clearcoatRoughness, sheen, sheenRoughness, iridescence, iridescenceIOR, iridescenceThickness } from '../core/PropertyNode.js';
 import { transformedNormalView, transformedClearcoatNormalView } from '../accessors/NormalNode.js';
 import { positionViewDirection } from '../accessors/PositionNode.js';
-import { float, vec3, mat3 } from '../shadernode/ShaderNode.js';
+import { tslFn, float, vec3, mat3 } from '../shadernode/ShaderNode.js';
 import { cond } from '../math/CondNode.js';
 import { mix, smoothstep } from '../math/MathNode.js';
 
@@ -56,11 +56,12 @@ const evalSensitivity = ( OPD, shift ) => {
 	xyz = vec3( xyz.x.add( x ), xyz.y, xyz.z ).div( 1.0685e-7 );
 
 	const rgb = XYZ_TO_REC709.mul( xyz );
+
 	return rgb;
 
 };
 
-const evalIridescence = ( outsideIOR, eta2, cosTheta1, thinFilmThickness, baseF0 ) => {
+const evalIridescence = tslFn( ( { outsideIOR, eta2, cosTheta1, thinFilmThickness, baseF0 } ) => {
 
 	// Force iridescenceIOR -> outsideIOR when thinFilmThickness -> 0.0
 	const iridescenceIOR = mix( outsideIOR, eta2, smoothstep( 0.0, 0.03, thinFilmThickness ) );
@@ -121,7 +122,17 @@ const evalIridescence = ( outsideIOR, eta2, cosTheta1, thinFilmThickness, baseF0
 	// Since out of gamut colors might be produced, negative color values are clamped to 0.
 	return I.max( vec3( 0.0 ) );
 
-};
+} ).setLayout( {
+	name: 'evalIridescence',
+	type: 'vec3',
+	inputs: [
+		{ name: 'outsideIOR', type: 'float' },
+		{ name: 'eta2', type: 'float' },
+		{ name: 'cosTheta1', type: 'float' },
+		{ name: 'thinFilmThickness', type: 'float' },
+		{ name: 'baseF0', type: 'vec3' }
+	]
+} );
 
 //
 //	Sheen
@@ -130,7 +141,7 @@ const evalIridescence = ( outsideIOR, eta2, cosTheta1, thinFilmThickness, baseF0
 // This is a curve-fit approxmation to the "Charlie sheen" BRDF integrated over the hemisphere from
 // Estevez and Kulla 2017, "Production Friendly Microfacet Sheen BRDF". The analysis can be found
 // in the Sheen section of https://drive.google.com/file/d/1T0D1VSyR4AllqIJTQAraEIzjlb5h4FKH/view?usp=sharing
-const IBLSheenBRDF = ( normal, viewDir, roughness ) => {
+const IBLSheenBRDF = tslFn( ( { normal, viewDir, roughness } ) => {
 
 	const dotNV = normal.dot( viewDir ).saturate();
 
@@ -152,7 +163,7 @@ const IBLSheenBRDF = ( normal, viewDir, roughness ) => {
 
 	return DG.mul( 1.0 / Math.PI ).saturate();
 
-};
+} );
 
 const clearcoatF0 = vec3( 0.04 );
 const clearcoatF90 = vec3( 1 );
@@ -196,7 +207,14 @@ class PhysicalLightingModel extends LightingModel {
 
 			const dotNVi = transformedNormalView.dot( positionViewDirection ).clamp();
 
-			this.iridescenceFresnel = evalIridescence( float( 1.0 ), iridescenceIOR, dotNVi, iridescenceThickness, specularColor );
+			this.iridescenceFresnel = evalIridescence( {
+				outsideIOR: float( 1.0 ),
+				eta2: iridescenceIOR,
+				cosTheta1: dotNVi,
+				thinFilmThickness: iridescenceThickness,
+				baseF0: specularColor
+			} );
+
 			this.iridescenceF0 = Schlick_to_F0( { f: this.iridescenceFresnel, f90: 1.0, dotVH: dotNVi } );
 
 		}
@@ -209,7 +227,9 @@ class PhysicalLightingModel extends LightingModel {
 
 	computeMultiscattering( stack, singleScatter, multiScatter, specularF90 = float( 1 ) ) {
 
-		const fab = DFGApprox( { roughness } );
+		const dotNV = transformedNormalView.dot( positionViewDirection ).clamp(); // @ TODO: Move to core dotNV
+
+		const fab = DFGApprox( { roughness, dotNV } );
 
 		const Fr = this.iridescenceF0 ? iridescence.mix( specularColor, this.iridescenceF0 ) : specularColor;
 
@@ -264,7 +284,11 @@ class PhysicalLightingModel extends LightingModel {
 
 			stack.addAssign( this.sheenSpecular, iblIrradiance.mul(
 				sheen,
-				IBLSheenBRDF( transformedNormalView, positionViewDirection, sheenRoughness )
+				IBLSheenBRDF( {
+					normal: transformedNormalView,
+					viewDir: positionViewDirection,
+					roughness: sheenRoughness
+				} )
 			) );
 
 		}

examples/jsm/nodes/shadernode/ShaderNode.js

@@ -97,6 +97,7 @@ const shaderNodeHandler = {
 };
 
 const nodeObjectsCacheMap = new WeakMap();
+const nodeBuilderFunctionsCacheMap = new WeakMap();
 
 const ShaderNodeObject = function ( obj, altType = null ) {
 
@@ -109,6 +110,7 @@ const ShaderNodeObject = function ( obj, altType = null ) {
 		if ( nodeObject === undefined ) {
 
 			nodeObject = new Proxy( obj, shaderNodeHandler );
+
 			nodeObjectsCacheMap.set( obj, nodeObject );
 			nodeObjectsCacheMap.set( nodeObject, nodeObject );
 
@@ -219,6 +221,32 @@ class ShaderCallNodeInternal extends Node {
 
 		const { shaderNode, inputNodes } = this;
 
+		if ( shaderNode.layout ) {
+
+			let functionNodesCacheMap = nodeBuilderFunctionsCacheMap.get( builder.constructor );
+
+			if ( functionNodesCacheMap === undefined ) {
+
+				functionNodesCacheMap = new WeakMap();
+
+				nodeBuilderFunctionsCacheMap.set( builder.constructor, functionNodesCacheMap );
+
+			}
+
+			let functionNode = functionNodesCacheMap.get( shaderNode );
+
+			if ( functionNode === undefined ) {
+
+				functionNode = nodeObject( builder.buildFunctionNode( shaderNode ) );
+
+				functionNodesCacheMap.set( shaderNode, functionNode );
+
+			}
+
+			return nodeObject( functionNode.call( nodeObjects( inputNodes ) ) );
+
+		}
+
 		const jsFunc = shaderNode.jsFunc;
 		const outputNode = inputNodes !== null ? jsFunc( nodeObjects( inputNodes ), builder.stack, builder ) : jsFunc( builder.stack, builder );
 
@@ -261,6 +289,15 @@ class ShaderNodeInternal extends Node {
 		super();
 
 		this.jsFunc = jsFunc;
+		this.layout = null;
+
+	}
+
+	setLayout( layout ) {
+
+		this.layout = layout;
+
+		return this;
 
 	}
 
@@ -395,7 +432,18 @@ export const tslFn = ( jsFunc ) => {
 
 	const shaderNode = new ShaderNode( jsFunc );
 
-	return ( inputs ) => shaderNode.call( inputs );
+	const fn = ( inputs ) => shaderNode.call( inputs );
+	fn.shaderNode = shaderNode;
+
+	fn.setLayout = ( layout ) => {
+
+		shaderNode.setLayout( layout );
+
+		return fn;
+
+	}
+
+	return fn;
 
 };
 

examples/jsm/renderers/webgl/nodes/GLSLNodeBuilder.js

@@ -1,4 +1,4 @@
-import { MathNode, GLSLNodeParser, NodeBuilder, NodeMaterial } from '../../../nodes/Nodes.js';
+import { MathNode, GLSLNodeParser, NodeBuilder, NodeMaterial, FunctionNode } from '../../../nodes/Nodes.js';
 
 import UniformsGroup from '../../common/UniformsGroup.js';
 import { NodeSampledTexture, NodeSampledCubeTexture } from '../../common/nodes/NodeSampledTexture.js';
@@ -42,6 +42,36 @@ class GLSLNodeBuilder extends NodeBuilder {
 
 	}
 
+	buildFunctionNode( shaderNode ) {
+
+		const layout = shaderNode.layout;
+		const flowData = this.flowShaderNode( shaderNode );
+
+		const parameters = [];
+
+		for ( const input of layout.inputs ) {
+
+			parameters.push( this.getType( input.type ) + ' ' + input.name );
+
+		}
+
+		//
+
+		const code = `${ this.getType( layout.type ) } ${ layout.name }( ${ parameters.join( ', ' ) } ) {
+
+	${ flowData.vars }
+
+${ flowData.code }
+	return ${ flowData.result };
+
+}`;
+
+		//
+
+		return new FunctionNode( code );
+
+	}
+
 	getTexture( texture, textureProperty, uvSnippet ) {
 
 		if ( texture.isTextureCube ) {
@@ -451,9 +481,6 @@ void main() {
 			this.vertexShader = this._getGLSLVertexCode( shadersData.vertex );
 			this.fragmentShader = this._getGLSLFragmentCode( shadersData.fragment );
 
-			//console.log( this.vertexShader );
-			//console.log( this.fragmentShader );
-
 		} else {
 
 			console.warn( 'GLSLNodeBuilder: compute shaders are not supported.' );

examples/jsm/renderers/webgpu/nodes/WGSLNodeBuilder.js

@@ -9,7 +9,7 @@ import UniformBuffer from '../../common/UniformBuffer.js';
 import StorageBuffer from '../../common/StorageBuffer.js';
 import { getVectorLength, getStrideLength } from '../../common/BufferUtils.js';
 
-import { NodeBuilder, CodeNode, NodeMaterial } from '../../../nodes/Nodes.js';
+import { NodeBuilder, CodeNode, NodeMaterial, FunctionNode } from '../../../nodes/Nodes.js';
 
 import { getFormat } from '../utils/WebGPUTextureUtils.js';
 
@@ -436,6 +436,34 @@ class WGSLNodeBuilder extends NodeBuilder {
 
 	}
 
+	buildFunctionNode( shaderNode ) {
+
+		const layout = shaderNode.layout;
+		const flowData = this.flowShaderNode( shaderNode );
+
+		const parameters = [];
+
+		for ( const input of layout.inputs ) {
+
+			parameters.push( input.name + ' : ' + this.getType( input.type ) );
+
+		}
+
+		//
+
+		const code = `fn ${ layout.name }( ${ parameters.join( ', ' ) } ) -> ${ this.getType( layout.type ) } {
+${ flowData.vars }
+${ flowData.code }
+	return ${ flowData.result };
+
+}`;
+
+		//
+
+		return new FunctionNode( code );
+
+	}
+
 	getInstanceIndex() {
 
 		if ( this.shaderStage === 'vertex' ) {
@@ -551,9 +579,13 @@ class WGSLNodeBuilder extends NodeBuilder {
 		const snippets = [];
 		const vars = this.vars[ shaderStage ];
 
-		for ( const variable of vars ) {
+		if ( vars !== undefined ) {
+
+			for ( const variable of vars ) {
 
-			snippets.push( `\t${ this.getVar( variable.type, variable.name ) };` );
+				snippets.push( `\t${ this.getVar( variable.type, variable.name ) };` );
+
+			}
 
 		}