From 67d0126cf4a7fb4faa8309a8f6312eb2a0624990 Mon Sep 17 00:00:00 2001
From: Balazs Komuves <bkomuves@gmail.com>
Date: Mon, 15 Jun 2026 19:25:45 +0200
Subject: [PATCH] row permutations (so that the final witness delta's image can
 be in a subgroup)

---
 groth16/bn128/arrays.nim    |  60 ++++++++++++
 groth16/dynamic/permute.nim | 176 ++++++++++++++++++++++++++++++++++++
 groth16/dynamic/setup.nim   |   6 ++
 groth16/files/r1cs.nim      |  15 +++
 groth16/math/domain.nim     |  26 +++++-
 groth16/misc.nim            |   4 +
 groth16/partial/finish.nim  |   1 +
 groth16/partial/precalc.nim |  41 +++++++++
 groth16/partial/types.nim   |   2 +
 9 files changed, 326 insertions(+), 5 deletions(-)
 create mode 100644 groth16/dynamic/permute.nim

diff --git a/groth16/bn128/arrays.nim b/groth16/bn128/arrays.nim
index eeb96d2..4856d77 100644
--- a/groth16/bn128/arrays.nim
+++ b/groth16/bn128/arrays.nim
@@ -29,6 +29,66 @@ func constSeq*[T]( N: int , y: T ): seq[T] =
     xs[i] = y
   return xs
 
+#-------------------------------------------------------------------------------
+# bool arrays
+
+func countTrues*( bs: seq[bool] ): int =
+  var cnt = 0 
+  for b in bs: 
+    if b: 
+      cnt += 1
+  return cnt
+
+func countFalses*( bs: seq[bool] ): int =
+  var cnt = 0 
+  for b in bs: 
+    if not b: 
+      cnt += 1
+  return cnt
+
+func trueIndices*( bs: seq[bool] ): seq[int] =
+  let k = countTrues(bs)
+  var idxs: seq[int] = newSeq[int]( k )
+  var j = 0
+  for (i,b) in bs.pairs:
+    if b:
+      idxs[j] = i
+      j += 1
+  return idxs
+  
+func falseIndices*( bs: seq[bool] ): seq[int] =
+  let k = countFalses(bs)
+  var idxs: seq[int] = newSeq[int]( k )
+  var j = 0
+  for (i,b) in bs.pairs:
+    if not b:
+      idxs[j] = i
+      j += 1
+  return idxs
+
+func notBoolSeq*( us: seq[bool]): seq[bool] =
+  let n = us.len
+  var ws: seq[bool] = newSeq[bool]( n )
+  for i in 0..<n:
+    ws[i] = not us[i]
+  return ws
+
+func andBoolSeqs*( us: seq[bool] , vs: seq[bool]): seq[bool] =
+  let n = us.len
+  assert( n == vs.len )
+  var ws: seq[bool] = newSeq[bool]( n )
+  for i in 0..<n:
+    ws[i] = us[i] and vs[i]
+  return ws
+
+func orBoolSeqs*( us: seq[bool] , vs: seq[bool]): seq[bool] =
+  let n = us.len
+  assert( n == vs.len )
+  var ws: seq[bool] = newSeq[bool]( n )
+  for i in 0..<n:
+    ws[i] = us[i] or vs[i]
+  return ws
+
 #-------------------------------------------------------------------------------
 # Fr arrays
 
diff --git a/groth16/dynamic/permute.nim b/groth16/dynamic/permute.nim
new file mode 100644
index 0000000..6181847
--- /dev/null
+++ b/groth16/dynamic/permute.nim
@@ -0,0 +1,176 @@
+
+#
+# Ideally we want to the image of update (delta) of the witness under 
+# both the A and B matrices to lie in a small subgroup
+# 
+# This can be always achieved simply by permuting the rows. 
+# Here we compute such a permutation
+# 
+
+import std/sequtils
+import std/algorithm
+import std/tables
+
+import constantine/math/arithmetic
+import constantine/named/properties_fields
+
+import groth16/bn128
+import groth16/bn128/arrays
+import groth16/math/domain
+import groth16/math/matrix
+import groth16/zkey_types
+import groth16/misc
+
+#-------------------------------------------------------------------------------
+
+type 
+
+  F = Fr[BN254_Snarks]
+
+  Permutation* = object
+    srcIndices* : seq[int]
+
+#-------------------------------------------------------------------------------
+
+func isPermutation*(input: seq[int]): bool =
+  let n = input.len
+  let tgt: seq[int] = toSeq(0..<n)
+  var tmp = input
+  tmp.sort()
+  return (tmp == tgt)
+
+func isValidPermutation*(perm: Permutation): bool =
+  return isPermutation(perm.srcIndices)
+
+func inversePermutation*(P: Permutation): Permutation = 
+  let N = P.srcIndices.len
+  let ps: seq[int] = P.srcIndices
+  var qs: seq[int] = newSeq[int]( N )
+  for i in 0..<N:
+    qs[ps[i]] = i
+  return Permutation(srcIndices: qs)
+
+func permuteSeq*[T](P: Permutation, xs: seq[T]) = 
+  let N  = xs.len
+  let ps = P.srcIndices
+  assert( N == ps.len )
+
+  var ys: seq[T] = newSeq[T]( N )
+  for i in 0..<N:
+    ys[i] = xs[ps[i]]
+
+#-------------------------------------------------------------------------------
+
+func permuteSparseColumn*[T]( perm: Permutation , column: SparseColumn[T] ): SparseColumn[T] =
+  let tgtIndices = inversePermutation(perm).srcIndices
+  var output: Table[int,T]
+  for (i,y) in column:
+    output[tgtIndices[i]] = y
+  return output
+
+func permuteMatrixRows*[T]( perm: Permutation , mat: SparseMatrix[T] ): SparseMatrix[T] =
+  assert( mat.dims.nrows == perm.srcIndices.len )
+  let tgtIndices = inversePermutation(perm).srcIndices
+  var outputs: seq[SparseColumn[T]] = newSeq[SparseColumn[T]]( mat.dims.ncols )
+  for j in 0..<mat.dims.ncols:
+    var column: Table[int,T]
+    for (i,y) in mat.columns[j].pairs:
+      column[tgtIndices[i]] = y
+    outputs[j] = column
+  return SparseMatrix[T]( dims: mat.dims, columns: outputs )
+
+#-------------------------------------------------------------------------------
+
+func liesInSubgroup*( sg: Subgroup, mask: seq[bool] ): bool = 
+  let N = sg.bigDomain.domainSize
+  let K = sg.smallDomain.domainSize
+  let J = N div K
+
+  assert( N == mask.len )
+
+  var ok = true
+  for (i,b) in mask.pairs:
+    if b:
+      ok = ok and ((i mod J) == 0)
+
+  return ok
+
+#-------------------------------------------------------------------------------
+
+proc findRowPermutation*(A: SparseMatrix[F], B: SparseMatrix[F], witness_delta_mask: seq[bool]): (int,Permutation) =
+
+  let N = A.dims.nrows
+  let M = A.dims.ncols
+
+  assert( isPowerOfTwo(N) )
+  assert( A.dims == B.dims )
+  assert( witness_delta_mask.len == M )
+
+  let deltaImgA  = sparseMatrixImage( A , witness_delta_mask )
+  let deltaImgB  = sparseMatrixImage( B , witness_delta_mask )
+  let deltaImgAB = orBoolSeqs( deltaImgA , deltaImgB )
+
+  let k0   = countTrues(deltaImgAB)
+  let logK = ceilingLog2(k0)
+  let K    = (1 shl logK)           # size of the subgroup
+
+  var sgIndex = N div K             # index of the subgroup
+
+  # echo "domain size    = " & $N
+  # echo "image size     = " & $k0
+  # echo "subgroup size  = " & $K
+  # echo "subgroup index = " & $sgIndex
+
+  var used: seq[bool] = constSeq(N, false)
+  var perm: seq[int]  = newSeq[int]( N )
+
+  # first we fill in the subgroup part
+  var k: int = 0                 
+  for (i,b) in deltaImgAB.pairs:
+    if b:
+      let j = k * sgIndex
+      perm[j] = i
+      used[j] = true
+      k += 1
+
+  # next we fill in the rest
+  var j: int = 0
+  for (i,b) in deltaImgAB.pairs:
+    if not b:
+      while(used[j]): 
+        j += 1
+      perm[j] = i
+      j += 1
+
+  # # debugging
+  # echo $trueIndices(deltaImgAB)
+  # for j in 0..<K:
+  #   echo $perm[sgIndex*j]
+
+  let P = Permutation(srcIndices: perm)
+
+  # sanity checks
+
+  assert( isValidPermutation(P) , "findRowPermutation: the result is not valid permutation :(" )
+
+  let A1 = permuteMatrixRows( P , A )
+  let B1 = permuteMatrixRows( P , B )
+
+  let newImgA  = sparseMatrixImage( A1 , witness_delta_mask )
+  let newImgB  = sparseMatrixImage( B1 , witness_delta_mask )
+  let newImgAB = orBoolSeqs( newImgA , newImgB )
+ 
+  let sg = createSubgroup( createDomain(N) , K )
+
+  assert( liesInSubgroup(sg , newImgAB) , "permuted matrices are not compatible with the subgroup...")
+
+  return (K,P)
+  
+#-------------------------------------------------------------------------------
+
+proc zkeyFindRowPermutation*(zkey: ZKey, witness_delta_mask: seq[bool]): (int,Permutation) =
+  let matrices = zkeyToSparseMatrices(zkey)
+  return findRowPermutation( matrices.A , matrices.B , witness_delta_mask )
+
+#-------------------------------------------------------------------------------
+
diff --git a/groth16/dynamic/setup.nim b/groth16/dynamic/setup.nim
index 5d5169f..bba1977 100644
--- a/groth16/dynamic/setup.nim
+++ b/groth16/dynamic/setup.nim
@@ -19,6 +19,12 @@ import groth16/dynamic/shared
 
 #-------------------------------------------------------------------------------
 
+#  let deltaImgA  = sparseMatrixImage( A , delta_mask )
+#  let deltaImgB  = sparseMatrixImage( B , delta_mask )
+#  let deltaImgAB = orBoolSeqs( deltaImgA , deltaImgB )
+
+#-------------------------------------------------------------------------------
+
 # does the setup from the ZKey (prover key)
 func dynaSetupV1FromZKey*(zkey: Zkey): DynaSetupV1 = 
 
diff --git a/groth16/files/r1cs.nim b/groth16/files/r1cs.nim
index ad3d4fa..3f59a17 100644
--- a/groth16/files/r1cs.nim
+++ b/groth16/files/r1cs.nim
@@ -82,6 +82,21 @@ type
 
 #-------------------------------------------------------------------------------
 
+proc printWitnessConfig*(cfg: WitnessConfig) = 
+  echo "R1CS witness config:"
+  echo " - nWires  = " & $cfg.nWires  
+  echo " - nPubOut = " & $cfg.nPubOut 
+  echo " - nPubIn  = " & $cfg.nPubIn  
+  echo " - nPrivIn = " & $cfg.nPrivIn 
+  echo " - nLabels = " & $cfg.nLabels 
+
+proc printR1CSMetaData*(r1cs: R1CS) = 
+  printWitnessConfig(r1cs.cfg)
+  echo "nConstraints = " & $r1cs.nConstr
+  assert (r1cs.nConstr == r1cs.constraints.len)
+
+#-------------------------------------------------------------------------------
+
 proc parseSection1_header( stream: Stream, user: var R1CS, sectionLen: int ) =
   # echo "\nparsing r1cs header"
   
diff --git a/groth16/math/domain.nim b/groth16/math/domain.nim
index 160b418..ff3b048 100644
--- a/groth16/math/domain.nim
+++ b/groth16/math/domain.nim
@@ -14,12 +14,17 @@ import groth16/misc
 #-------------------------------------------------------------------------------
 
 type 
+
   Domain* = object
-    domainSize*    : int              # `N = 2^n`
-    logDomainSize* : int              # `n = log2(N)`
-    domainGen*     : Fr[BN254_Snarks] # `g`
-    invDomainGen*  : Fr[BN254_Snarks] # `g^-1`
-    invDomainSize* : Fr[BN254_Snarks] # `1/n`
+    domainSize*    : int                # `N = 2^n`
+    logDomainSize* : int                # `n = log2(N)`
+    domainGen*     : Fr[BN254_Snarks]   # `g`
+    invDomainGen*  : Fr[BN254_Snarks]   # `g^-1`
+    invDomainSize* : Fr[BN254_Snarks]   # `1/n`
+
+  Subgroup* = object
+    bigDomain*   : Domain
+    smallDomain* : Domain
 
 #-------------------------------------------------------------------------------
 
@@ -58,3 +63,14 @@ func enumerateDomain*(D: Domain): seq[Fr[BN254_Snarks]] =
 
 #-------------------------------------------------------------------------------
 
+func subgroupIndex*(sg: Subgroup): int = 
+  return (sg.bigDomain.domainSize div sg.smallDomain.domainSize)
+
+func createSubgroup*(D: Domain, K: int): Subgroup =
+  let N = D.domainSize
+  assert( K >= 1 )
+  assert( K <= N )
+  assert( (N mod K) == 0 )
+  return Subgroup( bigDomain: D, smallDomain: createDomain(K) )
+
+#-------------------------------------------------------------------------------
diff --git a/groth16/misc.nim b/groth16/misc.nim
index 2426925..247d0e2 100644
--- a/groth16/misc.nim
+++ b/groth16/misc.nim
@@ -46,6 +46,10 @@ func ceilingLog2* (x : int) : int =
   else:
     return (floorLog2(x-1) + 1)
 
+func isPowerOfTwo*( N: int): bool =
+  let reN = (1 shl ceilingLog2(N))
+  return ( N == reN )
+
 #-------------------------------------------------------------------------------
 
 #[
diff --git a/groth16/partial/finish.nim b/groth16/partial/finish.nim
index 33857b4..e8676dd 100644
--- a/groth16/partial/finish.nim
+++ b/groth16/partial/finish.nim
@@ -14,6 +14,7 @@ import constantine/named/properties_fields
 import groth16/bn128
 import groth16/math/domain
 import groth16/math/poly
+import groth16/math/matrix
 import groth16/zkey_types
 import groth16/files/witness
 import groth16/misc
diff --git a/groth16/partial/precalc.nim b/groth16/partial/precalc.nim
index c1b1d83..9a179f1 100644
--- a/groth16/partial/precalc.nim
+++ b/groth16/partial/precalc.nim
@@ -18,15 +18,56 @@ import system
 import taskpools
 
 import groth16/bn128
+import groth16/bn128/arrays
 import groth16/zkey_types
 import groth16/misc
 
+import groth16/math/matrix
 import groth16/partial/types
 
 #import groth16/math/domain
 #import groth16/math/poly
 #import groth16/prover/shared
 
+#-------------------------------------------------------------------------------
+
+proc printPartialWitnessStats*(zkey: ZKey, pw: PartialWitness) = 
+  let n = zkey.header.domainSize
+  let m = zkey.header.nvars
+
+  echo ""
+  echo "N (domain)  = " & $n
+  echo "M (witness) = " & $m
+
+  let matrices = zkeyToSparseMatrices(zkey)
+  var partial_mask: seq[bool] = newSeq[bool] ( m )
+  var delta_mask:   seq[bool] = newSeq[bool] ( m )
+  for j in 0..<m:
+    partial_mask[j] = pw.values[j].isSome()
+    delta_mask[j]   = not partial_mask[j]
+
+  let partImgA   = sparseMatrixImage( matrices.A , partial_mask )
+  let partImgB   = sparseMatrixImage( matrices.B , partial_mask )
+
+  let deltaImgA  = sparseMatrixImage( matrices.A , delta_mask )
+  let deltaImgB  = sparseMatrixImage( matrices.B , delta_mask )
+  let deltaImgAB = orBoolSeqs( deltaImgA , deltaImgB )
+
+  let intersectA = andBoolSeqs( partImgA , deltaImgA )
+  let intersectB = andBoolSeqs( partImgB , deltaImgB )
+
+  echo ""
+  echo "A mtx partial witness image = " & $countTrues(partImgA  )
+  echo "A mtx delta image           = " & $countTrues(deltaImgA )
+  echo "A mtx images intersection   = " & $countTrues(intersectA)
+  echo ""
+  echo "B mtx partial witness image = " & $countTrues(partImgB  )
+  echo "B mtx delta image           = " & $countTrues(deltaImgB )
+  echo "B mtx images intersection   = " & $countTrues(intersectB)
+  echo ""
+  echo "union of A and B delta      = " & $countTrues(deltaImgAB)
+  echo ""
+  
 #-------------------------------------------------------------------------------
 # the prover
 #
diff --git a/groth16/partial/types.nim b/groth16/partial/types.nim
index c4b9f0f..b8aaf38 100644
--- a/groth16/partial/types.nim
+++ b/groth16/partial/types.nim
@@ -31,6 +31,7 @@ type
 func makePartialWitness*(vals: seq[Option[Fr[BN254_Snarks]]]): PartialWitness =
   return PartialWitness(values: vals)
 
+# true where it's filled
 func partialWitnessMask*(pw: PartialWitness): seq[bool] = 
   let N = pw.values.len
   var bs: seq[bool] = newSeq[bool]( N )
@@ -39,3 +40,4 @@ func partialWitnessMask*(pw: PartialWitness): seq[bool] =
   return bs 
 
 #-------------------------------------------------------------------------------
+