TSL: Support for split assignment

mrdoob · Feb 21, 2024 · 17734ce · 17734ce
1 parent 4abd4f2
commit 17734ce
Show file tree

Hide file tree

Showing 3 changed files with 56 additions and 13 deletions.
diff --git a/examples/jsm/nodes/core/AssignNode.js b/examples/jsm/nodes/core/AssignNode.js
@@ -1,6 +1,7 @@
 import { addNodeClass } from '../core/Node.js';
 import TempNode from '../core/TempNode.js';
 import { addNodeElement, nodeProxy } from '../shadernode/ShaderNode.js';
+import { vectorComponents } from '../core/constants.js';
 
 class AssignNode extends TempNode {
 
@@ -25,39 +26,82 @@ class AssignNode extends TempNode {
 
 	}
 
+	needsSplitAssign( builder ) {
+
+		const { targetNode } = this;
+
+		if ( builder.isAvailable( 'swizzleAssign' ) === false && targetNode.isSplitNode && targetNode.components.length > 1 ) {
+
+			const targetLength = builder.getTypeLength( targetNode.getNodeType( builder ) );
+			const assignDiferentVector = vectorComponents.join( '' ).slice( 0, targetLength ) !== targetNode.components;
+
+			return assignDiferentVector;
+
+		}
+
+		return false;
+
+	}
+
 	generate( builder, output ) {
 
 		const { targetNode, sourceNode } = this;
 
+		const needsSplitAssign = this.needsSplitAssign( builder );
+
+		if ( needsSplitAssign ) sourceNode.increaseUsage( builder ); // flag to cache system
+
 		const targetType = targetNode.getNodeType( builder );
 
 		const target = targetNode.context( { assign: true } ).build( builder );
 		const source = sourceNode.build( builder, targetType );
 
-		const snippet = `${ target } = ${ source }`;
+		const sourceType = sourceNode.getNodeType( builder );
 
-		if ( output === 'void' ) {
+		//
 
-			builder.addLineFlowCode( snippet );
+		let snippet;
 
-			return;
+		if ( needsSplitAssign ) {
 
-		} else {
+			const targetRoot = targetNode.node.context( { assign: true } ).build( builder );
+
+			for ( let i = 0; i < targetNode.components.length; i ++ ) {
 
-			const sourceType = sourceNode.getNodeType( builder );
+				const component = targetNode.components[ i ];
 
-			if ( sourceType === 'void' ) {
+				snippet = `${ targetRoot }.${ component } = ${ source }[ ${ i } ]`;
 
 				builder.addLineFlowCode( snippet );
 
-				return target;
+			}
+
+			if ( output !== 'void' ) {
+
+				snippet = target;
 
 			}
 
-			return builder.format( snippet, targetType, output );
+		} else {
+
+			snippet = `${ target } = ${ source }`;
+
+			if ( output === 'void' || sourceType === 'void' ) {
+
+				builder.addLineFlowCode( snippet );
+
+				if ( output !== 'void' ) {
+
+					snippet = target;
+
+				}
+
+			}
 
 		}
 
+		return builder.format( snippet, targetType, output );
+
 	}
 
 }

diff --git a/examples/jsm/renderers/webgl/nodes/GLSLNodeBuilder.js b/examples/jsm/renderers/webgl/nodes/GLSLNodeBuilder.js
@@ -19,7 +19,8 @@ const precisionLib = {
 };
 
 const supports = {
-	instance: true
+	instance: true,
+	swizzleAssign: true
 };
 
 const defaultPrecisions = `
@@ -540,7 +541,6 @@ ${ flowData.code }
 
 	}
 
-
 	isFlipY() {
 
 		return true;

diff --git a/examples/webgpu_compute_particles_rain.html b/examples/webgpu_compute_particles_rain.html
@@ -167,9 +167,8 @@
 
 						position.y = 25;
 
-						ripplePosition.x = position.x;
+						ripplePosition.xz = position.xz;
 						ripplePosition.y = floorPosition;
-						ripplePosition.z = position.z;
 
 						// reset hit time: x = time