Repository: celestiaorg/nmt (go)

Issue:
feat: compute leaf hashes immediately after data items are pushed
# Problem

When adding data items to the tree, their hash is not directly calculated. Instead, every time we need to access the leaf hashes, we have to compare the number of data items added so far (which can be found in the [`leaves`](url) field of [`NamespaceMerkleTree`](https://github.com/celestiaorg/nmt/blob/master/nmt.go#L85)) with the leaf hashes (which are stored in the [`leafHashes`](https://github.com/celestiaorg/nmt/blob/master/nmt.go#L98) field of the same struct). For example, see [the computeRoot() method](https://github.com/celestiaorg/nmt/blob/master/nmt.go#L427) and [computeLeafHashesIfNecessary()  function](https://github.com/celestiaorg/nmt/blob/fa98a59cbf582b3a0880ba0222a49e3bf5032f6d/nmt.go#L518).
This can make debugging and code maintenance more challenging.

# Acceptance criteria 
To simplify the process, it would be better to hash the data items immediately after they are pushed. We can address this issue by implementing this feature in the `Push` method and refactoring the other parts of the code as necessary.

Pull request description:
feat:  compute leaf hashes in the Push method
## Overview
Closes #108  and https://github.com/celestiaorg/nmt/issues/143

## Checklist

- [x] New and updated code has appropriate documentation
- [x] New and updated code has new and/or updated testing
- [x] Required CI checks are passing
- [x] Visual proof for any user facing features like CLI or documentation updates
- [x] Linked issues closed with keywords


Patch:
diff --git a/hasher.go b/hasher.go
index cbcbf2c..0c2bf59 100644
--- a/hasher.go
+++ b/hasher.go
@@ -183,6 +183,16 @@ func (n *Hasher) HashLeaf(ndata []byte) ([]byte, error) {
 	return nameSpacedHash, nil
 }
 
+// MustHashLeaf is a wrapper around HashLeaf that panics if an error is
+// encountered. The ndata must be a valid leaf node.
+func (n *Hasher) MustHashLeaf(ndata []byte) []byte {
+	res, err := n.HashLeaf(ndata)
+	if err != nil {
+		panic(err)
+	}
+	return res
+}
+
 // ValidateNodeFormat checks whether the supplied node conforms to the
 // namespaced hash format and returns an error if it does not. Specifically, it returns ErrInvalidNodeLen if the length of the node is less than the 2*namespace length which indicates it does not match the namespaced hash format.
 func (n *Hasher) ValidateNodeFormat(node []byte) (err error) {
diff --git a/nmt.go b/nmt.go
index 449eb8a..4b01b51 100644
--- a/nmt.go
+++ b/nmt.go
@@ -167,9 +167,6 @@ func (n *NamespacedMerkleTree) Prove(index int) (Proof, error) {
 // then ProveRange returns an ErrInvalidRange error. Any errors rather than ErrInvalidRange are irrecoverable and indicate an illegal state of the tree (n).
 func (n *NamespacedMerkleTree) ProveRange(start, end int) (Proof, error) {
 	isMaxNsIgnored := n.treeHasher.IsMaxNamespaceIDIgnored()
-	if err := n.computeLeafHashesIfNecessary(); err != nil {
-		return Proof{}, err // this never happens
-	}
 	// TODO: store nodes and re-use the hashes instead recomputing parts of the
 	// tree here
 	if start < 0 || start >= end || end > len(n.leafHashes) {
@@ -234,9 +231,6 @@ func (n *NamespacedMerkleTree) ProveNamespace(nID namespace.ID) (Proof, error) {
 	// case 3) At this point we either found leaves with the namespace nID in
 	// the tree or calculated the range it would be in (to generate a proof of
 	// absence and to return the corresponding leaf hashes).
-	if err := n.computeLeafHashesIfNecessary(); err != nil {
-		return Proof{}, err
-	}
 
 	proof, err := n.buildRangeProof(proofStart, proofEnd)
 	if err != nil {
@@ -418,8 +412,15 @@ func (n *NamespacedMerkleTree) Push(namespacedData namespace.PrefixedData) error
 		return err
 	}
 
+	// compute the leaf hash
+	res, err := n.treeHasher.HashLeaf(namespacedData)
+	if err != nil {
+		return err
+	}
+
 	// update relevant "caches":
 	n.leaves = append(n.leaves, namespacedData)
+	n.leafHashes = append(n.leafHashes, res)
 	n.updateNamespaceRanges()
 	n.updateMinMaxID(nID)
 	n.rawRoot = nil
@@ -472,13 +473,8 @@ func (n *NamespacedMerkleTree) computeRoot(start, end int) ([]byte, error) {
 		n.visit(rootHash)
 		return rootHash, nil
 	case 1:
-		leafHash, err := n.treeHasher.HashLeaf(n.leaves[start])
-		if err != nil { // this should never happen since leaves are added through the Push method, during which leaves formats are validated to make sure they are hashable.
-			return nil, fmt.Errorf("failed to hash leaf: %w", err)
-		}
-		if len(n.leafHashes) < len(n.leaves) {
-			n.leafHashes = append(n.leafHashes, leafHash)
-		}
+		leafHash := make([]byte, len(n.leafHashes[start]))
+		copy(leafHash, n.leafHashes[start])
 		n.visit(leafHash, n.leaves[start])
 		return leafHash, nil
 	default:
@@ -577,24 +573,6 @@ func (n *NamespacedMerkleTree) updateMinMaxID(id namespace.ID) {
 	}
 }
 
-// computes the leaf hashes if not already done in a previous call of
-// NamespacedMerkleTree.Root()
-// Any errors return by this method is irrecoverable and indicate an illegal state of the tree (n).
-func (n *NamespacedMerkleTree) computeLeafHashesIfNecessary() error {
-	// check whether all the hash of all the existing leaves are available
-	if len(n.leafHashes) < len(n.leaves) {
-		n.leafHashes = make([][]byte, len(n.leaves))
-		for i, leaf := range n.leaves {
-			res, err := n.treeHasher.HashLeaf(leaf)
-			if err != nil { // should never happen since the validity of leaves is checked in the Push method
-				return err
-			}
-			n.leafHashes[i] = res
-		}
-	}
-	return nil
-}
-
 type leafRange struct {
 	// start and end denote the indices of a leaf in the tree. start ranges from
 	// 0 up to the total number of leaves minus 1 end ranges from 1 up to the


Test patch:
diff --git a/hasher_test.go b/hasher_test.go
index 0e70d51..9c9a703 100644
--- a/hasher_test.go
+++ b/hasher_test.go
@@ -619,3 +619,29 @@ func TestValidateNodes(t *testing.T) {
 		})
 	}
 }
+
+// Test_MustHashLeaf_panic checks that the MustHashLeaf method panics only on invalid inputs.
+func Test_MustHashLeaf_Panic(t *testing.T) {
+	hasher := NewNmtHasher(sha256.New(), 2, false)
+	tests := []struct {
+		name      string
+		leaf      []byte
+		wantPanic bool
+	}{
+		{"valid leaf length", []byte{0, 0}, false},
+		{"invalid leaf length", []byte{0}, true},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if tt.wantPanic {
+				assert.Panics(t, func() {
+					hasher.MustHashLeaf(tt.leaf)
+				})
+			} else {
+				assert.NotPanics(t, func() {
+					hasher.MustHashLeaf(tt.leaf)
+				})
+			}
+		})
+	}
+}
diff --git a/nmt_test.go b/nmt_test.go
index 8f4820e..02fd9ef 100644
--- a/nmt_test.go
+++ b/nmt_test.go
@@ -892,8 +892,6 @@ func swap(slice [][]byte, i int, j int) {
 func Test_buildRangeProof_Err(t *testing.T) {
 	// create a nmt, 8 leaves namespaced sequentially from 1-8
 	treeWithCorruptLeafHash := exampleTreeWithEightLeaves()
-	err := treeWithCorruptLeafHash.computeLeafHashesIfNecessary()
-	require.NoError(t, err)
 	// corrupt a leaf hash
 	treeWithCorruptLeafHash.leafHashes[4] = treeWithCorruptLeafHash.leafHashes[4][:treeWithCorruptLeafHash.NamespaceSize()]
 
@@ -901,8 +899,7 @@ func Test_buildRangeProof_Err(t *testing.T) {
 	treeWithUnorderedLeafHashes := exampleTreeWithEightLeaves()
 	// swap the positions of the 4th and 5th leaves
 	swap(treeWithUnorderedLeafHashes.leaves, 4, 5)
-	err = treeWithUnorderedLeafHashes.computeLeafHashesIfNecessary()
-	require.NoError(t, err)
+	swap(treeWithUnorderedLeafHashes.leafHashes, 4, 5)
 
 	tests := []struct {
 		name                 string
@@ -931,15 +928,8 @@ func Test_buildRangeProof_Err(t *testing.T) {
 
 // Test_ProveRange_Err tests that ProveRange returns an error when the underlying tree has an invalid state e.g., leaves are not ordered by namespace ID or a leaf hash is corrupted.
 func Test_ProveRange_Err(t *testing.T) {
-	// create an NMT with 8 sequentially namespaced leaves, numbered from 1 to 8.
-	treeWithCorruptLeaf := exampleTreeWithEightLeaves()
-	// corrupt a leaf
-	treeWithCorruptLeaf.leaves[4] = treeWithCorruptLeaf.leaves[4][:treeWithCorruptLeaf.NamespaceSize()-1]
-
 	// create an NMT with 8 sequentially namespaced leaves, numbered from 1 to 8.
 	treeWithCorruptLeafHash := exampleTreeWithEightLeaves()
-	err := treeWithCorruptLeafHash.computeLeafHashesIfNecessary()
-	require.NoError(t, err)
 	// corrupt a leaf hash
 	treeWithCorruptLeafHash.leafHashes[4] = treeWithCorruptLeafHash.leafHashes[4][:treeWithCorruptLeafHash.NamespaceSize()]
 
@@ -947,8 +937,7 @@ func Test_ProveRange_Err(t *testing.T) {
 	treeWithUnorderedLeafHashes := exampleTreeWithEightLeaves()
 	// swap the positions of the 4th and 5th leaves
 	swap(treeWithUnorderedLeafHashes.leaves, 4, 5)
-	err = treeWithUnorderedLeafHashes.computeLeafHashesIfNecessary()
-	require.NoError(t, err)
+	swap(treeWithUnorderedLeafHashes.leafHashes, 4, 5)
 
 	tests := []struct {
 		name                 string
@@ -957,7 +946,6 @@ func Test_ProveRange_Err(t *testing.T) {
 		wantErr              bool
 		errType              error
 	}{
-		{"corrupt leaf", treeWithCorruptLeaf, 4, 5, true, ErrInvalidLeafLen},
 		{"corrupt leaf hash", treeWithCorruptLeafHash, 4, 5, true, ErrInvalidNodeLen},
 		{"unordered leaf hashes: the out of order leaf", treeWithUnorderedLeafHashes, 4, 5, true, ErrUnorderedSiblings},
 		{"unordered leaf hashes: first leaf", treeWithUnorderedLeafHashes, 1, 2, true, ErrUnorderedSiblings}, // for a tree with an unordered set of leaves, the ProveRange method  should produce an error for any input range,
@@ -978,15 +966,8 @@ func Test_ProveRange_Err(t *testing.T) {
 
 // The Test_ProveNamespace_Err function tests that ProveNamespace returns an error when the underlying tree is in an invalid state, such as when the leaves are not ordered by namespace ID or when a leaf hash is corrupt.
 func Test_ProveNamespace_Err(t *testing.T) {
-	// create an NMT with 8 sequentially namespaced leaves, numbered from 1 to 8.
-	treeWithCorruptLeaf := exampleTreeWithEightLeaves()
-	// corrupt a leaf
-	treeWithCorruptLeaf.leaves[4] = treeWithCorruptLeaf.leaves[4][:treeWithCorruptLeaf.NamespaceSize()-1]
-
 	// create an NMT with 8 sequentially namespaced leaves, numbered from 1 to 8.
 	treeWithCorruptLeafHash := exampleTreeWithEightLeaves()
-	err := treeWithCorruptLeafHash.computeLeafHashesIfNecessary()
-	require.NoError(t, err)
 	// corrupt a leaf hash
 	treeWithCorruptLeafHash.leafHashes[4] = treeWithCorruptLeafHash.leafHashes[4][:treeWithCorruptLeafHash.NamespaceSize()]
 
@@ -994,8 +975,7 @@ func Test_ProveNamespace_Err(t *testing.T) {
 	treeWithUnorderedLeafHashes := exampleTreeWithEightLeaves()
 	// swap the positions of the 4th and 5th leaves
 	swap(treeWithUnorderedLeafHashes.leaves, 4, 5)
-	err = treeWithUnorderedLeafHashes.computeLeafHashesIfNecessary()
-	require.NoError(t, err)
+	swap(treeWithUnorderedLeafHashes.leafHashes, 4, 5)
 
 	tests := []struct {
 		name    string
@@ -1004,7 +984,6 @@ func Test_ProveNamespace_Err(t *testing.T) {
 		wantErr bool
 		errType error
 	}{
-		{"corrupt leaf", treeWithCorruptLeaf, namespace.ID{5, 5}, true, ErrInvalidLeafLen},
 		{"corrupt leaf hash", treeWithCorruptLeafHash, namespace.ID{5, 5}, true, ErrInvalidNodeLen},
 		{"unordered leaf hashes: the queried namespace falls in the corrupted range", treeWithUnorderedLeafHashes, namespace.ID{5, 5}, true, ErrUnorderedSiblings},
 		{"unordered leaf hashes: query for the first namespace", treeWithUnorderedLeafHashes, namespace.ID{1, 1}, true, ErrUnorderedSiblings}, // for a tree with an unordered set of leaves,
@@ -1026,14 +1005,15 @@ func Test_ProveNamespace_Err(t *testing.T) {
 // Test_Root_Error tests that the Root method returns an error when the underlying tree is in an invalid state, such as when the leaves are not ordered by namespace ID or when a leaf is corrupt.
 func Test_Root_Error(t *testing.T) {
 	// create an NMT with 8 sequentially namespaced leaves, numbered from 1 to 8.
-	treeWithCorruptLeaf := exampleTreeWithEightLeaves()
-	// corrupt a leaf
-	treeWithCorruptLeaf.leaves[4] = treeWithCorruptLeaf.leaves[4][:treeWithCorruptLeaf.NamespaceSize()-1]
+	treeWithCorruptLeafHash := exampleTreeWithEightLeaves()
+	// corrupt a leaf hash
+	treeWithCorruptLeafHash.leafHashes[4] = treeWithCorruptLeafHash.leafHashes[4][:treeWithCorruptLeafHash.NamespaceSize()-1]
 
 	// create an NMT with 8 sequentially namespaced leaves, numbered from 1 to 8.
 	treeWithUnorderedLeaves := exampleTreeWithEightLeaves()
 	// swap the positions of the 4th and 5th leaves
 	swap(treeWithUnorderedLeaves.leaves, 4, 5)
+	swap(treeWithUnorderedLeaves.leafHashes, 4, 5)
 
 	tests := []struct {
 		name    string
@@ -1041,7 +1021,7 @@ func Test_Root_Error(t *testing.T) {
 		wantErr bool
 		errType error
 	}{
-		{"corrupt leaf hash", treeWithCorruptLeaf, true, ErrInvalidLeafLen},
+		{"corrupt leaf hash", treeWithCorruptLeafHash, true, ErrInvalidNodeLen},
 		{"unordered leaf hashes", treeWithUnorderedLeaves, true, ErrUnorderedSiblings},
 	}
 	for _, tt := range tests {
@@ -1058,14 +1038,15 @@ func Test_Root_Error(t *testing.T) {
 // Test_computeRoot_Error tests that the computeRoot method returns an error when the underlying tree is in an invalid state, such as when the leaves are not ordered by namespace ID or when a leaf is corrupt.
 func Test_computeRoot_Error(t *testing.T) {
 	// create an NMT with 8 sequentially namespaced leaves, numbered from 1 to 8.
-	treeWithCorruptLeaf := exampleTreeWithEightLeaves()
-	// corrupt a leaf
-	treeWithCorruptLeaf.leaves[4] = treeWithCorruptLeaf.leaves[4][:treeWithCorruptLeaf.NamespaceSize()-1]
+	treeWithCorruptLeafHash := exampleTreeWithEightLeaves()
+	// corrupt a leaf hash
+	treeWithCorruptLeafHash.leafHashes[4] = treeWithCorruptLeafHash.leafHashes[4][:treeWithCorruptLeafHash.NamespaceSize()-1]
 
 	// create an NMT with 8 sequentially namespaced leaves, numbered from 1 to 8.
 	treeWithUnorderedLeaves := exampleTreeWithEightLeaves()
 	// swap the positions of the 4th and 5th leaves
 	swap(treeWithUnorderedLeaves.leaves, 4, 5)
+	swap(treeWithUnorderedLeaves.leafHashes, 4, 5)
 
 	tests := []struct {
 		name       string
@@ -1074,10 +1055,9 @@ func Test_computeRoot_Error(t *testing.T) {
 		wantErr    bool
 		errType    error
 	}{
-		{"corrupt leaf: the entire tree", treeWithCorruptLeaf, 0, 7, true, ErrInvalidLeafLen},
-		{"corrupt leaf: the corrupt node", treeWithCorruptLeaf, 4, 5, true, ErrInvalidLeafLen},
-		{"corrupt leaf: from the corrupt node until the end of the tree", treeWithCorruptLeaf, 4, 7, true, ErrInvalidLeafLen},
-		{"corrupt leaf: the corrupt node and the node to its left", treeWithCorruptLeaf, 3, 5, true, ErrInvalidLeafLen},
+		{"corrupt leaf hash: the entire tree", treeWithCorruptLeafHash, 0, 7, true, ErrInvalidNodeLen},
+		{"corrupt leaf: from the corrupt node until the end of the tree", treeWithCorruptLeafHash, 4, 7, true, ErrInvalidNodeLen},
+		{"corrupt leaf: the corrupt node and the node to its left", treeWithCorruptLeafHash, 3, 5, true, ErrInvalidNodeLen},
 		{"unordered leaves: the entire tree", treeWithUnorderedLeaves, 0, 7, true, ErrUnorderedSiblings},
 		{"unordered leaves: the unordered portion", treeWithUnorderedLeaves, 4, 6, true, ErrUnorderedSiblings},
 		{"unordered leaves: a portion of the tree containing the unordered leaves", treeWithUnorderedLeaves, 3, 7, true, ErrUnorderedSiblings},
@@ -1096,14 +1076,15 @@ func Test_computeRoot_Error(t *testing.T) {
 // Test_MinMaxNamespace_Err tests that the MinNamespace and MaxNamespace methods return an error when the underlying tree is in an invalid state, such as when the leaves are not ordered by namespace ID or when a leaf is corrupt.
 func Test_MinMaxNamespace_Err(t *testing.T) {
 	// create an NMT with 8 sequentially namespaced leaves, numbered from 1 to 8.
-	treeWithCorruptLeaf := exampleTreeWithEightLeaves()
-	// corrupt a leaf
-	treeWithCorruptLeaf.leaves[4] = treeWithCorruptLeaf.leaves[4][:treeWithCorruptLeaf.NamespaceSize()-1]
+	treeWithCorruptLeafHash := exampleTreeWithEightLeaves()
+	// corrupt a leaf hash
+	treeWithCorruptLeafHash.leafHashes[4] = treeWithCorruptLeafHash.leafHashes[4][:treeWithCorruptLeafHash.NamespaceSize()-1]
 
 	// create an NMT with 8 sequentially namespaced leaves, numbered from 1 to 8.
 	treeWithUnorderedLeaves := exampleTreeWithEightLeaves()
 	// swap the positions of the 4th and 5th leaves
 	swap(treeWithUnorderedLeaves.leaves, 4, 5)
+	swap(treeWithUnorderedLeaves.leafHashes, 4, 5)
 
 	tests := []struct {
 		name    string
@@ -1111,7 +1092,7 @@ func Test_MinMaxNamespace_Err(t *testing.T) {
 		wantErr bool
 		errType error
 	}{
-		{"corrupt leaf", treeWithCorruptLeaf, true, ErrInvalidLeafLen},
+		{"corrupt leaf hash", treeWithCorruptLeafHash, true, ErrInvalidNodeLen},
 		{"unordered leaves", treeWithUnorderedLeaves, true, ErrUnorderedSiblings},
 	}
 	for _, tt := range tests {
@@ -1130,29 +1111,3 @@ func Test_MinMaxNamespace_Err(t *testing.T) {
 		})
 	}
 }
-
-// Test_computeLeafHashesIfNecessary_err tests that the computeLeafHashesIfNecessary method returns an error when the underlying tree is in an invalid state, such as when a leaf is corrupt.
-func Test_computeLeafHashesIfNecessary_err(t *testing.T) {
-	// create an NMT with 8 sequentially namespaced leaves, numbered from 1 to 8.
-	treeWithCorruptLeaf := exampleTreeWithEightLeaves()
-	// corrupt a leaf
-	treeWithCorruptLeaf.leaves[4] = treeWithCorruptLeaf.leaves[4][:treeWithCorruptLeaf.NamespaceSize()-1]
-
-	tests := []struct {
-		name    string
-		tree    *NamespacedMerkleTree
-		wantErr bool
-		errType error
-	}{
-		{"corrupt leaf", treeWithCorruptLeaf, true, ErrInvalidLeafLen},
-	}
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			err := tt.tree.computeLeafHashesIfNecessary()
-			assert.Equal(t, tt.wantErr, err != nil)
-			if tt.wantErr {
-				assert.True(t, errors.Is(err, tt.errType))
-			}
-		})
-	}
-}
diff --git a/proof_test.go b/proof_test.go
index a06c33d..8adb61c 100644
--- a/proof_test.go
+++ b/proof_test.go
@@ -48,7 +48,6 @@ func TestProof_VerifyNamespace_False(t *testing.T) {
 	leafIndex := 3
 	inclusionProofOfLeafIndex, err := n.buildRangeProof(leafIndex, leafIndex+1)
 	require.NoError(t, err)
-	require.NoError(t, n.computeLeafHashesIfNecessary())
 	leafHash := n.leafHashes[leafIndex] // the only data item with namespace ID = 2 in the constructed tree is at index 3
 	invalidAbsenceProof := NewAbsenceProof(leafIndex, leafIndex+1, inclusionProofOfLeafIndex, leafHash, false)