test: add benchmarks and share commitment test (#1726)

Closes: celestiaorg/celestia-app#1717

---------

Co-authored-by: Rootul P <rootulp@gmail.com>

pkg/square/builder.go

@@ -316,6 +316,29 @@ func (b *Builder) FindTxShareRange(txIndex int) (shares.ShareRange, error) {
 	return shares.ShareRange{Start: start, End: end}, nil
 }
 
+func (b *Builder) GetWrappedPFB(txIndex int) (*coretypes.IndexWrapper, error) {
+	if txIndex < 0 {
+		return nil, fmt.Errorf("txIndex %d must not be negative", txIndex)
+	}
+
+	if txIndex < len(b.txs) {
+		return nil, fmt.Errorf("txIndex %d does not match a pfb", txIndex)
+	}
+
+	if txIndex >= len(b.txs)+len(b.pfbs) {
+		return nil, fmt.Errorf("txIndex %d out of range", txIndex)
+	}
+
+	if !b.done {
+		_, err := b.Export()
+		if err != nil {
+			return nil, fmt.Errorf("building square: %w", err)
+		}
+	}
+
+	return b.pfbs[txIndex-len(b.txs)], nil
+}
+
 func (b *Builder) canFit(shareNum int) bool {
 	return b.currentSize+shareNum <= b.maxCapacity
 }

pkg/square/builder_test.go

@@ -52,6 +52,10 @@ func TestBuilderSquareSizeEstimation(t *testing.T) {
 }
 
 func generateMixedTxs(normalTxCount, pfbCount, pfbSize int) [][]byte {
+	return shuffle(generateOrderedTxs(normalTxCount, pfbCount, pfbSize))
+}
+
+func generateOrderedTxs(normalTxCount, pfbCount, pfbSize int) [][]byte {
 	encCfg := encoding.MakeConfig(app.ModuleEncodingRegisters...)
 	pfbTxs := blobfactory.RandBlobTxs(encCfg.TxConfig.TxEncoder(), pfbCount, pfbSize)
 	normieTxs := blobfactory.GenerateManyRawSendTxs(encCfg.TxConfig, normalTxCount)
@@ -60,7 +64,7 @@ func generateMixedTxs(normalTxCount, pfbCount, pfbSize int) [][]byte {
 		normieTxs...),
 		pfbTxs...,
 	)
-	return shuffle(coretypes.Txs(txs).ToSliceOfBytes())
+	return coretypes.Txs(txs).ToSliceOfBytes()
 }
 
 func shuffle(slice [][]byte) [][]byte {

pkg/square/square_benchmark_test.go

@@ -0,0 +1,34 @@
+package square_test
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/celestiaorg/celestia-app/pkg/appconsts"
+	"github.com/celestiaorg/celestia-app/pkg/square"
+	"github.com/stretchr/testify/require"
+)
+
+func BenchmarkSquareConstruct(b *testing.B) {
+	for _, txCount := range []int{10, 100, 1000} {
+		b.Run(fmt.Sprintf("txCount=%d", txCount), func(b *testing.B) {
+			txs := generateOrderedTxs(txCount/2, txCount/2, 1024)
+			for i := 0; i < b.N; i++ {
+				_, err := square.Construct(txs, appconsts.DefaultMaxSquareSize)
+				require.NoError(b, err)
+			}
+		})
+	}
+}
+
+func BenchmarkSquareBuild(b *testing.B) {
+	for _, txCount := range []int{10, 100, 1000, 10000} {
+		b.Run(fmt.Sprintf("txCount=%d", txCount), func(b *testing.B) {
+			txs := generateMixedTxs(txCount/2, txCount/2, 1024)
+			for i := 0; i < b.N; i++ {
+				_, _, err := square.Build(txs, appconsts.DefaultMaxSquareSize)
+				require.NoError(b, err)
+			}
+		})
+	}
+}

pkg/square/square_test.go

@@ -9,11 +9,14 @@ import (
 	"github.com/celestiaorg/celestia-app/app/encoding"
 	"github.com/celestiaorg/celestia-app/pkg/appconsts"
 	"github.com/celestiaorg/celestia-app/pkg/da"
+	"github.com/celestiaorg/celestia-app/pkg/inclusion"
 	ns "github.com/celestiaorg/celestia-app/pkg/namespace"
 	"github.com/celestiaorg/celestia-app/pkg/shares"
 	"github.com/celestiaorg/celestia-app/pkg/square"
 	"github.com/celestiaorg/celestia-app/test/util/blobfactory"
 	"github.com/celestiaorg/celestia-app/test/util/testfactory"
+	blob "github.com/celestiaorg/celestia-app/x/blob/types"
+	"github.com/celestiaorg/rsmt2d"
 	"github.com/stretchr/testify/require"
 	"github.com/tendermint/tendermint/types"
 	coretypes "github.com/tendermint/tendermint/types"
@@ -362,3 +365,35 @@ func TestSquareBlobShareRange(t *testing.T) {
 	_, err = square.BlobShareRange(txs, 0, 10)
 	require.Error(t, err)
 }
+
+func TestSquareShareCommitments(t *testing.T) {
+	const numTxs = 10
+	txs := generateOrderedTxs(numTxs, numTxs, 5)
+	builder, err := square.NewBuilder(appconsts.DefaultMaxSquareSize, txs...)
+	require.NoError(t, err)
+
+	dataSquare, err := builder.Export()
+	require.NoError(t, err)
+
+	cacher := inclusion.NewSubtreeCacher(dataSquare.Size())
+	eds, err := rsmt2d.ComputeExtendedDataSquare(shares.ToBytes(dataSquare), appconsts.DefaultCodec(), cacher.Constructor)
+	require.NoError(t, err)
+	dah := da.NewDataAvailabilityHeader(eds)
+	decoder := encoding.MakeConfig(app.ModuleEncodingRegisters...).TxConfig.TxDecoder()
+
+	for pfbIndex := 0; pfbIndex < numTxs; pfbIndex++ {
+		wpfb, err := builder.GetWrappedPFB(pfbIndex + numTxs)
+		require.NoError(t, err)
+		tx, err := decoder(wpfb.Tx)
+		require.NoError(t, err)
+
+		pfb, ok := tx.GetMsgs()[0].(*blob.MsgPayForBlobs)
+		require.True(t, ok)
+
+		for blobIndex, shareIndex := range wpfb.ShareIndexes {
+			commitment, err := inclusion.GetCommitment(cacher, dah, int(shareIndex), shares.SparseSharesNeeded(pfb.BlobSizes[blobIndex]))
+			require.NoError(t, err)
+			require.Equal(t, pfb.ShareCommitments[blobIndex], commitment)
+		}
+	}
+}