Fix computation of location values when processing the shader code

src/Engines/Processors/iShaderProcessor.ts

@@ -3,8 +3,8 @@ import { ShaderProcessingContext } from "./shaderProcessingOptions";
 
 /** @hidden */
 export interface IShaderProcessor {
-    attributeProcessor?: (attribute: string, processingContext: Nullable<ShaderProcessingContext>) => string;
-    varyingProcessor?: (varying: string, isFragment: boolean, processingContext: Nullable<ShaderProcessingContext>) => string;
+    attributeProcessor?: (attribute: string, preProcessors: { [key: string]: string }, processingContext: Nullable<ShaderProcessingContext>) => string;
+    varyingProcessor?: (varying: string, isFragment: boolean, preProcessors: { [key: string]: string }, processingContext: Nullable<ShaderProcessingContext>) => string;
     uniformProcessor?: (uniform: string, isFragment: boolean, preProcessors: { [key: string]: string }, processingContext: Nullable<ShaderProcessingContext>) => string;
     uniformBufferProcessor?: (uniformBuffer: string, isFragment: boolean, processingContext: Nullable<ShaderProcessingContext>) => string;
     endOfUniformBufferProcessor?: (closingBracketLine: string, isFragment: boolean, processingContext: Nullable<ShaderProcessingContext>) => string;

src/Engines/Processors/shaderCodeNode.ts

@@ -24,9 +24,9 @@ export class ShaderCodeNode {
                 }
 
                 if (processor.attributeProcessor && StringTools.StartsWith(this.line, "attribute")) {
-                    value = processor.attributeProcessor(this.line, options.processingContext);
+                    value = processor.attributeProcessor(this.line, preprocessors, options.processingContext);
                 } else if (processor.varyingProcessor && StringTools.StartsWith(this.line, "varying")) {
-                    value = processor.varyingProcessor(this.line, options.isFragment, options.processingContext);
+                    value = processor.varyingProcessor(this.line, options.isFragment, preprocessors, options.processingContext);
                 } else if ((processor.uniformProcessor || processor.uniformBufferProcessor) && StringTools.StartsWith(this.line, "uniform") && !options.lookForClosingBracketForUniformBuffer) {
                     let regex = /uniform\s+(?:(?:highp)?|(?:lowp)?)\s*(\S+)\s+(\S+)\s*;/;
 

src/Engines/WebGPU/webgpuShaderProcessingContext.ts

@@ -3,12 +3,17 @@ import { ShaderProcessingContext } from "../Processors/shaderProcessingOptions";
 const _maxSets = 4;
 const _maxBindingsPerSet = 16;
 
+// all types not listed are assumed to consume 1 location
+const _webGLTypeToLocationSize: { [key: string]: number } = {
+    "mat2": 2,
+    "mat3": 3,
+    "mat4": 4,
+};
+
 /**
  * @hidden
  */
 export class WebGPUShaderProcessingContext implements ShaderProcessingContext {
-    public attributeNextLocation: number;
-    public varyingNextLocation: number;
     public uboNextBindingIndex: number;
     public freeSetIndex: number;
     public freeBindingIndex: number;
@@ -23,9 +28,28 @@ export class WebGPUShaderProcessingContext implements ShaderProcessingContext {
     public orderedAttributes: string[];
     public orderedUBOsAndSamplers: { name: string, isSampler: boolean, isComparisonSampler?: boolean, textureDimension?: GPUTextureViewDimension }[][];
 
+    private _attributeNextLocation: number;
+    private _varyingNextLocation: number;
+
+    public getAttributeNextLocation(dataType: string, arrayLength: number = 0): number {
+        const index = this._attributeNextLocation;
+
+        this._attributeNextLocation += (_webGLTypeToLocationSize[dataType] ?? 1) * (arrayLength || 1);
+
+        return index;
+    }
+
+    public getVaryingNextLocation(dataType: string, arrayLength: number = 0): number {
+        const index = this._varyingNextLocation;
+
+        this._varyingNextLocation += (_webGLTypeToLocationSize[dataType] ?? 1) * (arrayLength || 1);
+
+        return index;
+    }
+
     constructor() {
-        this.attributeNextLocation = 0;
-        this.varyingNextLocation = 0;
+        this._attributeNextLocation = 0;
+        this._varyingNextLocation = 0;
         this.freeSetIndex = 2;
         this.freeBindingIndex = 0;
 

src/Engines/WebGPU/webgpuShaderProcessors.ts

@@ -58,16 +58,32 @@ export class WebGPUShaderProcessor implements IShaderProcessor {
 
     protected _missingVaryings: Array<string> = [];
 
+    private _getArraySize(name: string, preProcessors: { [key: string]: string }): [string, number] {
+        let length = 0;
+        const startArray = name.indexOf("[");
+        const endArray = name.indexOf("]");
+        if (startArray > 0 && endArray > 0) {
+            const lengthInString = name.substring(startArray + 1, endArray);
+            length = +(lengthInString);
+            if (isNaN(length)) {
+                length = +(preProcessors[lengthInString]);
+            }
+            name = name.substr(0, startArray);
+        }
+        return [name, length];
+    }
+
     public initializeShaders(processingContext: Nullable<ShaderProcessingContext>): void {
         this._missingVaryings.length = 0;
     }
 
-    public varyingProcessor(varying: string, isFragment: boolean, processingContext: Nullable<ShaderProcessingContext>) {
+    public varyingProcessor(varying: string, isFragment: boolean, preProcessors: { [key: string]: string }, processingContext: Nullable<ShaderProcessingContext>) {
         const webgpuProcessingContext = processingContext! as WebGPUShaderProcessingContext;
 
         const varyingRegex = new RegExp(/\s*varying\s+(\S+)\s+(\S+)\s*;/gm);
         const match = varyingRegex.exec(varying);
         if (match != null) {
+            const varyingType = match[1];
             const name = match[2];
             let location: number;
             if (isFragment) {
@@ -75,29 +91,30 @@ export class WebGPUShaderProcessor implements IShaderProcessor {
                 this._missingVaryings[location] = "";
             }
             else {
-                location = webgpuProcessingContext.varyingNextLocation++;
+                location = webgpuProcessingContext.getVaryingNextLocation(varyingType, this._getArraySize(name, preProcessors)[1]);
                 webgpuProcessingContext.availableVaryings[name] = location;
-                this._missingVaryings.push(`layout(location = ${location}) in ${match[1]} ${name};`);
+                this._missingVaryings[location] = `layout(location = ${location}) in ${varyingType} ${name};`;
             }
 
-            varying = varying.replace(match[0], `layout(location = ${location}) ${isFragment ? "in" : "out"} ${match[1]} ${name};`);
+            varying = varying.replace(match[0], `layout(location = ${location}) ${isFragment ? "in" : "out"} ${varyingType} ${name};`);
         }
         return varying;
     }
 
-    public attributeProcessor(attribute: string, processingContext: Nullable<ShaderProcessingContext>) {
+    public attributeProcessor(attribute: string, preProcessors: { [key: string]: string }, processingContext: Nullable<ShaderProcessingContext>) {
         const webgpuProcessingContext = processingContext! as WebGPUShaderProcessingContext;
 
         const attribRegex = new RegExp(/\s*attribute\s+(\S+)\s+(\S+)\s*;/gm);
         const match = attribRegex.exec(attribute);
         if (match != null) {
+            const varyingType = match[1];
             const name = match[2];
-            const location = webgpuProcessingContext.attributeNextLocation++;
+            const location = webgpuProcessingContext.getAttributeNextLocation(varyingType, this._getArraySize(name, preProcessors)[1]);
 
             webgpuProcessingContext.availableAttributes[name] = location;
             webgpuProcessingContext.orderedAttributes[location] = name;
 
-            attribute = attribute.replace(match[0], `layout(location = ${location}) in ${match[1]} ${name};`);
+            attribute = attribute.replace(match[0], `layout(location = ${location}) in ${varyingType} ${name};`);
         }
         return attribute;
     }
@@ -148,16 +165,8 @@ export class WebGPUShaderProcessor implements IShaderProcessor {
             else {
                 // Check the size of the uniform array in case of array.
                 let length = 0;
-                const startArray = name.indexOf("[");
-                const endArray = name.indexOf("]");
-                if (startArray > 0 && endArray > 0) {
-                    const lengthInString = name.substring(startArray + 1, endArray);
-                    length = +(lengthInString);
-                    if (isNaN(length)) {
-                        length = +(preProcessors[lengthInString]);
-                    }
-                    name = name.substr(0, startArray);
-                }
+
+                [name, length] = this._getArraySize(name, preProcessors);
 
                 for (let i = 0; i < webgpuProcessingContext.leftOverUniforms.length; i++) {
                     if (webgpuProcessingContext.leftOverUniforms[i].name === name) {
@@ -273,7 +282,7 @@ export class WebGPUShaderProcessor implements IShaderProcessor {
         // inject the missing varying in the fragment shader
         for (let i = 0; i < this._missingVaryings.length; ++i) {
             const decl = this._missingVaryings[i];
-            if (decl.length > 0) {
+            if (decl && decl.length > 0) {
                 fragmentCode = decl + "\n" + fragmentCode;
             }
         }