TSL: Introduce ShaderCallNode & tslFn improvements (mrdoob#26824)

* Add ShaderCallNode * cleanup * Use tslFn as default
aardgoose · Sep 23, 2023 · 9d2d7eb · 9d2d7eb
1 parent 45505ed
commit 9d2d7eb
Show file tree

Hide file tree

Showing 7 changed files with 82 additions and 45 deletions.
diff --git a/examples/jsm/nodes/procedural/CheckerNode.js b/examples/jsm/nodes/procedural/CheckerNode.js
@@ -25,9 +25,9 @@ class CheckerNode extends TempNode {
 
 	}
 
-	generate( builder ) {
+	construct() {
 
-		return checkerShaderNode( { uv: this.uvNode } ).build( builder );
+		return checkerShaderNode( { uv: this.uvNode } );
 
 	}
 

diff --git a/examples/jsm/nodes/shadernode/ShaderNode.js b/examples/jsm/nodes/shadernode/ShaderNode.js
@@ -196,42 +196,86 @@ const ShaderNodeImmutable = function ( NodeClass, ...params ) {
 
 };
 
-class ShaderNodeInternal extends Node {
+class ShaderCallNodeInternal extends Node {
 
-	constructor( jsFunc ) {
+	constructor( shaderNode, inputNodes ) {
 
 		super();
 
-		this._jsFunc = jsFunc;
+		this.shaderNode = shaderNode;
+		this.inputNodes = inputNodes;
 
 	}
 
-	call( inputs, stack, builder ) {
+	getNodeType( builder ) {
 
-		inputs = nodeObjects( inputs );
+		const { outputNode } = builder.getNodeProperties( this );
 
-		return nodeObject( this._jsFunc( inputs, stack, builder ) );
+		return outputNode ? outputNode.getNodeType( builder ) : super.getNodeType( builder );
 
 	}
 
-	getNodeType( builder ) {
+	call( builder ) {
 
-		const { outputNode } = builder.getNodeProperties( this );
+		const { shaderNode, inputNodes } = this;
 
-		return outputNode ? outputNode.getNodeType( builder ) : super.getNodeType( builder );
+		const jsFunc = shaderNode.jsFunc;
+		const outputNode = inputNodes !== null ? jsFunc( nodeObjects( inputNodes ), builder.stack, builder ) : jsFunc( builder.stack, builder );
+
+		return nodeObject( outputNode );
 
 	}
 
 	construct( builder ) {
 
 		builder.addStack();
 
-		builder.stack.outputNode = nodeObject( this._jsFunc( builder.stack, builder ) );
+		builder.stack.outputNode = this.call( builder );
 
 		return builder.removeStack();
 
 	}
 
+	generate( builder, output ) {
+
+		const { outputNode } = builder.getNodeProperties( this );
+
+		if ( outputNode === null ) {
+
+			// TSL: It's recommended to use `tslFn` in construct() pass.
+
+			return this.call( builder ).build( builder, output );
+
+		}
+
+		return super.generate( builder, output );
+
+	}
+
+}
+
+class ShaderNodeInternal extends Node {
+
+	constructor( jsFunc ) {
+
+		super();
+
+		this.jsFunc = jsFunc;
+
+	}
+
+	call( inputs = null ) {
+
+		return nodeObject( new ShaderCallNodeInternal( this, inputs ) );
+
+	}
+
+	construct() {
+
+		return this.call();
+
+	}
+
 }
 
 const bools = [ false, true ];
@@ -349,15 +393,9 @@ export const shader = ( jsFunc ) => { // @deprecated, r154
 
 export const tslFn = ( jsFunc ) => {
 
-	let shaderNode = null;
+	const shaderNode = new ShaderNode( jsFunc );
 
-	return ( ...params ) => {
-
-		if ( shaderNode === null ) shaderNode = new ShaderNode( jsFunc );
-
-		return shaderNode.call( ...params );
-
-	};
+	return ( inputs ) => shaderNode.call( inputs );
 
 };
 

diff --git a/examples/jsm/nodes/utils/LoopNode.js b/examples/jsm/nodes/utils/LoopNode.js
@@ -1,7 +1,7 @@
 import Node, { addNodeClass } from '../core/Node.js';
 import { expression } from '../code/ExpressionNode.js';
 import { bypass } from '../core/BypassNode.js';
-import { context as contextNode } from '../core/ContextNode.js';
+import { context } from '../core/ContextNode.js';
 import { addNodeElement, nodeObject, nodeArray } from '../shadernode/ShaderNode.js';
 
 class LoopNode extends Node {
@@ -65,13 +65,11 @@ class LoopNode extends Node {
 
 		const properties = this.getProperties( builder );
 
-		const context = { tempWrite: false };
+		const contextData = { tempWrite: false };
 
 		const params = this.params;
 		const stackNode = properties.stackNode;
 
-		const returnsSnippet = properties.returnsNode ? properties.returnsNode.build( builder ) : '';
-
 		for ( let i = 0, l = params.length - 1; i < l; i ++ ) {
 
 			const param = params[ i ];
@@ -82,7 +80,7 @@ class LoopNode extends Node {
 			if ( param.isNode ) {
 
 				start = '0';
-				end = param.generate( builder, 'int' );
+				end = param.build( builder, 'int' );
 				direction = 'forward';
 
 			} else {
@@ -92,10 +90,10 @@ class LoopNode extends Node {
 				direction = param.direction;
 
 				if ( typeof start === 'number' ) start = start.toString();
-				else if ( start && start.isNode ) start = start.generate( builder, 'int' );
+				else if ( start && start.isNode ) start = start.build( builder, 'int' );
 
 				if ( typeof end === 'number' ) end = end.toString();
-				else if ( end && end.isNode ) end = end.generate( builder, 'int' );
+				else if ( end && end.isNode ) end = end.build( builder, 'int' );
 
 				if ( start !== undefined && end === undefined ) {
 
@@ -159,7 +157,9 @@ class LoopNode extends Node {
 
 		}
 
-		const stackSnippet = contextNode( stackNode, context ).build( builder, 'void' );
+		const stackSnippet = context( stackNode, contextData ).build( builder, 'void' );
+
+		const returnsSnippet = properties.returnsNode ? properties.returnsNode.build( builder ) : '';
 
 		builder.removeFlowTab().addFlowCode( '\n' + builder.tab + stackSnippet );
 

diff --git a/examples/webgpu_audio_processing.html b/examples/webgpu_audio_processing.html
@@ -31,7 +31,7 @@
 		<script type="module">
 
 			import * as THREE from 'three';
-			import { ShaderNode, uniform, storage, instanceIndex, float, texture, viewportTopLeft, color } from 'three/nodes';
+			import { tslFn, uniform, storage, instanceIndex, float, texture, viewportTopLeft, color } from 'three/nodes';
 
 			import { GUI } from 'three/addons/libs/lil-gui.module.min.js';
 
@@ -136,11 +136,10 @@
 
 				// compute (shader-node)
 
-				const computeShaderNode = new ShaderNode( ( stack ) => {
+				const computeShaderFn = tslFn( ( stack ) => {
 
 					const index = float( instanceIndex );
 
-
 					// pitch
 
 					const time = index.mul( pitch );
@@ -171,7 +170,7 @@
 
 				// compute
 
-				computeNode = computeShaderNode.compute( waveBuffer.length );
+				computeNode = computeShaderFn().compute( waveBuffer.length );
 
 
 				// gui

diff --git a/examples/webgpu_compute.html b/examples/webgpu_compute.html
@@ -26,7 +26,7 @@
 		<script type="module">
 
 			import * as THREE from 'three';
-			import { ShaderNode, uniform, storage, attribute, float, vec2, vec3, color, instanceIndex, PointsNodeMaterial } from 'three/nodes';
+			import { tslFn, uniform, storage, attribute, float, vec2, vec3, color, instanceIndex, PointsNodeMaterial } from 'three/nodes';
 
 			import { GUI } from 'three/addons/libs/lil-gui.module.min.js';
 
@@ -74,7 +74,7 @@
 
 				// create function
 
-				const computeShaderNode = new ShaderNode( ( stack ) => {
+				const computeShaderFn = tslFn( ( stack ) => {
 
 					const particle = particleBufferNode.element( instanceIndex );
 					const velocity = velocityBufferNode.element( instanceIndex );
@@ -98,10 +98,10 @@
 
 				// compute
 
-				computeNode = computeShaderNode.compute( particleNum );
+				computeNode = computeShaderFn().compute( particleNum );
 				computeNode.onInit = ( { renderer } ) => {
 
-					const precomputeShaderNode = new ShaderNode( ( stack ) => {
+					const precomputeShaderNode = tslFn( ( stack ) => {
 
 						const particleIndex = float( instanceIndex );
 
@@ -117,7 +117,7 @@
 
 					} );
 
-					renderer.compute( precomputeShaderNode.compute( particleNum ) );
+					renderer.compute( precomputeShaderNode().compute( particleNum ) );
 
 				};
 

diff --git a/examples/webgpu_compute_particles.html b/examples/webgpu_compute_particles.html
@@ -26,7 +26,7 @@
 		<script type="module">
 
 			import * as THREE from 'three';
-			import { ShaderNode, uniform, texture, instanceIndex, float, vec3, storage, SpriteNodeMaterial } from 'three/nodes';
+			import { tslFn, uniform, texture, instanceIndex, float, vec3, storage, SpriteNodeMaterial } from 'three/nodes';
 
 			import WebGPU from 'three/addons/capabilities/WebGPU.js';
 			import WebGPURenderer from 'three/addons/renderers/webgpu/WebGPURenderer.js';
@@ -83,7 +83,7 @@
 
 				// compute
 
-				const computeInit = new ShaderNode( ( stack ) => {
+				const computeInit = tslFn( ( stack ) => {
 
 					const position = positionBuffer.element( instanceIndex );
 					const color = colorBuffer.element( instanceIndex );
@@ -98,11 +98,11 @@
 
 					stack.assign( color, vec3( randX, randY, randZ ) );
 
-				} ).compute( particleCount );
+				} )().compute( particleCount );
 
 				//
 
-				const computeUpdate = new ShaderNode( ( stack ) => {
+				const computeUpdate = tslFn( ( stack ) => {
 
 					const position = positionBuffer.element( instanceIndex );
 					const velocity = velocityBuffer.element( instanceIndex );
@@ -128,7 +128,7 @@
 
 				} );
 
-				computeParticles = computeUpdate.compute( particleCount );
+				computeParticles = computeUpdate().compute( particleCount );
 
 				// create nodes
 
@@ -179,7 +179,7 @@
 
 				// click event
 
-				const computeHit = new ShaderNode( ( stack ) => {
+				const computeHit = tslFn( ( stack ) => {
 
 					const position = positionBuffer.element( instanceIndex );
 					const velocity = velocityBuffer.element( instanceIndex );
@@ -193,7 +193,7 @@
 
 					stack.assign( velocity, velocity.add( direction.mul( relativePower ) ) );
 
-				} ).compute( particleCount );
+				} )().compute( particleCount );
 
 				//
 

diff --git a/examples/webgpu_materials.html b/examples/webgpu_materials.html
@@ -29,7 +29,7 @@
 			import * as THREE from 'three';
 			import * as Nodes from 'three/nodes';
 
-			import { tslFn, wgslFn, attribute, positionLocal, positionWorld, normalLocal, normalWorld, normalView, color, texture, uv, float, vec2, vec3, vec4, oscSine, triplanarTexture, viewportBottomLeft, js, string, global, loop, MeshBasicNodeMaterial, NodeObjectLoader } from 'three/nodes';
+			import { tslFn, wgslFn, positionLocal, positionWorld, normalLocal, normalWorld, normalView, color, texture, uv, float, vec2, vec3, vec4, oscSine, triplanarTexture, viewportBottomLeft, js, string, global, loop, MeshBasicNodeMaterial, NodeObjectLoader } from 'three/nodes';
 
 			import WebGPU from 'three/addons/capabilities/WebGPU.js';
 			import WebGPURenderer from 'three/addons/renderers/webgpu/WebGPURenderer.js';