From 7965dc5d6bc7f8e19fa17c29a13edd9acd6d5693 Mon Sep 17 00:00:00 2001
From: Jordi Baylina <jordi@baylina.cat>
Date: Mon, 8 Jun 2020 23:05:22 +0200
Subject: [PATCH] phase2 optimized

---
 src/binfileutils.js |   8 +-
 src/phase2_new.js   | 462 +++++++++++++++++++++++++++++---------------
 src/zkeyfile.js     |  16 +-
 3 files changed, 320 insertions(+), 166 deletions(-)

diff --git a/src/binfileutils.js b/src/binfileutils.js
index 9582e4a..725ec72 100644
--- a/src/binfileutils.js
+++ b/src/binfileutils.js
@@ -86,14 +86,14 @@ async function endReadSection(fd, noCheck) {
     delete fd.readingSection;
 }
 
-async function writeBigInt(fd, n, n8) {
+async function writeBigInt(fd, n, n8, pos) {
     const buff = new Uint8Array(n8);
     Scalar.toRprLE(buff, 0, n, n8);
-    await fd.write(buff);
+    await fd.write(buff, pos);
 }
 
-async function readBigInt(fd, n8) {
-    const buff = await fd.read(n8);
+async function readBigInt(fd, n8, pos) {
+    const buff = await fd.read(n8, pos);
     return Scalar.fromRprLE(buff, 0, n8);
 }
 
diff --git a/src/phase2_new.js b/src/phase2_new.js
index 4397efe..7dbf485 100644
--- a/src/phase2_new.js
+++ b/src/phase2_new.js
@@ -1,23 +1,26 @@
 
-const loadR1cs = require("r1csfile").load;
+const r1csFile = require("r1csfile");
 const utils = require("./powersoftau_utils");
 const binFileUtils = require("./binfileutils");
-const writeZKey = require("./zkeyfile").write;
 const assert = require("assert");
-
-
-function log2( V )
-{
-    return( ( ( V & 0xFFFF0000 ) !== 0 ? ( V &= 0xFFFF0000, 16 ) : 0 ) | ( ( V & 0xFF00FF00 ) !== 0 ? ( V &= 0xFF00FF00, 8 ) : 0 ) | ( ( V & 0xF0F0F0F0 ) !== 0 ? ( V &= 0xF0F0F0F0, 4 ) : 0 ) | ( ( V & 0xCCCCCCCC ) !== 0 ? ( V &= 0xCCCCCCCC, 2 ) : 0 ) | ( ( V & 0xAAAAAAAA ) !== 0 ) );
-}
+const {log2} = require("./misc");
+const Scalar = require("ffjavascript").Scalar;
 
 
 module.exports  = async function phase2new(r1csName, ptauName, zkeyName, verbose) {
 
-    const r1cs = await loadR1cs(r1csName, true);
+    const {fd: fdR1cs, sections: sectionsR1cs} = await binFileUtils.readBinFile(r1csName, "r1cs", 1);
+    const r1cs = await r1csFile.loadHeader(fdR1cs, sectionsR1cs);
 
-    const {fd: ptauFd, sections} = await binFileUtils.readBinFile(ptauName, "ptau", 1);
-    const {curve, power} = await utils.readPTauHeader(ptauFd, sections);
+    const {fd: fdPTau, sections: sectionsPTau} = await binFileUtils.readBinFile(ptauName, "ptau", 1);
+    const {curve, power} = await utils.readPTauHeader(fdPTau, sectionsPTau);
+
+    await curve.loadEngine();
+
+    const fdZKey = await binFileUtils.createBinFile(zkeyName, "zkey", 1, 9);
+
+    const sG1 = curve.G1.F.n8*2;
+    const sG2 = curve.G2.F.n8*2;
 
     if (r1cs.prime != curve.r) {
         console.log("r1cs curve does not match powers of tau ceremony curve");
@@ -31,89 +34,170 @@ module.exports  = async function phase2new(r1csName, ptauName, zkeyName, verbose
         return -1;
     }
 
-    if (!sections[12]) {
+    if (!sectionsPTau[12]) {
         console.log("Powers of tau is not prepared.");
         return -1;
     }
 
-    const zKey = {
+    const nPublic = r1cs.nOutputs + r1cs.nPubInputs;
+    const domainSize = 1 << cirPower;
 
-        nPublic: r1cs.nOutputs + r1cs.nPubInputs,
-        nVars: r1cs.nVars,
-        q: curve.q,
-        r: curve.r,
-        domainBits: cirPower,
-        domainSize: 1 << cirPower
+    // Write the header
+    ///////////
+    await binFileUtils.startWriteSection(fdZKey, 1);
+    await fdZKey.writeULE32(1); // Groth
+    await binFileUtils.endWriteSection(fdZKey);
 
-    };
+    // Write the Groth header section
+    ///////////
 
-    calculatePolinomials(curve, zKey,r1cs);
+    await binFileUtils.startWriteSection(fdZKey, 2);
+    const primeQ = curve.q;
+    const n8q = (Math.floor( (Scalar.bitLength(primeQ) - 1) / 64) +1)*8;
 
-    zKey.A = new Array(r1cs.nVars);
-    zKey.B1 = new Array(r1cs.nVars);
-    zKey.B2 = new Array(r1cs.nVars);
-    zKey.C = new Array(r1cs.nVars);
-    zKey.IC = new Array(zKey.nPublic+1);
-    for (let i=0; i<r1cs.nVars; i++) {
-        zKey.A[i] = curve.G1.zero;
-        zKey.B1[i] = curve.G1.zero;
-        zKey.B2[i] = curve.G2.zero;
-        if (i>zKey.nPublic) {
-            zKey.C[i] = curve.G1.zero;
-        } else {
-            zKey.IC[i] = curve.G1.zero;
+    const primeR = curve.r;
+    const n8r = (Math.floor( (Scalar.bitLength(primeR) - 1) / 64) +1)*8;
+    const Rr = Scalar.mod(Scalar.shl(1, n8r*8), primeR);
+    const R2r = Scalar.mod(Scalar.mul(Rr,Rr), primeR);
+
+    await fdZKey.writeULE32(n8q);
+    await binFileUtils.writeBigInt(fdZKey, primeQ, n8q);
+    await fdZKey.writeULE32(n8r);
+    await binFileUtils.writeBigInt(fdZKey, primeR, n8r);
+    await fdZKey.writeULE32(r1cs.nVars);                         // Total number of bars
+    await fdZKey.writeULE32(nPublic);                       // Total number of public vars (not including ONE)
+    await fdZKey.writeULE32(domainSize);                  // domainSize
+
+    const bAlpha1 = await fdPTau.read(sG1, sectionsPTau[4][0].p);
+    await fdZKey.write(bAlpha1);
+
+    const bBeta1 = await fdPTau.read(sG1, sectionsPTau[5][0].p);
+    await fdZKey.write(bBeta1);
+
+    const bg1 = new Uint8Array(sG1);
+    curve.G1.toRprLEM(bg1, 0, curve.G1.g);
+    await fdZKey.write(bg1);        // delta1
+
+    const bBeta2 = await fdPTau.read(sG2, sectionsPTau[6][0].p);
+    await fdZKey.write(bBeta2);
+
+    const bg2 = new Uint8Array(sG2);
+    curve.G2.toRprLEM(bg2, 0, curve.G2.g);
+    await fdZKey.write(bg2);        // gamma2
+    await fdZKey.write(bg2);        // delta2
+
+    await binFileUtils.endWriteSection(fdZKey);
+
+
+    const A = new Array(r1cs.nVars);
+    const B1 = new Array(r1cs.nVars);
+    const B2 = new Array(r1cs.nVars);
+    const C = new Array(r1cs.nVars- nPublic -1);
+    const IC = new Array(nPublic+1);
+
+    const lTauG1 = sectionsPTau[12][0].p + ((1 << cirPower) -1)*sG1;
+    const lTauG2 = sectionsPTau[13][0].p + ((1 << cirPower) -1)*sG2;
+    const lAlphaTauG1 = sectionsPTau[14][0].p + ((1 << cirPower) -1)*sG1;
+    const lBetaTauG1 = sectionsPTau[15][0].p + ((1 << cirPower) -1)*sG1;
+
+    await binFileUtils.startWriteSection(fdZKey, 4);
+    await binFileUtils.startReadUniqueSection(fdR1cs, sectionsR1cs, 2);
+
+    const pNCoefs =  fdZKey.pos;
+    let nCoefs = 0;
+    fdZKey.pos += 4;
+    for (let c=0; c<r1cs.nConstraints; c++) {
+        if (verbose && (c%1000 == 0) && (c >0)) console.log(`${c}/${r1cs.nConstraints}`);
+        const nA = await fdR1cs.readULE32();
+        for (let i=0; i<nA; i++) {
+            const s = await fdR1cs.readULE32();
+            const coef = await fdR1cs.read(r1cs.n8);
+
+            const l1 = lTauG1 + sG1*c;
+            const l2 = lBetaTauG1 + sG1*c;
+            if (typeof A[s] === "undefined") A[s] = [];
+            A[s].push([l1, coef]);
+
+            if (s <= nPublic) {
+                if (typeof IC[s] === "undefined") IC[s] = [];
+                IC[s].push([l2, coef]);
+            } else {
+                if (typeof C[s- nPublic -1] === "undefined") C[s- nPublic -1] = [];
+                C[s - nPublic -1].push([l2, coef]);
+            }
+            await fdZKey.writeULE32(0);
+            await fdZKey.writeULE32(c);
+            await fdZKey.writeULE32(s);
+            await writeFr2(coef);
+            nCoefs ++;
+        }
+
+        const nB = await fdR1cs.readULE32();
+        for (let i=0; i<nB; i++) {
+            const s = await fdR1cs.readULE32();
+            const coef = await fdR1cs.read(r1cs.n8);
+
+            const l1 = lTauG1 + sG1*c;
+            const l2 = lTauG2 + sG2*c;
+            const l3 = lAlphaTauG1 + sG1*c;
+            if (typeof B1[s] === "undefined") B1[s] = [];
+            B1[s].push([l1, coef]);
+            if (typeof B2[s] === "undefined") B2[s] = [];
+            B2[s].push([l2, coef]);
+
+            if (s <= nPublic) {
+                if (typeof IC[s] === "undefined") IC[s] = [];
+                IC[s].push([l3, coef]);
+            } else {
+                if (typeof C[s- nPublic -1] === "undefined") C[s- nPublic -1] = [];
+                C[s- nPublic -1].push([l3, coef]);
+            }
+            await fdZKey.writeULE32(1);
+            await fdZKey.writeULE32(c);
+            await fdZKey.writeULE32(s);
+            await writeFr2(coef);
+            nCoefs ++;
+        }
+
+        const nC = await fdR1cs.readULE32();
+        for (let i=0; i<nC; i++) {
+            const s = await fdR1cs.readULE32();
+            const coef = await fdR1cs.read(r1cs.n8);
+
+            const l1 = lTauG1 + sG1*c;
+            if (s <= nPublic) {
+                if (typeof IC[s] === "undefined") IC[s] = [];
+                IC[s].push([l1, coef]);
+            } else {
+                if (typeof C[s- nPublic -1] === "undefined") C[s- nPublic -1] = [];
+                C[s- nPublic -1].push([l1, coef]);
+            }
         }
     }
-    for (let i=0; i<zKey.ccoefs.length; i++) {
-        if (verbose && (i%1000 == 0) && (i >0)) console.log(`${i}/${zKey.ccoefs.length}`);
-        const c = zKey.ccoefs[i];
-        let CIC;
-        if (c.matrix == 0) {
-            const l1 = await readEvaluation("lTauG1", c.constraint);
-            const l2 = await readEvaluation("lBetaTauG1", c.constraint);
-            zKey.A[c.signal] =
-                curve.G1.add(
-                    zKey.A[c.signal],
-                    curve.G1.mulScalar(l1, c.value)
-                );
-            CIC = curve.G1.mulScalar(l2, c.value);
-        } else if (c.matrix == 1) {
-            const l1 = await readEvaluation("lTauG1", c.constraint);
-            const l2 = await readEvaluation("lTauG2", c.constraint);
-            const l3 = await readEvaluation("lAlphaTauG1", c.constraint);
-            zKey.B1[c.signal] =
-                curve.G1.add(
-                    zKey.B1[c.signal],
-                    curve.G1.mulScalar(l1, c.value)
-                );
-            zKey.B2[c.signal] =
-                curve.G2.add(
-                    zKey.B2[c.signal],
-                    curve.G2.mulScalar(l2, c.value)
-                );
-            CIC = curve.G1.mulScalar(l3, c.value);
-        } else if (c.matrix == 2) {
-            const l1 = await readEvaluation("lTauG1", c.constraint);
-            CIC = curve.G1.mulScalar(l1, c.value);
-        } else {
-            assert(false);
-        }
 
-        if (c.signal <= zKey.nPublic) {
-            zKey.IC[c.signal] =
-                curve.G1.add(
-                    zKey.IC[c.signal],
-                    CIC
-                );
-        } else {
-            zKey.C[c.signal] =
-                curve.G1.add(
-                    zKey.C[c.signal],
-                    CIC
-                );
-        }
+    const bOne = new Uint8Array(curve.Fr.n8);
+    curve.Fr.toRprLE(bOne, 0, curve.Fr.e(1));
+    for (let s = 0; s <= nPublic ; s++) {
+        const l1 = lTauG1 + sG1*(r1cs.nConstraints + s);
+        const l2 = lBetaTauG1 + sG1*(r1cs.nConstraints + s);
+        if (typeof A[s] === "undefined") A[s] = [];
+        A[s].push([l1, bOne]);
+        if (typeof IC[s] === "undefined") IC[s] = [];
+        IC[s].push([l2, bOne]);
+        await fdZKey.writeULE32(0);
+        await fdZKey.writeULE32(r1cs.nConstraints + s);
+        await fdZKey.writeULE32(s);
+        await writeFr2(bOne);
+        nCoefs ++;
     }
 
+    const oldPos = fdZKey.pos;
+    await fdZKey.writeULE32(nCoefs, pNCoefs);
+    fdZKey.pos = oldPos;
+
+    await binFileUtils.endWriteSection(fdZKey);
+    await binFileUtils.endReadSection(fdR1cs);
+
 /*
     zKey.hExps = new Array(zKey.domainSize-1);
     for (let i=0; i< zKey.domainSize; i++) {
@@ -123,86 +207,156 @@ module.exports  = async function phase2new(r1csName, ptauName, zkeyName, verbose
     }
 */
 
-    zKey.hExps = new Array(zKey.domainSize);
-    for (let i=0; i< zKey.domainSize; i++) {
-        zKey.hExps[i] = await readEvaluation("lTauG1", i*2+1, cirPower+1);
+    await composeAndWritePoints(3, "G1", IC, "IC");
+    await composeAndWritePoints(5, "G1", A, "A");
+    await composeAndWritePoints(6, "G1", B1, "B1");
+    await composeAndWritePoints(7, "G2", B2, "B2");
+    await composeAndWritePoints(8, "G1", C, "C");
+
+    // Write Hs
+    await binFileUtils.startWriteSection(fdZKey, 9);
+    const o = sectionsPTau[12][0].p + ((1 << (cirPower+1)) -1)*sG1;
+    for (let i=0; i< domainSize; i++) {
+        const buff = await fdPTau.read(sG1, o + (i*2+1)*sG1 );
+        await fdZKey.write(buff);
     }
+    await binFileUtils.endWriteSection(fdZKey);
 
-    zKey.vk_alfa_1 = await readEvaluation("alphaTauG1", 0);
-    zKey.vk_beta_1 = await readEvaluation("betaTauG1", 0);
-    zKey.vk_delta_1 = curve.G1.g;
-    zKey.vk_beta_2 = await readEvaluation("betaG2", 0);
-    zKey.vk_gamma_2 = curve.G2.g;
-    zKey.vk_delta_2 = curve.G2.g;
-
-
-    curve.G1.multiAffine(zKey.A);
-    curve.G1.multiAffine(zKey.B1);
-    curve.G2.multiAffine(zKey.B2);
-    curve.G1.multiAffine(zKey.C);
-    curve.G1.multiAffine(zKey.hExps);
-    curve.G1.multiAffine(zKey.IC);
-
-    await writeZKey(zkeyName, zKey);
+    await fdZKey.close();
+    await fdPTau.close();
+    await fdR1cs.close();
 
     return 0;
 
-    async function readEvaluation(sectionName, idx, p) {
-        p = p || cirPower;
-        let o;
-        let G;
-        switch (sectionName) {
-        case "tauG1": o = sections[2][0].p; G = curve.G1;  break;
-        case "tauG2": o = sections[3][0].p; G = curve.G2;  break;
-        case "alphaTauG1": o = sections[4][0].p; G = curve.G1;  break;
-        case "betaTauG1": o = sections[5][0].p; G = curve.G1;  break;
-        case "betaG2": o = sections[6][0].p; G = curve.G2;  break;
-        case "lTauG1": o = sections[12][0].p; G = curve.G1;  break;
-        case "lTauG2": o = sections[13][0].p; G = curve.G2; break;
-        case "lAlphaTauG1": o = sections[14][0].p; G = curve.G1; break;
-        case "lBetaTauG1": o = sections[15][0].p; G = curve.G1; break;
-        }
-        const sG = G.F.n8*2;
-        if (["lTauG1","lTauG2", "lAlphaTauG1", "lBetaTauG1" ].indexOf(sectionName) >= 0) {
-            o += ((1 << p)-1)*sG;
-        }
-        ptauFd.pos = o + sG*idx;
-        const buff = await ptauFd.read(sG);
-        return G.fromRprLEM(buff, 0);
+    async function writeFr2(buff) {
+        const n = curve.Fr.fromRprLE(buff, 0);
+        const nR2 = curve.Fr.mul(n, R2r);
+        const buff2 = new Uint8Array(curve.Fr.n8);
+        curve.Fr.toRprLE(buff2, 0, nR2);
+        await fdZKey.write(buff2);
     }
 
+
+    async function composeAndWritePoints(idSection, groupName, arr, sectionName) {
+        const CHUNK_SIZE= 1<<18;
+        await binFileUtils.startWriteSection(fdZKey, idSection);
+
+        for (let i=0; i<arr.length; i+= CHUNK_SIZE) {
+            if (verbose)  console.log(`${sectionName}: ${i}/${arr.length}`);
+            const n = Math.min(arr.length -i, CHUNK_SIZE);
+            const subArr = arr.slice(i, i + n);
+            await composeAndWritePointsChunk(groupName, subArr);
+        }
+        await binFileUtils.endWriteSection(fdZKey);
+    }
+
+    async function composeAndWritePointsChunk(groupName, arr) {
+        const concurrency= curve.engine.concurrency;
+        const nElementsPerThread = Math.floor(arr.length / concurrency);
+        const opPromises = [];
+        for (let i=0; i<concurrency; i++) {
+            let n;
+            if (i< concurrency-1) {
+                n = nElementsPerThread;
+            } else {
+                n = arr.length - i*nElementsPerThread;
+            }
+            if (n==0) continue;
+
+            const subArr = arr.slice(i*nElementsPerThread, i*nElementsPerThread + n);
+            opPromises.push(composeAndWritePointsThread(groupName, subArr));
+        }
+
+        const result = await Promise.all(opPromises);
+
+        for (let i=0; i<result.length; i++) {
+            await fdZKey.write(result[i][0]);
+        }
+    }
+
+    async function composeAndWritePointsThread(groupName, arr) {
+        const G = curve[groupName];
+        const sGin = G.F.n8*2;
+        const sGmid = G.F.n8*3;
+        const sGout = G.F.n8*2;
+        let fnExp, fnMultiExp, fnBatchToAffine, fnZero;
+        if (groupName == "G1") {
+            fnExp = "g1m_timesScalarAffine";
+            fnMultiExp = "g1m_multiexpAffine";
+            fnBatchToAffine = "g1m_batchToAffine";
+            fnZero = "g1m_zero";
+        } else if (groupName == "G2") {
+            fnExp = "g2m_timesScalarAffine";
+            fnMultiExp = "g2m_multiexpAffine";
+            fnBatchToAffine = "g2m_batchToAffine";
+            fnZero = "g2m_zero";
+        } else {
+            assert(false);
+        }
+        let acc =0;
+        for (let i=0; i<arr.length; i++) acc += arr[i] ? arr[i].length : 0;
+        const bBases = new Uint8Array(acc*sGin);
+        const bScalars = new Uint8Array(acc*curve.Fr.n8);
+        let pB =0;
+        let pS =0;
+        for (let i=0; i<arr.length; i++) {
+            if (!arr[i]) continue;
+            for (let j=0; j<arr[i].length; j++) {
+                const bBase = await fdPTau.read(sGin, arr[i][j][0]);
+                bBases.set(bBase, pB);
+                pB += sGin;
+                bScalars.set(arr[i][j][1], pS);
+                pS += curve.Fr.n8;
+            }
+        }
+        const task = [];
+        task.push({cmd: "ALLOCSET", var: 0, buff: bBases});
+        task.push({cmd: "ALLOCSET", var: 1, buff: bScalars});
+        task.push({cmd: "ALLOC", var: 2, len: arr.length*sGmid});
+        pB = 0;
+        pS = 0;
+        let pD =0;
+        for (let i=0; i<arr.length; i++) {
+            if (!arr[i]) {
+                task.push({cmd: "CALL", fnName: fnZero, params: [
+                    {var: 2, offset: pD}
+                ]});
+                pD += sGmid;
+                continue;
+            }
+            if (arr[i].length == 1) {
+                task.push({cmd: "CALL", fnName: fnExp, params: [
+                    {var: 0, offset: pB},
+                    {var: 1, offset: pS},
+                    {val: curve.Fr.n8},
+                    {var: 2, offset: pD}
+                ]});
+            } else {
+                task.push({cmd: "CALL", fnName: fnMultiExp, params: [
+                    {var: 0, offset: pB},
+                    {var: 1, offset: pS},
+                    {val: curve.Fr.n8},
+                    {val: arr[i].length},
+                    {var: 2, offset: pD}
+                ]});
+            }
+            pB += sGin*arr[i].length;
+            pS += curve.Fr.n8*arr[i].length;
+            pD += sGmid;
+        }
+        task.push({cmd: "CALL", fnName: fnBatchToAffine, params: [
+            {var: 2},
+            {val: arr.length},
+            {var: 2},
+        ]});
+        task.push({cmd: "GET", out: 0, var: 2, len: arr.length*sGout});
+
+        const res = await curve.engine.queueAction(task);
+
+        return res;
+    }
+
+
 };
 
 
-function calculatePolinomials(curve, zKey, r1cs) {
-
-    zKey.ccoefs = [];
-    for (let m=0; m<3; m++) {
-        for (let c=0; c<r1cs.nConstraints; c++) {
-            const signals = Object.keys(r1cs.constraints[c][m]);
-            signals.forEach( (s) => {
-                zKey.ccoefs.push({
-                    matrix: m,
-                    constraint: c,
-                    signal: Number(s),
-                    value: r1cs.constraints[c][m][s]
-                });
-            });
-        }
-    }
-
-    /**
-     * add and process the constraints
-     *     input_i * 0 = 0
-     * to ensure soundness of input consistency
-     */
-    for (let i = 0; i < r1cs.nPubInputs + r1cs.nOutputs + 1; ++i)
-    {
-        zKey.ccoefs.push({
-            matrix: 0,
-            constraint: r1cs.nConstraints + i,
-            signal: i,
-            value: curve.Fr.one
-        });
-    }
-}
diff --git a/src/zkeyfile.js b/src/zkeyfile.js
index 037bafa..321f39b 100644
--- a/src/zkeyfile.js
+++ b/src/zkeyfile.js
@@ -77,14 +77,6 @@ async function writeZKey(fileName, zkey) {
     await binFileUtils.endWriteSection(fd);
 
 
-    // Write IC Section
-    ///////////
-    await binFileUtils.startWriteSection(fd, 3);
-    for (let i=0; i<= zkey.nPublic; i++) {
-        await writePointG1(zkey.IC[i] );
-    }
-    await binFileUtils.endWriteSection(fd);
-
 
     // Write Pols (A and B (C can be ommited))
     ///////////
@@ -103,6 +95,14 @@ async function writeZKey(fileName, zkey) {
     await binFileUtils.endWriteSection(fd);
 
 
+    // Write IC Section
+    ///////////
+    await binFileUtils.startWriteSection(fd, 3);
+    for (let i=0; i<= zkey.nPublic; i++) {
+        await writePointG1(zkey.IC[i] );
+    }
+    await binFileUtils.endWriteSection(fd);
+
 
     // Write A
     ///////////