Merge pull request celestiaorg#671 from vgonkivs/add_tree_method_to_nmt

Author: Wondertan

CELESTIA_CHANGELOG.md

@@ -3,6 +3,7 @@
 ## Unreleased Changes
 - Add .github/workflows/pending-changelog-checker.yml to enforce updating this file (`CELESTIA_CHANGELOG.md`) on commits to `master` branch.
 - [proto: moves proto files to](https://github.com/celestiaorg/celestia-core/pull/608) [spec](https://github.com/celestiaorg/spec) @sweexordious
+- [nmt_wrapper: extend nmt wrapper with tree method](https://github.com/celestiaorg/celestia-core/pull/671) @vgonkivs
 
 ## v0.1
 

pkg/wrapper/nmt_wrapper.go

@@ -72,3 +72,8 @@ func (w *ErasuredNamespacedMerkleTree) Push(data []byte, idx rsmt2d.SquareIndex)
 func (w *ErasuredNamespacedMerkleTree) Root() []byte {
 	return w.tree.Root()
 }
+
+// Tree returns the underlying NamespacedMerkleTree
+func (w *ErasuredNamespacedMerkleTree) Tree() *nmt.NamespacedMerkleTree {
+	return w.tree
+}

pkg/wrapper/nmt_wrapper_test.go

@@ -108,6 +108,21 @@ func TestExtendedDataSquare(t *testing.T) {
 	assert.NoError(t, err)
 }
 
+func TestErasuredNamespacedMerkleTree(t *testing.T) {
+	// check that the Tree() returns exact underlying nmt tree
+	size := 8
+	data := generateRandNamespacedRawData(size, consts.NamespaceSize, consts.MsgShareSize)
+	n := NewErasuredNamespacedMerkleTree(uint64(size))
+	tree := n.Constructor()
+
+	for i, d := range data {
+		tree.Push(d, rsmt2d.SquareIndex{Axis: uint(0), Cell: uint(i)})
+	}
+
+	assert.Equal(t, n.Tree(), n.tree)
+	assert.Equal(t, n.Tree().Root(), n.tree.Root())
+}
+
 // generateErasuredData produces a slice that is twice as long as it erasures
 // the data
 func generateErasuredData(t *testing.T, numLeaves int, codec rsmt2d.Codec) [][]byte {