feat: add support for subtree root verification

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)
+			}
+		})
+	}
+}