From d04eff6c0d5f80149addb031cf15eba8c8e8e09d Mon Sep 17 00:00:00 2001
From: Jordi Baylina <jordi@baylina.cat>
Date: Wed, 18 Sep 2019 17:43:14 +0200
Subject: [PATCH] Optimize optimization and fix out<==in

---
 src/compiler.js            | 167 +++++++++++++++++++++++++------------
 src/exec.js                |  25 ++++--
 test/cases.js              |   5 ++
 test/circuits/circuit.json |  75 +++++++++++++++++
 test/circuits/inout.circom |  18 ++++
 5 files changed, 233 insertions(+), 57 deletions(-)
 create mode 100644 test/circuits/circuit.json
 create mode 100644 test/circuits/inout.circom

diff --git a/src/compiler.js b/src/compiler.js
index 0b53ade..91052c0 100644
--- a/src/compiler.js
+++ b/src/compiler.js
@@ -226,71 +226,134 @@ function reduceConstants(ctx) {
 }
 
 function reduceConstrains(ctx) {
-    const newConstraints = [];
-    for (let i=0; i<ctx.constraints.length; i++) {
-        const c = ctx.constraints[i];
+    indexVariables();
+    let possibleConstraints = Object.keys(ctx.constraints);
+    while (possibleConstraints.length>0) {
+        let nextPossibleConstraints = {};
+        for (let i in possibleConstraints) {
+            if (!ctx.constraints[i]) continue;
+            const c = ctx.constraints[i];
+
+            // Swap a and b if b has more variables.
+            if (Object.keys(c.b).length > Object.keys(c.a).length) {
+                const aux = c.a;
+                c.a=c.b;
+                c.b=aux;
+            }
 
-        // Swap a and b if b has more variables.
-        if (Object.keys(c.b).length > Object.keys(c.a).length) {
-            const aux = c.a;
-            c.a=c.b;
-            c.b=aux;
-        }
+            // Mov to C if possible.
+            if (isConstant(c.a)) {
+                const ct = {type: "NUMBER", value: c.a.values["one"]};
+                c.c = lc.add(lc.mul(c.b, ct), c.c);
+                c.a = { type: "LINEARCOMBINATION", values: {} };
+                c.b = { type: "LINEARCOMBINATION", values: {} };
+            }
+            if (isConstant(c.b)) {
+                const ct = {type: "NUMBER", value: c.b.values["one"]};
+                c.c = lc.add(lc.mul(c.a, ct), c.c);
+                c.a = { type: "LINEARCOMBINATION", values: {} };
+                c.b = { type: "LINEARCOMBINATION", values: {} };
+            }
+
+            if (lc.isZero(c.a) || lc.isZero(c.b)) {
+                const isolatedSignal = getFirstInternalSignal(ctx, c.c);
+                if (isolatedSignal) {
+
+                    let lSignal = ctx.signals[isolatedSignal];
+                    while (lSignal.equivalence) {
+                        lSignal = ctx.signals[lSignal.equivalence];
+                    }
 
-        // Mov to C if possible.
-        if (isConstant(c.a)) {
-            const ct = {type: "NUMBER", value: c.a.values["one"]};
-            c.c = lc.add(lc.mul(c.b, ct), c.c);
-            c.a = { type: "LINEARCOMBINATION", values: {} };
-            c.b = { type: "LINEARCOMBINATION", values: {} };
-        }
-        if (isConstant(c.b)) {
-            const ct = {type: "NUMBER", value: c.b.values["one"]};
-            c.c = lc.add(lc.mul(c.a, ct), c.c);
-            c.a = { type: "LINEARCOMBINATION", values: {} };
-            c.b = { type: "LINEARCOMBINATION", values: {} };
-        }
 
-        if (lc.isZero(c.a) || lc.isZero(c.b)) {
-            const isolatedSignal = getFirstInternalSignal(ctx, c.c);
-            if (isolatedSignal) {
-                const isolatedSignalEquivalence = {
-                    type: "LINEARCOMBINATION",
-                    values: {}
-                };
-                const invCoef = c.c.values[isolatedSignal].modInv(__P__);
-                for (const s in c.c.values) {
-                    if (s != isolatedSignal) {
-                        const v = __P__.minus(c.c.values[s]).times(invCoef).mod(__P__);
-                        if (!v.isZero()) {
-                            isolatedSignalEquivalence.values[s] = v;
+                    const isolatedSignalEquivalence = {
+                        type: "LINEARCOMBINATION",
+                        values: {}
+                    };
+                    const invCoef = c.c.values[isolatedSignal].modInv(__P__);
+                    for (const s in c.c.values) {
+                        if (s != isolatedSignal) {
+                            const v = __P__.minus(c.c.values[s]).times(invCoef).mod(__P__);
+                            if (!v.isZero()) {
+                                isolatedSignalEquivalence.values[s] = v;
+                            }
                         }
                     }
-                }
 
-                for (let j=0; j<newConstraints.length; j++) {
-                    newConstraints[j] = lc.substitute(newConstraints[j], isolatedSignal, isolatedSignalEquivalence);
-                }
-                for (let j=i+1; j<ctx.constraints.length; j++ ) {
-                    ctx.constraints[j] = lc.substitute(ctx.constraints[j], isolatedSignal, isolatedSignalEquivalence);
-                }
-                c.a={ type: "LINEARCOMBINATION", values: {} };
-                c.b={ type: "LINEARCOMBINATION", values: {} };
-                c.c={ type: "LINEARCOMBINATION", values: {} };
+                    for (let j in lSignal.inConstraints) {
+                        if ((j!=i)&&(ctx.constraints[j])) {
+                            ctx.constraints[j] = lc.substitute(ctx.constraints[j], isolatedSignal, isolatedSignalEquivalence);
+                            linkSignalsConstraint(j);
+                            if (j<i) {
+                                nextPossibleConstraints[j] = true;
+                            }
+                        }
+                    }
+
+                    ctx.constraints[i] = null;
 
-                let lSignal = ctx.signals[isolatedSignal];
-                while (lSignal.equivalence) {
-                    lSignal = ctx.signals[lSignal.equivalence];
+                    lSignal.category = "constant";
+                } else {
+                    if (lc.isZero(c.c)) ctx.constraints[i] = null;
                 }
-                lSignal.category = "constant";
             }
         }
+        possibleConstraints = Object.keys(nextPossibleConstraints);
+    }
+    unindexVariables();
 
-        if (!lc.isZero(c)) {
-            newConstraints.push(c);
+    // Pack the constraints
+    let o = 0;
+    for (let i=0; i<ctx.constraints.length; i++) {
+        if (ctx.constraints[i]) {
+            if (o != i) {
+                ctx.constraints[o] = ctx.constraints[i];
+            }
+            o++;
         }
     }
-    ctx.constraints = newConstraints;
+    ctx.constraints.length = o;
+
+    function indexVariables() {
+        for (let i=0; i<ctx.constraints.length; i++) linkSignalsConstraint(i);
+    }
+
+    function linkSignalsConstraint(cidx) {
+        const ct = ctx.constraints[cidx];
+        for (let k in ct.a.values) linkSignal(k, cidx);
+        for (let k in ct.b.values) linkSignal(k, cidx);
+        for (let k in ct.c.values) linkSignal(k, cidx);
+    }
+
+    function unindexVariables() {
+        for (let s in ctx.signals) {
+            let lSignal = ctx.signals[s];
+            while (lSignal.equivalence) {
+                lSignal = ctx.signals[lSignal.equivalence];
+            }
+            if (lSignal.inConstraints) delete lSignal.inConstraints;
+        }
+    }
+
+/*
+    function unlinkSignal(signalName, cidx) {
+        let lSignal = ctx.signals[signalName];
+        while (lSignal.equivalence) {
+            lSignal = ctx.signals[lSignal.equivalence];
+        }
+        if ((lSignal.inConstraints)&&(lSignal.inConstraints[cidx])) {
+            delete lSignal.inConstraints[cidx];
+        }
+    }
+*/
+
+    function linkSignal(signalName, cidx) {
+        let lSignal = ctx.signals[signalName];
+        while (lSignal.equivalence) {
+            lSignal = ctx.signals[lSignal.equivalence];
+        }
+        if (!lSignal.inConstraints) lSignal.inConstraints = {};
+        lSignal.inConstraints[cidx] = true;
+    }
 
     function getFirstInternalSignal(ctx, l) {
         for (let k in l.values) {
diff --git a/src/exec.js b/src/exec.js
index 35d7c24..1c38159 100644
--- a/src/exec.js
+++ b/src/exec.js
@@ -996,7 +996,12 @@ function execSignalAssign(ctx, ast) {
 
     let sDest=ctx.signals[dst.fullName];
     if (!sDest) return error(ctx, ast, "Invalid signal: "+dst.fullName);
-    while (sDest.equivalence) sDest=ctx.signals[sDest.equivalence];
+
+    let isOut = (sDest.component == "main")&&(sDest.direction=="OUT");
+    while (sDest.equivalence) {
+        sDest=ctx.signals[sDest.equivalence];
+        isOut = isOut || ((sDest.component == "main")&&(sDest.direction=="OUT"));
+    }
 
     if (sDest.value) return error(ctx, ast, "Signals cannot be assigned twice");
 
@@ -1024,10 +1029,20 @@ function execSignalAssign(ctx, ast) {
 
     let assignValue = true;
     if (src.type == "SIGNAL") {
-        sDest.equivalence = src.fullName;
-        sDest.alias = sDest.alias.concat(src.alias);
-        while (sDest.equivalence) sDest=ctx.signals[sDest.equivalence];
-        assignValue = false;
+        let sSrc = ctx.signals[src.fullName];
+        let isIn  = (sSrc.component == "main")&&(sSrc.direction == "IN");
+        while (sSrc.equivalence) {
+            sSrc=ctx.signals[sSrc.equivalence];
+            isIn = isIn || ((sSrc.component == "main")&&(sSrc.direction == "IN"));
+        }
+
+        // Skip if an out is assigned directly to an input.
+        if ((!isIn)||(!isOut)) {
+            sDest.equivalence = src.fullName;
+            sDest.alias = sDest.alias.concat(src.alias);
+            while (sDest.equivalence) sDest=ctx.signals[sDest.equivalence];
+            assignValue = false;
+        }
     }
 
     if (assignValue) {
diff --git a/test/cases.js b/test/cases.js
index a12a04d..7b98d83 100644
--- a/test/cases.js
+++ b/test/cases.js
@@ -59,4 +59,9 @@ describe("Sum test", () => {
         assert(witness[1].equals(bigInt(37)));
         assert(witness[2].equals(bigInt(6)));
     });
+    it("Should compile a code with compute", async () => {
+        const cirDef = await compiler(path.join(__dirname, "circuits", "inout.circom"));
+
+        assert.equal(cirDef.constraints.length, 1);
+    });
 });
diff --git a/test/circuits/circuit.json b/test/circuits/circuit.json
new file mode 100644
index 0000000..c0636c4
--- /dev/null
+++ b/test/circuits/circuit.json
@@ -0,0 +1,75 @@
+{
+ "mainCode": "{\n}\n",
+ "signalName2Idx": {
+  "one": 0,
+  "main.in": 2,
+  "main.out": 1,
+  "main.internal.in": 2,
+  "main.internal.out": 2
+ },
+ "components": [
+  {
+   "name": "main",
+   "params": {},
+   "template": "InOut",
+   "inputSignals": 1
+  },
+  {
+   "name": "main.internal",
+   "params": {},
+   "template": "Internal",
+   "inputSignals": 1
+  }
+ ],
+ "componentName2Idx": {
+  "main": 0,
+  "main.internal": 1
+ },
+ "signals": [
+  {
+   "names": [
+    "one"
+   ],
+   "triggerComponents": []
+  },
+  {
+   "names": [
+    "main.out"
+   ],
+   "triggerComponents": []
+  },
+  {
+   "names": [
+    "main.in",
+    "main.internal.in",
+    "main.internal.out"
+   ],
+   "triggerComponents": [
+    0,
+    1
+   ]
+  }
+ ],
+ "constraints": [
+  [
+   {},
+   {},
+   {
+    "1": "21888242871839275222246405745257275088548364400416034343698204186575808495616",
+    "2": "1"
+   }
+  ]
+ ],
+ "templates": {
+  "Internal": "function(ctx) {\n    ctx.setSignal(\"out\", [], ctx.getSignal(\"in\", []));\n}\n",
+  "InOut": "function(ctx) {\n    ctx.setPin(\"internal\", [], \"in\", [], ctx.getSignal(\"in\", []));\n    ctx.setSignal(\"out\", [], ctx.getPin(\"internal\", [], \"out\", []));\n}\n"
+ },
+ "functions": {},
+ "nPrvInputs": 0,
+ "nPubInputs": 1,
+ "nInputs": 1,
+ "nOutputs": 1,
+ "nVars": 3,
+ "nConstants": 0,
+ "nSignals": 3
+}
\ No newline at end of file
diff --git a/test/circuits/inout.circom b/test/circuits/inout.circom
new file mode 100644
index 0000000..cf42e36
--- /dev/null
+++ b/test/circuits/inout.circom
@@ -0,0 +1,18 @@
+template Internal() {
+    signal input in;
+    signal output out;
+
+    out <== in;
+}
+
+template InOut() {
+    signal input in;
+    signal output out;
+
+    component internal = Internal();
+
+    internal.in <== in;
+    internal.out ==> out;
+}
+
+component main = InOut();