From f55f88cf5fdb0ac5c46eb7ec8340499184e8dd41 Mon Sep 17 00:00:00 2001
From: Kyryl R <kyryl.ryabov@gmail.com>
Date: Thu, 31 Oct 2024 15:46:24 +0200
Subject: [PATCH] Added more tests

---
 test/circom-template-inputs-visitor.test.ts | 32 +++++++++++++++++++--
 test/mocks/CircomTemplateInputsVisitor.ts   |  8 ++++++
 2 files changed, 37 insertions(+), 3 deletions(-)

diff --git a/test/circom-template-inputs-visitor.test.ts b/test/circom-template-inputs-visitor.test.ts
index 08a8706..fe6ea86 100644
--- a/test/circom-template-inputs-visitor.test.ts
+++ b/test/circom-template-inputs-visitor.test.ts
@@ -38,9 +38,7 @@ describe("Circom Template Inputs Visitor", () => {
       mainComponentData,
     );
 
-    const parser = getCircomParser(`test/data/curve.circom`);
-
-    visitor.visit(parser.circuit());
+    visitor.startParse();
 
     expect(visitor.errors.length).to.equal(0);
 
@@ -91,4 +89,32 @@ describe("Circom Template Inputs Visitor", () => {
     expect(visitor.templateInputs.pkIdentityHash.type).to.equal("output");
     expect(visitor.templateInputs.pkIdentityHash.dimension).to.deep.equal([]);
   });
+
+  it("should analyse the MainComponent.circom circuit", () => {
+    const mainComponentData: VariableContext = {
+      p1: 2n,
+      p2: 8n,
+    };
+
+    const data = getData("MainComponent.circom");
+
+    const visitor = new CircomTemplateInputsVisitor(
+      "MainComponent.circom",
+      data["C"].context,
+      mainComponentData,
+    );
+
+    visitor.startParse();
+
+    expect(visitor.errors.length).to.equal(0);
+
+    expect(visitor.templateInputs.in1.type).to.equal("input");
+    expect(visitor.templateInputs.in1.dimension).to.deep.equal([]);
+
+    expect(visitor.templateInputs.in2.type).to.equal("input");
+    expect(visitor.templateInputs.in2.dimension).to.deep.equal([3n, 2n]);
+
+    expect(visitor.templateInputs.out.type).to.equal("output");
+    expect(visitor.templateInputs.out.dimension).to.deep.equal([]);
+  });
 });
diff --git a/test/mocks/CircomTemplateInputsVisitor.ts b/test/mocks/CircomTemplateInputsVisitor.ts
index a22feb3..ee29c93 100644
--- a/test/mocks/CircomTemplateInputsVisitor.ts
+++ b/test/mocks/CircomTemplateInputsVisitor.ts
@@ -55,6 +55,10 @@ export class CircomTemplateInputsVisitor extends CircomVisitor<void> {
     this._validateVariableContext();
   }
 
+  startParse = () => {
+    this.visit(this.templateContext);
+  };
+
   visitVarDeclaration = (ctx: VarDeclarationContext) => {
     if (ctx.LP() && ctx.RP()) {
       this.errors.push({
@@ -295,6 +299,10 @@ export class CircomTemplateInputsVisitor extends CircomVisitor<void> {
   };
 
   visitSubsLeftAssignment = (ctx: SubsLeftAssignmentContext) => {
+    if (ctx.LEFT_ASSIGNMENT() || ctx.LEFT_CONSTRAINT()) {
+      return;
+    }
+
     const primaryExpression = ctx._lhs.primaryExpression();
 
     if (!primaryExpression) {