From 2231d800c747e9f726efe4626ec07411ba4901ab Mon Sep 17 00:00:00 2001
From: Balazs Komuves <bkomuves@gmail.com>
Date: Sat, 13 Jun 2026 13:29:07 +0200
Subject: [PATCH] add group FFT

---
 groth16.nimble                |   2 +-
 groth16/bn128/curves.nim      |   7 +-
 groth16/dynamic/group_fft.nim | 188 ++++++++++++++++++++++++++++++++++
 tests/groth16/testFFT.nim     |  76 +++++++++-----
 4 files changed, 247 insertions(+), 26 deletions(-)
 create mode 100644 groth16/dynamic/group_fft.nim

diff --git a/groth16.nimble b/groth16.nimble
index 29b38d3..ec8762a 100644
--- a/groth16.nimble
+++ b/groth16.nimble
@@ -1,5 +1,5 @@
 
-version     = "0.1.1"
+version     = "0.1.2"
 author      = "Balazs Komuves"
 description = "Groth16 proof system"
 license     = "MIT OR Apache-2.0"
diff --git a/groth16/bn128/curves.nim b/groth16/bn128/curves.nim
index 76edcac..b4e0d1d 100644
--- a/groth16/bn128/curves.nim
+++ b/groth16/bn128/curves.nim
@@ -34,11 +34,14 @@ import groth16/bn128/fields
 
 #-------------------------------------------------------------------------------
 
-type G1*   = aff.EC_ShortW_Aff[Fp[BN254_Snarks] , aff.G1]
-type G2*   = aff.EC_ShortW_Aff[Fp2[BN254_Snarks], aff.G2]
+type AffG1*   = aff.EC_ShortW_Aff[Fp[BN254_Snarks] , aff.G1]
+type AffG2*   = aff.EC_ShortW_Aff[Fp2[BN254_Snarks], aff.G2]
 
 type ProjG1*  = prj.EC_ShortW_Prj[Fp[BN254_Snarks] , prj.G1]
 type ProjG2*  = prj.EC_ShortW_Prj[Fp2[BN254_Snarks], prj.G2]
+ 
+type G1* = AffG1
+type G2* = AffG2
 
 #-------------------------------------------------------------------------------
 # compressed points (supposedly compatible with arkworks-0.5)
diff --git a/groth16/dynamic/group_fft.nim b/groth16/dynamic/group_fft.nim
new file mode 100644
index 0000000..27a5b06
--- /dev/null
+++ b/groth16/dynamic/group_fft.nim
@@ -0,0 +1,188 @@
+
+#
+# FFT for groups elements
+#
+
+#-------------------------------------------------------------------------------
+
+import constantine/math/arithmetic
+import constantine/math/io/io_fields
+import constantine/named/properties_fields
+
+# import constantine/math/elliptic/ec_shortweierstrass_affine     as aff 
+import constantine/math/elliptic/ec_shortweierstrass_projective as prj 
+import constantine/math/elliptic/ec_scalar_mul_vartime          as scl 
+
+import groth16/bn128
+import groth16/bn128/curves
+import groth16/math/domain
+
+#-------------------------------------------------------------------------------
+
+func forward_FFT_worker( m: int
+                       , srcStride: int
+                       , gpows: seq[Fr[BN254_Snarks]]
+                       , src:     seq[ProjG1] , srcOfs: int
+                       , buf: var seq[ProjG1] , bufOfs: int
+                       , tgt: var seq[ProjG1] , tgtOfs: int ) =
+  case m 
+
+    of 0: 
+      tgt[tgtOfs] = src[srcOfs]
+  
+    of 1:
+      tgt[tgtOfs  ] = src[srcOfs] ; tgt[tgtOfs  ] += src[srcOfs+srcStride]
+      tgt[tgtOfs+1] = src[srcOfs] ; tgt[tgtOfs+1] -= src[srcOfs+srcStride]
+
+    else:
+      let N     : int =  1 shl  m  
+      let halfN : int =  1 shl (m-1)  
+      forward_FFT_worker( m-1
+                        , srcStride shl 1
+                        , gpows
+                        , src , srcOfs
+                        , buf , bufOfs + N
+                        , buf , bufOfs )
+      forward_FFT_worker( m-1
+                        , srcStride shl 1
+                        , gpows
+                        , src , srcOfs + srcStride
+                        , buf , bufOfs + N
+                        , buf , bufOfs + halfN )
+      for j in 0..<halfN:
+        var y = buf[bufOfs+j+halFN]
+        y.scalarMul_vartime( gpows[j*srcStride] ) 
+        tgt[tgtOfs+j      ] = buf[bufOfs+j] ; tgt[tgtOfs+j      ] += y
+        tgt[tgtOfs+j+halfN] = buf[bufOfs+j] ; tgt[tgtOfs+j+halfN] -= y
+
+#---------------------------------------
+
+# forward number-theoretical transform (corresponds to polynomial evaluation)
+func forwardGroupFFT*(affSrc: seq[AffG1], D: Domain): seq[AffG1] =
+  let N = D.domainSize
+  assert( N == (1 shl D.logDomainSize) , "domain must have a power-of-two size"        )
+  assert( N == affSrc.len              , "input must have the same size as the domain" )
+
+  var src = newSeq[ProjG1](     N )
+  var buf = newSeq[ProjG1]( 2 * N )
+  var tgt = newSeq[ProjG1](     N )
+
+  for i in 0..<N:
+    src[i].fromAffine(affSrc[i])
+
+  # precalc powers of gen
+  let halFN = N div 2
+  var gpows = newSeq[Fr[BN254_Snarks]]( halFN )
+  var x     = oneFr
+  let gen   = D.domainGen
+  for i in 0..<halfN:
+    gpows[i] = x
+    x *= gen
+
+  forward_FFT_worker( D.logDomainSize
+                    , 1
+                    , gpows
+                    , src , 0
+                    , buf , 0
+                    , tgt , 0 )
+
+  var affTgt = newSeq[AffG1]( N )
+  for i in 0..<N:
+    affTgt[i].affine(tgt[i])
+
+  return affTgt
+
+#-------------------------------------------------------------------------------
+
+# unscaled!
+func inverse_FFT_worker( m: int
+                       , tgtStride: int
+                       , gpows: seq[Fr[BN254_Snarks]]
+                       , src:     seq[ProjG1] , srcOfs: int
+                       , buf: var seq[ProjG1] , bufOfs: int
+                       , tgt: var seq[ProjG1] , tgtOfs: int ) =
+  case m 
+
+    of 0: 
+      tgt[tgtOfs] = src[srcOfs]
+  
+    of 1:
+      tgt[tgtOfs          ] = src[srcOfs] ; tgt[tgtOfs          ] += src[srcOfs+1] 
+      tgt[tgtOfs+tgtStride] = src[srcOfs] ; tgt[tgtOfs+tgtStride] -= src[srcOfs+1] 
+
+    else:
+      let N     : int =  1 shl  m  
+      let halfN : int =  1 shl (m-1)  
+
+      for j in 0..<halfN:
+        buf[bufOfs+j      ] = src[srcOfs+j] ; buf[bufOfs+j      ] += src[srcOfs+j+halfN] 
+        buf[bufOfs+j+halfN] = src[srcOfs+j] ; buf[bufOfs+j+halfN] -= src[srcOfs+j+halfN] 
+        buf[bufOfs+j+halfN].scalarMul_vartime( gpows[ j*tgtStride ] ) 
+
+      inverse_FFT_worker( m-1
+                        , tgtStride shl 1
+                        , gpows
+                        , buf , bufOfs
+                        , buf , bufOfs + N
+                        , tgt , tgtOfs )
+      inverse_FFT_worker( m-1
+                        , tgtStride shl 1
+                        , gpows
+                        , buf , bufOfs + halfN
+                        , buf , bufOfs + N
+                        , tgt , tgtOfs + tgtStride )
+
+#---------------------------------------
+
+# inverse number-theoretical transform (with the 1/N rescaling toggable)
+func toggableInverseGroupFFT(affSrc: seq[AffG1], D: Domain, do_rescale: bool): seq[AffG1] =
+  let N = D.domainSize
+  assert( N == (1 shl D.logDomainSize) , "domain must have a power-of-two size"        )
+  assert( N == affSrc.len              , "input must have the same size as the domain" )
+
+  var src = newSeq[ProjG1](     N )
+  var buf = newSeq[ProjG1]( 2 * N )
+  var tgt = newSeq[ProjG1](     N )
+
+  for i in 0..<N:
+    src[i].fromAffine(affSrc[i])
+
+  # precalc times powers of gen^-1
+  let halFN = N div 2
+  var gpows = newSeq[Fr[BN254_Snarks]]( halFN )
+  var x     = oneFr
+  let ginv  = invFr( D.domainGen )
+  for i in 0..<halfN:
+    gpows[i] = x
+    x *= ginv
+
+  inverse_FFT_worker( D.logDomainSize
+                    , 1
+                    , gpows
+                    , src , 0
+                    , buf , 0
+                    , tgt , 0 )
+
+  if do_rescale:
+    var invN : Fr[BN254_Snarks]
+    invN.fromInt(N)
+    invN.inv() 
+    for i in 0..<N:
+      tgt[i].scalarMul_vartime( invN )
+
+  var affTgt = newSeq[AffG1]( N )
+  for i in 0..<N:
+    affTgt[i].affine(tgt[i])
+
+  return affTgt
+
+#---------------------------------------
+
+func inverseGroupFFT*(src: seq[AffG1], D: Domain): seq[AffG1] =
+  toggableInverseGroupFFT(src , D , true)
+
+# inverse number-theoretical transform (without the 1/N rescaling)
+func unscaledInverseGroupFFT*(src: seq[AffG1], D: Domain): seq[AffG1] =
+  toggableInverseGroupFFT(src , D , false)
+
+#-------------------------------------------------------------------------------
diff --git a/tests/groth16/testFFT.nim b/tests/groth16/testFFT.nim
index d57bcf6..6971459 100644
--- a/tests/groth16/testFFT.nim
+++ b/tests/groth16/testFFT.nim
@@ -15,6 +15,8 @@ import groth16/bn128/debug
 import groth16/math/domain
 import groth16/math/ntt
 
+import groth16/dynamic/group_fft
+
 #-------------------------------------------------------------------------------
 
 proc randFrSeq(N: int) : seq[Fr[BN254_Snarks]] = 
@@ -28,7 +30,7 @@ proc scaleFrSeqInPlace(s: Fr[BN254_Snarks], arr: var seq[Fr[BN254_Snarks]] ) =
   for i in 0..<N:
     arr[i] *= s
 
-func isEqualFrSeq(xs : seq[Fr[BN254_Snarks]], ys: seq[Fr[BN254_Snarks]]): bool =
+func isEqualFrSeq(xs: seq[Fr[BN254_Snarks]], ys: seq[Fr[BN254_Snarks]]): bool =
   let N = xs.len
   let M = ys.len
 
@@ -37,7 +39,7 @@ func isEqualFrSeq(xs : seq[Fr[BN254_Snarks]], ys: seq[Fr[BN254_Snarks]]): bool =
   else:
     var ok: bool = true
     for i in 0..<N:
-      ok = ok and (xs === ys)
+      ok = ok and (xs[i] === ys[i])
     return ok
 
 #---------------------------------------
@@ -48,33 +50,61 @@ proc randG1Seq(N: int) : seq[G1] =
     xs[i] = randG1()
   return xs
 
+func isEqualG1Seq(xs: seq[G1], ys: seq[G1] ): bool =
+  let N = xs.len
+  let M = ys.len
+
+  if N != M:
+    return false
+  else:
+    var ok: bool = true
+    for i in 0..<N:
+      ok = ok and (xs[i] === ys[i])
+    return ok
+
 #-------------------------------------------------------------------------------
 
-suite "FFT checks":
+suite "field NTT checks":
 
-  suite "field NTT":
+  let D = createDomain(128)
+  let N = D.domainSize
+  let xs : seq[Fr[BN254_Snarks]] = randFrSeq(N)
 
-    let D = createDomain(128)
-    let N = D.domainSize
+  test "INTT(NTT(xs) == xs":
+    let ys = forwardNTT(xs, D)
+    let zs = inverseNTT(ys, D) 
+    check isEqualFrSeq( xs , zs )
   
-    let xs : seq[Fr[BN254_Snarks]] = randFrSeq(N)
-  
-    test "INTT(NTT(xs) == xs":
-      let ys = forwardNTT(xs, D)
-      let zs = inverseNTT(ys, D) 
-      check isEqualFrSeq( xs , zs )
-    
-    test "NTT(INTT(xs) == xs":
-      let ys = inverseNTT(xs, D)
-      let zs = forwardNTT(ys, D) 
-      check isEqualFrSeq( xs , zs )
+  test "NTT(INTT(xs) == xs":
+    let ys = inverseNTT(xs, D)
+    let zs = forwardNTT(ys, D) 
+    check isEqualFrSeq( xs , zs )
+
+  test "unscaledINTT(NTT(xs) == N * xs":
+    let ys = forwardNTT(xs, D)
+    let zs = unscaledInverseNTT(ys, D) 
+    var ws = xs
+    scaleFrSeqInPlace(intToFr(N) , ws)
+    check isEqualFrSeq( ws , zs )
+
+#-------------------------------------
+
+suite "group FFT checks (for the group G1)":
+
+  let D = createDomain(64)
+  let N = D.domainSize
+  let gs : seq[G1] = randG1Seq(N)
+
+  test "IFFT(FFT(gs) == gs":
+    let hs = forwardGroupFFT(gs, D)
+    let rs = inverseGroupFFT(hs, D) 
+    check isEqualG1Seq( gs , rs )
+  
+  test "NTT(INTT(xs) == xs":
+    let hs = inverseGroupFFT(gs, D)
+    let rs = forwardGroupFFT(hs, D) 
+    check isEqualG1Seq( gs , rs )
 
-    test "unscaledINTT(NTT(xs) == N * xs":
-      let ys = forwardNTT(xs, D)
-      let zs = unscaledInverseNTT(ys, D) 
-      var ws = xs
-      scaleFrSeqInPlace(intToFr(N) , ws)
-      check isEqualFrSeq( ws , zs )
 
 #-------------------------------------------------------------------------------