From 4cc8fdfa944181c4274c91061574bcf033315bee Mon Sep 17 00:00:00 2001
From: sunag <sunagbrasil@gmail.com>
Date: Fri, 26 Jul 2024 02:14:18 -0300
Subject: [PATCH] TSL: Add `spherizeUV` (#28976)

* TSL: Add `spherizeUV`

* update imports
---
 examples/webgpu_tsl_vfx_flames.html | 18 +++------------
 src/nodes/Nodes.js                  |  2 +-
 src/nodes/utils/RotateUVNode.js     | 35 -----------------------------
 src/nodes/utils/UVUtils.js          | 21 +++++++++++++++++
 4 files changed, 25 insertions(+), 51 deletions(-)
 delete mode 100644 src/nodes/utils/RotateUVNode.js
 create mode 100644 src/nodes/utils/UVUtils.js

diff --git a/examples/webgpu_tsl_vfx_flames.html b/examples/webgpu_tsl_vfx_flames.html
index 3d6581cbe860b9..c5fa0b106cf12e 100644
--- a/examples/webgpu_tsl_vfx_flames.html
+++ b/examples/webgpu_tsl_vfx_flames.html
@@ -27,7 +27,7 @@
 		<script type="module">
 
 			import * as THREE from 'three';
-			import { PI2, dot, sin, step, texture, timerLocal, tslFn, uv, vec2, vec3, vec4, mix } from 'three/tsl';
+			import { PI2, spherizeUV, sin, step, texture, timerLocal, tslFn, uv, vec2, vec3, vec4, mix } from 'three/tsl';
 
 			import { OrbitControls } from 'three/addons/controls/OrbitControls.js';
 
@@ -50,18 +50,6 @@
 				const cellularTexture = textureLoader.load( './textures/noises/voronoi/grayscale-256x256.png' );
 				const perlinTexture = textureLoader.load( './textures/noises/perlin/rgb-256x256.png' );
 
-				// TSL functions
-
-				const spherizeUv = tslFn( ( [ input, center, strength, offset ] ) => {
-
-					const delta = input.sub( center );
-					const delta2 = dot( delta, delta );
-					const delta4 = delta2.mul( delta2 );
-					const deltaOffset = delta4.mul( strength );
-					return input.add( delta.mul( deltaOffset ) ).add( offset );
-
-				} );
-
 				// gradient canvas
 
 				const gradient = {};
@@ -112,7 +100,7 @@
 
 					// main UV
 					const mainUv = uv().toVar();
-					mainUv.assign( spherizeUv( mainUv, vec2( 0.5 ), 10, vec2( 0 ) ).mul( 0.6 ).add( 0.2 ) ); // spherize
+					mainUv.assign( spherizeUV( mainUv, 10 ).mul( 0.6 ).add( 0.2 ) ); // spherize
 					mainUv.assign( mainUv.pow( vec2( 1, 2 ) ) ); // stretch
 					mainUv.assign( mainUv.mul( 2, 1 ).sub( vec2( 0.5, 0 ) ) ); // scale
 
@@ -150,7 +138,7 @@
 
 					// main UV
 					const mainUv = uv().toVar();
-					mainUv.assign( spherizeUv( mainUv, vec2( 0.5 ), 10, vec2( 0 ) ).mul( 0.6 ).add( 0.2 ) ); // spherize
+					mainUv.assign( spherizeUV( mainUv, 10 ).mul( 0.6 ).add( 0.2 ) ); // spherize
 					mainUv.assign( mainUv.pow( vec2( 1, 3 ) ) ); // stretch
 					mainUv.assign( mainUv.mul( 2, 1 ).sub( vec2( 0.5, 0 ) ) ); // scale
 
diff --git a/src/nodes/Nodes.js b/src/nodes/Nodes.js
index b7439ab67a7bce..e8cf03d535c434 100644
--- a/src/nodes/Nodes.js
+++ b/src/nodes/Nodes.js
@@ -62,7 +62,7 @@ export { default as MaxMipLevelNode, maxMipLevel } from './utils/MaxMipLevelNode
 export { default as OscNode, oscSine, oscSquare, oscTriangle, oscSawtooth } from './utils/OscNode.js';
 export { default as PackingNode, directionToColor, colorToDirection } from './utils/PackingNode.js';
 export { default as RemapNode, remap, remapClamp } from './utils/RemapNode.js';
-export { default as RotateUVNode, rotateUV } from './utils/RotateUVNode.js';
+export * from './utils/UVUtils.js';
 export { default as RotateNode, rotate } from './utils/RotateNode.js';
 export { default as SetNode } from './utils/SetNode.js';
 export { default as SplitNode } from './utils/SplitNode.js';
diff --git a/src/nodes/utils/RotateUVNode.js b/src/nodes/utils/RotateUVNode.js
deleted file mode 100644
index 4f432cd09d7873..00000000000000
--- a/src/nodes/utils/RotateUVNode.js
+++ /dev/null
@@ -1,35 +0,0 @@
-import TempNode from '../core/TempNode.js';
-import { addNodeClass } from '../core/Node.js';
-import { addNodeElement, nodeProxy, vec2 } from '../shadernode/ShaderNode.js';
-
-class RotateUVNode extends TempNode {
-
-	constructor( uvNode, rotationNode, centerNode = vec2( 0.5 ) ) {
-
-		super( 'vec2' );
-
-		this.uvNode = uvNode;
-		this.rotationNode = rotationNode;
-		this.centerNode = centerNode;
-
-	}
-
-	setup() {
-
-		const { uvNode, rotationNode, centerNode } = this;
-
-		const vector = uvNode.sub( centerNode );
-
-		return vector.rotate( rotationNode ).add( centerNode );
-
-	}
-
-}
-
-export default RotateUVNode;
-
-export const rotateUV = nodeProxy( RotateUVNode );
-
-addNodeElement( 'rotateUV', rotateUV );
-
-addNodeClass( 'RotateUVNode', RotateUVNode );
diff --git a/src/nodes/utils/UVUtils.js b/src/nodes/utils/UVUtils.js
new file mode 100644
index 00000000000000..f7eaa7206d0881
--- /dev/null
+++ b/src/nodes/utils/UVUtils.js
@@ -0,0 +1,21 @@
+import { addNodeElement, tslFn, vec2 } from '../shadernode/ShaderNode.js';
+
+export const rotateUV = tslFn( ( [ uv, rotation, center = vec2( 0.5 ) ] ) => {
+
+	return uv.sub( center ).rotate( rotation ).add( center );
+
+} );
+
+export const spherizeUV = tslFn( ( [ uv, strength, center = vec2( 0.5 ) ] ) => {
+
+	const delta = uv.sub( center );
+	const delta2 = delta.dot( delta );
+	const delta4 = delta2.mul( delta2 );
+	const deltaOffset = delta4.mul( strength );
+
+	return uv.add( delta.mul( deltaOffset ) );
+
+} );
+
+addNodeElement( 'rotateUV', rotateUV );
+addNodeElement( 'spherizeUV', spherizeUV );