chore: update kzg module's abstraction (#89)

* update kzg module

* refactor code to match

* lint

api.go

@@ -6,6 +6,7 @@ import (
 	"github.com/crate-crypto/go-eth-kzg/internal/domain"
 	"github.com/crate-crypto/go-eth-kzg/internal/erasure_code"
 	"github.com/crate-crypto/go-eth-kzg/internal/kzg"
+	kzgmulti "github.com/crate-crypto/go-eth-kzg/internal/kzg_multi"
 	"github.com/crate-crypto/go-eth-kzg/internal/kzg_multi/fk20"
 )
 
@@ -18,7 +19,8 @@ type Context struct {
 	domainExtended    *domain.Domain
 	commitKeyLagrange *kzg.CommitKey
 	commitKeyMonomial *kzg.CommitKey
-	openKey           *kzg.OpeningKey
+	openKey4844       *kzg.OpeningKey
+	openKey7594       *kzgmulti.OpeningKey
 
 	fk20 *fk20.FK20
 
@@ -114,14 +116,14 @@ func NewContext4096(trustedSetup *JSONTrustedSetup) (*Context, error) {
 		panic("The number of G2 points in the trusted setup is less than the number of scalars per blob")
 	}
 
-	openingKey := kzg.OpeningKey{
+	openingKey4844 := kzg.OpeningKey{
 		GenG1:   genG1,
 		GenG2:   genG2,
 		AlphaG2: alphaGenG2,
-		G1:      setupMonomialG1Points[:len(setupG2Points)],
-		G2:      setupG2Points,
 	}
 
+	openingKey7594 := kzgmulti.NewOpeningKey(setupMonomialG1Points[:len(setupG2Points)], setupG2Points, ScalarsPerBlob, scalarsPerExtBlob, scalarsPerCell)
+
 	domainBlobLen := domain.NewDomain(ScalarsPerBlob)
 	// Bit-Reverse the roots and the trusted setup according to the specs
 	// The bit reversal is not needed for simple KZG however it was
@@ -139,12 +141,9 @@ func NewContext4096(trustedSetup *JSONTrustedSetup) (*Context, error) {
 		domainExtended:    domainExtended,
 		commitKeyLagrange: &commitKeyLagrange,
 		commitKeyMonomial: &commitKeyMonomial,
-		openKey:           &openingKey,
+		openKey4844:       &openingKey4844,
+		openKey7594:       openingKey7594,
 		fk20:              &fk20,
-		// TODO: We compute the extendedDomain again in here.
-		// TODO: We could pass it in, but it breaks the API.
-		// TODO: And although its not an issue now because fft uses just the primitiveGenerator, the extended domain
-		// TODO: that recovery takes is not bit reversed.
-		dataRecovery: erasure_code.NewDataRecovery(scalarsPerCell, ScalarsPerBlob, expansionFactor),
+		dataRecovery:      erasure_code.NewDataRecovery(scalarsPerCell, ScalarsPerBlob, expansionFactor),
 	}, nil
 }

api_eip7594.go

@@ -156,7 +156,7 @@ func (ctx *Context) VerifyCellKZGProof(commitment KZGCommitment, cellID uint64,
 	// partition the extended roots to form cosets
 	cosets := partition(ctx.domainExtended.Roots, scalarsPerCell)
 
-	return kzgmulti.VerifyMultiPointKZGProof(commitmentG1, proofG1, cosetEvals, cosets[cellID], ctx.openKey)
+	return kzgmulti.VerifyMultiPointKZGProof(commitmentG1, proofG1, cosetEvals, cosets[cellID], ctx.openKey7594)
 }
 
 func (ctx *Context) VerifyCellKZGProofBatch(rowCommitments []KZGCommitment, rowIndices, columnIndices []uint64, cells []*Cell, proofs []KZGProof) error {

internal/kzg/kzg.go

@@ -6,6 +6,10 @@ import (
 )
 
 // A polynomial in lagrange form
+//
+// Note: This is intentionally not in the `poly` package as
+// all methods, we want to do on the lagrange form as `kzg`
+// related.
 type Polynomial = []fr.Element
 
 // A commitment to a polynomial

internal/kzg/kzg_prove.go

@@ -28,7 +28,7 @@ func Open(domain *domain.Domain, p Polynomial, evaluationPoint fr.Element, ck *C
 	}
 
 	// Commit to Quotient polynomial
-	quotientCommit, err := Commit(quotientPoly, ck, numGoRoutines)
+	quotientCommit, err := ck.Commit(quotientPoly, numGoRoutines)
 	if err != nil {
 		return OpeningProof{}, err
 	}

internal/kzg/kzg_test.go

@@ -17,7 +17,7 @@ func TestProofVerifySmoke(t *testing.T) {
 	// polynomial in lagrange form
 	poly := Polynomial{fr.NewElement(2), fr.NewElement(3), fr.NewElement(4), fr.NewElement(5)}
 
-	comm, _ := Commit(poly, &srs.CommitKey, 0)
+	comm, _ := srs.CommitKey.Commit(poly, 0)
 	point := samplePointOutsideDomain(*domain)
 	proof, _ := Open(domain, poly, *point, &srs.CommitKey, 0)
 
@@ -155,7 +155,7 @@ func computeQuotientEvalWithinDomain(domain domain.Domain, z fr.Element, polynom
 func randValidOpeningProof(t *testing.T, domain domain.Domain, srs SRS) (OpeningProof, Commitment) {
 	t.Helper()
 	poly := randPoly(t, domain)
-	comm, _ := Commit(poly, &srs.CommitKey, 0)
+	comm, _ := srs.CommitKey.Commit(poly, 0)
 	point := samplePointOutsideDomain(domain)
 	proof, _ := Open(&domain, poly, *point, &srs.CommitKey, 0)
 	return proof, *comm

internal/kzg/srs.go

@@ -1,10 +1,7 @@
 package kzg
 
 import (
-	"errors"
-
 	bls12381 "github.com/consensys/gnark-crypto/ecc/bls12-381"
-	"github.com/consensys/gnark-crypto/ecc/bls12-381/fr"
 	"github.com/crate-crypto/go-eth-kzg/internal/domain"
 	"github.com/crate-crypto/go-eth-kzg/internal/multiexp"
 )
@@ -20,11 +17,6 @@ type OpeningKey struct {
 	// This is the degree-1 G_2 element in the trusted setup.
 	// In the specs, this is denoted as `KZG_SETUP_G2[1]`
 	AlphaG2 bls12381.G2Affine
-	// These are the G1 elements in monomial form from the trusted setup
-	G1 []bls12381.G1Affine
-	// These are the G2 elements in monomial form from the trusted setup
-	// Note: the length of this list is the same as the length of the G1 list
-	G2 []bls12381.G2Affine
 }
 
 // CommitKey holds the data needed to commit to polynomials and by proxy make opening proofs
@@ -61,29 +53,10 @@ type SRS struct {
 //
 // numGoRoutines is used to configure the amount of concurrency needed. Setting this
 // value to a negative number or 0 will make it default to the number of CPUs.
-// TODO: Move this to a method on CommitKey
-func Commit(p Polynomial, ck *CommitKey, numGoRoutines int) (*Commitment, error) {
-	if len(p) == 0 || len(p) > len(ck.G1) {
+func (c *CommitKey) Commit(p Polynomial, numGoRoutines int) (*Commitment, error) {
+	if len(p) == 0 || len(p) > len(c.G1) {
 		return nil, ErrInvalidPolynomialSize
 	}
 
-	return multiexp.MultiExpG1(p, ck.G1[:len(p)], numGoRoutines)
-}
-
-// TODO: Move this to a method on OpeningKey
-func CommitG1(scalars []fr.Element, ok *OpeningKey) (*bls12381.G1Affine, error) {
-	if len(scalars) == 0 || len(scalars) > len(ok.G1) {
-		return nil, errors.New("invalid vector size for G1 commitment")
-	}
-
-	return multiexp.MultiExpG1(scalars, ok.G1[:len(scalars)], 0)
-}
-
-// TODO: Move this to a method on OpeningKey
-func CommitG2(scalars []fr.Element, ok *OpeningKey) (*bls12381.G2Affine, error) {
-	if len(scalars) == 0 || len(scalars) > len(ok.G2) {
-		return nil, errors.New("invalid vector size for G2 commitment")
-	}
-
-	return multiexp.MultiExpG2(scalars, ok.G2[:len(scalars)], 0)
+	return multiexp.MultiExpG1(p, c.G1[:len(p)], numGoRoutines)
 }

internal/kzg/srs_test.go

@@ -50,8 +50,8 @@ func TestLagrangeSRSSmoke(t *testing.T) {
 	}
 	polyLagrange := Polynomial{f(domain.Roots[0]), f(domain.Roots[1]), f(domain.Roots[2]), f(domain.Roots[3])}
 
-	commitmentLagrange, _ := Commit(polyLagrange, &srsLagrange.CommitKey, 0)
-	commitmentMonomial, _ := Commit(polyMonomial, &srsMonomial.CommitKey, 0)
+	commitmentLagrange, _ := srsLagrange.CommitKey.Commit(polyLagrange, 0)
+	commitmentMonomial, _ := srsMonomial.CommitKey.Commit(polyMonomial, 0)
 	require.Equal(t, commitmentLagrange, commitmentMonomial)
 }
 
@@ -60,7 +60,7 @@ func TestCommitRegression(t *testing.T) {
 	srsLagrange, _ := newLagrangeSRSInsecure(*domain, big.NewInt(100))
 
 	poly := Polynomial{fr.NewElement(12345), fr.NewElement(123456), fr.NewElement(1234567), fr.NewElement(12345678)}
-	cLagrange, _ := Commit(poly, &srsLagrange.CommitKey, 0)
+	cLagrange, _ := srsLagrange.CommitKey.Commit(poly, 0)
 	cLagrangeBytes := cLagrange.Bytes()
 	gotCommitment := hex.EncodeToString(cLagrangeBytes[:])
 	expectedCommitment := "85bdf872da5b8561d23055d32db3fc86c672b0be7543b8c1e48634af07231bf7ab6385b765750921017cbcdbcd14f8e0"

internal/kzg_multi/errors.go

@@ -0,0 +1,5 @@
+package kzgmulti
+
+import "errors"
+
+var ErrMinSRSSize = errors.New("minimum srs size is 2")

internal/kzg_multi/kzg_prove_test.go

@@ -38,7 +38,7 @@ func computeMultiPointKZGProof(polyCoeff poly.PolynomialCoeff, inputPoints []fr.
 	}
 
 	// Commit to the quotient polynomial
-	proof, err := kzg.Commit(quotient, ck, 0)
+	proof, err := ck.Commit(quotient, 0)
 	if err != nil {
 		return bls12381.G1Affine{}, nil, err
 	}

internal/kzg_multi/kzg_verify.go

@@ -23,18 +23,18 @@ func vanishingPolyCoeff(xs []fr.Element) poly.PolynomialCoeff {
 	return result
 }
 
-func VerifyMultiPointKZGProof(commitment, proof bls12381.G1Affine, outputPoints, inputPoints []fr.Element, openKey *kzg.OpeningKey) error {
+func VerifyMultiPointKZGProof(commitment, proof bls12381.G1Affine, outputPoints, inputPoints []fr.Element, openKey *OpeningKey) error {
 	// Compute the following pairing check:
 	// e([Q(X)]_1, [Z(X)]_2) == e([f(X)]_1 - [I(X)]_1, [1]_2)
 
 	zeroPoly := vanishingPolyCoeff(inputPoints)
-	zeroPolyComm, err := kzg.CommitG2(zeroPoly, openKey)
+	zeroPolyComm, err := openKey.CommitG2(zeroPoly)
 	if err != nil {
 		return err
 	}
 
 	interpolatedPoly := poly.LagrangeInterpolate(inputPoints, outputPoints)
-	interpolatedPolyComm, err := kzg.CommitG1(interpolatedPoly, openKey)
+	interpolatedPolyComm, err := openKey.CommitG1(interpolatedPoly)
 	if err != nil {
 		return err
 	}
@@ -45,7 +45,7 @@ func VerifyMultiPointKZGProof(commitment, proof bls12381.G1Affine, outputPoints,
 	fMinusIx.Add(&fMinusIx, &commitment)
 
 	var negG2Gen bls12381.G2Affine
-	negG2Gen.Neg(&openKey.GenG2)
+	negG2Gen.Neg(openKey.genG2())
 
 	check, err := bls12381.PairingCheck(
 		[]bls12381.G1Affine{proof, fMinusIx},

internal/kzg_multi/srs.go

@@ -0,0 +1,125 @@
+package kzgmulti
+
+import (
+	"errors"
+	"math/big"
+
+	bls12381 "github.com/consensys/gnark-crypto/ecc/bls12-381"
+	"github.com/consensys/gnark-crypto/ecc/bls12-381/fr"
+	"github.com/crate-crypto/go-eth-kzg/internal/domain"
+	"github.com/crate-crypto/go-eth-kzg/internal/kzg"
+	"github.com/crate-crypto/go-eth-kzg/internal/multiexp"
+)
+
+// The commit key stays the same between the kzg single opening
+// use case and the multi opening use case
+type CommitKey = kzg.CommitKey
+
+type SRS struct {
+	OpeningKey OpeningKey
+	CommitKey  CommitKey
+}
+
+type OpeningKey struct {
+	// These are the G1 elements in monomial form from the trusted setup
+	G1 []bls12381.G1Affine
+	// These are the G2 elements in monomial form from the trusted setup
+	// Note: the length of this list is the same as the length of the G1 list
+	G2 []bls12381.G2Affine
+	// Number of points that the prover will open to in a single multi-point proof
+	//
+	// The points that the prover can open to are not arbitrary sets, but are cosets.
+	CosetSize uint64
+	// A bound on the polynomial length that we can commit, create and
+	// verify proofs about.
+	//
+	// Note: This is not the degree of the polynomial.
+	// One can compute the degree of the polynomial by doing `PolySize-1``
+	PolySize uint64
+	// The total number of points that the prover will create proofs for.
+	//
+	// Example, we could have f(x) = x^3 + x^2 + x + 1
+	// and a protocol may require that this polynomial which has PolySize = 4
+	// to be opened up at 32 points. The number 32 is the NumPointsToOpen constant.
+	NumPointsToOpen uint64
+	// CosetShiftsPowCosetSize contains precomputed powers of coset shifts.
+	// For each coset k, it stores (ω_k)^n where ω_k is the k-th coset shift
+	// and n is the coset size. These values are used for verifying multiple
+	// multi-point KZG opening proofs.
+	CosetShiftsPowCosetSize []fr.Element
+	// cosetDomains is a slice of CosetDomain objects, one for each coset.
+	//
+	// Each CosetDomain encapsulates the necessary information and methods
+	// to perform FFT operations over a specific coset of the evaluation domain.
+	//
+	// Note: This should not be confused with the cosets that we are creating
+	// and verifying opening proofs for.
+	cosetDomains []*domain.CosetDomain
+}
+
+func NewOpeningKey(g1s []bls12381.G1Affine, g2s []bls12381.G2Affine, polySize, numPointsToOpen, cosetSize uint64) *OpeningKey {
+	cosetDomain := domain.NewDomain(cosetSize)
+
+	extDomain := domain.NewDomain(numPointsToOpen)
+	domain.BitReverse(extDomain.Roots)
+
+	numCosets := numPointsToOpen / cosetSize
+	cosetShifts := make([]fr.Element, numCosets)
+	for k := 0; k < int(numCosets); k++ {
+		cosetShifts[k] = extDomain.Roots[k*int(cosetSize)]
+	}
+
+	invCosetShifts := make([]fr.Element, numCosets)
+	for k := 0; k < int(numCosets); k++ {
+		// Note: This is safe because the coset shifts are roots of unity
+		// and zero is not a root of unity.
+		invCosetShifts[k].Inverse(&cosetShifts[k])
+	}
+
+	cosetShiftsPowCosetSize := make([]fr.Element, numCosets)
+	cosetSizeBigInt := big.NewInt(int64(cosetSize))
+	for k := 0; k < int(numCosets); k++ {
+		cosetShiftsPowCosetSize[k].Exp(cosetShifts[k], cosetSizeBigInt)
+	}
+
+	cosetDomains := make([]*domain.CosetDomain, numCosets)
+	for k := 0; k < int(numCosets); k++ {
+		fftCoset := domain.FFTCoset{
+			CosetGen:    cosetShifts[k],
+			InvCosetGen: invCosetShifts[k],
+		}
+		cosetDomains[k] = domain.NewCosetDomain(cosetDomain, fftCoset)
+	}
+
+	return &OpeningKey{
+		G1:                      g1s,
+		G2:                      g2s,
+		CosetSize:               cosetSize,
+		PolySize:                polySize,
+		NumPointsToOpen:         numPointsToOpen,
+		CosetShiftsPowCosetSize: cosetShiftsPowCosetSize,
+		cosetDomains:            cosetDomains,
+	}
+}
+
+// This is the degree-0 G_2 element in the trusted setup.
+// In the specs, this is denoted as `KZG_SETUP_G2[0]`
+func (o *OpeningKey) genG2() *bls12381.G2Affine {
+	return &o.G2[0]
+}
+
+func (ok *OpeningKey) CommitG1(scalars []fr.Element) (*bls12381.G1Affine, error) {
+	if len(scalars) == 0 || len(scalars) > len(ok.G1) {
+		return nil, errors.New("invalid vector size for G1 commitment")
+	}
+
+	return multiexp.MultiExpG1(scalars, ok.G1[:len(scalars)], 0)
+}
+
+func (ok *OpeningKey) CommitG2(scalars []fr.Element) (*bls12381.G2Affine, error) {
+	if len(scalars) == 0 || len(scalars) > len(ok.G2) {
+		return nil, errors.New("invalid vector size for G2 commitment")
+	}
+
+	return multiexp.MultiExpG2(scalars, ok.G2[:len(scalars)], 0)
+}

prove.go

@@ -20,7 +20,7 @@ func (c *Context) BlobToKZGCommitment(blob *Blob, numGoRoutines int) (KZGCommitm
 	}
 
 	// 2. Commit to polynomial
-	commitment, err := kzg.Commit(polynomial, c.commitKeyLagrange, numGoRoutines)
+	commitment, err := c.commitKeyLagrange.Commit(polynomial, numGoRoutines)
 	if err != nil {
 		return KZGCommitment{}, err
 	}

verify.go

@@ -39,7 +39,7 @@ func (c *Context) VerifyKZGProof(blobCommitment KZGCommitment, inputPointBytes,
 		ClaimedValue:       claimedValue,
 	}
 
-	return kzg.Verify(&polynomialCommitment, &proof, c.openKey)
+	return kzg.Verify(&polynomialCommitment, &proof, c.openKey4844)
 }
 
 // VerifyBlobKZGProof implements [verify_blob_kzg_proof].
@@ -79,7 +79,7 @@ func (c *Context) VerifyBlobKZGProof(blob *Blob, blobCommitment KZGCommitment, k
 		ClaimedValue:       *outputPoint,
 	}
 
-	return kzg.Verify(&polynomialCommitment, &openingProof, c.openKey)
+	return kzg.Verify(&polynomialCommitment, &openingProof, c.openKey4844)
 }
 
 // VerifyBlobKZGProofBatch implements [verify_blob_kzg_proof_batch].
@@ -140,7 +140,7 @@ func (c *Context) VerifyBlobKZGProofBatch(blobs []Blob, polynomialCommitments []
 	}
 
 	// 3. Verify opening proofs
-	return kzg.BatchVerifyMultiPoints(commitments, openingProofs, c.openKey)
+	return kzg.BatchVerifyMultiPoints(commitments, openingProofs, c.openKey4844)
 }
 
 // VerifyBlobKZGProofBatchPar implements [verify_blob_kzg_proof_batch]. This is the parallelized version of