From 159f8b78c401ff21275432c68ea931a8275634cc Mon Sep 17 00:00:00 2001
From: Balazs Komuves <bkomuves@gmail.com>
Date: Mon, 15 Jun 2026 14:13:56 +0200
Subject: [PATCH] some sparse matrix stuff

---
 cli/cli_main.nim               |   1 +
 groth16/dynamic/preprocess.nim |   4 +-
 groth16/dynamic/types.nim      |   8 ++-
 groth16/fake_setup.nim         |  47 ++++----------
 groth16/math/matrix.nim        | 115 +++++++++++++++++++++++++++++++++
 groth16/misc.nim               |  11 ----
 groth16/zkey_types.nim         |  55 +++++++++++++---
 tests/groth16/testMisc.nim     |  35 ++++++++++
 tests/test.nim                 |   1 +
 9 files changed, 216 insertions(+), 61 deletions(-)
 create mode 100644 groth16/math/matrix.nim
 create mode 100644 tests/groth16/testMisc.nim

diff --git a/cli/cli_main.nim b/cli/cli_main.nim
index e1dc300..273f3a8 100644
--- a/cli/cli_main.nim
+++ b/cli/cli_main.nim
@@ -175,6 +175,7 @@ proc cliMain(cfg: Config) =
   
   if cfg.debug:
     printGrothHeader(zkey.header)
+    printR1csStats(zkey)
     # debugPrintCoeffs(zkey.coeffs)
 
   if cfg.partial_sanity:
diff --git a/groth16/dynamic/preprocess.nim b/groth16/dynamic/preprocess.nim
index ce51ea6..1776f9f 100644
--- a/groth16/dynamic/preprocess.nim
+++ b/groth16/dynamic/preprocess.nim
@@ -115,7 +115,7 @@ proc testProjectionElementsV1*(N: int, tau: F, delta: F ): bool =
 
 #---------------------------------------
 
-func dynaPreprocessV1*(zkey: ZKey, setup: DynaSetupV1, partialWitness: PartialWitness): DynaPreprocess =
+func dynaPreprocessV1*(zkey: ZKey, setup: DynaSetupV1, partialWitness: PartialWitness): DynaPreprocessV1 =
   let N = zkey.header.domainSize
   let D = createDomain(N)
 
@@ -126,7 +126,7 @@ func dynaPreprocessV1*(zkey: ZKey, setup: DynaSetupV1, partialWitness: PartialWi
   let projA = projectionElementsV1( setup , D , partialAB.valuesAz , partialAB.complImageA )
   let projB = projectionElementsV1( setup , D , partialAB.valuesBz , partialAB.complImageB )
   
-  return DynaPreprocess( projA0: projA, projB0: projB )
+  return DynaPreprocessV1( projA0: projA, projB0: projB )
 
 #-------------------------------------------------------------------------------
 
diff --git a/groth16/dynamic/types.nim b/groth16/dynamic/types.nim
index 07d5226..67e1b06 100644
--- a/groth16/dynamic/types.nim
+++ b/groth16/dynamic/types.nim
@@ -36,18 +36,20 @@ type
     wConvDeltaLZ*  : seq[G1]      # the convolution of `W` and `pointsDeltaLZ`
 
   # things we can compute from the partial witness
-  DynaPreprocess* = object
+  DynaPreprocessV1* = object
     projA0*       : seq[G1]      # the points `delta^-1 * A0(tau) * L_i(tau) * g1` (well, "modulo Z(tau)")
     projB0*       : seq[G1]      # the same for B0
 
-  DynaPreProof* = object
+  DynaPreProofV1* = object
     partialProof*   : PartialProof
-    dyanPreprocess* : DynaPreprocess
+    dyanPreprocess* : DynaPreprocessV1
 
+  # this is used in the "cross-term"
   OnlyAB* = object
     valuesAz*    : seq[F]
     valuesBz*    : seq[F]
 
+  # this one is used in the "projection-term"
   PartialAB* = object
     valuesAz*    : seq[F]
     valuesBz*    : seq[F]
diff --git a/groth16/fake_setup.nim b/groth16/fake_setup.nim
index 95b6c2b..92e65fb 100644
--- a/groth16/fake_setup.nim
+++ b/groth16/fake_setup.nim
@@ -15,6 +15,7 @@ import constantine/named/properties_fields
 import groth16/bn128
 #import groth16/bn128/arrays
 import groth16/math/domain
+import groth16/math/matrix
 import groth16/math/poly
 import groth16/zkey_types
 import groth16/files/r1cs
@@ -69,16 +70,6 @@ func r1csToCoeffs*(r1cs: R1CS): seq[Coeff] =
 #-------------------------------------------------------------------------------
 # Note: dense matrices can be very big, this is only feasible for small circuits
 
-type DenseColumn*[T] = seq[T]
-
-type DenseMatrix*[T] = seq[DenseColumn[T]]
-
-type 
-  DenseMatrices* = object
-    A* : DenseMatrix[Fr[BN254_Snarks]]
-    B* : DenseMatrix[Fr[BN254_Snarks]]
-    C* : DenseMatrix[Fr[BN254_Snarks]]
-
 #[
 
 func r1csToDenseMatrices*(r1cs: R1CS): DenseMatrices =
@@ -137,36 +128,17 @@ func denseMatricesToCoeffs*(matrices: DenseMatrices): seq[Coeff] =
 
 #-------------------------------------------------------------------------------
 
-type SparseColumn*[T] = Table[int,T]
-
-proc columnInsertWithAddFr( col: var SparseColumn[Fr[BN254_Snarks]] , i: int,  y: Fr[BN254_Snarks] ) =
-  var x = getOrDefault( col, i, zeroFr )
-  x += y
-  col[i] = x
-
-proc sparseDenseDotProdFr( U: SparseColumn[Fr[BN254_Snarks]], V: DenseColumn[Fr[BN254_Snarks]] ): Fr[BN254_Snarks] =
-  var acc : Fr[BN254_Snarks] = zeroFr
-  for i,x in U.pairs:
-    acc += x * V[i]
-  return acc
-
-type SparseMatrix*[T] = seq[SparseColumn[T]]
-
-type 
-  SparseMatrices* = object
-    A* : SparseMatrix[Fr[BN254_Snarks]]
-    B* : SparseMatrix[Fr[BN254_Snarks]]
-    C* : SparseMatrix[Fr[BN254_Snarks]]
-
 func r1csToSparseMatrices*(r1cs: R1CS): SparseMatrices =
   let n = r1cs.constraints.len
   let m = r1cs.cfg.nWires
   let p = r1cs.cfg.nPubIn + r1cs.cfg.nPubOut
 
+  let dims = MatrixDims( nrows: n , ncols: m )
+
   let logDomSize = ceilingLog2(n+p+1)
   let domSize    = 1 shl logDomSize
 
-  var matA, matB, matC: SparseMatrix[Fr[BN254_Snarks]]
+  var matA, matB, matC: SparseMatrixColumns[Fr[BN254_Snarks]]
   for i in 0..<m:
     var colA : SparseColumn[Fr[BN254_Snarks]] = initTable[int,Fr[BN254_Snarks]]()
     var colB : SparseColumn[Fr[BN254_Snarks]] = initTable[int,Fr[BN254_Snarks]]()
@@ -186,7 +158,10 @@ func r1csToSparseMatrices*(r1cs: R1CS): SparseMatrices =
   for i in n..n+p:
     columnInsertWithAddFr( matA[i-n] , i , oneFr )
 
-  return SparseMatrices(A:matA, B:matB, C:matC)
+  let mA = SparseMatrix[Fr[BN254_Snarks]]( dims: dims , columns: matA )
+  let mB = SparseMatrix[Fr[BN254_Snarks]]( dims: dims , columns: matB )
+  let mC = SparseMatrix[Fr[BN254_Snarks]]( dims: dims , columns: matC )
+  return SparseMatrices(A:mA, B:mB, C:mC)
 
 #-------------------------------------------------------------------------------
 
@@ -253,9 +228,9 @@ func fakeCircuitSetup*(r1cs: R1CS, toxic: ToxicWaste, flavour=Snarkjs): ZKey =
   let columnTausC  : seq[Fr] = collect( newSeq, (for col in matrices.C: dotProdFr(col,lagrangeTaus) ))
 ]#
 
-  let columnTausA  : seq[Fr[BN254_Snarks]] = collect( newSeq, (for col in matrices.A: sparseDenseDotProdFr(col,lagrangeTaus) ))
-  let columnTausB  : seq[Fr[BN254_Snarks]] = collect( newSeq, (for col in matrices.B: sparseDenseDotProdFr(col,lagrangeTaus) ))
-  let columnTausC  : seq[Fr[BN254_Snarks]] = collect( newSeq, (for col in matrices.C: sparseDenseDotProdFr(col,lagrangeTaus) ))
+  let columnTausA  : seq[Fr[BN254_Snarks]] = collect( newSeq, (for col in matrices.A.columns: sparseDenseDotProdFr(col,lagrangeTaus) ))
+  let columnTausB  : seq[Fr[BN254_Snarks]] = collect( newSeq, (for col in matrices.B.columns: sparseDenseDotProdFr(col,lagrangeTaus) ))
+  let columnTausC  : seq[Fr[BN254_Snarks]] = collect( newSeq, (for col in matrices.C.columns: sparseDenseDotProdFr(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) ))
diff --git a/groth16/math/matrix.nim b/groth16/math/matrix.nim
new file mode 100644
index 0000000..2b09a37
--- /dev/null
+++ b/groth16/math/matrix.nim
@@ -0,0 +1,115 @@
+
+import std/tables
+
+import constantine/math/arithmetic
+import constantine/named/properties_fields
+
+import groth16/bn128
+import groth16/bn128/arrays
+
+#-------------------------------------------------------------------------------
+# dimensions
+
+type
+  MatrixDims* = object 
+    nrows* : int
+    ncols* : int
+
+#-------------------------------------------------------------------------------
+# Dense matrices
+#
+# Note: dense matrices can be very big, this is only feasible for small circuits
+
+type 
+
+  DenseColumn*[T] = seq[T]
+
+  DenseMatrixColumns*[T] = seq[DenseColumn[T]]
+
+  DenseMatrix*[T] = object
+    dims*    : MatrixDims
+    columns* : seq[DenseColumn[T]]
+
+  DenseMatrices* = object
+    A* : DenseMatrix[Fr[BN254_Snarks]]
+    B* : DenseMatrix[Fr[BN254_Snarks]]
+    C* : DenseMatrix[Fr[BN254_Snarks]]
+
+#-------------------------------------------------------------------------------
+# Sparse matrices
+
+type 
+
+  SparseColumn*[T] = Table[int,T]
+
+  SparseMatrixColumns*[T] = seq[SparseColumn[T]]
+
+  SparseMatrix*[T] = object
+    dims*    : MatrixDims
+    columns* : seq[SparseColumn[T]]
+
+  SparseMatrices* = object
+    A* : SparseMatrix[Fr[BN254_Snarks]]
+    B* : SparseMatrix[Fr[BN254_Snarks]]
+    C* : SparseMatrix[Fr[BN254_Snarks]]
+
+proc columnInsertWithAddFr*( column: var SparseColumn[Fr[BN254_Snarks]] , row: int,  y: Fr[BN254_Snarks] ) =
+  var x = getOrDefault( column, row, zeroFr )
+  x += y
+  column[row] = x
+
+proc sparseDenseDotProdFr*( U: SparseColumn[Fr[BN254_Snarks]], V: DenseColumn[Fr[BN254_Snarks]] ): Fr[BN254_Snarks] =
+  var acc : Fr[BN254_Snarks] = zeroFr
+  for i,x in U.pairs:
+    acc += x * V[i]
+  return acc
+
+#-------------------------------------------------------------------------------
+# densities
+
+# counts the non-zero elements in each row
+func sparseMatrixRowCounts*( A : SparseMatrix[Fr[BN254_Snarks]] ): seq[int] =
+  var rowCounts: seq[int] = newSeq[int]( A.dims.nrows )
+  for j,column in A.columns.pairs:
+    for i,value in column.pairs:
+      if not isZeroFr(value):
+        rowCounts[i] += 1
+  return rowCounts
+
+# counts the non-zero elements in each column
+func sparseMatrixColumnCounts*( A : SparseMatrix[Fr[BN254_Snarks]] ): seq[int] =
+  var colCounts: seq[int] = newSeq[int]( A.dims.ncols )
+  for j,column in A.columns.pairs:
+    for i,value in column.pairs:
+      if not isZeroFr(value):
+        colCounts[j] += 1
+  return colCounts
+
+# average count of non-zero elements in the rows
+func sparseMatrixAvgRowDensity*( A : SparseMatrix[Fr[BN254_Snarks]] ): float64 =
+  let rowCounts = sparseMatrixRowCounts( A )
+  var s: float64 = 0
+  for x in rowCounts: s += x.float64
+  return (s / rowCounts.len.float64)
+
+# average count of non-zero elements in the columns
+func sparseMatrixAvgColumnDensity*( A : SparseMatrix[Fr[BN254_Snarks]] ): float64 =
+  let colCounts = sparseMatrixColumnCounts( A )
+  var s: float64 = 0
+  for x in colCounts: s += x.float64
+  return (s / colCounts.len.float64)
+
+#-------------------------------------------------------------------------------
+
+# image of subspace (of the witness space) 
+func sparseMatrixImage*( A : SparseMatrix[Fr[BN254_Snarks]] , subspace: seq[bool] ): seq[bool] =
+  assert( A.dims.ncols == subspace.len )
+  var image: seq[bool] = newSeq[bool]( A.dims.nrows )
+  for j,column in A.columns.pairs:
+    if subspace[j]:
+      for i,value in column.pairs:
+        if not isZeroFr(value):
+          image[i] = true
+  return image
+
+#-------------------------------------------------------------------------------
diff --git a/groth16/misc.nim b/groth16/misc.nim
index 003f5e2..2426925 100644
--- a/groth16/misc.nim
+++ b/groth16/misc.nim
@@ -46,17 +46,6 @@ func ceilingLog2* (x : int) : int =
   else:
     return (floorLog2(x-1) + 1)
 
-#-------------------
-
-#[
-import std/math
-
-proc sanityCheckLog2* () =
-  for i in 0..18:
-    let x = float64(i)
-    echo( i," | ",floorLog2(i),"=",floor(log2(x))," | ",ceilingLog2(i),"=",ceil(log2(x)) )
-]#
-
 #-------------------------------------------------------------------------------
 
 #[
diff --git a/groth16/zkey_types.nim b/groth16/zkey_types.nim
index 1ce1d0b..41aa95e 100644
--- a/groth16/zkey_types.nim
+++ b/groth16/zkey_types.nim
@@ -4,6 +4,7 @@ import constantine/named/properties_fields
 import constantine/math/extension_fields/towers
 
 import groth16/bn128
+import groth16/math/matrix
 
 #-------------------------------------------------------------------------------
 
@@ -68,6 +69,7 @@ type
 
 #-------------------------------------------------------------------------------
 
+# extract the verification key
 func extractVKey*(zkey: Zkey): VKey = 
   let curve = zkey.header.curve
   let spec  = zkey.specPoints
@@ -75,19 +77,53 @@ func extractVKey*(zkey: Zkey): VKey =
   return VKey(curve:curve, spec:spec, vpoints:vpts)
 
 #-------------------------------------------------------------------------------
+# extract the three matrices from the ZKey
 
-proc printGrothHeader*(hdr: GrothHeader) = 
-  echo("curve        = " & ($hdr.curve        ) ) 
-  echo("flavour      = " & ($hdr.flavour      ) ) 
-  echo("|Fp|         = " & (toDecimalBig(hdr.p)) ) 
-  echo("|Fr|         = " & (toDecimalBig(hdr.r)) ) 
-  echo("nvars        = " & ($hdr.nvars        ) ) 
-  echo("npubs        = " & ($hdr.npubs        ) ) 
-  echo("domainSize   = " & ($hdr.domainSize   ) ) 
-  echo("logDomainSize= " & ($hdr.logDomainSize) ) 
+func coeffsToSparseMatrices*( dims: MatrixDims, coeffs: seq[Coeff]): SparseMatrices = 
+  var A: SparseMatrixColumns[Fr[BN254_Snarks]] = newSeq[SparseColumn[Fr[BN254_Snarks]]] ( dims.ncols )
+  var B: SparseMatrixColumns[Fr[BN254_Snarks]] = newSeq[SparseColumn[Fr[BN254_Snarks]]] ( dims.ncols )
+  var C: SparseMatrixColumns[Fr[BN254_Snarks]] = newSeq[SparseColumn[Fr[BN254_Snarks]]] ( dims.ncols )
+ 
+  for cf in coeffs:
+    case cf.matrix
+      of MatrixA: columnInsertWithAddFr( A[cf.col] , cf.row , cf.coeff )
+      of MatrixB: columnInsertWithAddFr( B[cf.col] , cf.row , cf.coeff ) 
+      of MatrixC: columnInsertWithAddFr( C[cf.col] , cf.row , cf.coeff )
+ 
+  let mA = SparseMatrix[Fr[BN254_Snarks]]( dims: dims , columns: A )
+  let mB = SparseMatrix[Fr[BN254_Snarks]]( dims: dims , columns: B )
+  let mC = SparseMatrix[Fr[BN254_Snarks]]( dims: dims , columns: C )
+  return SparseMatrices( A: mA, B: mB, C: mC )
+
+func zkeyToSparseMatrices*(zkey: ZKey): SparseMatrices = 
+  let dims = MatrixDims( nrows: zkey.header.domainSize , ncols: zkey.header.nvars )
+  return coeffsToSparseMatrices( dims , zkey.coeffs )
 
 #-------------------------------------------------------------------------------
 
+proc printGrothHeader*(hdr: GrothHeader) = 
+  echo("")
+  echo("curve         = " & ($hdr.curve        ) ) 
+  echo("flavour       = " & ($hdr.flavour      ) ) 
+  echo("|Fp|          = " & (toDecimalBig(hdr.p)) ) 
+  echo("|Fr|          = " & (toDecimalBig(hdr.r)) ) 
+  echo("nvars         = " & ($hdr.nvars        ) ) 
+  echo("npubs         = " & ($hdr.npubs        ) ) 
+  echo("domainSize    = " & ($hdr.domainSize   ) ) 
+  echo("logDomainSize = " & ($hdr.logDomainSize) ) 
+
+#-------------------------------------------------------------------------------
+
+proc printR1csStats*(zkey: ZKey) = 
+  let matrices = zkeyToSparseMatrices(zkey)
+  echo "average row density of A = " & $(sparseMatrixAvgRowDensity(matrices.A))
+  echo "average row density of B = " & $(sparseMatrixAvgRowDensity(matrices.B))
+  # echo "average row density of C = " & $(sparseMatrixAvgRowDensity(matrices.C))
+  # the matrix C is not included in ZKey
+
+#-------------------------------------------------------------------------------
+# debugging
+
 func matrixSelToString(sel: MatrixSel): string = 
   case sel 
     of MatrixA: return "A"
@@ -105,3 +141,4 @@ proc debugPrintCoeffs*(cfs: seq[Coeff]) =
   for cf in cfs: debugPrintCoeff(cf)
 
 #-------------------------------------------------------------------------------
+
diff --git a/tests/groth16/testMisc.nim b/tests/groth16/testMisc.nim
new file mode 100644
index 0000000..206838f
--- /dev/null
+++ b/tests/groth16/testMisc.nim
@@ -0,0 +1,35 @@
+
+{.used.} 
+
+import std/unittest
+import std/math
+
+import groth16/misc
+
+#-------------------------------------------------------------------------------
+
+suite "misc":
+
+  test "floorLog2":
+
+    var ok: bool = true
+    for i in 1..100:
+      let x = float64(i)
+      let lhs = floorLog2(i) 
+      let rhs = floor(log2(x)).int
+      # echo $lhs & " = " & $rhs 
+      ok = ok and ( lhs == rhs )
+    check ok
+
+  test "ceilingLog2":
+
+    var ok: bool = true
+    for i in 1..100:
+      let x = float64(i)
+      let lhs = ceilingLog2(i) 
+      let rhs = ceil(log2(x)).int
+      # echo $lhs & " = " & $rhs 
+      ok = ok and ( lhs == rhs )
+    check ok
+
+#-------------------------------------------------------------------------------
diff --git a/tests/test.nim b/tests/test.nim
index f46e1a6..085bf14 100644
--- a/tests/test.nim
+++ b/tests/test.nim
@@ -1,4 +1,5 @@
 
+import ./groth16/testMisc
 import ./groth16/testField
 import ./groth16/testCurve
 import ./groth16/testPoly