feat: evaluate at zeta wo FFTs

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.

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),

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) {

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))

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