NodeMaterial: Introduce `.geometryNode` and jelly example (#29551)

* fix `compute()` during `render()` call

* NodeMaterial: Introduce `.geometryNode`

* update compute geometry example

* rename jellyNode -> jelly

* cleanup

examples/webgpu_compute_geometry.html

@@ -25,83 +25,139 @@
 		<script type="module">
 
 			import * as THREE from 'three';
-			import { vec3, cos, sin, mat3, storage, Fn, instanceIndex, timerLocal } from 'three/tsl';
+			import { vec3, vec4, storage, Fn, If, uniform, instanceIndex, objectWorldMatrix, color, screenUV, attribute } from 'three/tsl';
 
 			import { GLTFLoader } from 'three/addons/loaders/GLTFLoader.js';
 
 			import { OrbitControls } from 'three/addons/controls/OrbitControls.js';
 
+			import { GUI } from 'three/addons/libs/lil-gui.module.min.js';
+
+			import Stats from 'three/addons/libs/stats.module.js';
+
 			let camera, scene, renderer;
-			let computeUpdate;
+			let raycaster, pointer;
+			let stats;
+
+			const pointerPosition = uniform( vec4( 0 ) );
+			const elasticity = uniform( .4 ); // elasticity ( how "strong" the spring is )
+			const damping = uniform( .94 ); // damping factor ( energy loss )
+			const brushSize = uniform( .25 );
+			const brushStrength = uniform( .22 );
 
 			init();
 
-			function init() {
+			const jelly = Fn( ( { renderer, geometry, object } ) => {
 
-				camera = new THREE.PerspectiveCamera( 50, window.innerWidth / window.innerHeight, 0.1, 10 );
-				camera.position.set( 0, 0, 1 );
+				const count = geometry.attributes.position.count;
 
-				scene = new THREE.Scene();
-				scene.background = new THREE.Color( 0x333333 );
+				const speedBufferAttribute = new THREE.StorageBufferAttribute( count, 4 );
 
-				new GLTFLoader().load( 'models/gltf/LeePerrySmith/LeePerrySmith.glb', function ( gltf ) {
+				// replace geometry attributes for storage buffer attributes
 
-					const mesh = gltf.scene.children[ 0 ];
-					mesh.scale.setScalar( .1 );
-					mesh.material = new THREE.MeshNormalMaterial();
-					scene.add( mesh );
+				const positionBaseAttribute = geometry.attributes.position;
+				const positionStorageBufferAttribute = new THREE.StorageBufferAttribute( count, 4 );
+
+				geometry.setAttribute( 'storagePosition', positionStorageBufferAttribute );
 
-					//
+				// compute ( jelly )
 
-					const positionBaseAttribute = mesh.geometry.attributes.position;
-					const normalBaseAttribute = mesh.geometry.attributes.normal;
+				const positionAttribute = storage( positionBaseAttribute, 'vec3', count ).toReadOnly();
+				const positionStorageAttribute = storage( positionStorageBufferAttribute, 'vec4', count );
 
-					// replace geometry attributes for storage buffer attributes
+				const speedAttribute = storage( speedBufferAttribute, 'vec4', count );
 
-					const positionStorageBufferAttribute = new THREE.StorageBufferAttribute( positionBaseAttribute.count, 4 );
-					const normalStorageBufferAttribute = new THREE.StorageBufferAttribute( normalBaseAttribute.count, 4 );
+				// vectors
 
-					mesh.geometry.setAttribute( 'position', positionStorageBufferAttribute );
-					mesh.geometry.setAttribute( 'normal', normalStorageBufferAttribute );
+				const basePosition = vec3( positionAttribute.element( instanceIndex ) );
+				const currentPosition = positionStorageAttribute.element( instanceIndex );
+				const currentSpeed = speedAttribute.element( instanceIndex );
 
-					// compute shader
+				//
 
-					const computeFn = Fn( () => {
+				const computeInit = Fn( () => {
 
-						const positionAttribute = storage( positionBaseAttribute, 'vec3', positionBaseAttribute.count ).toReadOnly();
-						const normalAttribute = storage( normalBaseAttribute, 'vec3', normalBaseAttribute.count ).toReadOnly();
+					// copy position to storage
 
-						const positionStorageAttribute = storage( positionStorageBufferAttribute, 'vec4', positionStorageBufferAttribute.count );
-						const normalStorageAttribute = storage( normalStorageBufferAttribute, 'vec4', normalStorageBufferAttribute.count );
+					currentPosition.assign( basePosition );
 
-						const time = timerLocal( 1 );
-						const scale = 0.3;
+				} )().compute( count );
 
-						//
+				//
 
-						const position = vec3( positionAttribute.element( instanceIndex ) );
-						const normal = vec3( normalAttribute.element( instanceIndex ) );
+				const computeUpdate = Fn( () => {
 
-						const theta = sin( time.add( position.y ) ).mul( scale );
+					// pinch
 
-						const c = cos( theta );
-						const s = sin( theta );
+					If( pointerPosition.w.equal( 1 ), () => {
 
-						const m = mat3(
-							c, 0, s,
-							0, 1, 0,
-							s.negate(), 0, c
-						);
+						const worldPosition = objectWorldMatrix( object ).mul( currentPosition.xyz );
 
-						const transformed = position.mul( m );
-						const transformedNormal = normal.mul( m );
+						const dist = worldPosition.distance( pointerPosition.xyz );
+						const direction = pointerPosition.xyz.sub( worldPosition ).normalize();
 
-						positionStorageAttribute.element( instanceIndex ).assign( transformed );
-						normalStorageAttribute.element( instanceIndex ).assign( transformedNormal );
+						const power = brushSize.sub( dist ).max( 0 ).mul( brushStrength );
+
+						currentPosition.xyz.addAssign( direction.mul( power ) );
 
 					} );
 
-					computeUpdate = computeFn().compute( positionBaseAttribute.count );
+					// update
+
+					const distance = basePosition.distance( currentPosition );
+					const force = elasticity.mul( distance ).mul( basePosition.sub( currentPosition ) );
+
+					currentSpeed.addAssign( force );
+					currentSpeed.mulAssign( damping );
+
+					currentPosition.addAssign( currentSpeed );
+
+				} )().compute( count );
+
+				// initialize the storage buffer with the base position
+
+				computeUpdate.onInit( () => renderer.compute( computeInit ) );
+
+				//
+
+				return computeUpdate;
+
+			} );
+
+			function init() {
+
+				camera = new THREE.PerspectiveCamera( 50, window.innerWidth / window.innerHeight, 0.1, 10 );
+				camera.position.set( 0, 0, 1 );
+
+				scene = new THREE.Scene();
+
+				raycaster = new THREE.Raycaster();
+				pointer = new THREE.Vector2();
+
+				// background
+
+				const bgColor = screenUV.y.mix( color( 0x9f87f7 ), color( 0xf2cdcd ) );
+				const bgVignet = screenUV.distance( .5 ).remapClamp( 0.3, .8 ).oneMinus();
+				const bgIntensity = 4;
+
+				scene.backgroundNode = bgColor.mul( bgVignet.mul( color( 0xa78ff6 ).mul( bgIntensity ) ) );
+
+				// model
+
+				new GLTFLoader().load( 'models/gltf/LeePerrySmith/LeePerrySmith.glb', function ( gltf ) {
+
+					// create jelly effect material
+
+					const material = new THREE.MeshNormalNodeMaterial();
+					material.geometryNode = jelly();
+					material.positionNode = attribute( 'storagePosition' );
+
+					// apply the material to the mesh
+
+					const mesh = gltf.scene.children[ 0 ];
+					mesh.scale.setScalar( .1 );
+					mesh.material = material;
+					scene.add( mesh );
 
 				} );
 
@@ -117,7 +173,40 @@
 				controls.minDistance = .7;
 				controls.maxDistance = 2;
 
+				const gui = new GUI();
+				gui.add( elasticity, 'value', 0, .5 ).name( 'elasticity' );
+				gui.add( damping, 'value', .9, .98 ).name( 'damping' );
+				gui.add( brushSize, 'value', .1, .5 ).name( 'brush size' );
+				gui.add( brushStrength, 'value', .1, .3 ).name( 'brush strength' );
+
+				stats = new Stats();
+				document.body.appendChild( stats.dom );
+
 				window.addEventListener( 'resize', onWindowResize );
+				window.addEventListener( 'pointermove', onPointerMove );
+
+			}
+
+			function onPointerMove( event ) {
+
+				pointer.set( ( event.clientX / window.innerWidth ) * 2 - 1, - ( event.clientY / window.innerHeight ) * 2 + 1 );
+
+				raycaster.setFromCamera( pointer, camera );
+
+				const intersects = raycaster.intersectObject( scene );
+
+				if ( intersects.length > 0 ) {
+
+					const intersect = intersects[ 0 ];
+
+					pointerPosition.value.copy( intersect.point );
+					pointerPosition.value.w = 1; // enable
+
+				} else {
+
+					pointerPosition.value.w = 0; // disable
+
+				}
 
 			}
 
@@ -132,7 +221,7 @@
 
 			async function animate() {
 
-				if ( computeUpdate ) renderer.compute( computeUpdate );
+				stats.update();
 
 				renderer.render( scene, camera );
 

src/materials/nodes/NodeMaterial.js

@@ -57,6 +57,7 @@ class NodeMaterial extends Material {
 		this.alphaTestNode = null;
 
 		this.positionNode = null;
+		this.geometryNode = null;
 
 		this.depthNode = null;
 		this.shadowNode = null;
@@ -98,6 +99,12 @@ class NodeMaterial extends Material {
 
 		builder.stack.outputNode = this.vertexNode || this.setupPosition( builder );
 
+		if ( this.geometryNode !== null ) {
+
+			builder.stack.outputNode = builder.stack.outputNode.bypass( this.geometryNode );
+
+		}
+
 		builder.addFlow( 'vertex', builder.removeStack() );
 
 		// < FRAGMENT STAGE >
@@ -634,6 +641,7 @@ class NodeMaterial extends Material {
 		this.alphaTestNode = source.alphaTestNode;
 
 		this.positionNode = source.positionNode;
+		this.geometryNode = source.geometryNode;
 
 		this.depthNode = source.depthNode;
 		this.shadowNode = source.shadowNode;

src/renderers/webgl-fallback/WebGLBackend.js

@@ -605,6 +605,12 @@ class WebGLBackend extends Backend {
 
 		this.prepareTimestampBuffer( computeGroup );
 
+		if ( this._currentContext ) {
+
+			this._setFramebuffer( this._currentContext );
+
+		}
+
 	}
 
 	draw( renderObject/*, info*/ ) {