Fix for resources used as array index

src/wgsl_ast.ts

@@ -1110,6 +1110,9 @@ export class VariableExpr extends Expression {
 
   search(callback: (node: Node) => void) {
     callback(this);
+    if (this.postfix) {
+      this.postfix.search(callback);
+    }
   }
 
   evaluate(context: ParseContext): number {
@@ -1262,6 +1265,23 @@ export class GroupingExpr extends Expression {
   }
 }
 
+/**
+ * @class ArrayIndex
+ * @extends Expression
+ * @category AST
+ */
+export class ArrayIndex extends Expression {
+  index: Expression;
+  constructor(index: Expression) {
+    super();
+    this.index = index;
+  }
+
+  search(callback: (node: Node) => void): void {
+    this.index.search(callback);
+  }
+}
+
 /**
  * @class Operator
  * @extends Expression

src/wgsl_parser.ts

@@ -1052,7 +1052,9 @@ export class WgslParser {
     // primary_expression postfix_expression ?
     const expr = this._primary_expression();
     const p = this._postfix_expression();
-    if (p) expr.postfix = p;
+    if (p) {
+      expr.postfix = p;
+    }
     return expr;
   }
 
@@ -1061,9 +1063,12 @@ export class WgslParser {
     if (this._match(TokenTypes.tokens.bracket_left)) {
       const expr = this._short_circuit_or_expression();
       this._consume(TokenTypes.tokens.bracket_right, "Expected ']'.");
+      const arrayIndex = new AST.ArrayIndex(expr);
       const p = this._postfix_expression();
-      if (p) expr.postfix = p;
-      return expr;
+      if (p) {
+        arrayIndex.postfix = p;
+      }
+      return arrayIndex;
     }
 
     // period ident postfix_expression?
@@ -1074,7 +1079,9 @@ export class WgslParser {
       );
       const p = this._postfix_expression();
       const expr = new AST.StringExpr(name.lexeme);
-      if (p) expr.postfix = p;
+      if (p) {
+        expr.postfix = p;
+      }
       return expr;
     }
 

test/tests/test_reflect.js

@@ -2,6 +2,21 @@ import { test, group } from "../test.js";
 import { WgslReflect, ResourceType } from "../../../wgsl_reflect.module.js";
 
 group("Reflect", function () {
+  test("uniform index", function (test) {
+    const reflect = new WgslReflect(`
+      @group(0) @binding(2) var<uniform>  batchIndex: u32;
+      @group(0) @binding(3) var<storage, read> batchOffsets: array<u32>;
+      @vertex
+      fn main() {
+        let batchOffset = batchOffsets[batchIndex];
+      }`);
+      test.equals(reflect.uniforms.length, 1);
+      test.equals(reflect.uniforms[0].name, "batchIndex");
+      test.equals(reflect.entry.vertex[0].resources.length, 2);
+      test.equals(reflect.entry.vertex[0].resources[0].name, "batchOffsets");
+      test.equals(reflect.entry.vertex[0].resources[1].name, "batchIndex");
+  });
+
   test("enable", function (test) {
     const reflect = new WgslReflect(`enable chromium_experimental_subgroups;
       @compute @workgroup_size(64) fn main(

types/wgsl_ast.d.ts

@@ -529,6 +529,16 @@ export declare class GroupingExpr extends Expression {
     evaluate(context: ParseContext): number;
     search(callback: (node: Node) => void): void;
 }
+/**
+ * @class ArrayIndex
+ * @extends Expression
+ * @category AST
+ */
+export declare class ArrayIndex extends Expression {
+    index: Expression;
+    constructor(index: Expression);
+    search(callback: (node: Node) => void): void;
+}
 /**
  * @class Operator
  * @extends Expression