From b433cbf7adca85d11d8213d7c44ed76d5ff90edc Mon Sep 17 00:00:00 2001
From: Balazs Komuves <bkomuves@gmail.com>
Date: Thu, 23 Nov 2023 13:25:26 +0100
Subject: [PATCH] speed up fake setup generation

---
 groth16/fake_setup.nim | 29 ++++++++++++++++++++++++++++-
 groth16/test_proof.nim | 12 +++++++++---
 2 files changed, 37 insertions(+), 4 deletions(-)

diff --git a/groth16/fake_setup.nim b/groth16/fake_setup.nim
index f4a19ec..f39f67a 100644
--- a/groth16/fake_setup.nim
+++ b/groth16/fake_setup.nim
@@ -32,7 +32,7 @@ proc randomToxicWaste(): ToxicWaste =
   let b = randFr()
   let c = randFr()
   let d = randFr()
-  let t = randFr()
+  let t = randFr()  # intToFr(106)   
   return 
     ToxicWaste( alpha: a
               , beta:  b
@@ -129,6 +129,16 @@ func matricesToCoeffs*(matrices: Matrices): seq[Coeff] =
 
 #-------------------------------------------------------------------------------
 
+func dotProdFr(xs, ys: seq[Fr]): Fr = 
+  let n = xs.len
+  assert( n == ys.len, "dotProdFr: incompatible vector lengths" )
+  var s : Fr = zeroFr
+  for i in 0..<n:
+    s += xs[i] * ys[i]
+  return s
+
+#-------------------------------------------------------------------------------
+
 func fakeCircuitSetup*(r1cs: R1CS, toxic: ToxicWaste, flavour=Snarkjs): ZKey = 
  
   let neqs       = r1cs.constraints.len
@@ -171,6 +181,9 @@ func fakeCircuitSetup*(r1cs: R1CS, toxic: ToxicWaste, flavour=Snarkjs): ZKey =
 
   let D : Domain = createDomain(domSize) 
 
+#[
+  # this approach is very inefficient
+
   let polyAs : seq[Poly] = collect( newSeq , (for col in matrices.A: polyInverseNTT(col, D) ))
   let polyBs : seq[Poly] = collect( newSeq , (for col in matrices.B: polyInverseNTT(col, D) ))
   let polyCs : seq[Poly] = collect( newSeq , (for col in matrices.C: polyInverseNTT(col, D) ))
@@ -179,6 +192,20 @@ func fakeCircuitSetup*(r1cs: R1CS, toxic: ToxicWaste, flavour=Snarkjs): ZKey =
   let pointsB1 : seq[G1] = collect( newSeq , (for p in polyBs: polyEvalAt(p, toxic.tau) ** gen1) )
   let pointsB2 : seq[G2] = collect( newSeq , (for p in polyBs: polyEvalAt(p, toxic.tau) ** gen2) )
   let pointsC  : seq[G1] = collect( newSeq , (for p in polyCs: polyEvalAt(p, toxic.tau) ** gen1) )
+]#
+
+  # the Lagrange polynomials L_k(x) evaluated at x=tau
+  # we can then simply take the dot product of these with the column vectors to compute the points A,B1,B2,C
+  let lagrangeTaus : seq[Fr] = collect( newSeq, (for k in 0..<domSize: evalLagrangePolyAt(D, k, toxic.tau) ))
+  
+  let columnTausA  : seq[Fr] = collect( newSeq, (for col in matrices.A: dotProdFr(col,lagrangeTaus) ))
+  let columnTausB  : seq[Fr] = collect( newSeq, (for col in matrices.B: dotProdFr(col,lagrangeTaus) ))
+  let columnTausC  : seq[Fr] = collect( newSeq, (for col in matrices.C: dotProdFr(col,lagrangeTaus) ))
+
+  let pointsA  : seq[G1] = collect( newSeq , (for y in columnTausA: (y ** gen1) ))
+  let pointsB1 : seq[G1] = collect( newSeq , (for y in columnTausB: (y ** gen1) ))
+  let pointsB2 : seq[G2] = collect( newSeq , (for y in columnTausB: (y ** gen2) ))
+  let pointsC  : seq[G1] = collect( newSeq , (for y in columnTausC: (y ** gen1) ))
 
   let gammaInv : Fr = invFr(toxic.gamma)
   let deltaInv : Fr = invFr(toxic.delta)
diff --git a/groth16/test_proof.nim b/groth16/test_proof.nim
index cf3d603..bdb8474 100644
--- a/groth16/test_proof.nim
+++ b/groth16/test_proof.nim
@@ -1,6 +1,7 @@
 
 
-import std/[times, os]
+import std/[times,os]
+import strformat
 
 import groth16/prover
 import groth16/verifier
@@ -10,6 +11,8 @@ import groth16/files/zkey
 import groth16/zkey_types
 import groth16/fake_setup
 
+func seconds(x: float): string = fmt"{x:.4f} seconds"
+
 #-------------------------------------------------------------------------------
 
 proc testProveAndVerify*( zkey_fname, wtns_fname: string): (VKey,Proof) = 
@@ -22,7 +25,7 @@ proc testProveAndVerify*( zkey_fname, wtns_fname: string): (VKey,Proof) =
   let start = cpuTime()
   let proof = generateProof( zkey, witness )
   let elapsed = cpuTime() - start
-  echo("proving took ",elapsed)
+  echo("proving took ",seconds(elapsed))
 
   echo("verifying the proof...")
   let vkey = extractVKey( zkey)
@@ -41,7 +44,10 @@ proc testFakeSetupAndVerify*( r1cs_fname, wtns_fname: string, flavour=Snarkjs):
   let r1cs    = parseR1CS( r1cs_fname)
 
   echo("performing fake trusted setup...")
+  let start1 = cpuTime()
   let zkey = createFakeCircuitSetup( r1cs, flavour=flavour )
+  let elapsed1 = cpuTime() - start1
+  echo("fake setup took ",seconds(elapsed1))
 
   # printCoeffs(zkey.coeffs)
 
@@ -51,7 +57,7 @@ proc testFakeSetupAndVerify*( r1cs_fname, wtns_fname: string, flavour=Snarkjs):
   let start = cpuTime()
   let proof = generateProof( zkey, witness )
   let elapsed = cpuTime() - start
-  echo("proving took ",elapsed)
+  echo("proving took ",seconds(elapsed))
 
   echo("verifying the proof...")
   let ok = verifyProof( vkey, proof )