From be34a776973b6ee0b592cb02cfd1f7deb34697a1 Mon Sep 17 00:00:00 2001 From: Thomas Piellard Date: Thu, 28 May 2026 15:31:22 +0200 Subject: [PATCH] feat: evaluate at zeta wo FFTs --- internal/poly/ext.go | 24 +++++ internal/poly/ext_test.go | 36 +++++++ internal/poly/polynomial_test.go | 17 ++++ internal/poly/utils.go | 42 +++++++++ prover/prover.go | 157 +++++++++++++++++++------------ 5 files changed, 216 insertions(+), 60 deletions(-) diff --git a/internal/poly/ext.go b/internal/poly/ext.go index 2ac5887..0ca9af2 100644 --- a/internal/poly/ext.go +++ b/internal/poly/ext.go @@ -78,6 +78,30 @@ func MulExt(P1, P2 ExtPolynomial) (ExtPolynomial, error) { return res, nil } +// EvaluateLagrangeAtExt lagrangeZeta = [L_0(zeta), L_1(zeta), ..., L_n-1(zeta)] (n = power of two), +// p polynomial in lagrange form, of size n +func EvaluateLagrangeAtExt(p Polynomial, lagrangeZetaCache []ext.E4, shift int) ext.E4 { + n := len(lagrangeZetaCache) + var tmp, res ext.E4 + for i := 0; i < len(p); i++ { + tmp.Set(&lagrangeZetaCache[(n+i-shift)%n]).MulByElement(&tmp, &p[i]) + res.Add(&res, &tmp) + } + return res +} + +// ExtEvaluateLagrangeAtExt lagrangeZeta = [L_0(zeta), L_1(zeta), ..., L_n-1(zeta)] (n = power of two), +// p polynomial in lagrange form, of size n +func ExtEvaluateLagrangeAtExt(p ExtPolynomial, lagrangeZetaCache []ext.E4, shift int) ext.E4 { + n := len(lagrangeZetaCache) + var tmp, res ext.E4 + for i := 0; i < len(p); i++ { + tmp.Set(&lagrangeZetaCache[(n+i-shift)%n]).Mul(&tmp, &p[i]) + res.Add(&res, &tmp) + } + return res +} + // EvaluateAtExt evaluates a base-field polynomial p, stored in Lagrange normal // form over d, at the extension-field point zeta. Base coefficients are lifted // during Horner evaluation. diff --git a/internal/poly/ext_test.go b/internal/poly/ext_test.go index 826fbd7..0438815 100644 --- a/internal/poly/ext_test.go +++ b/internal/poly/ext_test.go @@ -14,6 +14,7 @@ package poly import ( + "fmt" "math/big" "testing" @@ -133,6 +134,41 @@ func TestEvaluateAtExt(t *testing.T) { } } +func TestEvaluateLagrangeAtExt(t *testing.T) { + const domainSize = 32 + zeta := e4FromU64(3, 5, 7, 11) + lagrangeZetaCache := LagrangesAtZeta(zeta, domainSize) + d := fft.NewDomain(domainSize) + w := d.Generator + + t.Run(fmt.Sprintf("size_%d_no_shift", domainSize), func(t *testing.T) { + p := make(Polynomial, domainSize) + for i := 0; i < domainSize; i++ { + p[i].SetUint64(uint64(3*i + 5)) + } + + got := EvaluateLagrangeAtExt(p, lagrangeZetaCache, 0) + want := EvaluateAtExt(p, d, zeta) + if !got.Equal(&want) { + t.Fatalf("EvaluateLagrangeAtExt = %s, want %s", got.String(), want.String()) + } + }) + + t.Run(fmt.Sprintf("size_%d_with_shift", domainSize), func(t *testing.T) { + p := make(Polynomial, domainSize) + for i := 0; i < domainSize; i++ { + p[i].SetUint64(uint64(3*i + 5)) + } + zetaShifted := zeta.MulByElement(&zeta, &w) + got := EvaluateLagrangeAtExt(p, lagrangeZetaCache, 1) + want := EvaluateAtExt(p, d, *zetaShifted) + if !got.Equal(&want) { + t.Fatalf("EvaluateLagrangeAtExt = %s, want %s", got.String(), want.String()) + } + }) + +} + func TestExtEvaluateAtExt(t *testing.T) { coeffs := ExtPolynomial{ e4FromU64(1, 2, 3, 4), diff --git a/internal/poly/polynomial_test.go b/internal/poly/polynomial_test.go index 58d38ec..b3adbda 100644 --- a/internal/poly/polynomial_test.go +++ b/internal/poly/polynomial_test.go @@ -98,6 +98,23 @@ func TestLagrangeAtZeta(t *testing.T) { } +func TestLagrangesAtZetaExt(t *testing.T) { + n := 16 + zeta := e4FromU64(3, 5, 7, 11) + + computed := LagrangesAtZeta(zeta, n) + if len(computed) != n { + t.Fatalf("expected %d lagrange evaluations, got %d", n, len(computed)) + } + + for i := 0; i < n; i++ { + expected := LagrangeAtZetaExt(zeta, n, i) + if !expected.Equal(&computed[i]) { + t.Errorf("L_%d(zeta): expected %s, got %s", i, expected.String(), computed[i].String()) + } + } +} + func TestEvalPointWise(t *testing.T) { t.Run("NonHomogeneous_x0sq_plus_x1", func(t *testing.T) { diff --git a/internal/poly/utils.go b/internal/poly/utils.go index 43c12eb..599db7b 100644 --- a/internal/poly/utils.go +++ b/internal/poly/utils.go @@ -61,6 +61,48 @@ func NextPowerOfTwo(n int) int { return n + 1 } +// LagrangesAtZeta computes [L_0(zeta), L_1(zeta), ..., L_n-1(zeta)] (n = power of two) +func LagrangesAtZeta(zeta ext.E4, n int) []ext.E4 { + res := make([]ext.E4, n) + if n == 0 { + return res + } + + omega, err := koalabear.Generator(uint64(n)) + if err != nil { + panic(err) + } + + denominators := make([]ext.E4, n) + var omegaI koalabear.Element + omegaI.SetOne() + for i := 0; i < n; i++ { + var omegaIExt ext.E4 + omegaIExt.Lift(&omegaI) + denominators[i].Sub(&zeta, &omegaIExt) + omegaI.Mul(&omegaI, &omega) + } + invDenominators := ext.BatchInvertE4(denominators) + + var zetaN, one, common ext.E4 + zetaN.Exp(zeta, big.NewInt(int64(n))) + one.SetOne() + common.Sub(&zetaN, &one) + + var nElt, nInv koalabear.Element + nElt.SetUint64(uint64(n)) + nInv.Inverse(&nElt) + common.MulByElement(&common, &nInv) + + omegaI.SetOne() + for i := 0; i < n; i++ { + res[i].MulByElement(&common, &omegaI) + res[i].Mul(&res[i], &invDenominators[i]) + omegaI.Mul(&omegaI, &omega) + } + return res +} + func LagrangeAtZeta(zeta koalabear.Element, N, i int) koalabear.Element { var zetan, omegai, one, Nk koalabear.Element Nk.SetUint64(uint64(N)) diff --git a/prover/prover.go b/prover/prover.go index 155590a..5d5c134 100644 --- a/prover/prover.go +++ b/prover/prover.go @@ -615,81 +615,118 @@ func (pr *proverRuntime) ComputeEvaluationsAtZeta() error { expr.WithoutPublicColumns(), ) - type zetaEval struct { - key string - val ext.E4 - } - moduleNames := make([]string, 0, len(pr.program.Modules)) for name := range pr.program.Modules { moduleNames = append(moduleNames, name) } sort.Strings(moduleNames) - results := make([][]zetaEval, len(moduleNames)) - errs := make([]error, len(moduleNames)) - parallel.Execute(len(moduleNames), func(start, end int) { - for idx := start; idx < end; idx++ { + lagrangesByN := make(map[int][]ext.E4) + for _, moduleName := range moduleNames { + N := pr.program.Modules[moduleName].N + if _, ok := lagrangesByN[N]; !ok { + lagrangesByN[N] = poly.LagrangesAtZeta(pr.zeta, N) + } + } - // trace polynomials - module := pr.program.Modules[moduleNames[idx]] - leaves := module.VanishingRelation.LeavesFull(config) - local := make([]zetaEval, 0, len(leaves)) - for _, leaf := range leaves { - evalPoint := pr.zeta - if leaf.Type == expr.RotatedColumn { - shift := ((leaf.Shift % module.N) + module.N) % module.N - var omegaPow koalabear.Element - omegaPow.SetOne() - for k := 0; k < shift; k++ { - omegaPow.Mul(&omegaPow, &module.D.Generator) - } - evalPoint.MulByElement(&evalPoint, &omegaPow) - } + type evalTask struct { + key string + field field.Kind + shift int + lagranges []ext.E4 + base poly.Polynomial + ext poly.ExtPolynomial + } - if p, ok := pr.t.Ext[leaf.Name]; ok { - val := poly.ExtEvaluateAtExt(p, module.D, evalPoint) - local = append(local, zetaEval{key: leaf.String(), val: val}) - continue - } - p, ok := pr.t.Base[leaf.Name] - if !ok { - errs[idx] = fmt.Errorf("ComputeEvaluationsAtZeta: column %q not found in trace", leaf.Name) - return + tasks := make([]evalTask, 0) + addBaseTask := func(key string, p poly.Polynomial, moduleN, shift int) error { + if len(p) != 1 && len(p) != moduleN { + return fmt.Errorf("ComputeEvaluationsAtZeta: base polynomial %q has size %d, module has size %d", key, len(p), moduleN) + } + tasks = append(tasks, evalTask{ + key: key, + field: field.Base, + shift: shift, + lagranges: lagrangesByN[moduleN], + base: p, + }) + return nil + } + addExtTask := func(key string, p poly.ExtPolynomial, moduleN, shift int) error { + if len(p) != 1 && len(p) != moduleN { + return fmt.Errorf("ComputeEvaluationsAtZeta: extension polynomial %q has size %d, module has size %d", key, len(p), moduleN) + } + tasks = append(tasks, evalTask{ + key: key, + field: field.Ext, + shift: shift, + lagranges: lagrangesByN[moduleN], + ext: p, + }) + return nil + } + + for _, moduleName := range moduleNames { + module := pr.program.Modules[moduleName] + for _, leaf := range module.VanishingRelation.LeavesFull(config) { + shift := 0 + if leaf.Type == expr.RotatedColumn { + shift = ((leaf.Shift % module.N) + module.N) % module.N + } + if p, ok := pr.t.Ext[leaf.Name]; ok { + if err := addExtTask(leaf.String(), p, module.N, shift); err != nil { + return err } - val := poly.EvaluateAtExt(p, module.D, evalPoint) - local = append(local, zetaEval{key: leaf.String(), val: val}) - } - - // air chunks - evalPoint := pr.zeta - for i := 0; ; i++ { - chunkName := constants.QuotientChunkName(moduleNames[idx], i) - if chunk, ok := pr.airTrace.Base[chunkName]; ok { - val := poly.EvaluateAtExt(chunk, module.D, evalPoint) - local = append(local, zetaEval{key: chunkName, val: val}) - continue + continue + } + p, ok := pr.t.Base[leaf.Name] + if !ok { + return fmt.Errorf("ComputeEvaluationsAtZeta: column %q not found in trace", leaf.Name) + } + if err := addBaseTask(leaf.String(), p, module.N, shift); err != nil { + return err + } + } + + for i := 0; ; i++ { + chunkName := constants.QuotientChunkName(moduleName, i) + if chunk, ok := pr.airTrace.Base[chunkName]; ok { + if err := addBaseTask(chunkName, chunk, module.N, 0); err != nil { + return err } - if chunk, ok := pr.airTrace.Ext[chunkName]; ok { - val := poly.ExtEvaluateAtExt(chunk, module.D, evalPoint) - local = append(local, zetaEval{key: chunkName, val: val}) - continue + continue + } + if chunk, ok := pr.airTrace.Ext[chunkName]; ok { + if err := addExtTask(chunkName, chunk, module.N, 0); err != nil { + return err } - break + continue } - - results[idx] = local - } - }) - for _, err := range errs { - if err != nil { - return err + break } } - for _, local := range results { - for _, ev := range local { - pr.Proof.SetValueAtZetaExt(ev.key, ev.val) + + values := make([]ext.E4, len(tasks)) + parallel.Execute(len(tasks), func(start, end int) { + for i := start; i < end; i++ { + task := tasks[i] + if task.field == field.Ext { + if len(task.ext) == 1 { + values[i].Set(&task.ext[0]) + } else { + values[i] = poly.ExtEvaluateLagrangeAtExt(task.ext, task.lagranges, task.shift) + } + continue + } + if len(task.base) == 1 { + values[i] = liftBaseToExt(task.base[0]) + } else { + values[i] = poly.EvaluateLagrangeAtExt(task.base, task.lagranges, task.shift) + } } + }) + for i, task := range tasks { + pr.Proof.SetValueAtZetaExt(task.key, values[i]) } return nil