feat: addition of Trace constructor of plonk on koalabear

backend/plonk/koalabear/setup.go

@@ -0,0 +1,214 @@
+package plonk
+
+import (
+	"github.com/consensys/gnark-crypto/field/koalabear"
+	"github.com/consensys/gnark-crypto/field/koalabear/fft"
+	"github.com/consensys/gnark-crypto/field/koalabear/iop"
+	"github.com/consensys/gnark/constraint"
+	cs "github.com/consensys/gnark/constraint/koalabear"
+)
+
+// Trace stores a plonk trace as columns
+type Trace struct {
+	// Constants describing a plonk circuit. The first entries
+	// of LQk (whose index correspond to the public inputs) are set to 0, and are to be
+	// completed by the prover. At those indices i (so from 0 to nb_public_variables), LQl[i]=-1
+	// so the first nb_public_variables constraints look like this:
+	// -1*Wire[i] + 0* + 0 . It is zero when the constant coefficient is replaced by Wire[i].
+	Ql, Qr, Qm, Qo, Qk *iop.Polynomial
+	Qcp                []*iop.Polynomial
+
+	// Polynomials representing the splitted permutation. The full permutation's support is 3*N where N=nb wires.
+	// The set of interpolation is <g> of size N, so to represent the permutation S we let S acts on the
+	// set A=(<g>, u*<g>, u^{2}*<g>) of size 3*N, where u is outside <g> (its use is to shift the set <g>).
+	// We obtain a permutation of A, A'. We split A' in 3 (A'_{1}, A'_{2}, A'_{3}), and S1, S2, S3 are
+	// respectively the interpolation of A'_{1}, A'_{2}, A'_{3} on <g>.
+	S1, S2, S3 *iop.Polynomial
+
+	// S full permutation, i -> S[i]
+	S []int64
+}
+
+// NewTrace returns a new Trace object from the constraint system.
+// It fills the constant columns ql, qr, qm, qo, qk, and qcp with the
+// coefficients of the constraints.
+// Size is the size of the system that is next power of 2 (nb_constraints+nb_public_variables)
+// The permutation is also computed and stored in the Trace.
+func NewTrace(spr *cs.SparseR1CS, domain *fft.Domain) *Trace {
+	var trace Trace
+
+	size := int(domain.Cardinality)
+	commitmentInfo := spr.CommitmentInfo.(constraint.PlonkCommitments)
+
+	ql := make([]koalabear.Element, size)
+	qr := make([]koalabear.Element, size)
+	qm := make([]koalabear.Element, size)
+	qo := make([]koalabear.Element, size)
+	qk := make([]koalabear.Element, size)
+	qcp := make([][]koalabear.Element, len(commitmentInfo))
+
+	for i := 0; i < len(spr.Public); i++ { // placeholders (-PUB_INPUT_i + qk_i = 0) TODO should return error if size is inconsistent
+		ql[i].SetOne().Neg(&ql[i])
+		qr[i].SetZero()
+		qm[i].SetZero()
+		qo[i].SetZero()
+		qk[i].SetZero() // → to be completed by the prover
+	}
+	offset := len(spr.Public)
+
+	j := 0
+	it := spr.GetSparseR1CIterator()
+	for c := it.Next(); c != nil; c = it.Next() {
+		ql[offset+j].Set(&spr.Coefficients[c.QL])
+		qr[offset+j].Set(&spr.Coefficients[c.QR])
+		qm[offset+j].Set(&spr.Coefficients[c.QM])
+		qo[offset+j].Set(&spr.Coefficients[c.QO])
+		qk[offset+j].Set(&spr.Coefficients[c.QC])
+		j++
+	}
+
+	lagReg := iop.Form{Basis: iop.Lagrange, Layout: iop.Regular}
+
+	trace.Ql = iop.NewPolynomial(&ql, lagReg)
+	trace.Qr = iop.NewPolynomial(&qr, lagReg)
+	trace.Qm = iop.NewPolynomial(&qm, lagReg)
+	trace.Qo = iop.NewPolynomial(&qo, lagReg)
+	trace.Qk = iop.NewPolynomial(&qk, lagReg)
+	trace.Qcp = make([]*iop.Polynomial, len(qcp))
+
+	for i := range commitmentInfo {
+		qcp[i] = make([]koalabear.Element, size)
+		for _, committed := range commitmentInfo[i].Committed {
+			qcp[i][offset+committed].SetOne()
+		}
+		trace.Qcp[i] = iop.NewPolynomial(&qcp[i], lagReg)
+	}
+
+	// build the permutation and build the polynomials S1, S2, S3 to encode the permutation.
+	// Note: at this stage, the permutation takes in account the placeholders
+	nbVariables := spr.NbInternalVariables + len(spr.Public) + len(spr.Secret)
+	buildPermutation(spr, &trace, nbVariables)
+	s := computePermutationPolynomials(&trace, domain)
+	trace.S1 = s[0]
+	trace.S2 = s[1]
+	trace.S3 = s[2]
+
+	return &trace
+}
+
+// buildPermutation builds the Permutation associated with a circuit.
+//
+// The permutation s is composed of cycles of maximum length such that
+//
+//	s. (l∥r∥o) = (l∥r∥o)
+//
+// , where l∥r∥o is the concatenation of the indices of l, r, o in
+// ql.l+qr.r+qm.l.r+qo.O+k = 0.
+//
+// The permutation is encoded as a slice s of size 3*size(l), where the
+// i-th entry of l∥r∥o is sent to the s[i]-th entry, so it acts on a tab
+// like this: for i in tab: tab[i] = tab[permutation[i]]
+func buildPermutation(spr *cs.SparseR1CS, trace *Trace, nbVariables int) {
+
+	// nbVariables := spr.NbInternalVariables + len(spr.Public) + len(spr.Secret)
+	sizeSolution := len(trace.Ql.Coefficients())
+	sizePermutation := 3 * sizeSolution
+
+	// init permutation
+	permutation := make([]int64, sizePermutation)
+	for i := 0; i < len(permutation); i++ {
+		permutation[i] = -1
+	}
+
+	// init LRO position -> variable_ID
+	lro := make([]int, sizePermutation) // position -> variable_ID
+	for i := 0; i < len(spr.Public); i++ {
+		lro[i] = i // IDs of LRO associated to placeholders (only L needs to be taken care of)
+	}
+
+	offset := len(spr.Public)
+
+	j := 0
+	it := spr.GetSparseR1CIterator()
+	for c := it.Next(); c != nil; c = it.Next() {
+		lro[offset+j] = int(c.XA)
+		lro[sizeSolution+offset+j] = int(c.XB)
+		lro[2*sizeSolution+offset+j] = int(c.XC)
+
+		j++
+	}
+
+	// init cycle:
+	// map ID -> last position the ID was seen
+	cycle := make([]int64, nbVariables)
+	for i := 0; i < len(cycle); i++ {
+		cycle[i] = -1
+	}
+
+	for i := 0; i < len(lro); i++ {
+		if cycle[lro[i]] != -1 {
+			// if != -1, it means we already encountered this value
+			// so we need to set the corresponding permutation index.
+			permutation[i] = cycle[lro[i]]
+		}
+		cycle[lro[i]] = int64(i)
+	}
+
+	// complete the Permutation by filling the first IDs encountered
+	for i := 0; i < sizePermutation; i++ {
+		if permutation[i] == -1 {
+			permutation[i] = cycle[lro[i]]
+		}
+	}
+
+	trace.S = permutation
+}
+
+// computePermutationPolynomials computes the LDE (Lagrange basis) of the permutation.
+// We let the permutation act on <g> || u<g> || u^{2}<g>, split the result in 3 parts,
+// and interpolate each of the 3 parts on <g>.
+func computePermutationPolynomials(trace *Trace, domain *fft.Domain) [3]*iop.Polynomial {
+
+	nbElmts := int(domain.Cardinality)
+
+	var res [3]*iop.Polynomial
+
+	// Lagrange form of ID
+	evaluationIDSmallDomain := getSupportPermutation(domain)
+
+	// Lagrange form of S1, S2, S3
+	s1Canonical := make([]koalabear.Element, nbElmts)
+	s2Canonical := make([]koalabear.Element, nbElmts)
+	s3Canonical := make([]koalabear.Element, nbElmts)
+	for i := 0; i < nbElmts; i++ {
+		s1Canonical[i].Set(&evaluationIDSmallDomain[trace.S[i]])
+		s2Canonical[i].Set(&evaluationIDSmallDomain[trace.S[nbElmts+i]])
+		s3Canonical[i].Set(&evaluationIDSmallDomain[trace.S[2*nbElmts+i]])
+	}
+
+	lagReg := iop.Form{Basis: iop.Lagrange, Layout: iop.Regular}
+	res[0] = iop.NewPolynomial(&s1Canonical, lagReg)
+	res[1] = iop.NewPolynomial(&s2Canonical, lagReg)
+	res[2] = iop.NewPolynomial(&s3Canonical, lagReg)
+
+	return res
+}
+
+// getSupportPermutation returns the support on which the permutation acts, it is
+// <g> || u<g> || u^{2}<g>
+func getSupportPermutation(domain *fft.Domain) []koalabear.Element {
+
+	res := make([]koalabear.Element, 3*domain.Cardinality)
+
+	res[0].SetOne()
+	res[domain.Cardinality].Set(&domain.FrMultiplicativeGen)
+	res[2*domain.Cardinality].Square(&domain.FrMultiplicativeGen)
+
+	for i := uint64(1); i < domain.Cardinality; i++ {
+		res[i].Mul(&res[i-1], &domain.Generator)
+		res[domain.Cardinality+i].Mul(&res[domain.Cardinality+i-1], &domain.Generator)
+		res[2*domain.Cardinality+i].Mul(&res[2*domain.Cardinality+i-1], &domain.Generator)
+	}
+
+	return res
+}

go.mod

@@ -9,7 +9,7 @@ require (
 	github.com/blang/semver/v4 v4.0.0
 	github.com/consensys/bavard v0.1.31-0.20250406004941-2db259e4b582
 	github.com/consensys/compress v0.2.5
-	github.com/consensys/gnark-crypto v0.18.0
+	github.com/consensys/gnark-crypto v0.18.1-0.20250613145137-bf7ac9d06da2
 	github.com/fxamacker/cbor/v2 v2.8.0
 	github.com/google/go-cmp v0.7.0
 	github.com/google/pprof v0.0.0-20250607225305-033d6d78b36a

go.sum

@@ -63,6 +63,8 @@ github.com/consensys/compress v0.2.5 h1:gJr1hKzbOD36JFsF1AN8lfXz1yevnJi1YolffY19
 github.com/consensys/compress v0.2.5/go.mod h1:pyM+ZXiNUh7/0+AUjUf9RKUM6vSH7T/fsn5LLS0j1Tk=
 github.com/consensys/gnark-crypto v0.18.0 h1:vIye/FqI50VeAr0B3dx+YjeIvmc3LWz4yEfbWBpTUf0=
 github.com/consensys/gnark-crypto v0.18.0/go.mod h1:L3mXGFTe1ZN+RSJ+CLjUt9x7PNdx8ubaYfDROyp2Z8c=
+github.com/consensys/gnark-crypto v0.18.1-0.20250613145137-bf7ac9d06da2 h1:LZlarR2EgBB7NDuCf6A9OS2UUa1RAjV2l7T9C0Wvn1I=
+github.com/consensys/gnark-crypto v0.18.1-0.20250613145137-bf7ac9d06da2/go.mod h1:L3mXGFTe1ZN+RSJ+CLjUt9x7PNdx8ubaYfDROyp2Z8c=
 github.com/coreos/go-semver v0.3.0/go.mod h1:nnelYz7RCh+5ahJtPPxZlU+153eP4D4r3EedlOD2RNk=
 github.com/coreos/go-systemd/v22 v22.3.2/go.mod h1:Y58oyj3AT4RCenI/lSvhwexgC+NSVTIJ3seZv2GcEnc=
 github.com/coreos/go-systemd/v22 v22.5.0/go.mod h1:Y58oyj3AT4RCenI/lSvhwexgC+NSVTIJ3seZv2GcEnc=