add overrides to FunctionInfo

src/wgsl_reflect.ts

@@ -272,12 +272,32 @@ export class OutputInfo {
   }
 }
 
+export class OverrideInfo {
+  name: string;
+  type: TypeInfo | null;
+  attributes: Array<AST.Attribute> | null;
+  id: number;
+
+  constructor(
+    name: string,
+    type: TypeInfo | null,
+    attributes: Array<AST.Attribute> | null,
+    id: number
+  ) {
+    this.name = name;
+    this.type = type;
+    this.attributes = attributes;
+    this.id = id;
+  }
+}
+
 export class FunctionInfo {
   name: string;
   stage: string | null = null;
   inputs: Array<InputInfo> = [];
   outputs: Array<OutputInfo> = [];
   resources: Array<VariableInfo> = [];
+  overrides: Array<OverrideInfo> = [];
   startLine: number = -1;
   endLine: number = -1;
   inUse: boolean = false;
@@ -295,25 +315,6 @@ export class EntryFunctions {
   compute: Array<FunctionInfo> = [];
 }
 
-export class OverrideInfo {
-  name: string;
-  type: TypeInfo | null;
-  attributes: Array<AST.Attribute> | null;
-  id: number;
-
-  constructor(
-    name: string,
-    type: TypeInfo | null,
-    attributes: Array<AST.Attribute> | null,
-    id: number
-  ) {
-    this.name = name;
-    this.type = type;
-    this.attributes = attributes;
-    this.id = id;
-  }
-}
-
 class _FunctionResources {
   node: AST.Function;
   resources: Array<VariableInfo> | null = null;
@@ -505,6 +506,19 @@ export class WgslReflect {
       }
     }
 
+    for (const fn of this._functions.values()) {
+      fn.node.search((node) => {
+        if (node.astNodeType === "varExpr") {
+          const v = node as AST.VariableExpr;
+          for (const override of this.overrides) {
+            if (v.name == override.name) {
+              fn.info?.overrides.push(override);
+            }
+          }
+        }
+      });
+    }
+
     for (const u of this.uniforms) {
       this._markStructsInUse(u.type);
     }

test/tests/test_reflect.js

@@ -2,6 +2,18 @@ import { test, group } from "../test.js";
 import { WgslReflect, ResourceType } from "../../../wgsl_reflect.module.js";
 
 group("Reflect", function () {
+  test("override", function (test) {
+    const reflect = new WgslReflect(`override red = 0.0;
+      override green = 0.0;
+      override blue = 0.0;
+      @fragment fn fs() -> @location(0) vec4f {
+        return vec4f(red, green, blue, 1.0);
+      }`);
+      test.equals(reflect.overrides.length, 3);
+      test.equals(reflect.entry.fragment[0].name, "fs");
+      test.equals(reflect.entry.fragment[0].overrides.length, 3);
+  });
+
   test("f16", function (test) {
     const reflect = new WgslReflect(`
       @binding(0) @group(0) var<uniform> a1: f16;            // This is correctly sized at 2 bytes!
@@ -37,11 +49,14 @@ group("Reflect", function () {
 
   test("uniform index", function (test) {
     const reflect = new WgslReflect(`
-      @group(0) @binding(2) var<uniform>  batchIndex: u32;
+      struct BatchIndex {
+        index: u32,
+      }
+      @group(0) @binding(2) var<uniform>  batchIndex: BatchIndex;
       @group(0) @binding(3) var<storage, read> batchOffsets: array<u32>;
       @vertex
       fn main() {
-        let batchOffset = batchOffsets[batchIndex];
+        let batchOffset = batchOffsets[batchIndex.index];
       }`);
       test.equals(reflect.uniforms.length, 1);
       test.equals(reflect.uniforms[0].name, "batchIndex");

types/wgsl_reflect.d.ts

@@ -97,12 +97,20 @@ export declare class OutputInfo {
     location: number | string;
     constructor(name: string, type: TypeInfo | null, locationType: string, location: number | string);
 }
+export declare class OverrideInfo {
+    name: string;
+    type: TypeInfo | null;
+    attributes: Array<AST.Attribute> | null;
+    id: number;
+    constructor(name: string, type: TypeInfo | null, attributes: Array<AST.Attribute> | null, id: number);
+}
 export declare class FunctionInfo {
     name: string;
     stage: string | null;
     inputs: Array<InputInfo>;
     outputs: Array<OutputInfo>;
     resources: Array<VariableInfo>;
+    overrides: Array<OverrideInfo>;
     startLine: number;
     endLine: number;
     inUse: boolean;
@@ -114,13 +122,6 @@ export declare class EntryFunctions {
     fragment: Array<FunctionInfo>;
     compute: Array<FunctionInfo>;
 }
-export declare class OverrideInfo {
-    name: string;
-    type: TypeInfo | null;
-    attributes: Array<AST.Attribute> | null;
-    id: number;
-    constructor(name: string, type: TypeInfo | null, attributes: Array<AST.Attribute> | null, id: number);
-}
 declare class _FunctionResources {
     node: AST.Function;
     resources: Array<VariableInfo> | null;