Introduce storageStruct

src/nodes/TSL.js

@@ -11,6 +11,7 @@ export * from './core/IndexNode.js';
 export * from './core/ParameterNode.js';
 export * from './core/PropertyNode.js';
 export * from './core/StackNode.js';
+export * from './core/StructTypeNode.js';
 export * from './core/UniformGroupNode.js';
 export * from './core/UniformNode.js';
 export * from './core/VaryingNode.js';

src/nodes/accessors/StorageBufferNode.js

@@ -22,6 +22,7 @@ class StorageBufferNode extends BufferNode {
 		this.isAtomic = false;
 
 		this.bufferObject = false;
+		this.bufferStruct = false;
 		this.bufferCount = bufferCount;
 
 		this._attribute = null;
@@ -84,6 +85,14 @@ class StorageBufferNode extends BufferNode {
 
 	}
 
+	setBufferStruct( value ) {
+
+		this.bufferStruct = value;
+
+		return this;
+
+	}
+
 	setAccess( value ) {
 
 		this.access = value;
@@ -166,4 +175,5 @@ export default StorageBufferNode;
 
 // Read-Write Storage
 export const storage = ( value, type, count ) => nodeObject( new StorageBufferNode( value, type, count ) );
+export const storageStruct = ( value, type, count ) => nodeObject( new StorageBufferNode( value, type, count ).setBufferStruct( true ) );
 export const storageObject = ( value, type, count ) => nodeObject( new StorageBufferNode( value, type, count ).setBufferObject( true ) );

src/nodes/core/StructTypeNode.js

@@ -26,3 +26,19 @@ class StructTypeNode extends Node {
 }
 
 export default StructTypeNode;
+
+export const struct = ( members ) => {
+
+	return Object.entries( members ).map( ( [ name, value ] ) => {
+
+		if ( typeof value === 'string' ) {
+
+			return { name, type: value, isAtomic: false };
+
+		}
+
+		return { name, type: value.type, isAtomic: value.atomic || false };
+
+	} );
+
+};

src/renderers/webgpu/nodes/WGSLNodeBuilder.js

@@ -964,6 +964,88 @@
 
 	}
 
+	getTypeFromCustomStruct( shaderStage ) {
+
+		const uniforms = this.uniforms[ shaderStage ];
+
+		const bufferStructMap = new Map();
+
+		uniforms.forEach( ( uniform ) => {
+
+			const { name, node } = uniform;
+			const hasBufferStruct = node.bufferStruct === true;
+			bufferStructMap.set( name, hasBufferStruct );
+
+		} );
+
+		return bufferStructMap;
+
+	}
+
+	getMembersFromCustomStruct( members ) {
+
+		const structMembers = members.map( ( { name, type, isAtomic } ) => {
+
+			let finalType = wgslTypeLib[ type ];
+
+			if ( ! finalType ) {
+
+				console.warn( `Unrecognized type: ${type}` );
+				finalType = 'vec4<f32>';
+
+			}
+
+			return `${name}: ${isAtomic ? `atomic<${finalType}>` : finalType}`;
+
+		} );
+
+		return `\t${structMembers.join( ',\n\t' )}`;
+
+	}
+
+	getCustomStructNameFromShader( source ) {
+
+		const functionRegex = /fn\s+\w+\s*\(([\s\S]*?)\)/g; // filter shader header
+		const parameterRegex = /(\w+)\s*:\s*(ptr<\s*([\w]+),\s*(?:array<([\w<>]+)>|(\w+))[^>]*>|[\w<>,]+)/g; // filter parameters
+
+		const results = [];
+
+		let match;
+
+		while ( ( match = functionRegex.exec( source ) ) !== null ) {
+
+			const parameterString = match[ 1 ];
+
+			let paramMatch;
+
+			while ( ( paramMatch = parameterRegex.exec( parameterString ) ) !== null ) {
+
+				const [ fullMatch, name, fullType, ptrType, arrayType, directStructName ] = paramMatch;
+
+				const structName = arrayType || directStructName || null;
+
+				const type = ptrType || fullType;
+
+				if (Object.values(wgslTypeLib).includes(structName) || structName === null) {
+
+					continue;
+
+				}
+
+				results.push( {
+					name,
+					type,
+					structName
+				} );
+
+			}
+
+		}
+
+		return results;
+
+	}
+
 	getStructMembers( struct ) {
 
 		const snippets = [];
@@ -1171,12 +1253,14 @@
 				const bufferType = this.getType( bufferNode.bufferType );
 				const bufferCount = bufferNode.bufferCount;
 
+
+				const isArray = bufferNode.value.array.length !== bufferNode.value.itemSize;
 				const bufferCountSnippet = bufferCount > 0 && uniform.type === 'buffer' ? ', ' + bufferCount : '';
 				const bufferTypeSnippet = bufferNode.isAtomic ? `atomic<${bufferType}>` : `${bufferType}`;
-				const bufferSnippet = `\t${ uniform.name } : array< ${ bufferTypeSnippet }${ bufferCountSnippet } >\n`;
+				const bufferSnippet = bufferNode.bufferStruct ? this.getMembersFromCustomStruct( bufferType ) : `\t${ uniform.name } : array< ${ bufferTypeSnippet }${ bufferCountSnippet } >\n`;
 				const bufferAccessMode = bufferNode.isStorageBufferNode ? `storage, ${ this.getStorageAccess( bufferNode ) }` : 'uniform';
 
-				bufferSnippets.push( this._getWGSLStructBinding( 'NodeBuffer_' + bufferNode.id, bufferSnippet, bufferAccessMode, uniformIndexes.binding ++, uniformIndexes.group ) );
+				bufferSnippets.push( this._getWGSLStructBinding( bufferNode.bufferStruct, isArray, 'NodeBuffer_' + bufferNode.id, bufferSnippet, bufferAccessMode, uniformIndexes.binding ++, uniformIndexes.group ) );
 
 			} else {
 
@@ -1199,7 +1283,7 @@
 
 			const group = uniformGroups[ name ];
 
-			structSnippets.push( this._getWGSLStructBinding( name, group.snippets.join( ',\n' ), 'uniform', group.index, group.id ) );
+			structSnippets.push( this._getWGSLStructBinding( false, false, name, group.snippets.join( ',\n' ), 'uniform', group.index, group.id ) );
 
 		}
 
@@ -1228,9 +1312,69 @@
 			stageData.codes = this.getCodes( shaderStage );
 			stageData.directives = this.getDirectives( shaderStage );
 			stageData.scopedArrays = this.getScopedArrays( shaderStage );
+			stageData.isBufferStruct = this.getTypeFromCustomStruct( shaderStage );
+			stageData.customStructNames = this.getCustomStructNameFromShader( stageData.codes );
 
 			//
 
+			const reduceFlow = ( flow ) => {
+
+				return flow.replace( /&(\w+)\.(\w+)/g, ( match, bufferName, uniformName ) =>
+
+					stageData.isBufferStruct.get( uniformName ) === true ? `&${bufferName}` : match
+
+				);
+
+			};
+
+			const extractPointerNames = ( source ) => {
+
+				const match = source.match( /\(([^)]+)\)/ );
+				if ( ! match ) return [];
+
+				const content = match[ 1 ];
+
+				return content
+					.split( /\s*,\s*/ )
+					.map( part => part.trim() )
+					.filter( part => part.includes( '&' ) )
+					.map( part => part.replace( '&', '' ) )
+					.filter( part => ! part.includes( '.' ) );
+
+			};
+
+			const createStructNameMapping = ( nodeBuffers, structs ) => {
+
+				const resultMap = new Map();
+
+				for (let i = 0; i < nodeBuffers.length; i++) {
+
+					const bufferName = nodeBuffers[i];
+					const struct = structs[i];
+
+					resultMap.set(bufferName, struct.structName);
+
+				}
+
+				return resultMap;
+
+			};
+
+			const replaceStructNamesInUniforms = ( shaderCode, map ) => {
+
+				for ( const [ key, value ] of map.entries() ) {
+
+					const regex = new RegExp( `\\b${key}Struct\\b`, 'g' );
+					shaderCode = shaderCode.replace( regex, value );
+
+				}
+
+				return shaderCode;
+
+			};
+
+
+			let pointerNames, structnameMapping;
 			let flow = '// code\n\n';
 			flow += this.flowCode[ shaderStage ];
 
@@ -1255,6 +1399,68 @@
 
 				flow += `${ flowSlotData.code }\n\t`;
 
+
+				/*
+				REVIEW COMMENT remove after review
+
+				before reduceFlow:
+				compute( &NodeBuffer_554.nodeUniform0, &NodeBuffer_558.nodeUniform1, &NodeBuffer_555.nodeUniform2, &NodeBuffer_556.nodeUniform3, &NodeBuffer_557.nodeUniform4, &NodeBuffer_559.nodeUniform5, instanceIndex, object.nodeUniform6 );
+
+				after reduceFlow: reduceFlow checks whether there is a storageStruct and if so then the 
+				postfix is ​​removed so that the pointer points to the struct and not to a content in the struct
+				compute( &NodeBuffer_554, &NodeBuffer_558, &NodeBuffer_555.nodeUniform2, &NodeBuffer_556.nodeUniform3, &NodeBuffer_557.nodeUniform4, &NodeBuffer_559.nodeUniform5, instanceIndex, object.nodeUniform6 );
+
+				extractPointerNames reads the names of the reduced pointers and stores them in an array
+				Array(2)
+				0: "NodeBuffer_554"
+				1: "NodeBuffer_558"
+
+				getCustomStructNameFromShader at the beginning reads the structNames from the shader header
+				Array(2)
+				0: {name: 'drawBuffer', type: 'storage', structName: 'DrawBuffer'}
+				1: {name: 'meshletInfo', type: 'storage', structName: 'MeshletInfo'}
+
+
+				createStructNameMapping links the automatic generated WGSLNodeBuilder for each struct with the struct name specified by the user. 
+				This is necessary because in wgslFn the user can choose any name in the shader in ptr for structs.
+
+				Map(2)
+				[[Entries]]
+				0: {"NodeBuffer_554" => "DrawBuffer"}
+				1: {"NodeBuffer_558" => "MeshletInfo"}
+
+				replaceStructNames then replaces the names in the uniforms in the custom structs that the WGSLNodeBuilder 
+				created with the name chosen by the user.
+
+				before replaceStructNames:
+
+				struct NodeBuffer_554Struct {
+					vertexCount: u32,
+					instanceCount: u32,
+					firstVertex: u32,
+					firstInstance: u32
+				};
+				@binding( 0 ) @group( 0 )
+				var<storage, read_write> NodeBuffer_554 : NodeBuffer_554Struct;
+
+				after replaceStructNames:
+
+				struct DrawBuffer {
+					vertexCount: u32,
+					instanceCount: u32,
+					firstVertex: u32,
+					firstInstance: u32
+				};
+				@binding( 0 ) @group( 0 )
+				var<storage, read_write> NodeBuffer_554 : DrawBuffer;
+				*/
+
+
+				flow = reduceFlow( flow );
+				pointerNames = extractPointerNames(flow);
+				structnameMapping = createStructNameMapping(pointerNames, stageData.customStructNames);
+				stageData.uniforms = replaceStructNamesInUniforms(stageData.uniforms, structnameMapping);
+
 				if ( node === mainNode && shaderStage !== 'compute' ) {
 
 					flow += '// result\n\n\t';
@@ -1289,6 +1495,11 @@
 
 					}
 
+					flow = reduceFlow( flow );
+					pointerNames = extractPointerNames(flow);
+					structnameMapping = createStructNameMapping(pointerNames, stageData.customStructNames);
+					stageData.uniforms = replaceStructNamesInUniforms(stageData.uniforms, structnameMapping);
+
 				}
 
 			}
@@ -1493,15 +1704,27 @@
 
 	}
 
-	_getWGSLStructBinding( name, vars, access, binding = 0, group = 0 ) {
+	_getWGSLStructBinding( isBufferStruct, isArray, name, vars, access, binding = 0, group = 0 ) {
 
 		const structName = name + 'Struct';
 		const structSnippet = this._getWGSLStruct( structName, vars );
 
-		return `${structSnippet}
+		if ( ! isBufferStruct ) {
+
+			return `${structSnippet}
 @binding( ${binding} ) @group( ${group} )
 var<${access}> ${name} : ${structName};`;
 
+		} else {
+
+			const StructName = isArray ? `array<${structName}>` : structName;
+
+			return `${structSnippet}
+@binding( ${binding} ) @group( ${group} )
+var<${access}> ${name} : ${StructName};`;
+
+		}
+
 	}
 
 }