feat: Cache and traverse nmt sub tree roots

go.mod

@@ -22,7 +22,10 @@ require (
 	google.golang.org/grpc v1.47.0
 )
 
-require github.com/cosmos/cosmos-sdk/errors v1.0.0-beta.3
+require (
+	github.com/celestiaorg/rsmt2d v0.6.0
+	github.com/cosmos/cosmos-sdk/errors v1.0.0-beta.3
+)
 
 require (
 	cloud.google.com/go v0.100.2 // indirect
@@ -40,7 +43,6 @@ require (
 	github.com/btcsuite/btcd v0.22.1 // indirect
 	github.com/celestiaorg/go-leopard v0.1.0 // indirect
 	github.com/celestiaorg/merkletree v0.0.0-20210714075610-a84dc3ddbbe4 // indirect
-	github.com/celestiaorg/rsmt2d v0.6.0 // indirect
 	github.com/cenkalti/backoff/v4 v4.1.1 // indirect
 	github.com/cespare/xxhash v1.1.0 // indirect
 	github.com/cespare/xxhash/v2 v2.1.2 // indirect

pkg/inclusion/sub_tree_root.go

@@ -0,0 +1,125 @@
+package inclusion
+
+import (
+	"fmt"
+
+	"github.com/celestiaorg/nmt"
+	"github.com/celestiaorg/rsmt2d"
+	"github.com/tendermint/tendermint/pkg/da"
+	"github.com/tendermint/tendermint/pkg/wrapper"
+)
+
+// TODO optimize https://github.com/celestiaorg/nmt/blob/e679661c6776d8a694f4f7c423c2e2eccb6c5aaa/subrootpaths.go#L15-L28
+
+// subTreeRootCacher keep track of all the inner nodes of an nmt using a simple
+// map. Note: this cacher does not cache individual leaves or their hashes, only
+// inner nodes.
+type subTreeRootCacher struct {
+	cache map[string][2]string
+}
+
+func newSubTreeRootCacher() *subTreeRootCacher {
+	return &subTreeRootCacher{cache: make(map[string][2]string)}
+}
+
+// Visit fullfills the nmt.VisitorNode function definition. It stores each inner
+// node in a simple map, which can later be used to walk the tree. This function
+// is called by the nmt when calculating the root.
+func (strc *subTreeRootCacher) Visit(hash []byte, children ...[]byte) {
+	switch len(children) {
+	case 2:
+		strc.cache[string(hash)] = [2]string{string(children[0]), string(children[1])}
+	case 1:
+		return
+	default:
+		panic("unexpected visit")
+	}
+}
+
+// walk recursively traverses the subTreeRootCacher's internal tree by using the
+// provided sub tree root and path. The provided path should be a []bool, false
+// indicating that the first child node (left most node) should be used to find
+// the next path, and the true indicating that the second (right) should be used.
+// walk throws an error if the sub tree cannot be found.
+func (strc subTreeRootCacher) walk(root []byte, path []bool) ([]byte, error) {
+	// return if we've reached the end of the path
+	if len(path) == 0 {
+		return root, nil
+	}
+	// try to lookup the provided sub root
+	children, has := strc.cache[string(root)]
+	if !has {
+		// note: we might want to consider panicing here
+		return nil, fmt.Errorf("did not find sub tree root: %v", root)
+	}
+
+	// continue to traverse the tree by recursively calling this function on the next root
+	switch path[0] {
+	// walk left
+	case false:
+		return strc.walk([]byte(children[0]), path[1:])
+	// walk right
+	case true:
+		return strc.walk([]byte(children[1]), path[1:])
+	default:
+		// this is unreachable code, but the compiler doesn't recognize this somehow
+		panic("bool other than true or false, computers were a mistake, everything is a lie, math is fake.")
+	}
+}
+
+// EDSSubTreeRootCacher caches the inner nodes for each row so that we can
+// traverse it later to check for message inclusion. NOTE: Currently this has to
+// use a leaky abstraction (see docs on counter field below), and is not
+// threadsafe, but with a future refactor, we could simply read from rsmt2d and
+// not use the tree constructor which would fix both of these issues.
+type EDSSubTreeRootCacher struct {
+	caches     []*subTreeRootCacher
+	squareSize uint64
+	// counter is used to ignore columns NOTE: this is a leaky abstraction that
+	// we make because rsmt2d is used to generate the roots for us, so we have
+	// to assume that it will generate a row root every other tree contructed.
+	// This is also one of the reasons this implementation is not thread safe.
+	// Please see note above on a better refactor.
+	counter int
+}
+
+func NewCachedSubtreeCacher(squareSize uint64) *EDSSubTreeRootCacher {
+	return &EDSSubTreeRootCacher{
+		caches:     []*subTreeRootCacher{},
+		squareSize: squareSize,
+	}
+}
+
+// Constructor fullfills the rsmt2d.TreeCreatorFn by keeping a pointer to the
+// cache and embedding it as a nmt.NodeVisitor into a new wrapped nmt. I only
+func (stc *EDSSubTreeRootCacher) Constructor() rsmt2d.Tree {
+	// see docs of counter field for more
+	// info. if the counter is even or == 0, then we make the assumption that we
+	// are creating a tree for a row
+	var newTree wrapper.ErasuredNamespacedMerkleTree
+	switch stc.counter % 2 {
+	case 0:
+		strc := newSubTreeRootCacher()
+		stc.caches = append(stc.caches, strc)
+		newTree = wrapper.NewErasuredNamespacedMerkleTree(stc.squareSize, nmt.NodeVisitor(strc.Visit))
+	default:
+		newTree = wrapper.NewErasuredNamespacedMerkleTree(stc.squareSize)
+	}
+
+	stc.counter++
+	return &newTree
+}
+
+// GetSubTreeRoot traverses the nmt of the selected row and returns the
+// subtree root. An error is thrown if the subtree cannot be found.
+func (stc *EDSSubTreeRootCacher) GetSubTreeRoot(dah da.DataAvailabilityHeader, row int, path []bool) ([]byte, error) {
+	const unexpectedDAHErr = "data availability header has unexpected number of row roots: expected %d got %d"
+	if len(stc.caches) != len(dah.RowsRoots) {
+		return nil, fmt.Errorf(unexpectedDAHErr, len(stc.caches), len(dah.RowsRoots))
+	}
+	const rowOutOfBoundsErr = "row exceeds range of cache: max %d got %d"
+	if row >= len(stc.caches) {
+		return nil, fmt.Errorf(rowOutOfBoundsErr, len(stc.caches), row)
+	}
+	return stc.caches[row].walk(dah.RowsRoots[row], path)
+}

pkg/inclusion/sub_tree_root_test.go

@@ -0,0 +1,178 @@
+package inclusion
+
+import (
+	"testing"
+
+	"github.com/celestiaorg/celestia-app/testutil/coretestutil"
+	"github.com/celestiaorg/nmt"
+	"github.com/celestiaorg/rsmt2d"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"github.com/tendermint/tendermint/pkg/consts"
+	"github.com/tendermint/tendermint/pkg/da"
+	"github.com/tendermint/tendermint/pkg/wrapper"
+)
+
+func TestWalkCachedSubTreeRoot(t *testing.T) {
+	// create the first main tree
+	strc := newSubTreeRootCacher()
+	oss := uint64(8)
+	tr := wrapper.NewErasuredNamespacedMerkleTree(oss, nmt.NodeVisitor(strc.Visit))
+	d := []byte{0, 0, 0, 0, 0, 0, 0, 1, 1, 2, 3, 4, 5, 6, 7, 8}
+	for i := 0; i < 8; i++ {
+		tr.Push(d, rsmt2d.SquareIndex{
+			Axis: uint(rsmt2d.Row),
+			Cell: uint(i),
+		})
+	}
+	highestRoot := tr.Root()
+
+	// create a small sub tree
+	subTree1 := wrapper.NewErasuredNamespacedMerkleTree(oss)
+	for i := 0; i < 2; i++ {
+		subTree1.Push(d, rsmt2d.SquareIndex{
+			Axis: uint(rsmt2d.Row),
+			Cell: uint(i),
+		})
+	}
+	shortSTR := subTree1.Root()
+
+	// create a larger sub tree root
+	subTree2 := wrapper.NewErasuredNamespacedMerkleTree(oss)
+	for i := 0; i < 4; i++ {
+		subTree2.Push(d, rsmt2d.SquareIndex{
+			Axis: uint(rsmt2d.Row),
+			Cell: uint(i),
+		})
+	}
+	tallSTR := subTree2.Root()
+
+	type test struct {
+		name        string
+		path        []bool
+		subTreeRoot []byte
+		error       string
+	}
+
+	tests := []test{
+		{
+			"left most short sub tree",
+			[]bool{false, false},
+			shortSTR,
+			"",
+		},
+		{
+			"left middle short sub tree",
+			[]bool{false, true},
+			shortSTR,
+			"",
+		},
+		{
+			"right middle short sub tree",
+			[]bool{true, false},
+			shortSTR,
+			"",
+		},
+		{
+			"right most short sub tree",
+			[]bool{true, true},
+			shortSTR,
+			"",
+		},
+		{
+			"left most tall sub tree",
+			[]bool{false},
+			tallSTR,
+			"",
+		},
+		{
+			"right most tall sub tree",
+			[]bool{true},
+			tallSTR,
+			"",
+		},
+	}
+
+	for _, tt := range tests {
+		foundSubRoot, err := strc.walk(highestRoot, tt.path)
+		if tt.error != "" {
+			require.Error(t, err, tt.name)
+			assert.Contains(t, err.Error(), tt.error, tt.name)
+			continue
+		}
+
+		require.NoError(t, err)
+		require.Equal(t, tt.subTreeRoot, foundSubRoot, tt.name)
+	}
+}
+
+func TestEDSSubRootCacher(t *testing.T) {
+	oss := uint64(8)
+	d := coretestutil.GenerateRandNamespacedRawData(uint32(oss*oss), consts.NamespaceSize, consts.ShareSize-consts.NamespaceSize)
+	stc := NewCachedSubtreeCacher(oss)
+
+	eds, err := rsmt2d.ComputeExtendedDataSquare(d, consts.DefaultCodec(), stc.Constructor)
+	require.NoError(t, err)
+
+	dah := da.NewDataAvailabilityHeader(eds)
+
+	for i := range dah.RowsRoots[:oss] {
+		expectedSubTreeRoots := calculateSubTreeRoots(eds.Row(uint(i))[:oss], 2)
+		require.NotNil(t, expectedSubTreeRoots)
+		// note: the depth is one greater than expected because we're dividing
+		// the row in half when we calculate the expected roots.
+		result, err := stc.GetSubTreeRoot(dah, i, []bool{false, false, false})
+		require.NoError(t, err)
+		assert.Equal(t, expectedSubTreeRoots[0], result)
+	}
+}
+
+// calculateSubTreeRoots generates the subtree roots for a given row. If the
+// selected depth is too deep for the tree, nil is returned. It relies on
+// passing a row whose length is a power of 2 and assumes that the row is
+// **NOT** extended since calculating subtree root for erasure data using the
+// nmt wrapper makes this difficult.
+func calculateSubTreeRoots(row [][]byte, depth int) [][]byte {
+	subLeafRange := len(row)
+	for i := 0; i < depth; i++ {
+		subLeafRange = subLeafRange / 2
+	}
+
+	if subLeafRange == 0 || subLeafRange%2 != 0 {
+		return nil
+	}
+
+	count := len(row) / subLeafRange
+	subTreeRoots := make([][]byte, count)
+	chunks := chunkSlice(row, subLeafRange)
+	for i, rowChunk := range chunks {
+		tr := wrapper.NewErasuredNamespacedMerkleTree(uint64(len(row)))
+		for j, r := range rowChunk {
+			c := (i * subLeafRange) + j
+			tr.Push(r, rsmt2d.SquareIndex{
+				Axis: uint(rsmt2d.Row),
+				Cell: uint(c),
+			})
+		}
+		subTreeRoots[i] = tr.Root()
+	}
+
+	return subTreeRoots
+}
+
+func chunkSlice(slice [][]byte, chunkSize int) [][][]byte {
+	var chunks [][][]byte
+	for i := 0; i < len(slice); i += chunkSize {
+		end := i + chunkSize
+
+		// necessary check to avoid slicing beyond
+		// slice capacity
+		if end > len(slice) {
+			end = len(slice)
+		}
+
+		chunks = append(chunks, slice[i:end])
+	}
+
+	return chunks
+}

testutil/coretestutil/core.go

@@ -0,0 +1,28 @@
+package coretestutil
+
+import (
+	"bytes"
+	"math/rand"
+	"sort"
+)
+
+func GenerateRandNamespacedRawData(total, nidSize, leafSize uint32) [][]byte {
+	data := make([][]byte, total)
+	for i := uint32(0); i < total; i++ {
+		nid := make([]byte, nidSize)
+		rand.Read(nid)
+		data[i] = nid
+	}
+	sortByteArrays(data)
+	for i := uint32(0); i < total; i++ {
+		d := make([]byte, leafSize)
+		rand.Read(d)
+		data[i] = append(data[i], d...)
+	}
+
+	return data
+}
+
+func sortByteArrays(src [][]byte) {
+	sort.Slice(src, func(i, j int) bool { return bytes.Compare(src[i], src[j]) < 0 })
+}

x/payment/types/payfordata.go

@@ -1,17 +1,17 @@
 package types
 
 import (
-	"crypto/sha256"
 	"fmt"
 	"math/bits"
 
-	"github.com/celestiaorg/nmt"
+	"github.com/celestiaorg/rsmt2d"
 	sdkclient "github.com/cosmos/cosmos-sdk/client"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/cosmos/cosmos-sdk/types/tx/signing"
 	authsigning "github.com/cosmos/cosmos-sdk/x/auth/signing"
 	"github.com/tendermint/tendermint/crypto/merkle"
 	"github.com/tendermint/tendermint/pkg/consts"
+	"github.com/tendermint/tendermint/pkg/wrapper"
 	coretypes "github.com/tendermint/tendermint/types"
 )
 
@@ -145,13 +145,13 @@ func CreateCommitment(k uint64, namespace, message []byte) ([]byte, error) {
 	subTreeRoots := make([][]byte, len(leafSets))
 	for i, set := range leafSets {
 		// create the nmt todo(evan) use nmt wrapper
-		tree := nmt.New(sha256.New(), nmt.NamespaceIDSize(NamespaceIDSize))
+		tree := wrapper.NewErasuredNamespacedMerkleTree(k)
 		for _, leaf := range set {
 			nsLeaf := append(make([]byte, 0), append(namespace, leaf...)...)
-			err := tree.Push(nsLeaf)
-			if err != nil {
-				return nil, err
-			}
+			// note: we're not concerned about adding the correct namespace to
+			// erasure data since we're only dealing with original square data,
+			// so we can push to the wrapped nmt using Axis and Cell == 0
+			tree.Push(nsLeaf, rsmt2d.SquareIndex{Axis: 0, Cell: 0})
 		}
 		// add the root
 		subTreeRoots[i] = tree.Root()