feat(shwap): Add axis roots to Accessor (#3586)

share/eds/edstest/testing.go

@@ -1,6 +1,7 @@
 package edstest
 
 import (
+	"crypto/rand"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -47,3 +48,21 @@ func RandEDSWithNamespace(
 	require.NoError(t, err)
 	return eds, roots
 }
+
+// RandomAxisRoots generates random share.AxisRoots for the given eds size.
+func RandomAxisRoots(t testing.TB, edsSize int) *share.AxisRoots {
+	roots := make([][]byte, edsSize*2)
+	for i := range roots {
+		root := make([]byte, edsSize)
+		_, err := rand.Read(root)
+		require.NoError(t, err)
+		roots[i] = root
+	}
+
+	rows := roots[:edsSize]
+	cols := roots[edsSize:]
+	return &share.AxisRoots{
+		RowRoots:    rows,
+		ColumnRoots: cols,
+	}
+}

share/new_eds/accessor.go

@@ -16,6 +16,8 @@ type Accessor interface {
 	Size(ctx context.Context) int
 	// DataHash returns data hash of the Accessor.
 	DataHash(ctx context.Context) (share.DataHash, error)
+	// AxisRoots returns share.AxisRoots (DataAvailabilityHeader) of the Accessor.
+	AxisRoots(ctx context.Context) (*share.AxisRoots, error)
 	// Sample returns share and corresponding proof for row and column indices. Implementation can
 	// choose which axis to use for proof. Chosen axis for proof should be indicated in the returned
 	// Sample.

share/new_eds/close_once.go

@@ -49,6 +49,13 @@ func (c *closeOnce) DataHash(ctx context.Context) (share.DataHash, error) {
 	return c.f.DataHash(ctx)
 }
 
+func (c *closeOnce) AxisRoots(ctx context.Context) (*share.AxisRoots, error) {
+	if c.closed.Load() {
+		return nil, errAccessorClosed
+	}
+	return c.f.AxisRoots(ctx)
+}
+
 func (c *closeOnce) Sample(ctx context.Context, rowIdx, colIdx int) (shwap.Sample, error) {
 	if c.closed.Load() {
 		return shwap.Sample{}, errAccessorClosed

share/new_eds/close_once_test.go

@@ -51,7 +51,11 @@ func (s *stubEdsAccessorCloser) Size(context.Context) int {
 }
 
 func (s *stubEdsAccessorCloser) DataHash(context.Context) (share.DataHash, error) {
-	return nil, nil
+	return share.DataHash{}, nil
+}
+
+func (s *stubEdsAccessorCloser) AxisRoots(context.Context) (*share.AxisRoots, error) {
+	return &share.AxisRoots{}, nil
 }
 
 func (s *stubEdsAccessorCloser) Sample(context.Context, int, int) (shwap.Sample, error) {

share/new_eds/nd.go

@@ -13,13 +13,15 @@ import (
 // avoiding the need to recalculate the row roots for each row.
 func NamespacedData(
 	ctx context.Context,
-	root *share.AxisRoots,
 	eds Accessor,
 	namespace share.Namespace,
 ) (shwap.NamespacedData, error) {
-	rowIdxs := share.RowsWithNamespace(root, namespace)
+	roots, err := eds.AxisRoots(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get AxisRoots: %w", err)
+	}
+	rowIdxs := share.RowsWithNamespace(roots, namespace)
 	rows := make(shwap.NamespacedData, len(rowIdxs))
-	var err error
 	for i, idx := range rowIdxs {
 		rows[i], err = eds.RowNamespaceData(ctx, namespace, idx)
 		if err != nil {

share/new_eds/nd_test.go

@@ -21,7 +21,7 @@ func TestNamespacedData(t *testing.T) {
 	for amount := 1; amount < sharesAmount; amount++ {
 		eds, root := edstest.RandEDSWithNamespace(t, namespace, amount, odsSize)
 		rsmt2d := &Rsmt2D{ExtendedDataSquare: eds}
-		nd, err := NamespacedData(ctx, root, rsmt2d, namespace)
+		nd, err := NamespacedData(ctx, rsmt2d, namespace)
 		require.NoError(t, err)
 		require.True(t, len(nd) > 0)
 		require.Len(t, nd.Flatten(), amount)

share/new_eds/proofs_cache.go

@@ -25,17 +25,23 @@ var _ AccessorStreamer = (*proofsCache)(nil)
 
 // proofsCache is eds accessor that caches proofs for rows and columns. It also caches extended
 // axis Shares. It is used to speed up the process of building proofs for rows and columns,
-// reducing the number of reads from the underlying accessor.
+// reducing the number of reads from the underlying accessor. Cache does not synchronize access
+// to the underlying accessor.
 type proofsCache struct {
 	inner AccessorStreamer
 
-	// lock protects axisCache
-	lock sync.RWMutex
+	// size caches the size of the data square
+	size atomic.Int32
+	// dataHash caches the data hash
+	dataHash atomic.Pointer[share.DataHash]
+	// rootsCache caches the axis roots
+	rootsCache atomic.Pointer[share.AxisRoots]
+	// axisCacheLock protects proofCache
+	axisCacheLock sync.RWMutex
 	// axisCache caches the axis Shares and proofs. Index in the slice corresponds to the axis type.
 	// The map key is the index of the axis.
 	axisCache []map[int]axisWithProofs
-	// size caches the size of the data square
-	size atomic.Int32
+
 	// disableCache disables caching of rows for testing purposes
 	disableCache bool
 }
@@ -78,7 +84,31 @@ func (c *proofsCache) Size(ctx context.Context) int {
 }
 
 func (c *proofsCache) DataHash(ctx context.Context) (share.DataHash, error) {
-	return c.inner.DataHash(ctx)
+	dataHash := c.dataHash.Load()
+	if dataHash != nil {
+		return *dataHash, nil
+	}
+	loaded, err := c.inner.DataHash(ctx)
+	if err != nil {
+		return nil, err
+	}
+	c.dataHash.Store(&loaded)
+	return loaded, nil
+}
+
+func (c *proofsCache) AxisRoots(ctx context.Context) (*share.AxisRoots, error) {
+	roots := c.rootsCache.Load()
+	if roots != nil {
+		return roots, nil
+	}
+
+	// if roots are not in cache, read them from the inner accessor
+	roots, err := c.inner.AxisRoots(ctx)
+	if err != nil {
+		return nil, err
+	}
+	c.rootsCache.Store(roots)
+	return roots, nil
 }
 
 func (c *proofsCache) Sample(ctx context.Context, rowIdx, colIdx int) (shwap.Sample, error) {
@@ -250,14 +280,14 @@ func (c *proofsCache) axisShares(ctx context.Context, axisType rsmt2d.Axis, axis
 }
 
 func (c *proofsCache) storeAxisInCache(axisType rsmt2d.Axis, axisIdx int, axis axisWithProofs) {
-	c.lock.Lock()
-	defer c.lock.Unlock()
+	c.axisCacheLock.Lock()
+	defer c.axisCacheLock.Unlock()
 	c.axisCache[axisType][axisIdx] = axis
 }
 
 func (c *proofsCache) getAxisFromCache(axisType rsmt2d.Axis, axisIdx int) (axisWithProofs, bool) {
-	c.lock.RLock()
-	defer c.lock.RUnlock()
+	c.axisCacheLock.RLock()
+	defer c.axisCacheLock.RUnlock()
 	ax, ok := c.axisCache[axisType][axisIdx]
 	return ax, ok
 }

share/new_eds/rsmt2d.go

@@ -33,6 +33,15 @@ func (eds *Rsmt2D) DataHash(context.Context) (share.DataHash, error) {
 	return roots.Hash(), nil
 }
 
+// AxisRoots returns AxisRoots of the Accessor.
+func (eds *Rsmt2D) AxisRoots(context.Context) (*share.AxisRoots, error) {
+	roots, err := share.NewAxisRoots(eds.ExtendedDataSquare)
+	if err != nil {
+		return nil, fmt.Errorf("while creating axis roots: %w", err)
+	}
+	return roots, nil
+}
+
 // Sample returns share and corresponding proof for row and column indices.
 func (eds *Rsmt2D) Sample(
 	_ context.Context,

share/new_eds/testing.go

@@ -35,6 +35,14 @@ func TestSuiteAccessor(
 		t.Errorf("minSize must be power of 2: %v", maxSize)
 	}
 	for size := minSize; size <= maxSize; size *= 2 {
+		t.Run(fmt.Sprintf("DataHash:%d", size), func(t *testing.T) {
+			testAccessorDataHash(ctx, t, createAccessor, size)
+		})
+
+		t.Run(fmt.Sprintf("AxisRoots:%d", size), func(t *testing.T) {
+			testAccessorAxisRoots(ctx, t, createAccessor, size)
+		})
+
 		t.Run(fmt.Sprintf("Sample:%d", size), func(t *testing.T) {
 			testAccessorSample(ctx, t, createAccessor, size)
 		})
@@ -64,6 +72,40 @@ func TestStreamer(
 	})
 }
 
+func testAccessorDataHash(
+	ctx context.Context,
+	t *testing.T,
+	createAccessor createAccessor,
+	odsSize int,
+) {
+	eds := edstest.RandEDS(t, odsSize)
+	fl := createAccessor(t, eds)
+
+	roots, err := share.NewAxisRoots(eds)
+	require.NoError(t, err)
+
+	datahash, err := fl.DataHash(ctx)
+	require.NoError(t, err)
+	require.Equal(t, share.DataHash(roots.Hash()), datahash)
+}
+
+func testAccessorAxisRoots(
+	ctx context.Context,
+	t *testing.T,
+	createAccessor createAccessor,
+	odsSize int,
+) {
+	eds := edstest.RandEDS(t, odsSize)
+	fl := createAccessor(t, eds)
+
+	roots, err := share.NewAxisRoots(eds)
+	require.NoError(t, err)
+
+	axisRoots, err := fl.AxisRoots(ctx)
+	require.NoError(t, err)
+	require.True(t, roots.Equals(axisRoots))
+}
+
 func testAccessorSample(
 	ctx context.Context,
 	t *testing.T,

share/root.go

@@ -11,6 +11,13 @@ import (
 	"github.com/celestiaorg/rsmt2d"
 )
 
+const (
+	// DataHashSize is the size of the DataHash.
+	DataHashSize = 32
+	// AxisRootSize is the size of the single root in AxisRoots.
+	AxisRootSize = 90
+)
+
 // AxisRoots represents root commitment to multiple Shares.
 // In practice, it is a commitment to all the Data in a square.
 type AxisRoots = da.DataAvailabilityHeader
@@ -19,7 +26,7 @@ type AxisRoots = da.DataAvailabilityHeader
 type DataHash []byte
 
 func (dh DataHash) Validate() error {
-	if len(dh) != 32 {
+	if len(dh) != DataHashSize {
 		return fmt.Errorf("invalid hash size, expected 32, got %d", len(dh))
 	}
 	return nil

share/shwap/row_namespace_data_test.go

@@ -66,7 +66,7 @@ func TestValidateNamespacedRow(t *testing.T) {
 	for amount := 1; amount < sharesAmount; amount++ {
 		randEDS, root := edstest.RandEDSWithNamespace(t, namespace, amount, odsSize)
 		rsmt2d := &eds.Rsmt2D{ExtendedDataSquare: randEDS}
-		nd, err := eds.NamespacedData(ctx, root, rsmt2d, namespace)
+		nd, err := eds.NamespacedData(ctx, rsmt2d, namespace)
 		require.NoError(t, err)
 		require.True(t, len(nd) > 0)
 
@@ -86,9 +86,9 @@ func TestNamespacedRowProtoEncoding(t *testing.T) {
 
 	const odsSize = 8
 	namespace := sharetest.RandV0Namespace()
-	randEDS, root := edstest.RandEDSWithNamespace(t, namespace, odsSize, odsSize)
+	randEDS, _ := edstest.RandEDSWithNamespace(t, namespace, odsSize, odsSize)
 	rsmt2d := &eds.Rsmt2D{ExtendedDataSquare: randEDS}
-	nd, err := eds.NamespacedData(ctx, root, rsmt2d, namespace)
+	nd, err := eds.NamespacedData(ctx, rsmt2d, namespace)
 	require.NoError(t, err)
 	require.True(t, len(nd) > 0)
 

store/cache/accessor_cache_test.go

@@ -304,6 +304,10 @@ func (m *mockAccessor) DataHash(context.Context) (share.DataHash, error) {
 	panic("implement me")
 }
 
+func (m *mockAccessor) AxisRoots(context.Context) (*share.AxisRoots, error) {
+	panic("implement me")
+}
+
 func (m *mockAccessor) Sample(context.Context, int, int) (shwap.Sample, error) {
 	panic("implement me")
 }

store/cache/noop.go

@@ -46,9 +46,12 @@ func (n NoopFile) Size(context.Context) int {
 	return 0
 }
 
-// DataHash returns root hash of Accessor's underlying EDS.
 func (n NoopFile) DataHash(context.Context) (share.DataHash, error) {
-	return nil, nil
+	return share.DataHash{}, nil
+}
+
+func (n NoopFile) AxisRoots(context.Context) (*share.AxisRoots, error) {
+	return &share.AxisRoots{}, nil
 }
 
 func (n NoopFile) Sample(context.Context, int, int) (shwap.Sample, error) {