chore: update uses of rsmt2d

celestiaorg · Apr 26, 2022 · 07bad81 · 07bad81
1 parent 33e33b1
commit 07bad81
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 5 deletions.
diff --git a/ipld/read.go b/ipld/read.go
@@ -42,12 +42,16 @@ func RetrieveData(
 	}
 	batchAdder := NewNmtNodeAdder(parentCtx, ipld.NewBatch(parentCtx, dag, ipld.MaxSizeBatchOption(batchSize(edsWidth))))
 	tree := wrapper.NewErasuredNamespacedMerkleTree(uint64(edsWidth)/2, nmt.NodeVisitor(batchAdder.Visit))
-	extended, err := rsmt2d.RepairExtendedDataSquare(dah.RowsRoots, dah.ColumnRoots, dataSquare, codec, tree.Constructor)
+	err = rsmt2d.RepairExtendedDataSquare(dah.RowsRoots, dah.ColumnRoots, dataSquare, codec, tree.Constructor)
 	if err != nil {
 		return nil, err
 	}
 
-	return extended, batchAdder.Commit()
+	eds, err := rsmt2d.ImportExtendedDataSquare(dataSquare, codec, tree.Constructor)
+	if err != nil {
+		return nil, err
+	}
+	return eds, batchAdder.Commit()
 }
 
 type quadrant struct {

diff --git a/ipld/read_test.go b/ipld/read_test.go
@@ -92,11 +92,13 @@ func TestBlockRecovery(t *testing.T) {
 
 			flat := flatten(eds)
 
+			dataSquare := removeRandShares(flat, tc.d)
+
 			// recover a partially complete square
-			reds, err := rsmt2d.RepairExtendedDataSquare(
+			err = rsmt2d.RepairExtendedDataSquare(
 				rowRoots,
 				colRoots,
-				removeRandShares(flat, tc.d),
+				dataSquare,
 				rsmt2d.NewRSGF8Codec(),
 				recoverTree.Constructor,
 			)
@@ -108,6 +110,13 @@ func TestBlockRecovery(t *testing.T) {
 			}
 			assert.NoError(t, err)
 
+			reds, err := rsmt2d.ImportExtendedDataSquare(dataSquare, rsmt2d.NewRSGF8Codec(), recoverTree.Constructor)
+			// if an error is expected during attempted EDS reparation,
+			// then an error should be expected during import
+			if tc.expectErr {
+				require.Error(t, err)
+			}
+			assert.NoError(t, err)
 			// check that the squares are equal
 			assert.Equal(t, flatten(eds), flatten(reds))
 		})

diff --git a/service/header/header.go b/service/header/header.go
@@ -45,7 +45,10 @@ func MakeExtendedHeader(
 ) (*ExtendedHeader, error) {
 	var dah DataAvailabilityHeader
 	if len(b.Txs) > 0 {
-		namespacedShares, _ := b.Data.ComputeShares()
+		namespacedShares, _, err := b.Data.ComputeShares(b.Data.OriginalSquareSize)
+		if err != nil {
+			return nil, err
+		}
 		extended, err := ipld.PutData(ctx, namespacedShares.RawShares(), dag)
 		if err != nil {
 			return nil, err