From 52858cb80fc412532fc81c09ea9a6e8d705be96c Mon Sep 17 00:00:00 2001
From: Balazs Komuves <bkomuves@gmail.com>
Date: Tue, 16 Jun 2026 13:50:28 +0200
Subject: [PATCH] merge the two witness MSM-s into a single one

---
 groth16/dynamic/v2/finish.nim     | 79 ++++++++++++++++++++++++++++++-
 groth16/dynamic/v2/preprocess.nim | 43 +++++++++++++++--
 2 files changed, 116 insertions(+), 6 deletions(-)

diff --git a/groth16/dynamic/v2/finish.nim b/groth16/dynamic/v2/finish.nim
index fbb7e93..7ba2acf 100644
--- a/groth16/dynamic/v2/finish.nim
+++ b/groth16/dynamic/v2/finish.nim
@@ -40,6 +40,83 @@ import groth16/dynamic/v2/setup
 
 #-------------------------------------------------------------------------------
 
+# you can turn on/off the computation of the nonlinear term (useful for the "dynark" version)
+proc finishLinearTerms( zkey: ZKey, wtns: Witness, partialProof: PartialProof, mask: Mask, pool: Taskpool, printTimings: bool): Proof =
+
+  assert( zkey.header.curve == wtns.curve )
+
+  let witness = wtns.values
+
+  let hdr  : GrothHeader  = zkey.header
+  let spec : SpecPoints   = zkey.specPoints
+  let pts  : ProverPoints = zkey.pPoints     
+
+  let nvars = hdr.nvars
+  let npubs = hdr.npubs
+
+  assert( nvars == witness.len , "wrong witness length" )
+
+  let partial_mask = partialProof.partial_mask
+
+  # remark: with the special variable "1" we actuall have (npub+1) public IO variables
+  var pubIO = newSeq[Fr[BN254_Snarks]]( npubs + 1)
+  for i in 0..npubs: pubIO[i] = witness[i]             
+
+  # note: we have to ignore public inputs (plus 1 for the special first entry)
+  let startIdx : int = npubs + 1
+  var count    : int = countFalses(partial_mask)       # number of NEW witness elements 
+
+  # compactify the stuff so we can call normal vector MSM
+  var compact_witness:  seq[F]  = newSeq[F](count)
+  var compact_pointsA1: seq[G1] = newSeq[G1](count)
+  var compact_pointsB1: seq[G1] = newSeq[G1](count)
+  var compact_pointsB2: seq[G2] = newSeq[G2](count)
+  block:
+    var j: int = 0;
+    for i in 0..<nvars:
+      if (not partial_mask[i]):
+        compact_witness[j ] = witness[i]
+        compact_pointsA1[j] = pts.pointsA1[i]
+        compact_pointsB1[j] = pts.pointsB1[i]
+        compact_pointsB2[j] = pts.pointsB2[i]
+        j += 1
+
+  # masking coeffs
+  let r = mask.r
+  let s = mask.s
+
+  assert( witness.len == pts.pointsA1.len )
+  assert( witness.len == pts.pointsB1.len )
+  assert( witness.len == pts.pointsB2.len )
+  assert( hdr.domainSize == pts.pointsH1.len )
+
+  # echo "count = " & $count
+  # for x in compact_witness:
+  #   echo toDecimalFr(x)
+
+  var pi_a : G1 = partialProof.partial_pi_a
+  withMeasureTime(printTimings,"computing pi_A (G1 MSM)"):
+    pi_a += r ** spec.delta1
+    pi_a += msmMultiThreadedG1( compact_witness , compact_pointsA1, pool )
+
+  var rho : G1 = partialProof.partial_rho
+  withMeasureTime(printTimings,"computing rho (G1 MSM)"):
+    # rho += s ** spec.delta1
+    rho += msmMultiThreadedG1( compact_witness , compact_pointsB1, pool )
+
+  var pi_b : G2 = partialProof.partial_pi_b
+  withMeasureTime(printTimings,"computing pi_B (G2 MSM)"):
+    pi_b += s ** spec.delta2
+    pi_b += msmMultiThreadedG2( compact_witness , compact_pointsB2, pool )
+
+  var pi_c : G1 = partialProof.partial_pi_c
+  pi_c += s ** pi_a
+  pi_c += r ** rho
+
+  return Proof( curve:"bn128", publicIO:pubIO, pi_a:pi_a, pi_b:pi_b, pi_c:pi_c )
+
+#-------------------------------------------------------------------------------
+
 proc finishDynaProofWithMaskV2*( zkey: ZKey, wtns: Witness, dynaPreProof: DynaPreProofV2, mask: Mask, pool: Taskpool, printTimings: bool): Proof =
 
   let N = zkey.header.domainSize
@@ -63,7 +140,7 @@ proc finishDynaProofWithMaskV2*( zkey: ZKey, wtns: Witness, dynaPreProof: DynaPr
 
   var proof: Proof
   withMeasureTime(printTimings,"finishing the linear terms"):
-    proof = finishPartialProofWithMaskGeneric( false, zkey, wtns, dynaPreProof.partialProof, mask, pool, false )
+    proof = finishLinearTerms( zkey, wtns, dynaPreProof.partialProof, mask, pool, false )
 
   var nonlin: G1 = infG1
 
diff --git a/groth16/dynamic/v2/preprocess.nim b/groth16/dynamic/v2/preprocess.nim
index 06816eb..7337b85 100644
--- a/groth16/dynamic/v2/preprocess.nim
+++ b/groth16/dynamic/v2/preprocess.nim
@@ -53,9 +53,39 @@ func projectionElementsV2*(setup: DynaSetupV2, D: Domain, As: seq[F], complImage
 
   return Us
 
-#---------------------------------------
+#-------------------------------------------------------------------------------
 
-proc dynaPreprocessV2*(zkey: ZKey, setup: DynaSetupV2, partialAB: PartialAB, zdeltaMask: seq[bool] ): DynaPreprocessV2 =
+proc compactifyPointsC1( zkey: ZKey, partial_mask: seq[bool] ): seq[G1] =
+
+  let hdr  : GrothHeader  = zkey.header
+  let spec : SpecPoints   = zkey.specPoints
+  let pts  : ProverPoints = zkey.pPoints     
+
+  let nvars = hdr.nvars
+  let npubs = hdr.npubs
+
+  # note: we have to ignore public inputs (plus 1 for the special first entry)
+  let startIdx : int = npubs + 1
+  let count    : int = countFalses(partial_mask) 
+
+  var compact_pointsC1: seq[G1] = newSeq[G1](count)
+
+  var j: int = 0;
+  for i in 0..<nvars:
+    if (not partial_mask[i]):
+      if i >= startIdx:
+        compact_pointsC1[j] = pts.pointsC1[ i - startIdx ]
+      else:
+        compact_pointsC1[j] = infG1
+      j += 1
+
+  assert( nvars - npubs - 1 == pts.pointsC1.len )
+
+  return compact_pointsC1
+
+#-------------------------------------------------------------------------------
+
+proc dynaPreprocessV2*(zkey: ZKey, setup: DynaSetupV2, partialAB: PartialAB, partial_mask: seq[bool], zdeltaMask: seq[bool] ): DynaPreprocessV2 =
   let N = zkey.header.domainSize
   let D = createDomain(N)
 
@@ -69,6 +99,8 @@ proc dynaPreprocessV2*(zkey: ZKey, setup: DynaSetupV2, partialAB: PartialAB, zde
   let m = countTrues(zdeltaMask)
   var Xs: seq[G1] = newSeq[G1]( m )
 
+  let points_C1 = compactifyPointsC1( zkey, partial_mask )
+
   withMeasureTime(true," + the \"unified\" points X_j"):
 
     let matABC = zkeyToSparseMatrices(zkey)
@@ -76,7 +108,7 @@ proc dynaPreprocessV2*(zkey: ZKey, setup: DynaSetupV2, partialAB: PartialAB, zde
     let colsB = selectTrues( zdeltaMask , matABC.B.columns )
 
     for j in 0..<m:
-      var h : G1 = infG1
+      var h : G1 = points_C1[j]
       for (k,a) in colsA[j].pairs: h += (a ** projB_full[k])
       for (k,b) in colsB[j].pairs: h += (b ** projA_full[k])
       Xs[j] = h
@@ -113,7 +145,8 @@ proc dynaPreProofV2*(zkey: ZKey, setup: DynaSetupV2, partialWitness: PartialWitn
   let N = zkey.header.domainSize
   let D = createDomain(N)
 
-  let zdeltaMask = notBoolSeq( partialWitnessMask( partialWitness ) )
+  let partialMask = partialWitnessMask( partialWitness )
+  let zdeltaMask  = notBoolSeq( partialMask )
 
   var partialAB : PartialAB
   withMeasureTime(printTimings,"build partial AB"):
@@ -137,7 +170,7 @@ proc dynaPreProofV2*(zkey: ZKey, setup: DynaSetupV2, partialWitness: PartialWitn
 
   var preprocess : DynaPreprocessV2
   withMeasureTime(printTimings,"precomputing the \"projection\" points"):
-    preprocess = dynaPreprocessV2( zkey, setup, partialAB, zdeltaMask )
+    preprocess = dynaPreprocessV2( zkey, setup, partialAB, partialMask, zdeltaMask )
 
   return DynaPreProofV2( partialProof   : partialProof ,
                          deltaImages    : deltaImages  ,