Merge pull request #63 from mhawryluk/function-arguments

Add argument and return type info to non-entry FunctionInfo

README.md

@@ -139,6 +139,8 @@ class FunctionInfo {
   stage: string | null;
   inputs: Array<InputInfo>;
   outputs: Array<OutputInfo>;
+  arguments: Array<ArgumentInfo>; // only for non-entry functions
+  returnType: TypeInfo | null;
   resources: Array<VariableInfo>;
   startLine: number;
   endLine: number;
@@ -166,6 +168,11 @@ class OverrideInfo {
   type: TypeInfo | null;
   id: number;
 }
+
+class ArgumentInfo {
+  name: string;
+  type: TypeInfo;
+}
 ```
 
 ## Examples

src/wgsl_reflect.ts

@@ -291,11 +291,26 @@ export class OverrideInfo {
   }
 }
 
+export class ArgumentInfo {
+  name: string;
+  type: TypeInfo;
+
+  constructor(
+    name: string,
+    type: TypeInfo,
+  ) {
+    this.name = name;
+    this.type = type;
+  }
+}
+
 export class FunctionInfo {
   name: string;
   stage: string | null = null;
   inputs: Array<InputInfo> = [];
   outputs: Array<OutputInfo> = [];
+  arguments: Array<ArgumentInfo> = [];
+  returnType: TypeInfo | null = null;
   resources: Array<VariableInfo> = [];
   overrides: Array<OverrideInfo> = [];
   startLine: number = -1;
@@ -494,6 +509,17 @@ export class WgslReflect {
           fn.inputs = this._getInputs(node.args);
           fn.outputs = this._getOutputs(node.returnType);
           this.entry[stage.name].push(fn);
+        } else {
+          fn.arguments = node.args.map(
+            (arg) =>
+              new ArgumentInfo(
+                arg.name,
+                this._getTypeInfo(arg.type, arg.attributes)
+              )
+          );
+          fn.returnType = node.returnType
+            ? this._getTypeInfo(node.returnType, node.attributes)
+            : null;
         }
         continue;
       }

test/tests/test_reflect.js

@@ -1225,4 +1225,28 @@ fn shuffler() { }
     test.equals(reflect.entry.vertex[5].resources[0].name, "u1");
     test.equals(reflect.entry.vertex[5].resources[1].name, "u2");
   });
+
+  test("function arguments and return types", function (test) {
+    const reflect = new WgslReflect(`
+      fn rotate(v: vec2<f32>, angle: f32) -> vec2f {
+        let pos = vec2(
+          (v.x * cos(angle)) - (v.y * sin(angle)),
+          (v.x * sin(angle)) + (v.y * cos(angle))
+        );
+        return pos;
+      }`);
+
+    const args = reflect.functions[0].arguments;
+    test.equals(args[0].name, "v");
+    test.equals(args[0].type.name, "vec2");
+    test.equals(args[0].type.format.name, "f32");
+    test.equals(args[1].name, "angle");
+    test.equals(args[1].type.name, "f32");
+
+    test.equals(reflect.functions[0].returnType.name, "vec2f");
+
+    const logFunc = new WgslReflect("fn log() {}").functions[0];
+    test.equals(logFunc.arguments.length, 0);
+    test.equals(logFunc.returnType, null);
+  });
 });

types/wgsl_reflect.d.ts

@@ -104,11 +104,18 @@ export declare class OverrideInfo {
     id: number;
     constructor(name: string, type: TypeInfo | null, attributes: Array<AST.Attribute> | null, id: number);
 }
+export declare class ArgumentInfo {
+    name: string;
+    type: TypeInfo;
+    constructor(name: string, type: TypeInfo);
+}
 export declare class FunctionInfo {
     name: string;
     stage: string | null;
     inputs: Array<InputInfo>;
     outputs: Array<OutputInfo>;
+    arguments: Array<ArgumentInfo>;
+    returnType: TypeInfo | null;
     resources: Array<VariableInfo>;
     overrides: Array<OverrideInfo>;
     startLine: number;