TSL: loop() (#25967)

* TSL: loop()

* TSL: Auto shader() function

* TSL: Add loop() example

* cleanup

* cleanup

* Update examples/webgpu_materials.html

Co-authored-by: Levi Pesin <[email protected]>

---------

Co-authored-by: Levi Pesin <[email protected]>

examples/jsm/nodes/Nodes.js

@@ -45,6 +45,7 @@ export { default as ConvertNode } from './utils/ConvertNode.js';
 export { default as DiscardNode, discard } from './utils/DiscardNode.js';
 export { default as EquirectUVNode, equirectUV } from './utils/EquirectUVNode.js';
 export { default as JoinNode } from './utils/JoinNode.js';
+export { default as LoopNode, loop } from './utils/LoopNode.js';
 export { default as MatcapUVNode, matcapUV } from './utils/MatcapUVNode.js';
 export { default as MaxMipLevelNode, maxMipLevel } from './utils/MaxMipLevelNode.js';
 export { default as OscNode, oscSine, oscSquare, oscTriangle, oscSawtooth } from './utils/OscNode.js';

examples/jsm/nodes/core/NodeBuilder.js

@@ -744,6 +744,12 @@ class NodeBuilder {
 
 	}
 
+	getVar( type, name ) {
+
+		return `${type} ${name}`;
+
+	}
+
 	getVars( shaderStage ) {
 
 		let snippet = '';
@@ -752,7 +758,7 @@ class NodeBuilder {
 
 		for ( const variable of vars ) {
 
-			snippet += `${variable.type} ${variable.name}; `;
+			snippet += `${ this.getVar( variable.type, variable.name ) }; `;
 
 		}
 

examples/jsm/nodes/core/NodeUtils.js

@@ -95,6 +95,10 @@ export function getValueType( value ) {
 
 		return 'string';
 
+	} else if ( typeOf === 'function' ) {
+
+		return 'shader';
+
 	} else if ( value.isVector2 === true ) {
 
 		return 'vec2';

examples/jsm/nodes/core/PropertyNode.js

@@ -42,7 +42,7 @@ class PropertyNode extends Node {
 
 export default PropertyNode;
 
-export const property = ( name, nodeOrType ) => nodeObject( new PropertyNode( name, getConstNodeType( nodeOrType ) ) );
+export const property = ( type, name ) => nodeObject( new PropertyNode( getConstNodeType( type ), name ) );
 
 export const diffuseColor = nodeImmutable( PropertyNode, 'vec4', 'DiffuseColor' );
 export const roughness = nodeImmutable( PropertyNode, 'float', 'Roughness' );

examples/jsm/nodes/core/StackNode.js

@@ -3,6 +3,7 @@ import { assign } from '../math/OperatorNode.js';
 import { bypass } from '../core/BypassNode.js';
 import { expression } from '../code/ExpressionNode.js';
 import { cond } from '../math/CondNode.js';
+import { loop } from '../utils/LoopNode.js';
 import { nodeProxy, shader } from '../shadernode/ShaderNode.js';
 
 class StackNode extends Node {
@@ -71,6 +72,12 @@ class StackNode extends Node {
 
 	}
 
+	loop( ...params ) {
+
+		return this.add( loop( ...params ) );
+
+	}
+
 	build( builder, ...params ) {
 
 		for ( const node of this.nodes ) {

examples/jsm/nodes/shadernode/ShaderNode.js

@@ -79,7 +79,7 @@ const shaderNodeHandler = {
 
 const nodeObjectsCacheMap = new WeakMap();
 
-const ShaderNodeObject = function ( obj ) {
+const ShaderNodeObject = function ( obj, altType = null ) {
 
 	const type = getValueType( obj );
 
@@ -97,13 +97,17 @@ const ShaderNodeObject = function ( obj ) {
 
 		return nodeObject;
 
-	} else if ( ( type === 'float' ) || ( type === 'boolean' ) ) {
+	} else if ( ( altType === null ) && ( ( type === 'float' ) || ( type === 'boolean' ) ) ) {
 
 		return nodeObject( getAutoTypedConstNode( obj ) );
 
+	} else if ( type === 'shader' ) {
+
+		return shader( obj );
+
 	} else if ( type && type !== 'string' ) {
 
-		return nodeObject( new ConstNode( obj ) );
+		return nodeObject( new ConstNode( obj, altType ) );
 
 	}
 
@@ -111,11 +115,11 @@ const ShaderNodeObject = function ( obj ) {
 
 };
 
-const ShaderNodeObjects = function ( objects ) {
+const ShaderNodeObjects = function ( objects, altType = null ) {
 
 	for ( const name in objects ) {
 
-		objects[ name ] = nodeObject( objects[ name ] );
+		objects[ name ] = nodeObject( objects[ name ], altType );
 
 	}
 
@@ -123,13 +127,13 @@ const ShaderNodeObjects = function ( objects ) {
 
 };
 
-const ShaderNodeArray = function ( array ) {
+const ShaderNodeArray = function ( array, altType = null ) {
 
 	const len = array.length;
 
 	for ( let i = 0; i < len; i ++ ) {
 
-		array[ i ] = nodeObject( array[ i ] );
+		array[ i ] = nodeObject( array[ i ], altType );
 
 	}
 
@@ -307,9 +311,9 @@ export function ShaderNode( jsFunc ) {
 
 }
 
-export const nodeObject = ( val ) => /* new */ ShaderNodeObject( val );
-export const nodeObjects = ( val ) => new ShaderNodeObjects( val );
-export const nodeArray = ( val ) => new ShaderNodeArray( val );
+export const nodeObject = ( val, altType = null ) => /* new */ ShaderNodeObject( val, altType );
+export const nodeObjects = ( val, altType = null ) => new ShaderNodeObjects( val, altType );
+export const nodeArray = ( val, altType = null ) => new ShaderNodeArray( val, altType );
 export const nodeProxy = ( ...val ) => new ShaderNodeProxy( ...val );
 export const nodeImmutable = ( ...val ) => new ShaderNodeImmutable( ...val );
 

examples/jsm/nodes/utils/LoopNode.js

@@ -0,0 +1,186 @@
+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 { addNodeElement, nodeObject, nodeArray } from '../shadernode/ShaderNode.js';
+
+class LoopNode extends Node {
+
+	constructor( params = [] ) {
+
+		super();
+
+		this.params = params;
+
+	}
+
+	getVarName( index ) {
+
+		return String.fromCharCode( 'i'.charCodeAt() + index );
+
+	}
+
+	getProperties( builder ) {
+
+		const properties = builder.getNodeProperties( this );
+
+		if ( properties.stackNode !== undefined ) return properties;
+
+		//
+
+		const inputs = {};
+
+		for ( let i = 0, l = this.params.length - 1; i < l; i ++ ) {
+
+			const prop = this.getVarName( i );
+
+			inputs[ prop ] = expression( prop, 'int' );
+
+		}
+
+		properties.returnsNode = this.params[ this.params.length - 1 ].call( inputs, builder.addStack(), builder );
+		properties.stackNode = builder.removeStack();
+
+		return properties;
+
+	}
+
+	getNodeType( builder ) {
+
+		const { returnsNode } = this.getProperties( builder );
+
+		return returnsNode ? returnsNode.getNodeType( builder ) : 'void';
+
+	}
+
+	construct( builder ) {
+
+		// construct properties
+
+		this.getProperties( builder );
+
+	}
+
+	generate( builder ) {
+
+		const properties = this.getProperties( builder );
+
+		const context = { 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 ];
+			const property = this.getVarName( i );
+
+			let start = null, end = null, direction = null;
+
+			if ( param.isNode ) {
+
+				start = '0';
+				end = param.generate( builder, 'int' );
+				direction = 'forward';
+
+			} else {
+
+				start = param.start;
+				end = param.end;
+				direction = param.direction;
+
+				if ( typeof start === 'number' ) start = start.toString();
+				else if ( start && start.isNode ) start = start.generate( builder, 'int' );
+
+				if ( typeof end === 'number' ) end = end.toString();
+				else if ( end && end.isNode ) end = end.generate( builder, 'int' );
+
+				if ( start !== undefined && end === undefined ) {
+
+					start = start + ' - 1';
+					end = '0';
+					direction = 'backwards';
+
+				} else if ( end !== undefined && start === undefined ) {
+
+					start = '0';
+					direction = 'forward';
+
+				}
+
+				if ( direction === undefined ) {
+
+					if ( Number( start ) > Number( end ) ) {
+
+						direction = 'backwards';
+
+					} else {
+
+						direction = 'forward';
+
+					}
+
+				}
+
+			}
+
+			const internalParam = { start, end, direction };
+
+			//
+
+			const startSnippet = internalParam.start;
+			const endSnippet = internalParam.end;
+
+			let declarationSnippet = '';
+			let conditionalSnippet = '';
+			let updateSnippet = '';
+
+			declarationSnippet += builder.getVar( 'int', property ) + ' = ' + startSnippet;
+
+			if ( internalParam.direction === 'backwards' ) {
+
+				conditionalSnippet += property + ' >= ' + endSnippet;
+				updateSnippet += property + ' --';
+
+			} else {
+
+				// forward
+
+				conditionalSnippet += property + ' < ' + endSnippet;
+				updateSnippet += property + ' ++';
+
+			}
+
+			const forSnippet = `for ( ${ declarationSnippet }; ${ conditionalSnippet }; ${ updateSnippet } )`;
+
+			builder.addFlowCode( ( i === 0 ? '\n' : '' ) + builder.tab + forSnippet + ' {\n\n' ).addFlowTab();
+
+		}
+
+		const stackSnippet = contextNode( stackNode, context ).build( builder, 'void' );
+
+		builder.removeFlowTab().addFlowCode( '\n' + builder.tab + stackSnippet );
+
+		for ( let i = 0, l = this.params.length - 1; i < l; i ++ ) {
+
+			builder.addFlowCode( ( i === 0 ? '' : builder.tab ) + '}\n\n' ).removeFlowTab();
+
+		}
+
+		builder.addFlowTab();
+
+		return returnsSnippet;
+
+	}
+
+}
+
+export default LoopNode;
+
+export const loop = ( ...params ) => nodeObject( new LoopNode( nodeArray( params, 'int' ) ) );
+
+addNodeElement( 'loop', ( returns, ...params ) => bypass( returns, loop( ...params ) ) );
+
+addNodeClass( LoopNode );

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

@@ -465,6 +465,12 @@ class WebGPUNodeBuilder extends NodeBuilder {
 
 	}
 
+	getVar( type, name ) {
+
+		return `var ${ name } : ${ this.getType( type ) }`;
+
+	}
+
 	getVars( shaderStage ) {
 
 		const snippets = [];
@@ -472,10 +478,7 @@ class WebGPUNodeBuilder extends NodeBuilder {
 
 		for ( const variable of vars ) {
 
-			const name = variable.name;
-			const type = this.getType( variable.type );
-
-			snippets.push( `\tvar ${name} : ${type};` );
+			snippets.push( `\t${ this.getVar( variable.type, variable.name ) };` );
 
 		}
 

examples/webgpu_materials.html

@@ -29,7 +29,7 @@
 			import * as THREE from 'three';
 			import * as Nodes from 'three/nodes';
 
-			import { attribute, positionLocal, positionWorld, normalLocal, normalWorld, normalView, color, texture, ShaderNode, func, uv, vec3, triplanarTexture, viewportBottomLeft, js, string, global, MeshBasicNodeMaterial, NodeObjectLoader } from 'three/nodes';
+			import { attribute, positionLocal, positionWorld, normalLocal, normalWorld, normalView, color, texture, ShaderNode, func, uv, 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';
@@ -210,6 +210,33 @@
 				material.colorNode = texture( uvTexture, viewportBottomLeft );
 				materials.push( material );
 
+				// Loop
+				material = new MeshBasicNodeMaterial();
+				materials.push( material );
+
+				const loopCount = 10;
+				material.colorNode = loop( loopCount, ( { i }, stack ) => {
+
+					const output = vec4().temp();
+					const scale = oscSine().mul( .09 ); // just a value to test
+
+					const scaleI = scale.mul( i );
+					const scaleINeg = scaleI.negate();
+
+					const leftUV = uv().add( vec2( scaleI, 0 ) );
+					const rightUV = uv().add( vec2( scaleINeg, 0 ) );
+					const topUV = uv().add( vec2( 0, scaleI ) );
+					const bottomUV = uv().add( vec2( 0, scaleINeg ) );
+
+					stack.assign( output, output.add( texture( uvTexture, leftUV ) ) );
+					stack.assign( output, output.add( texture( uvTexture, rightUV ) ) );
+					stack.assign( output, output.add( texture( uvTexture, topUV ) ) );
+					stack.assign( output, output.add( texture( uvTexture, bottomUV ) ) );
+
+					return output.div( loopCount * 4 );
+
+				} );
+
 				// Scriptable
 
 				global.set( 'THREE', THREE );