From ec0e7f421ba1611c519b2c1d8c175cd10c403124 Mon Sep 17 00:00:00 2001
From: Jordi Baylina <jordi@baylina.cat>
Date: Tue, 10 Dec 2019 18:46:10 +0100
Subject: [PATCH] roll unrolled loops with code

---
 c/zqfield.cpp                       |  4 ++++
 c/zqfield.h                         |  1 +
 src/c_build.js                      | 20 +++++++++++++++-----
 src/c_gen.js                        | 22 +++++++++++++++++-----
 src/exec.js                         | 26 +++++++++++++++++---------
 src/utils.js                        |  2 +-
 src/zqfield.js                      |  4 ++--
 test/basiccases.js                  | 11 +++++++++++
 test/circuits/componentarray.circom | 28 ++++++++++++++++++++++++++++
 9 files changed, 96 insertions(+), 22 deletions(-)
 create mode 100644 test/circuits/componentarray.circom

diff --git a/c/zqfield.cpp b/c/zqfield.cpp
index 3acf4e2..50551f4 100644
--- a/c/zqfield.cpp
+++ b/c/zqfield.cpp
@@ -189,3 +189,7 @@ void ZqField::shr(PBigInt r, PBigInt a, PBigInt b) {
     }
 }
 
+int ZqField::toInt(PBigInt a) {
+     return mpz_get_si (*a);
+}
+
diff --git a/c/zqfield.h b/c/zqfield.h
index e80c8a5..bac9b2a 100644
--- a/c/zqfield.h
+++ b/c/zqfield.h
@@ -45,6 +45,7 @@ public:
     void shr(PBigInt r, PBigInt a, PBigInt b);
 
     int isTrue(PBigInt a);
+    int toInt(PBigInt a);
 };
 
 #endif // ZQFIELD_H
diff --git a/src/c_build.js b/src/c_build.js
index 44ac485..711cefa 100644
--- a/src/c_build.js
+++ b/src/c_build.js
@@ -150,6 +150,8 @@ function buildCode(ctx) {
             ctx.codeFooter = "// Footer\n";
             ctx.uniqueNames = Object.assign({},ctx.globalNames);
             ctx.refs = [];
+            ctx.fileName = ctx.templates[ctx.components[i].template].fileName;
+            ctx.filePath = ctx.templates[ctx.components[i].template].filePath;
 
             for (let p in ctx.components[i].params) {
                 if (ctx.scopes[0][p]) return ctx.throwError(`Repeated parameter at ${ctx.components[i].template}: ${p}`);
@@ -245,7 +247,7 @@ function buildMapIsInput(ctx) {
             line += (i>31) ? "," : " ";
             line += toHex(acc);
             acc = 0;
-            if ( i % (32*64) == 0) {
+            if ( (i+1) % (32*64) == 0) {
                 arr.push(line);
                 line = "";
             }
@@ -299,9 +301,12 @@ function buildWit2Sig(ctx) {
         code += (i>0) ? ",": " ";
         code += arr[i];
         if ((i>0)&&(i%64 == 0)) {
-            if (code != "") codes.push(code + "\n");
-            codes.push(code);
-            code =0;
+            if (code != "") {
+                codes.push(code + "\n");
+            } else {
+                codes.push(code);
+            }
+            code ="";
         }
     }
     if (code != "") codes.push(code + "\n");
@@ -338,7 +343,6 @@ function addSizes(_sizes) {
     if (name=="sizes") name="sizes_0";
 
     if (this.definedSizes[name]) return this.definedSizes[name];
-    name = "_" + name;
     const labelName = this.getUniqueName(name);
     this.definedSizes[name] = labelName;
 
@@ -371,6 +375,8 @@ function buildFunction(name, paramValues) {
     const oldCodeHeader = ctx.codeHeader;
     const oldCodeFooter = ctx.codeFooter;
     const oldUniqueNames = ctx.uniqueNames;
+    const oldFileName = ctx.fileName;
+    const oldFilePath = ctx.oldFilePath;
 
 
     ctx.scopes = [{}];
@@ -382,6 +388,8 @@ function buildFunction(name, paramValues) {
     ctx.uniqueNames = Object.assign({},ctx.globalNames);
     ctx.returnValue = null;
     ctx.returnSizes = null;
+    ctx.fileName = ctx.functions[name].fileName;
+    ctx.filePath = ctx.functions[name].filePath;
 
     let paramsStr = "";
 
@@ -443,6 +451,8 @@ function buildFunction(name, paramValues) {
     ctx.codeHeader = oldCodeHeader;
     ctx.codeFooter = oldCodeFooter;
     ctx.uniqueNames = oldUniqueNames;
+    ctx.fileName = oldFileName;
+    ctx.filePath = oldFilePath;
 
     ctx.definedFunctions[h] = res;
 
diff --git a/src/c_gen.js b/src/c_gen.js
index bb845de..db3b1bb 100644
--- a/src/c_gen.js
+++ b/src/c_gen.js
@@ -40,6 +40,7 @@ function newRef(ctx, type, name, value, sizes) {
 function instantiateRef(ctx, refId, initValue) {
     const v = ctx.refs[refId];
     if (v.used) return;
+    if (!v.sizes) v.sizes = [1,0];
 
     if (v.type=="BIGINT") {
         ctx.codeHeader += `PBigInt ${v.label} = ctx->allocBigInts(${v.sizes[0]});\n`;
@@ -344,7 +345,7 @@ function genDeclareVariable(ctx, ast) {
 }
 
 function genNumber(ctx, ast) {
-    return newRef(ctx, "BIGINT", "_num", ast.value);
+    return newRef(ctx, "BIGINT", "_num", bigInt(ast.value));
 }
 
 
@@ -391,9 +392,9 @@ function genGetOffset(ctx, refOffset, vSizes, sels) {
 
                 if (rStr != "") rStr += " + ";
                 if (iIdx.used) {
-                    rStr += iIdx.label;
+                    rStr += `ctx->field->toInt(${iIdx.label})`;
                 } else {
-                    rStr += iIdx.value;
+                    rStr += iIdx.value[0].toString();
                 }
                 rStr += "*";
                 if (iSizes.used) {
@@ -741,6 +742,8 @@ function genConstraint(ctx, ast) {
     const b = ctx.refs[bRef];
     if (ctx.error) return;
     const strErr = ast.fileName + ":" + ast.first_line + ":" + ast.first_column;
+    instantiateRef(ctx, aRef);
+    instantiateRef(ctx, bRef);
     ctx.code += `ctx->checkConstraint(${a.label}, ${b.label}, "${strErr}");`;
 }
 
@@ -897,13 +900,18 @@ function genLoop(ctx, ast) {
 
         const cond2 = ctx.refs[condRef2];
 
+        if ((!cond2.used) &&(haveCode(ctx.code))) {
+            instantiateRef(ctx, condRef2, cond2.value);
+        }
+
         if (!inLoop) {
             if (cond2.used) {
                 ctx.code = oldCode + ctx.code;
                 inLoop = true;
                 enterConditionalCode(ctx, ast);
-                condVar = newRef(ctx, "INT", "_cond");
-                instantiateRef(ctx, condVar);
+                condVarRef = newRef(ctx, "INT", "_cond");
+                condVar = ctx.refs[condVarRef];
+                instantiateRef(ctx, condVarRef);
 
                 ctx.code =
                     oldCode +
@@ -928,6 +936,10 @@ function genLoop(ctx, ast) {
         leaveConditionalCode(ctx);
     }
     ctx.scopes.pop();
+
+    function haveCode(c) {
+        return c.replace(/\/\*[\s\S]*?\*\/|\/\/.*/g, "").trim() != "";
+    }
 }
 
 function genIf(ctx, ast) {
diff --git a/src/exec.js b/src/exec.js
index 74f8b9b..e065f3a 100644
--- a/src/exec.js
+++ b/src/exec.js
@@ -320,7 +320,9 @@ function execTemplateDef(ctx, ast) {
     };
     ctx.templates[ast.name] = {
         block: ast.block,
-        params: ast.params
+        params: ast.params,
+        fileName: ctx.fileName,
+        filePath: ctx.filePath,
     };
 }
 
@@ -344,7 +346,9 @@ function execFunctionDef(ctx, ast) {
     };
     ctx.functions[ast.name] = {
         block: ast.block,
-        params: ast.params
+        params: ast.params,
+        fileName: ctx.fileName,
+        filePath: ctx.filePath
     };
 }
 
@@ -383,7 +387,7 @@ function execDeclareComponent(ctx, ast) {
     };
 }
 
-function execInstantiateComponet(ctx, vr, fn) {
+function execInstantiateComponet(ctx, vr, fn, sels) {
 
     if (vr.type != "VARIABLE") return error(ctx, fn, "Left hand instatiate component must be a variable");
     if (fn.type != "FUNCTIONCALL") return error(ctx, fn, "Right type of instantiate component must be a function call");
@@ -408,13 +412,16 @@ function execInstantiateComponet(ctx, vr, fn) {
     if (template.params.length != paramValues.length) return error(ctx, fn, "Invalid Number of parameters");
 
 
-    const vComp = getScope(ctx, componentName, vr.selectors);
+    const vComp = getScope(ctx, componentName);
     if (vComp.type != "COMPONENT") return error(ctx, fn, "Assigning to a non component");
     const cIdx = vComp.cIdx;
     if (cIdx == -1) return error(ctx, fn, "Component not defined");
-    let l=1;
-    for (let i=0; i<vComp.sizes.length; i++) l = l*vComp.sizes[i];
-    for (let i=0; i<l; i++) {
+
+
+    const aSizes = utils.accSizes(vComp.sizes);
+    let o=0;
+    for (let i=0; i<sels.length; i++) o += sels[i] * aSizes[i+1];
+    for (let i=o; i<o+aSizes[sels.length]; i++) {
         instantiateComponent(cIdx+i);
     }
 
@@ -762,11 +769,12 @@ function execVarAssignement(ctx, ast) {
     const [num, sels, typ] = getScopeRef(ctx, v.name, v.selectors);
     if (ctx.error) return;
 
+    if (typ == "COMPONENT") return execInstantiateComponet(ctx, v, ast.values[1], sels);
+
     if ((typeof(num) != "object")||(num == null)) return  error(ctx, ast, "Variable not defined");
 
-    if (typ == "COMPONENT") return execInstantiateComponet(ctx, v, ast.values[1]);
     if (ctx.error) return;
-    if ((typ == "SIGNAL")&&(ast.op == "=")) return error(ctx, ast, "Cannot assign to a signal with `=` use <-- or <== ops");
+    if ((typ == "SIGNAL")&&(ast.op.indexOf("=")>=0)) return error(ctx, ast, "Cannot assign to a signal with `=` use <-- or <== ops");
     if ((["NUMBER", "COMPONENT"].indexOf(typ) >= 0 )&&(ast.op != "=")) return error(ctx, ast, `Cannot assign to a var with ${ast.op}. use = op`);
 
     const res = exec(ctx, ast.values[1]);
diff --git a/src/utils.js b/src/utils.js
index 847e33e..b922181 100644
--- a/src/utils.js
+++ b/src/utils.js
@@ -38,7 +38,7 @@ function flatArray(a) {
                 fillArray(res, a[i]);
             }
         } else {
-            res.push(a);
+            res.push(bigInt(a));
         }
     }
 }
diff --git a/src/zqfield.js b/src/zqfield.js
index 565788f..c134c14 100644
--- a/src/zqfield.js
+++ b/src/zqfield.js
@@ -30,7 +30,7 @@ module.exports = class ZqField {
     }
 
     mul(a, b) {
-        return a.mul(b).mod(this.p);
+        return a.times(b).mod(this.p);
     }
 
     lt(a, b) {
@@ -59,7 +59,7 @@ module.exports = class ZqField {
 
     div(a, b) {
         assert(!b.isZero(), "Division by zero");
-        return a.mul(b.modInv(this.p)).mod(this.p);
+        return a.times(b.modInv(this.p)).mod(this.p);
     }
 
     idiv(a, b) {
diff --git a/test/basiccases.js b/test/basiccases.js
index 3baeecf..5db5d68 100644
--- a/test/basiccases.js
+++ b/test/basiccases.js
@@ -287,4 +287,15 @@ describe("basic cases", function () {
             ]
         );
     });
+    it("Component array ", async () => {
+        await doTest(
+            "componentarray.circom",
+            [
+                [{in: 1}, {out:    1}],
+                [{in: 2}, {out:  256}],
+                [{in: 3}, {out: 6561}],
+                [{in:-1}, {out:    1}],
+            ]
+        );
+    });
 });
diff --git a/test/circuits/componentarray.circom b/test/circuits/componentarray.circom
new file mode 100644
index 0000000..62706da
--- /dev/null
+++ b/test/circuits/componentarray.circom
@@ -0,0 +1,28 @@
+template Square() {
+    signal input in;
+    signal output out;
+
+    out <== in*in;
+}
+
+template Main(n) {
+    signal input in;
+    signal output out;
+
+    component squares[n];
+
+    var i;
+
+    for (i=0; i<n; i++) {
+        squares[i] = Square();
+        if (i==0) {
+            squares[i].in <== in;
+        } else {
+            squares[i].in <== squares[i-1].out;
+        }
+    }
+
+    squares[n-1].out ==> out;
+}
+
+component main = Main(3);