Add ability to parse var declarations

package.json

@@ -1,6 +1,6 @@
 {
   "name": "@solarity/hardhat-zkit",
-  "version": "0.4.7",
+  "version": "0.4.8",
   "description": "The ultimate TypeScript environment for Circom development",
   "main": "dist/src/index.js",
   "types": "dist/src/index.d.ts",
@@ -53,7 +53,7 @@
     "node": ">=16"
   },
   "dependencies": {
-    "@distributedlab/circom-parser": "0.1.4",
+    "@distributedlab/circom-parser": "0.1.5",
     "@solarity/zkit": "0.2.6",
     "@solarity/zktype": "0.3.1",
     "@wasmer/wasi": "0.12.0",

src/core/dependencies/parser/CircomTemplateInputsVisitor.ts

@@ -6,6 +6,10 @@ import {
   BigIntOrNestedArray,
   IdentifierContext,
   Variables,
+  VarDeclarationContext,
+  VarDefinitionContext,
+  RhsValueContext,
+  TemplateStmtContext,
 } from "@distributedlab/circom-parser";
 
 import { InputData } from "../../../types/core";
@@ -26,7 +30,8 @@ import { HardhatZKitError } from "../../../errors";
  */
 export class CircomTemplateInputsVisitor extends CircomVisitor<void> {
   templateInputs: Record<string, InputData>;
-  expressionVisitor: CircomExpressionVisitor;
+
+  vars: Variables = {};
 
   constructor(
     private readonly _templateName: string,
@@ -36,15 +41,11 @@ export class CircomTemplateInputsVisitor extends CircomVisitor<void> {
 
     this.templateInputs = {};
 
-    const vars: Variables = {};
-
     for (const key of Object.keys(this._parameterValues)) {
-      vars[key] = {
+      this.vars[key] = {
         value: this._parameterValues[key],
       };
     }
-
-    this.expressionVisitor = new CircomExpressionVisitor(true, vars);
   }
 
   visitTemplateDeclaration = (ctx: TemplateDeclarationContext) => {
@@ -53,11 +54,115 @@ export class CircomTemplateInputsVisitor extends CircomVisitor<void> {
         .templateBlock()
         .templateStmt_list()
         .forEach((stmt) => {
-          if (stmt.signalDeclaration()) {
-            this.visit(stmt.signalDeclaration());
+          this.visitTemplateStmt(stmt);
+        });
+    }
+  };
+
+  visitTemplateStmt = (ctx: TemplateStmtContext) => {
+    if (ctx.signalDeclaration()) {
+      this.visit(ctx.signalDeclaration());
+
+      return;
+    }
+
+    if (ctx.identifier() && ctx.ASSIGNMENT() && ctx.expression(0)) {
+      const id = ctx.identifier().ID(0).getText();
+      const value = new CircomExpressionVisitor(true, this.vars).visitExpression(ctx.expression(0));
+
+      if (Array.isArray(value)) {
+        throw new HardhatZKitError(`Currently, only single value assignment is supported - ${value}`);
+      }
+
+      this.vars[id] = {
+        value: value,
+      };
+
+      return;
+    }
+
+    if (!ctx.IF()) {
+      this.visitChildren(ctx);
+
+      return;
+    }
+
+    const result = new CircomExpressionVisitor(true, this.vars).visitExpression(ctx.parExpression().expression());
+
+    if (Array.isArray(result)) {
+      throw new HardhatZKitError(
+        `Currently, only single value assignment is supported as a result inside if statement - ${result}`,
+      );
+    }
+
+    if (result === 1n) {
+      this.visitTemplateStmt(ctx.templateStmt(0));
+
+      return;
+    }
+
+    if (ctx.ELSE()) {
+      this.visitTemplateStmt(ctx.templateStmt(1));
+
+      return;
+    }
+  };
+
+  visitVarDeclaration = (ctx: VarDeclarationContext) => {
+    const vars = this._parseVarDefinition(ctx.varDefinition());
+
+    if (!ctx.ASSIGNMENT()) return;
+
+    const results = this._parseRHSValue(ctx.rhsValue());
+
+    if (vars.length !== results.length) {
+      throw new HardhatZKitError(`Mismatch between variable definitions and values - ${ctx.getText()}`);
+    }
+
+    vars.forEach((varName, index) => {
+      this.vars[varName] = {
+        value: results[index],
+      };
+    });
+  };
+
+  _parseVarDefinition = (ctx: VarDefinitionContext): string[] => {
+    return ctx.identifier_list().map((identifier) => identifier.ID(0).getText());
+  };
+
+  _parseRHSValue = (ctx: RhsValueContext): bigint[] => {
+    const expressionVisitor = new CircomExpressionVisitor(true, this.vars);
+
+    if (ctx.expression()) {
+      const expressionResult = expressionVisitor.visitExpression(ctx.expression());
+
+      if (Array.isArray(expressionResult)) {
+        throw new HardhatZKitError(`Currently, only single value assignment is supported - ${expressionResult}`);
+      }
+
+      return [expressionResult];
+    }
+
+    if (ctx.expressionList()) {
+      const expressionsResult: bigint[] = [];
+
+      ctx
+        .expressionList()
+        .expression_list()
+        .forEach((expression) => {
+          const expressionResult = expressionVisitor.visitExpression(expression);
+
+          if (Array.isArray(expressionResult)) {
+            throw new HardhatZKitError(`Currently, only single value assignment is supported - ${expressionResult}`);
           }
+
+          expressionsResult.push(expressionResult);
         });
+
+      return expressionsResult;
     }
+
+    throw new HardhatZKitError(`RHS value as function call is not supported - ${ctx.getText()}`);
   };
 
   visitSignalDeclaration = (ctx: SignalDeclarationContext) => {
@@ -87,7 +192,8 @@ export class CircomTemplateInputsVisitor extends CircomVisitor<void> {
     const inputDimension: string[] = [];
 
     identifier.arrayDimension_list().forEach((dimension) => {
-      const expressionResult = this.expressionVisitor.visitExpression(dimension.expression());
+      const expressionVisitor = new CircomExpressionVisitor(true, this.vars);
+      const expressionResult = expressionVisitor.visitExpression(dimension.expression());
 
       if (Array.isArray(expressionResult)) {
         throw new HardhatZKitError(