From 17734ce446a96227ffe23b2dd385ba84f36da581 Mon Sep 17 00:00:00 2001
From: sunag <sunagbrasil@gmail.com>
Date: Wed, 21 Feb 2024 16:52:46 -0300
Subject: [PATCH] TSL: Support for split assignment

---
 examples/jsm/nodes/core/AssignNode.js         | 62 ++++++++++++++++---
 .../renderers/webgl/nodes/GLSLNodeBuilder.js  |  4 +-
 examples/webgpu_compute_particles_rain.html   |  3 +-
 3 files changed, 56 insertions(+), 13 deletions(-)

diff --git a/examples/jsm/nodes/core/AssignNode.js b/examples/jsm/nodes/core/AssignNode.js
index 971135541e8b1..b54c1d830e04b 100644
--- 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
index a773a9927f99f..95ce76592a372 100644
--- 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
index f0bbcfe60dc34..fc4a5dd1a7f4f 100644
--- 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