test: extends computeRoots benchmark to cover larger square sizes and nmt tree

datasquare_test.go

@@ -3,10 +3,12 @@ package rsmt2d
 import (
 	"crypto/sha256"
 	"fmt"
+	"math"
 	"reflect"
 	"testing"
 
 	"github.com/celestiaorg/merkletree"
+	"github.com/celestiaorg/nmt"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -401,7 +403,7 @@ func Test_setColSlice(t *testing.T) {
 	}
 }
 
-func BenchmarkEDSRoots(b *testing.B) {
+func BenchmarkEDSRootsWithDefaultTree(b *testing.B) {
 	for i := 32; i < 513; i *= 2 {
 		square, err := newDataSquare(genRandDS(i*2, int(shareSize)), NewDefaultTree, shareSize)
 		if err != nil {
@@ -420,6 +422,54 @@ func BenchmarkEDSRoots(b *testing.B) {
 	}
 }
 
+func BenchmarkEDSRootsWithErasuredNMT(b *testing.B) {
+	ODSSizeByteUpperBound := 512 * 1024 * 1024 // converting 512 MiB to bytes
+	totalNumberOfShares := float64(ODSSizeByteUpperBound) / shareSize
+	// the closest power of 2 of the square root of
+	// the total number of shares
+	oDSShareSizeUpperBound := math.Pow(2, math.Ceil(math.Log2(math.Sqrt(
+		totalNumberOfShares))))
+	namespaceIDSize := 29
+
+	for i := 32; i <= int(oDSShareSizeUpperBound); i *= 2 {
+		// number of shares in the original data square's row/column
+		odsSize := i
+		// number of shares in the extended data square's row/column
+		edsSize := 2 * odsSize
+		// generate an EDS with edsSize X edsSize dimensions in terms of shares.
+		// the generated EDS does not conform to celestia-app specs in terms
+		// of namespace version, also no erasure encoding takes place
+		// yet none of these should impact the benchmarking
+		ds := genRandSortedDS(edsSize, shareSize, namespaceIDSize)
+
+		// a tree constructor for erasured nmt
+		treeConstructor := newErasuredNamespacedMerkleTreeConstructor(uint64(i*2),
+			nmt.NamespaceIDSize(namespaceIDSize), nmt.IgnoreMaxNamespace(true),
+			nmt.InitialCapacity(i*2))
+
+		square, err := newDataSquare(ds, treeConstructor, shareSize)
+		if err != nil {
+			b.Errorf("Failure to create square of size %d: %s", i, err)
+		}
+		// the total size of the ODS in MiB
+		odsSizeMiBytes := i * i * shareSize / (1024 * 1024)
+		// the total size of the EDS in MiB
+		edsSizeMiBytes := 4 * odsSizeMiBytes
+		b.Run(
+			fmt.Sprintf("%dx%dx%d ODS=%dMB, EDS=%dMB", i, i,
+				int(square.chunkSize),
+				odsSizeMiBytes, edsSizeMiBytes),
+			func(b *testing.B) {
+				for n := 0; n < b.N; n++ {
+					square.resetRoots()
+					err := square.computeRoots()
+					assert.NoError(b, err)
+				}
+			},
+		)
+	}
+}
+
 func computeRowProof(ds *dataSquare, x uint, y uint) ([]byte, [][]byte, uint, uint, error) {
 	tree := ds.createTreeFn(Row, x)
 	data := ds.row(x)

extendeddatacrossword_test.go

@@ -439,7 +439,7 @@ func createTestEdsWithNMT(t *testing.T, codec Codec, shareSize, namespaceSize in
 	edsWidth := 4            // number of shares per row/column in the extended data square
 	odsWidth := edsWidth / 2 // number of shares per row/column in the original data square
 
-	eds, err := ComputeExtendedDataSquare(shares, codec, newConstructor(uint64(odsWidth), nmt.NamespaceIDSize(namespaceSize)))
+	eds, err := ComputeExtendedDataSquare(shares, codec, newErasuredNamespacedMerkleTreeConstructor(uint64(odsWidth), nmt.NamespaceIDSize(namespaceSize)))
 	require.NoError(t, err)
 
 	return eds

extendeddatasquare_test.go

@@ -6,6 +6,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"reflect"
+	"sort"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -347,6 +348,18 @@ func genRandDS(width int, chunkSize int) [][]byte {
 	return ds
 }
 
+func genRandSortedDS(width int, chunkSize int, namespaceSize int) [][]byte {
+	ds := genRandDS(width, chunkSize)
+
+	// Sort the shares in the square based on their namespace
+	sort.Slice(ds, func(i, j int) bool {
+		// Compare only the first  namespaceSize bytes
+		return bytes.Compare(ds[i][:namespaceSize], ds[j][:namespaceSize]) < 0
+	})
+
+	return ds
+}
+
 // TestFlattened_EDS tests that eds.Flattened() returns all the shares in the
 // EDS. This function has the `_EDS` suffix to avoid a name collision with the
 // TestFlattened.

nmtwrapper_test.go

@@ -54,7 +54,9 @@ type nmtTree interface {
 // with an underlying NMT of namespace size `NamespaceSize` and with
 // `ignoreMaxNamespace=true`. axisIndex is the index of the row or column that
 // this tree is committing to. squareSize must be greater than zero.
-func newErasuredNamespacedMerkleTree(squareSize uint64, axisIndex uint, options ...nmt.Option) erasuredNamespacedMerkleTree {
+func newErasuredNamespacedMerkleTree(squareSize uint64, axisIndex uint,
+	options ...nmt.Option,
+) erasuredNamespacedMerkleTree {
 	if squareSize == 0 {
 		panic("cannot create a erasuredNamespacedMerkleTree of squareSize == 0")
 	}
@@ -73,10 +75,10 @@ type constructor struct {
 	opts       []nmt.Option
 }
 
-// newConstructor creates a tree constructor function as required by rsmt2d to
+// newErasuredNamespacedMerkleTreeConstructor creates a tree constructor function as required by rsmt2d to
 // calculate the data root. It creates that tree using the
 // erasuredNamespacedMerkleTree.
-func newConstructor(squareSize uint64, opts ...nmt.Option) TreeConstructorFn {
+func newErasuredNamespacedMerkleTreeConstructor(squareSize uint64, opts ...nmt.Option) TreeConstructorFn {
 	return constructor{
 		squareSize: squareSize,
 		opts:       opts,