extract edsFromRows and write a test

share/shwap/p2p/bitswap/getter.go

@@ -4,14 +4,16 @@ import (
 	"context"
 	"fmt"
 
-	"github.com/celestiaorg/celestia-app/pkg/wrapper"
-	"github.com/celestiaorg/rsmt2d"
 	"github.com/ipfs/boxo/blockstore"
 	"github.com/ipfs/boxo/exchange"
 
+	"github.com/celestiaorg/celestia-app/pkg/wrapper"
+	"github.com/celestiaorg/rsmt2d"
+
 	"github.com/celestiaorg/celestia-node/header"
 	"github.com/celestiaorg/celestia-node/pruner"
 	"github.com/celestiaorg/celestia-node/share"
+	"github.com/celestiaorg/celestia-node/share/shwap"
 )
 
 // Getter implements share.Getter.
@@ -139,25 +141,14 @@ func (g *Getter) GetEDS(
 		return nil, err
 	}
 
-	shrs := make([]share.Share, 0, sqrLn/2*sqrLn/2)
-	for i, row := range blks {
-		rowShrs, err := row.(*RowBlock).Container.Shares()
-		if err != nil {
-			return nil, fmt.Errorf("decoding Shares out of Row: %w", err)
-		}
-
-		for j, shr := range rowShrs {
-			shrs[i*j] = shr
-		}
+	rows := make([]shwap.Row, len(blks))
+	for i, blk := range blks {
+		rows[i] = blk.(*RowBlock).Container
 	}
 
-	square, err := rsmt2d.ImportExtendedDataSquare(
-		shrs,
-		share.DefaultRSMT2DCodec(),
-		wrapper.NewConstructor(uint64(sqrLn/2)),
-	)
+	square, err := edsFromRows(hdr.DAH, rows)
 	if err != nil {
-		return nil, fmt.Errorf("computing EDS: %w", err)
+		return nil, err
 	}
 
 	return square, nil
@@ -212,3 +203,34 @@ func (g *Getter) session(hdr *header.ExtendedHeader) exchange.Fetcher {
 
 	return g.archivalSession
 }
+
+// edsFromRows imports given Rows and computes EDS out of them, assuming enough were Rows provided.
+func edsFromRows(roots *share.Root, rows []shwap.Row) (*rsmt2d.ExtendedDataSquare, error) {
+	shrs := make([]share.Share, len(roots.RowRoots)*len(roots.RowRoots))
+	for i, row := range rows {
+		rowShrs, err := row.Shares()
+		if err != nil {
+			return nil, fmt.Errorf("decoding Shares out of Row: %w", err)
+		}
+
+		for j, shr := range rowShrs {
+			shrs[j+(i*len(roots.RowRoots))] = shr
+		}
+	}
+
+	square, err := rsmt2d.ImportExtendedDataSquare(
+		shrs,
+		share.DefaultRSMT2DCodec(),
+		wrapper.NewConstructor(uint64(len(roots.RowRoots)/2)),
+	)
+	if err != nil {
+		return nil, fmt.Errorf("importing EDS: %w", err)
+	}
+
+	err = square.Repair(roots.RowRoots, roots.ColumnRoots)
+	if err != nil {
+		return nil, fmt.Errorf("repairing EDS: %w", err)
+	}
+
+	return square, nil
+}

share/shwap/p2p/bitswap/getter_test.go

@@ -0,0 +1,27 @@
+package bitswap
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/celestiaorg/celestia-node/share"
+	"github.com/celestiaorg/celestia-node/share/eds/edstest"
+	"github.com/celestiaorg/celestia-node/share/shwap"
+)
+
+func TestEDSFromRows(t *testing.T) {
+	edsIn := edstest.RandEDS(t, 8)
+	roots, err := share.NewRoot(edsIn)
+	require.NoError(t, err)
+
+	rows := make([]shwap.Row, edsIn.Width()/2)
+	for i := range edsIn.Width() / 2 {
+		rowShrs := edsIn.Row(i)[:edsIn.Width()/2]
+		rows[i] = shwap.NewRow(rowShrs, shwap.Left)
+	}
+
+	edsOut, err := edsFromRows(roots, rows)
+	require.NoError(t, err)
+	require.True(t, edsIn.Equals(edsOut))
+}