From 8c3d85151441d5ba6d57b38f7d1a65552057381b Mon Sep 17 00:00:00 2001
From: kevaundray <kevtheappdev@gmail.com>
Date: Tue, 20 Aug 2024 12:42:49 +0100
Subject: [PATCH] add erasure-code (#87)

---
 api.go                                        |  6 ++--
 .../erasure_code.go}                          | 30 ++++++++++++++++++-
 internal/erasure_code/erasure_code_test.go    | 19 ++++++++++++
 3 files changed, 51 insertions(+), 4 deletions(-)
 rename internal/{kzg_multi/recovery.go => erasure_code/erasure_code.go} (84%)
 create mode 100644 internal/erasure_code/erasure_code_test.go

diff --git a/api.go b/api.go
index f2ec265..73e3b68 100644
--- a/api.go
+++ b/api.go
@@ -4,8 +4,8 @@ import (
 	"encoding/json"
 
 	"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"
 )
 
@@ -22,7 +22,7 @@ type Context struct {
 
 	fk20 *fk20.FK20
 
-	dataRecovery *kzgmulti.DataRecovery
+	dataRecovery *erasure_code.DataRecovery
 }
 
 // BlsModulus is the bytes representation of the bls12-381 scalar field modulus.
@@ -145,6 +145,6 @@ func NewContext4096(trustedSetup *JSONTrustedSetup) (*Context, error) {
 		// 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: kzgmulti.NewDataRecovery(scalarsPerCell, ScalarsPerBlob, expansionFactor),
+		dataRecovery: erasure_code.NewDataRecovery(scalarsPerCell, ScalarsPerBlob, expansionFactor),
 	}, nil
 }
diff --git a/internal/kzg_multi/recovery.go b/internal/erasure_code/erasure_code.go
similarity index 84%
rename from internal/kzg_multi/recovery.go
rename to internal/erasure_code/erasure_code.go
index 8703704..2cf3611 100644
--- a/internal/kzg_multi/recovery.go
+++ b/internal/erasure_code/erasure_code.go
@@ -1,10 +1,11 @@
-package kzgmulti
+package erasure_code
 
 import (
 	"errors"
 
 	"github.com/consensys/gnark-crypto/ecc/bls12-381/fr"
 	"github.com/crate-crypto/go-eth-kzg/internal/domain"
+	kzgmulti "github.com/crate-crypto/go-eth-kzg/internal/kzg_multi"
 )
 
 // BlockErasureIndex is used to indicate the index of the block erasure that is missing
@@ -88,6 +89,17 @@ func (dr *DataRecovery) constructVanishingPolyOnIndices(missingBlockErasureIndic
 	return zeroPolyCoeff
 }
 
+// Encode the polynomial by evaluating it on the extended domain.
+//
+// Note: `polyCoeff` is mutated in-place, ie it should be seen as mutable reference.
+func (dr *DataRecovery) Encode(polyCoeff []fr.Element) []fr.Element {
+	// Pad to the correct length
+	for i := len(polyCoeff); i < len(dr.domainExtended.Roots); i++ {
+		polyCoeff = append(polyCoeff, fr.Element{})
+	}
+	return dr.domainExtended.FftFr(polyCoeff)
+}
+
 // NumBlocksNeededToReconstruct returns the number of blocks that are needed to reconstruct
 // the original data word.
 func (dr *DataRecovery) NumBlocksNeededToReconstruct() int {
@@ -126,3 +138,19 @@ func (dr *DataRecovery) RecoverPolynomialCoefficients(data []fr.Element, missing
 	polyCoeff = polyCoeff[:dr.numScalarsInDataWord]
 	return polyCoeff, nil
 }
+
+// vanishingPolyCoeff returns the polynomial that has roots at the given points
+func vanishingPolyCoeff(xs []fr.Element) kzgmulti.PolynomialCoeff {
+	result := []fr.Element{fr.One()}
+
+	for _, x := range xs {
+		// This is to silence: G601: Implicit memory aliasing in for loop.
+		x := x
+
+		negX := fr.Element{}
+		negX.Neg(&x)
+		result = kzgmulti.PolyMul(result, []fr.Element{negX, fr.One()})
+	}
+
+	return result
+}
diff --git a/internal/erasure_code/erasure_code_test.go b/internal/erasure_code/erasure_code_test.go
new file mode 100644
index 0000000..f8929b3
--- /dev/null
+++ b/internal/erasure_code/erasure_code_test.go
@@ -0,0 +1,19 @@
+package erasure_code
+
+import (
+	"testing"
+
+	"github.com/consensys/gnark-crypto/ecc/bls12-381/fr"
+	kzgmulti "github.com/crate-crypto/go-eth-kzg/internal/kzg_multi"
+)
+
+func TestVanishingPoly(t *testing.T) {
+	points := []fr.Element{fr.NewElement(1), fr.NewElement(2), fr.NewElement(3), fr.NewElement(4)}
+	vanishingPoly := vanishingPolyCoeff(points)
+	for _, point := range points {
+		eval := kzgmulti.PolyEval(vanishingPoly, point)
+		if !eval.IsZero() {
+			t.Fatalf("expected evaluation at the vanishing polynomial to be zero")
+		}
+	}
+}