From 8dd031eb4124891e1c59fdf8c1808c72e5ead970 Mon Sep 17 00:00:00 2001
From: sunag <sunagbrasil@gmail.com>
Date: Sun, 4 Feb 2024 01:23:34 -0300
Subject: [PATCH] Add StorageArrayElementNode

---
 examples/jsm/nodes/Nodes.js                   |  1 +
 .../jsm/nodes/accessors/StorageBufferNode.js  |  7 +++
 .../nodes/utils/StorageArrayElementNode.js    | 53 +++++++++++++++++++
 3 files changed, 61 insertions(+)
 create mode 100644 examples/jsm/nodes/utils/StorageArrayElementNode.js

diff --git a/examples/jsm/nodes/Nodes.js b/examples/jsm/nodes/Nodes.js
index d0fedfa71acfab..6633c481fb6f3c 100644
--- a/examples/jsm/nodes/Nodes.js
+++ b/examples/jsm/nodes/Nodes.js
@@ -68,6 +68,7 @@ export { default as SetNode } from './utils/SetNode.js';
 export { default as SpecularMIPLevelNode, specularMIPLevel } from './utils/SpecularMIPLevelNode.js';
 export { default as SplitNode } from './utils/SplitNode.js';
 export { default as SpriteSheetUVNode, spritesheetUV } from './utils/SpriteSheetUVNode.js';
+export { default as StorageArrayElementNode } from './utils/StorageArrayElementNode.js';
 export { default as TimerNode, timerLocal, timerGlobal, timerDelta, frameId } from './utils/TimerNode.js';
 export { default as TriplanarTexturesNode, triplanarTextures, triplanarTexture } from './utils/TriplanarTexturesNode.js';
 export { default as ReflectorNode, reflector } from './utils/ReflectorNode.js';
diff --git a/examples/jsm/nodes/accessors/StorageBufferNode.js b/examples/jsm/nodes/accessors/StorageBufferNode.js
index b7599db4017fe5..1e3595b44aa3c0 100644
--- a/examples/jsm/nodes/accessors/StorageBufferNode.js
+++ b/examples/jsm/nodes/accessors/StorageBufferNode.js
@@ -3,6 +3,7 @@ import { bufferAttribute } from './BufferAttributeNode.js';
 import { addNodeClass } from '../core/Node.js';
 import { nodeObject } from '../shadernode/ShaderNode.js';
 import { varying } from '../core/VaryingNode.js';
+import { storageElement } from '../utils/StorageArrayElementNode.js';
 
 class StorageBufferNode extends BufferNode {
 
@@ -23,6 +24,12 @@ class StorageBufferNode extends BufferNode {
 
 	}
 
+	element( indexNode ) {
+
+		return storageElement( this, indexNode );
+
+	}
+
 	generate( builder ) {
 
 		if ( builder.isAvailable( 'storageBuffer' ) ) return super.generate( builder );
diff --git a/examples/jsm/nodes/utils/StorageArrayElementNode.js b/examples/jsm/nodes/utils/StorageArrayElementNode.js
new file mode 100644
index 00000000000000..9f8ad009ccdca0
--- /dev/null
+++ b/examples/jsm/nodes/utils/StorageArrayElementNode.js
@@ -0,0 +1,53 @@
+import { addNodeClass } from '../core/Node.js';
+import { nodeProxy, addNodeElement } from '../shadernode/ShaderNode.js';
+import ArrayElementNode from './ArrayElementNode.js';
+
+class StorageArrayElementNode extends ArrayElementNode {
+
+	constructor( storageBufferNode, indexNode ) {
+
+		super( storageBufferNode, indexNode );
+
+		this.isStorageArrayElementNode = true;
+
+	}
+
+	set storageBufferNode( value ) {
+
+		this.node = value;
+
+	}
+
+	get storageBufferNode() {
+
+		return this.node;
+
+	}
+
+	generate( builder ) {
+
+		let snippet;
+
+		if ( builder.isAvailable( 'storageBuffer' ) === false ) {
+
+			snippet = this.node.build( builder );
+
+		} else {
+
+			snippet = super.generate( builder );
+
+		}
+
+		return snippet;
+
+	}
+
+}
+
+export default StorageArrayElementNode;
+
+export const storageElement = nodeProxy( StorageArrayElementNode );
+
+addNodeElement( 'storageElement', storageElement );
+
+addNodeClass( 'StorageArrayElementNode', StorageArrayElementNode );