README.md

@@ -53,11 +53,16 @@ func main() {
     flattened[8], flattened[9], flattened[10] = nil, nil, nil
     flattened[12], flattened[13] = nil, nil
 
+    // Re-import the data square.
+    eds, err = rsmt2d.ImportExtendedDataSquare(flattened, codec, rsmt2d.NewDefaultTree)
+    if err != nil {
+        // ImportExtendedDataSquare failed
+    }
+
     // Repair square.
-    err := rsmt2d.RepairExtendedDataSquare(
+    err := eds.Repair(
         eds.RowRoots(),
         eds.ColRoots(),
-        flattened,
         codec,
         rsmt2d.NewDefaultTree,
     )

extendeddatacrossword.go

@@ -6,38 +6,36 @@ import (
 	"fmt"
 )
 
+// Axis represents which of a row or col.
+type Axis int
+
 const (
-	row = iota
-	col
+	Row Axis = iota
+	Col
 )
 
 // ErrUnrepairableDataSquare is thrown when there is insufficient chunks to repair the square.
 var ErrUnrepairableDataSquare = errors.New("failed to solve data square")
 
-// ErrNoChunksAvailable is thrown when dataSquare is empty.
-var ErrNoChunksAvailable = errors.New("no chunks available")
-
-// ErrByzantineRow is thrown when a repaired row does not match the expected row Merkle root.
-type ErrByzantineRow struct {
-	RowNumber uint     // Row index
-	Shares    [][]byte // Pre-repaired row shares. Missing shares are nil.
-}
-
-func (e *ErrByzantineRow) Error() string {
-	return fmt.Sprintf("byzantine row: %d", e.RowNumber)
+// ErrByzantineData is thrown when a repaired row or column does not match the expected row or column Merkle root.
+type ErrByzantineData struct {
+	Axis   Axis     // Axis of the data.
+	Index  uint     // Row/Col index.
+	Shares [][]byte // Pre-repaired shares. Missing shares are nil.
 }
 
-// ErrByzantineCol is thrown when a repaired column does not match the expected column Merkle root.
-type ErrByzantineCol struct {
-	ColNumber uint     // Column index
-	Shares    [][]byte // Pre-repaired column shares. Missing shares are nil.
-}
-
-func (e *ErrByzantineCol) Error() string {
-	return fmt.Sprintf("byzantine column: %d", e.ColNumber)
+func (e *ErrByzantineData) Error() string {
+	switch e.Axis {
+	case Row:
+		return fmt.Sprintf("byzantine row: %d", e.Index)
+	case Col:
+		return fmt.Sprintf("byzantine column: %d", e.Index)
+	default:
+		panic(fmt.Sprintf("invalid axis: %d", e.Axis))
+	}
 }
 
-// RepairExtendedDataSquare attempts to repair an incomplete extended data
+// Repair attempts to repair an incomplete extended data
 // square (EDS), comparing repaired rows and columns against expected Merkle
 // roots.
 //
@@ -52,29 +50,13 @@ func (e *ErrByzantineCol) Error() string {
 // prior to the Byzantine row or column being repaired, and the Byzantine row
 // or column prior to repair is returned in the error with missing shares as
 // nil.
-func RepairExtendedDataSquare(
+func (eds *ExtendedDataSquare) Repair(
 	rowRoots [][]byte,
 	colRoots [][]byte,
-	data [][]byte,
 	codec Codec,
 	treeCreatorFn TreeConstructorFn,
 ) error {
-	var isNotEmpty bool
-	for _, d := range data {
-		if d != nil {
-			isNotEmpty = true
-		}
-	}
-	if !isNotEmpty {
-		return ErrNoChunksAvailable
-	}
-
-	eds, err := ImportExtendedDataSquare(data, codec, treeCreatorFn)
-	if err != nil {
-		return err
-	}
-
-	err = eds.prerepairSanityCheck(rowRoots, colRoots, codec)
+	err := eds.prerepairSanityCheck(rowRoots, colRoots, codec)
 	if err != nil {
 		return err
 	}
@@ -281,7 +263,7 @@ func (eds *ExtendedDataSquare) verifyAgainstRowRoots(
 	root := eds.computeSharesRoot(shares, r)
 
 	if !bytes.Equal(root, rowRoots[r]) {
-		return &ErrByzantineRow{r, shares}
+		return &ErrByzantineData{Row, r, shares}
 	}
 
 	return nil
@@ -295,7 +277,7 @@ func (eds *ExtendedDataSquare) verifyAgainstColRoots(
 	root := eds.computeSharesRoot(shares, c)
 
 	if !bytes.Equal(root, colRoots[c]) {
-		return &ErrByzantineCol{c, shares}
+		return &ErrByzantineData{Col, c, shares}
 	}
 
 	return nil
@@ -332,7 +314,7 @@ func (eds *ExtendedDataSquare) prerepairSanityCheck(
 				return err
 			}
 			if !bytes.Equal(flattenChunks(parityShares), flattenChunks(eds.rowSlice(i, eds.originalDataWidth, eds.originalDataWidth))) {
-				return &ErrByzantineRow{i, eds.row(i)}
+				return &ErrByzantineData{Row, i, eds.row(i)}
 			}
 		}
 
@@ -342,7 +324,7 @@ func (eds *ExtendedDataSquare) prerepairSanityCheck(
 				return err
 			}
 			if !bytes.Equal(flattenChunks(parityShares), flattenChunks(eds.colSlice(eds.originalDataWidth, i, eds.originalDataWidth))) {
-				return &ErrByzantineCol{i, eds.col(i)}
+				return &ErrByzantineData{Col, i, eds.col(i)}
 			}
 		}
 	}

extendeddatacrossword_test.go

@@ -35,12 +35,22 @@ func TestRepairExtendedDataSquare(t *testing.T) {
 			panic(err)
 		}
 
+		rowRoots := original.RowRoots()
+		colRoots := original.ColRoots()
+
 		flattened := original.flattened()
 		flattened[0], flattened[2], flattened[3] = nil, nil, nil
 		flattened[4], flattened[5], flattened[6], flattened[7] = nil, nil, nil, nil
 		flattened[8], flattened[9], flattened[10] = nil, nil, nil
 		flattened[12], flattened[13] = nil, nil
-		err = RepairExtendedDataSquare(original.getRowRoots(), original.getColRoots(), flattened, codec, NewDefaultTree)
+
+		// Re-import the data square.
+		eds, err := ImportExtendedDataSquare(flattened, codec, NewDefaultTree)
+		if err != nil {
+			t.Errorf("ImportExtendedDataSquare failed: %v", err)
+		}
+
+		err = eds.Repair(rowRoots, colRoots, codec, NewDefaultTree)
 		if err != nil {
 			t.Errorf("unexpected err while repairing data square: %v, codec: :%s", err, codecName)
 		} else {
@@ -56,7 +66,13 @@ func TestRepairExtendedDataSquare(t *testing.T) {
 		flattened[8], flattened[9], flattened[10] = nil, nil, nil
 		flattened[12], flattened[13], flattened[14] = nil, nil, nil
 
-		err = RepairExtendedDataSquare(original.getRowRoots(), original.getColRoots(), flattened, codec, NewDefaultTree)
+		// Re-import the data square.
+		eds, err = ImportExtendedDataSquare(flattened, codec, NewDefaultTree)
+		if err != nil {
+			t.Errorf("ImportExtendedDataSquare failed: %v", err)
+		}
+
+		err = eds.Repair(rowRoots, colRoots, codec, NewDefaultTree)
 		if err == nil {
 			t.Errorf("did not return an error on trying to repair an unrepairable square")
 		}
@@ -67,7 +83,7 @@ func TestRepairExtendedDataSquare(t *testing.T) {
 		}
 		corruptChunk := bytes.Repeat([]byte{66}, bufferSize)
 		corrupted.setCell(0, 0, corruptChunk)
-		err = RepairExtendedDataSquare(original.getRowRoots(), original.getColRoots(), corrupted.flattened(), codec, NewDefaultTree)
+		err = corrupted.Repair(rowRoots, colRoots, codec, NewDefaultTree)
 		if err == nil {
 			t.Errorf("did not return an error on trying to repair a square with bad roots")
 		}
@@ -77,13 +93,13 @@ func TestRepairExtendedDataSquare(t *testing.T) {
 			t.Fatalf("unexpected err while copying original data: %v, codec: :%s", err, codecName)
 		}
 		corrupted.setCell(0, 0, corruptChunk)
-		err = RepairExtendedDataSquare(corrupted.getRowRoots(), corrupted.getColRoots(), corrupted.flattened(), codec, NewDefaultTree)
-		var byzRow *ErrByzantineRow
-		if !errors.As(err, &byzRow) {
-			t.Errorf("did not return a ErrByzantineRow for a bad row; got: %v", err)
+		err = corrupted.Repair(corrupted.getRowRoots(), corrupted.getColRoots(), codec, NewDefaultTree)
+		var byzData *ErrByzantineData
+		if !errors.As(err, &byzData) || byzData.Axis != Row {
+			t.Errorf("did not return a ErrByzantineData for a bad row; got: %v", err)
 		}
 		// Construct the fraud proof
-		fraudProof := PseudoFraudProof{row, byzRow.RowNumber, byzRow.Shares}
+		fraudProof := PseudoFraudProof{0, byzData.Index, byzData.Shares}
 		// Verify the fraud proof
 		// TODO in a real fraud proof, also verify Merkle proof for each non-nil share.
 		rebuiltShares, err := codec.Decode(fraudProof.Shares)
@@ -108,33 +124,34 @@ func TestRepairExtendedDataSquare(t *testing.T) {
 			t.Fatalf("unexpected err while copying original data: %v, codec: :%s", err, codecName)
 		}
 		corrupted.setCell(0, 3, corruptChunk)
-		err = RepairExtendedDataSquare(corrupted.getRowRoots(), corrupted.getColRoots(), corrupted.flattened(), codec, NewDefaultTree)
-		if !errors.As(err, &byzRow) {
-			t.Errorf("did not return a ErrByzantineRow for a bad row; got %v", err)
+		err = corrupted.Repair(corrupted.getRowRoots(), corrupted.getColRoots(), codec, NewDefaultTree)
+		if !errors.As(err, &byzData) || byzData.Axis != Row {
+			t.Errorf("did not return a ErrByzantineData for a bad row; got %v", err)
 		}
 
 		corrupted, err = original.deepCopy(codec)
 		if err != nil {
 			t.Fatalf("unexpected err while copying original data: %v, codec: :%s", err, codecName)
 		}
 		corrupted.setCell(0, 0, corruptChunk)
-		flattened = corrupted.flattened()
-		flattened[1], flattened[2], flattened[3] = nil, nil, nil
-		err = RepairExtendedDataSquare(corrupted.getRowRoots(), corrupted.getColRoots(), flattened, codec, NewDefaultTree)
-		var byzCol *ErrByzantineCol
-		if !errors.As(err, &byzCol) {
-			t.Errorf("did not return a ErrByzantineCol for a bad column; got %v", err)
+		corrupted.setCell(0, 1, nil)
+		corrupted.setCell(0, 2, nil)
+		corrupted.setCell(0, 3, nil)
+		err = corrupted.Repair(corrupted.getRowRoots(), corrupted.getColRoots(), codec, NewDefaultTree)
+		if !errors.As(err, &byzData) || byzData.Axis != Col {
+			t.Errorf("did not return a ErrByzantineData for a bad column; got %v", err)
 		}
 		corrupted, err = original.deepCopy(codec)
 		if err != nil {
 			t.Fatalf("unexpected err while copying original data: %v, codec: :%s", err, codecName)
 		}
 		corrupted.setCell(3, 0, corruptChunk)
-		flattened = corrupted.flattened()
-		flattened[1], flattened[2], flattened[3] = nil, nil, nil
-		err = RepairExtendedDataSquare(corrupted.getRowRoots(), corrupted.getColRoots(), flattened, codec, NewDefaultTree)
-		if !errors.As(err, &byzCol) {
-			t.Errorf("did not return a ErrByzantineCol for a bad column; got %v", err)
+		corrupted.setCell(0, 1, nil)
+		corrupted.setCell(0, 2, nil)
+		corrupted.setCell(0, 3, nil)
+		err = corrupted.Repair(corrupted.getRowRoots(), corrupted.getColRoots(), codec, NewDefaultTree)
+		if !errors.As(err, &byzData) || byzData.Axis != Col {
+			t.Errorf("did not return a ErrByzantineData for a bad column; got %v", err)
 		}
 	}
 }
@@ -151,6 +168,8 @@ func BenchmarkRepair(b *testing.B) {
 			}
 
 			extendedDataWidth := originalDataWidth * 2
+			rowRoots := eds.RowRoots()
+			colRoots := eds.ColRoots()
 
 			b.Run(
 				fmt.Sprintf(
@@ -176,12 +195,14 @@ func BenchmarkRepair(b *testing.B) {
 							}
 						}
 
+						// Re-import the data square.
+						eds, _ = ImportExtendedDataSquare(flattened, codec, NewDefaultTree)
+
 						b.StartTimer()
 
-						err := RepairExtendedDataSquare(
-							eds.getRowRoots(),
-							eds.getColRoots(),
-							flattened,
+						err := eds.Repair(
+							rowRoots,
+							colRoots,
 							codec,
 							NewDefaultTree,
 						)

rsmt2d_test.go

@@ -43,6 +43,9 @@ func TestEdsRepairRoundtripSimple(t *testing.T) {
 				t.Errorf("ComputeExtendedDataSquare failed: %v", err)
 			}
 
+			rowRoots := eds.RowRoots()
+			colRoots := eds.ColRoots()
+
 			// Save all shares in flattened form.
 			flattened := make([][]byte, 0, eds.Width()*eds.Width())
 			for i := uint(0); i < eds.Width(); i++ {
@@ -55,11 +58,16 @@ func TestEdsRepairRoundtripSimple(t *testing.T) {
 			flattened[8], flattened[9], flattened[10] = nil, nil, nil
 			flattened[12], flattened[13] = nil, nil
 
+			// Re-import the data square.
+			eds, err = rsmt2d.ImportExtendedDataSquare(flattened, tt.codec, rsmt2d.NewDefaultTree)
+			if err != nil {
+				t.Errorf("ImportExtendedDataSquare failed: %v", err)
+			}
+
 			// Repair square.
-			err = rsmt2d.RepairExtendedDataSquare(
-				eds.RowRoots(),
-				eds.ColRoots(),
-				flattened,
+			err = eds.Repair(
+				rowRoots,
+				colRoots,
 				tt.codec,
 				rsmt2d.NewDefaultTree,
 			)
@@ -105,6 +113,9 @@ func TestEdsRepairTwice(t *testing.T) {
 				t.Errorf("ComputeExtendedDataSquare failed: %v", err)
 			}
 
+			rowRoots := eds.RowRoots()
+			colRoots := eds.ColRoots()
+
 			// Save all shares in flattened form.
 			flattened := make([][]byte, 0, eds.Width()*eds.Width())
 			for i := uint(0); i < eds.Width(); i++ {
@@ -119,11 +130,16 @@ func TestEdsRepairTwice(t *testing.T) {
 			flattened[8], flattened[9], flattened[10] = nil, nil, nil
 			flattened[12], flattened[13] = nil, nil
 
+			// Re-import the data square.
+			eds, err = rsmt2d.ImportExtendedDataSquare(flattened, tt.codec, rsmt2d.NewDefaultTree)
+			if err != nil {
+				t.Errorf("ImportExtendedDataSquare failed: %v", err)
+			}
+
 			// Repair square.
-			err = rsmt2d.RepairExtendedDataSquare(
-				eds.RowRoots(),
-				eds.ColRoots(),
-				flattened,
+			err = eds.Repair(
+				rowRoots,
+				colRoots,
 				tt.codec,
 				rsmt2d.NewDefaultTree,
 			)
@@ -134,10 +150,16 @@ func TestEdsRepairTwice(t *testing.T) {
 			// Re-insert missing share and try again.
 			flattened[1] = make([]byte, bufferSize)
 			copy(flattened[1], missing)
-			err = rsmt2d.RepairExtendedDataSquare(
-				eds.RowRoots(),
-				eds.ColRoots(),
-				flattened,
+
+			// Re-import the data square.
+			eds, err = rsmt2d.ImportExtendedDataSquare(flattened, tt.codec, rsmt2d.NewDefaultTree)
+			if err != nil {
+				t.Errorf("ImportExtendedDataSquare failed: %v", err)
+			}
+
+			err = eds.Repair(
+				rowRoots,
+				colRoots,
 				tt.codec,
 				rsmt2d.NewDefaultTree,
 			)