fix uniform-buffer-layout size calculation and add support for arrays (#2398)

Author: georgioskarnas

modules/core/src/portable/uniform-buffer-layout.ts

@@ -29,14 +29,19 @@ export class UniformBufferLayout {
   readonly byteLength: number;
 
   /** Create a new UniformBufferLayout given a map of attributes. */
-  constructor(uniformTypes: Record<string, VariableShaderType>) {
+  constructor(
+    uniformTypes: Record<string, VariableShaderType>,
+    uniformSizes: Record<string, number> = {}
+  ) {
     /** number of 4 byte slots taken */
     let size: number = 0;
 
     // Add layout (type, size and offset) definitions for each uniform in the layout
     for (const [key, uniformType] of Object.entries(uniformTypes)) {
       const typeAndComponents = getVariableShaderTypeInfo(uniformType);
-      const {type, components: count} = typeAndComponents;
+      const {type, components} = typeAndComponents;
+      // Calculate total count for uniform arrays.
+      const count = components * (uniformSizes?.[key] ?? 1);
       // First, align (bump) current offset to an even multiple of current object (1, 2, 4)
       size = alignTo(size, count);
       // Use the aligned size as the offset of the current uniform.
@@ -53,19 +58,15 @@ export class UniformBufferLayout {
 
   /** Get the data for the complete buffer */
   getData(uniformValues: Record<string, UniformValue>): Uint8Array {
-    const bufferSize = Math.max(this.byteLength, minBufferSize);
-
     // Allocate three typed arrays pointing at same memory
-    const arrayBuffer = getScratchArrayBuffer(bufferSize);
+    const arrayBuffer = getScratchArrayBuffer(this.byteLength);
     const typedArrays = {
       i32: new Int32Array(arrayBuffer),
       u32: new Uint32Array(arrayBuffer),
       f32: new Float32Array(arrayBuffer),
       // TODO not implemented
       f16: new Uint16Array(arrayBuffer)
     };
-    // TODO is this needed?
-    // typedArrays.i32.fill(0);
 
     for (const [name, value] of Object.entries(uniformValues)) {
       const uniformLayout = this.layout[name];
@@ -101,7 +102,7 @@ export class UniformBufferLayout {
       }
     }
 
-    return new Uint8Array(arrayBuffer);
+    return new Uint8Array(arrayBuffer, 0, this.byteLength);
   }
 
   /** Does this layout have a field with specified name */

modules/core/src/portable/uniform-store.ts

@@ -39,6 +39,7 @@ export class UniformStore<
       keyof TPropGroups,
       {
         uniformTypes?: Record<string, VariableShaderType>;
+        uniformSizes?: Record<string, number>;
         defaultProps?: Record<string, unknown>;
         defaultUniforms?: Record<string, UniformValue>;
       }
@@ -48,7 +49,10 @@ export class UniformStore<
       const uniformBufferName = bufferName as keyof TPropGroups;
 
       // Create a layout object to help us generate correctly formatted binary uniform buffers
-      const uniformBufferLayout = new UniformBufferLayout(block.uniformTypes || {});
+      const uniformBufferLayout = new UniformBufferLayout(
+        block.uniformTypes ?? {},
+        block.uniformSizes ?? {}
+      );
       this.uniformBufferLayouts.set(uniformBufferName, uniformBufferLayout);
 
       // Create a Uniform block to store the uniforms for each buffer.

modules/shadertools/src/lib/shader-module/shader-module.ts

@@ -9,7 +9,7 @@ import {
   makePropValidators,
   getValidatedProperties
 } from '../filters/prop-types';
-import type {UniformTypes, UniformValue} from '../utils/uniform-types';
+import type {UniformSizes, UniformTypes, UniformValue} from '../utils/uniform-types';
 import {ShaderInjection, normalizeInjections} from '../shader-assembly/shader-injections';
 
 // To avoid dependency on core module, do not import `Binding` type.
@@ -58,6 +58,8 @@ export type ShaderModule<
 
   /** Uniform shader types @note: Both order and types MUST match uniform block declarations in shader */
   uniformTypes?: Required<UniformTypes<UniformsT>>; // Record<keyof UniformsT, UniformFormat>;
+  /** Uniform shader array sizes (default 1) */
+  uniformSizes?: Required<UniformSizes<UniformsT>>;
   /** Uniform JS prop types  */
   propTypes?: Record<keyof UniformsT, UniformInfo>;
   /** Default uniform values */

modules/shadertools/src/lib/utils/uniform-types.ts

@@ -64,3 +64,7 @@ type UniformProps = {
 export type UniformTypes<PropsT extends UniformProps> = {
   [name in keyof PropsT]: UniformType<PropsT[name]>;
 };
+
+export type UniformSizes<PropsT extends UniformProps> = {
+  [name in keyof PropsT]: number;
+};