feat: add support for subtree root verification (#260)

<!--
Please read and fill out this form before submitting your PR.

Please make sure you have reviewed our contributors guide before
submitting your
first PR.
-->

## Overview

Closes #256

I added this change here so that we have a reference implementation of
the algorithm that we will implement in Solidity.

Also, adds a method to generate the subtree roots, which didn't exist
before and will be needed during proof generation in Celestia-node.

The codecov missing coverage complaints are for conditions that are
checked twice in different contexts. So there is no way to bypass the
first check to arrive at the second check. So, I guess they're fine.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

- **New Features**
- Added subtree root computation functionality to the Namespaced Merkle
Tree (NMT).
  - Introduced new validation methods for subtree root inclusion in NMT.
  
- **Tests**
- Added comprehensive tests for subtree root computation and
verification in the Namespaced Merkle Tree.
  - Introduced helper functions for enhanced verification capabilities.
- Added edge case handling for various scenarios in the NMT proof
system.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

nmt.go

@@ -107,9 +107,9 @@ type NamespacedMerkleTree struct {
 
 	// namespaceRanges can be used to efficiently look up the range for an
 	// existing namespace without iterating through the leaves. The map key is
-	// the string representation of a namespace.ID  and the leafRange indicates
+	// the string representation of a namespace.ID  and the LeafRange indicates
 	// the range of the leaves matching that namespace ID in the tree
-	namespaceRanges map[string]leafRange
+	namespaceRanges map[string]LeafRange
 	// minNID is the minimum namespace ID of the leaves
 	minNID namespace.ID
 	// maxNID is the maximum namespace ID of the leaves
@@ -151,7 +151,7 @@ func New(h hash.Hash, setters ...Option) *NamespacedMerkleTree {
 		visit:           opts.NodeVisitor,
 		leaves:          make([][]byte, 0, opts.InitialCapacity),
 		leafHashes:      make([][]byte, 0, opts.InitialCapacity),
-		namespaceRanges: make(map[string]leafRange),
+		namespaceRanges: make(map[string]LeafRange),
 		minNID:          bytes.Repeat([]byte{0xFF}, int(opts.NamespaceIDSize)),
 		maxNID:          bytes.Repeat([]byte{0x00}, int(opts.NamespaceIDSize)),
 	}
@@ -436,7 +436,7 @@ func (n *NamespacedMerkleTree) foundInRange(nID namespace.ID) (found bool, start
 	// This is a faster version of this code snippet:
 	// https://github.com/celestiaorg/celestiaorg-prototype/blob/2aeca6f55ad389b9d68034a0a7038f80a8d2982e/simpleblock.go#L106-L117
 	foundRng, found := n.namespaceRanges[string(nID)]
-	return found, foundRng.start, foundRng.end
+	return found, foundRng.Start, foundRng.End
 }
 
 // NamespaceSize returns the underlying namespace size. Note that all namespaced
@@ -590,14 +590,14 @@ func (n *NamespacedMerkleTree) updateNamespaceRanges() {
 		lastNsStr := string(lastPushed[:n.treeHasher.NamespaceSize()])
 		lastRange, found := n.namespaceRanges[lastNsStr]
 		if !found {
-			n.namespaceRanges[lastNsStr] = leafRange{
-				start: lastIndex,
-				end:   lastIndex + 1,
+			n.namespaceRanges[lastNsStr] = LeafRange{
+				Start: lastIndex,
+				End:   lastIndex + 1,
 			}
 		} else {
-			n.namespaceRanges[lastNsStr] = leafRange{
-				start: lastRange.start,
-				end:   lastRange.end + 1,
+			n.namespaceRanges[lastNsStr] = LeafRange{
+				Start: lastRange.Start,
+				End:   lastRange.End + 1,
 			}
 		}
 	}
@@ -644,11 +644,38 @@ func (n *NamespacedMerkleTree) updateMinMaxID(id namespace.ID) {
 	}
 }
 
-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
-	// total number of leaves end is non-inclusive
-	start, end int
+// ComputeSubtreeRoot takes a leaf range and returns the corresponding subtree root.
+// Also, it requires the start and end range to correctly reference an inner node.
+// The provided range, defined by start and end, is end-exclusive.
+func (n *NamespacedMerkleTree) ComputeSubtreeRoot(start, end int) ([]byte, error) {
+	if start < 0 {
+		return nil, fmt.Errorf("start %d shouldn't be strictly negative", start)
+	}
+	if end <= start {
+		return nil, fmt.Errorf("end %d should be stricly bigger than start %d", end, start)
+	}
+	uStart, err := safeIntToUint(start)
+	if err != nil {
+		return nil, err
+	}
+	uEnd, err := safeIntToUint(end)
+	if err != nil {
+		return nil, err
+	}
+	// check if the provided range correctly references an inner node.
+	// calculates the ideal tree from the provided range, and verifies if it is the same as the range
+	if idealTreeRange := nextSubtreeSize(uint64(uStart), uint64(uEnd)); end-start != idealTreeRange {
+		return nil, fmt.Errorf("the provided range [%d, %d) does not construct a valid subtree root range", start, end)
+	}
+	return n.computeRoot(start, end)
+}
+
+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 total number of leaves.
+	// End is non-inclusive
+	Start, End int
 }
 
 // MinNamespace extracts the minimum namespace ID from a given namespace hash,

nmt_test.go

@@ -862,6 +862,20 @@ func exampleNMT(nidSize int, ignoreMaxNamespace bool, leavesNIDs ...byte) *Names
 	return tree
 }
 
+// exampleNMT2 Replica of exampleNMT except that it uses the namespace IDs in the
+// leaves instead of the index.
+func exampleNMT2(nidSize int, ignoreMaxNamespace bool, leavesNIDs ...byte) *NamespacedMerkleTree {
+	tree := New(sha256.New(), NamespaceIDSize(nidSize), IgnoreMaxNamespace(ignoreMaxNamespace))
+	for _, nid := range leavesNIDs {
+		namespace := bytes.Repeat([]byte{nid}, nidSize)
+		d := append(namespace, []byte(fmt.Sprintf("leaf_%d", nid))...)
+		if err := tree.Push(d); err != nil {
+			panic(fmt.Sprintf("unexpected error: %v", err))
+		}
+	}
+	return tree
+}
+
 func swap(slice [][]byte, i int, j int) {
 	temp := slice[i]
 	slice[i] = slice[j]
@@ -1175,3 +1189,126 @@ func TestForcedOutOfOrderNamespacedMerkleTree(t *testing.T) {
 		assert.NoError(t, err)
 	}
 }
+
+func TestComputeSubtreeRoot(t *testing.T) {
+	n := exampleNMT2(1, true, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15)
+	tests := []struct {
+		start, end   int
+		tree         *NamespacedMerkleTree
+		expectedRoot []byte
+		expectError  bool
+	}{
+		{
+			start: 0,
+			end:   16,
+			tree:  n,
+			expectedRoot: func() []byte {
+				root, err := n.Root()
+				require.NoError(t, err)
+				return root
+			}(),
+		},
+		{
+			start: 0,
+			end:   8,
+			tree:  n,
+			expectedRoot: func() []byte {
+				// because the root of the range [0,8) coincides with the root of this tree
+				root, err := exampleNMT2(1, true, 0, 1, 2, 3, 4, 5, 6, 7).Root()
+				require.NoError(t, err)
+				return root
+			}(),
+		},
+		{
+			start: 8,
+			end:   16,
+			tree:  n,
+			expectedRoot: func() []byte {
+				// because the root of the range [8,16) coincides with the root of this tree
+				root, err := exampleNMT2(1, true, 8, 9, 10, 11, 12, 13, 14, 15).Root()
+				require.NoError(t, err)
+				return root
+			}(),
+		},
+		{
+			start: 8,
+			end:   12,
+			tree:  n,
+			expectedRoot: func() []byte {
+				// because the root of the range [8,12) coincides with the root of this tree
+				root, err := exampleNMT2(1, true, 8, 9, 10, 11).Root()
+				require.NoError(t, err)
+				return root
+			}(),
+		},
+		{
+			start: 4,
+			end:   8,
+			tree:  n,
+			expectedRoot: func() []byte {
+				// because the root of the range [4,8) coincides with the root of this tree
+				root, err := exampleNMT2(1, true, 4, 5, 6, 7).Root()
+				require.NoError(t, err)
+				return root
+			}(),
+		},
+		{
+			start: 4,
+			end:   6,
+			tree:  n,
+			expectedRoot: func() []byte {
+				// because the root of the range [4,6) coincides with the root of this tree
+				root, err := exampleNMT2(1, true, 4, 5).Root()
+				require.NoError(t, err)
+				return root
+			}(),
+		},
+		{
+			start: 4,
+			end:   5,
+			tree:  n,
+			expectedRoot: func() []byte {
+				// because the root of the range [4,5) coincides with the root of this tree
+				root, err := exampleNMT2(1, true, 4).Root()
+				require.NoError(t, err)
+				return root
+			}(),
+		},
+		{ // doesn't correctly reference an inner node
+			start:       2,
+			end:         6,
+			tree:        n,
+			expectError: true,
+		},
+		{
+			start:       -1, // invalid start
+			end:         4,
+			tree:        n,
+			expectError: true,
+		},
+		{
+			start:       4,
+			end:         4, // start == end
+			tree:        n,
+			expectError: true,
+		},
+		{
+			start:       5, // start >= end
+			end:         4,
+			tree:        n,
+			expectError: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(fmt.Sprintf("treeSize=%d,start=%d,end=%d", tt.tree.Size(), tt.start, tt.end), func(t *testing.T) {
+			root, err := tt.tree.ComputeSubtreeRoot(tt.start, tt.end)
+			if tt.expectError {
+				assert.Error(t, err)
+			} else {
+				assert.NoError(t, err)
+				assert.Equal(t, tt.expectedRoot, root)
+			}
+		})
+	}
+}