Nodes: Allow JS-like assigns, implement AssignNode (#26795)

* Introduce AssignNode Signed-off-by: Levi Pesin <35454228+LeviPesin@users.noreply.github.com> * Move assign and bypass Signed-off-by: Levi Pesin <35454228+LeviPesin@users.noreply.github.com> * Fix Signed-off-by: Levi Pesin <35454228+LeviPesin@users.noreply.github.com> * Same for temp Signed-off-by: Levi Pesin <35454228+LeviPesin@users.noreply.github.com> * Implement SplitNode.assign() Signed-off-by: Levi Pesin <35454228+LeviPesin@users.noreply.github.com> * Revert bypass and temp change, fix in another way Signed-off-by: Levi Pesin <35454228+LeviPesin@users.noreply.github.com> * Implement ArrayElementNode.assign() Signed-off-by: Levi Pesin <35454228+LeviPesin@users.noreply.github.com> * Add the bonus usecase support Signed-off-by: Levi Pesin <35454228+LeviPesin@users.noreply.github.com> * Fix code issues Signed-off-by: Levi Pesin <35454228+LeviPesin@users.noreply.github.com> * Fixes * Fix, simplify Signed-off-by: Levi Pesin <35454228+LeviPesin@users.noreply.github.com> * Small fix --------- Signed-off-by: Levi Pesin <35454228+LeviPesin@users.noreply.github.com>
mrdoob · Oct 8, 2023 · e41196a · e41196a
1 parent e7a9554
commit e41196a
Show file tree

Hide file tree

Showing 5 changed files with 101 additions and 32 deletions.
diff --git a/examples/jsm/nodes/Nodes.js b/examples/jsm/nodes/Nodes.js
@@ -6,6 +6,7 @@ export * from './core/constants.js';
 
 // core
 export { default as ArrayUniformNode /* @TODO: arrayUniform */ } from './core/ArrayUniformNode.js';
+export { default as AssignNode, assign } from './core/AssignNode.js';
 export { default as AttributeNode, attribute } from './core/AttributeNode.js';
 export { default as BypassNode, bypass } from './core/BypassNode.js';
 export { default as CacheNode, cache } from './core/CacheNode.js';
@@ -37,7 +38,7 @@ export { NodeUtils };
 
 // math
 export { default as MathNode, EPSILON, INFINITY, radians, degrees, exp, exp2, log, log2, sqrt, inverseSqrt, floor, ceil, normalize, fract, sin, cos, tan, asin, acos, atan, abs, sign, length, negate, oneMinus, dFdx, dFdy, round, reciprocal, trunc, fwidth, atan2, min, max, mod, step, reflect, distance, difference, dot, cross, pow, pow2, pow3, pow4, transformDirection, mix, clamp, saturate, refract, smoothstep, faceForward } from './math/MathNode.js';
-export { default as OperatorNode, add, sub, mul, div, remainder, equal, assign, lessThan, greaterThan, lessThanEqual, greaterThanEqual, and, or, xor, bitAnd, bitOr, bitXor, shiftLeft, shiftRight } from './math/OperatorNode.js';
+export { default as OperatorNode, add, sub, mul, div, remainder, equal, lessThan, greaterThan, lessThanEqual, greaterThanEqual, and, or, xor, bitAnd, bitOr, bitXor, shiftLeft, shiftRight } from './math/OperatorNode.js';
 export { default as CondNode, cond } from './math/CondNode.js';
 export { default as HashNode, hash } from './math/HashNode.js';
 

diff --git a/examples/jsm/nodes/core/AssignNode.js b/examples/jsm/nodes/core/AssignNode.js
@@ -0,0 +1,65 @@
+import { addNodeClass } from '../core/Node.js';
+import TempNode from '../core/TempNode.js';
+import { addNodeElement, nodeProxy } from '../shadernode/ShaderNode.js';
+
+class AssignNode extends TempNode {
+
+	constructor( aNode, bNode ) {
+
+		super();
+
+		this.aNode = aNode;
+		this.bNode = bNode;
+
+	}
+
+	hasDependencies( builder ) {
+
+		return false;
+
+	}
+
+	getNodeType( builder, output ) {
+
+		const aNode = this.aNode;
+		const bNode = this.bNode;
+
+		const typeA = aNode.getNodeType( builder );
+		const typeB = bNode.getNodeType( builder );
+
+		return typeB === 'void' ? 'void' : typeA;
+
+	}
+
+	generate( builder, output ) {
+
+		const aNode = this.aNode;
+		const bNode = this.bNode;
+
+		const type = this.getNodeType( builder, output );
+
+		const a = aNode.build( builder, type );
+		const b = bNode.build( builder, type );
+
+		if ( output !== 'void' ) {
+
+			builder.addLineFlowCode( `${a} = ${b}` );
+			return a;
+
+		} else if ( type !== 'void' ) {
+
+			return builder.format( `${a} = ${b}`, type, output );
+
+		}
+
+	}
+
+}
+
+export default AssignNode;
+
+export const assign = nodeProxy( AssignNode );
+
+addNodeClass( 'AssignNode', AssignNode );
+
+addNodeElement( 'assign', assign );
diff --git a/examples/jsm/nodes/math/OperatorNode.js b/examples/jsm/nodes/math/OperatorNode.js
@@ -29,12 +29,6 @@ class OperatorNode extends TempNode {
 
 	}
 
-	hasDependencies( builder ) {
-
-		return this.op !== '=' ? super.hasDependencies( builder ) : false;
-
-	}
-
 	getNodeType( builder, output ) {
 
 		const op = this.op;
@@ -49,7 +43,7 @@ class OperatorNode extends TempNode {
 
 			return 'void';
 
-		} else if ( op === '=' || op === '%' ) {
+		} else if ( op === '%' ) {
 
 			return typeA;
 
@@ -116,11 +110,7 @@ class OperatorNode extends TempNode {
 			typeA = aNode.getNodeType( builder );
 			typeB = bNode.getNodeType( builder );
 
-			if ( op === '=' ) {
-
-				typeB = typeA;
-
-			} else if ( op === '<' || op === '>' || op === '<=' || op === '>=' || op === '==' ) {
+			if ( op === '<' || op === '>' || op === '<=' || op === '>=' || op === '==' ) {
 
 				if ( builder.isVector( typeA ) ) {
 
@@ -170,13 +160,7 @@ class OperatorNode extends TempNode {
 
 		if ( output !== 'void' ) {
 
-			if ( op === '=' ) {
-
-				builder.addLineFlowCode( `${a} ${this.op} ${b}` );
-
-				return a;
-
-			} else if ( op === '<' && outputLength > 1 ) {
+			if ( op === '<' && outputLength > 1 ) {
 
 				return builder.format( `${ builder.getMethod( 'lessThan' ) }( ${a}, ${b} )`, type, output );
 
@@ -232,7 +216,6 @@ export const mul = nodeProxy( OperatorNode, '*' );
 export const div = nodeProxy( OperatorNode, '/' );
 export const remainder = nodeProxy( OperatorNode, '%' );
 export const equal = nodeProxy( OperatorNode, '==' );
-export const assign = nodeProxy( OperatorNode, '=' );
 export const lessThan = nodeProxy( OperatorNode, '<' );
 export const greaterThan = nodeProxy( OperatorNode, '>' );
 export const lessThanEqual = nodeProxy( OperatorNode, '<=' );
@@ -252,7 +235,6 @@ addNodeElement( 'mul', mul );
 addNodeElement( 'div', div );
 addNodeElement( 'remainder', remainder );
 addNodeElement( 'equal', equal );
-addNodeElement( 'assign', assign );
 addNodeElement( 'lessThan', lessThan );
 addNodeElement( 'greaterThan', greaterThan );
 addNodeElement( 'lessThanEqual', lessThanEqual );

diff --git a/examples/jsm/nodes/shadernode/ShaderNode.js b/examples/jsm/nodes/shadernode/ShaderNode.js
@@ -34,13 +34,13 @@ const shaderNodeHandler = {
 
 	},
 
-	get: function ( node, prop, nodeObj ) {
+	get( node, prop, nodeObj ) {
 
 		if ( typeof prop === 'string' && node[ prop ] === undefined ) {
 
 			if ( node.isStackNode !== true && prop === 'assign' ) {
 
-				return ( ...params ) => assign( node, ...params );
+				return ( ...params ) => currentStack.assign( nodeObj, ...params );
 
 			} else if ( NodeElements.has( prop ) ) {
 
@@ -64,7 +64,7 @@ const shaderNodeHandler = {
 
 				prop = parseSwizzle( prop );
 
-				return nodeObject( new SplitNode( node, prop ) );
+				return nodeObject( new SplitNode( nodeObj, prop ) );
 
 			} else if ( /^set[XYZWRGBASTPQ]{1,4}$/.test( prop ) === true ) {
 
@@ -92,13 +92,33 @@ const shaderNodeHandler = {
 
 				// accessing array
 
-				return nodeObject( new ArrayElementNode( node, new ConstNode( Number( prop ), 'uint' ) ) );
+				return nodeObject( new ArrayElementNode( nodeObj, new ConstNode( Number( prop ), 'uint' ) ) );
+
+			}
+
+		}
+
+		return Reflect.get( node, prop, nodeObj );
+
+	},
+
+	set( node, prop, value, nodeObj ) {
+
+		if ( typeof prop === 'string' && node[ prop ] === undefined ) {
+
+			// setting properties
+
+			if ( /^[xyzwrgbastpq]{1,4}$/.test( prop ) === true || prop === 'width' || prop === 'height' || prop === 'depth' || /^\d+$/.test( prop ) === true ) {
+
+				nodeObj[ prop ].assign( value );
+
+				return true;
 
 			}
 
 		}
 
-		return node[ prop ];
+		return Reflect.set( node, prop, value, nodeObj );
 
 	}
 
@@ -465,8 +485,6 @@ export const getCurrentStack = () => currentStack;
 export const If = ( ...params ) => currentStack.if( ...params );
 export const append = ( ...params ) => currentStack.add( ...params );
 
-const assign = ( ...params ) => currentStack.assign( ...params );
-
 addNodeElement( 'append', append );
 
 // types

diff --git a/examples/jsm/renderers/common/Background.js b/examples/jsm/renderers/common/Background.js
@@ -1,6 +1,6 @@
 import DataMap from './DataMap.js';
 import { Color, Mesh, SphereGeometry, BackSide } from 'three';
-import { context, normalWorld, backgroundBlurriness, backgroundIntensity, NodeMaterial, modelViewProjection } from '../../nodes/Nodes.js';
+import { context, normalWorld, backgroundBlurriness, backgroundIntensity, NodeMaterial, modelViewProjection, tslFn } from '../../nodes/Nodes.js';
 
 let _clearAlpha;
 const _clearColor = new Color();
@@ -59,8 +59,11 @@ class Background extends DataMap {
 					getSamplerLevelNode: () => backgroundBlurriness
 				} ).mul( backgroundIntensity );
 
-				let viewProj = modelViewProjection();
-				viewProj = viewProj.setZ( viewProj.w );
+				const viewProj = tslFn( () => {
+					const matrix = modelViewProjection();
+					matrix.z = matrix.w;
+					return matrix;
+				} )();
 
 				const nodeMaterial = new NodeMaterial();
 				nodeMaterial.outputNode = this.backgroundMeshNode;