Add FindSubtrees (transparency-dev#250)

Add a FindSubtrees func which implements an extended version of the MTC algorithm for determining the one or two subtrees which efficiently cover a range of entries.

Author: AlCutter

proof/proof.go

@@ -268,6 +268,45 @@ func reverse(ids []compact.NodeID) {
 	}
 }
 
+// Subtree represents a valid Merkle subtree [Start, End).
+type Subtree struct {
+	Start, End uint64
+}
+
+// FindSubtrees returns one or two subtrees that efficiently cover [start, end).
+//
+// If the provided [start, end) range is already a valid subtree, then that subtree is returned directly.
+// Otherwise, this function continues by applying the "Selecting Two Subtrees" procedure
+// from Section 4.5.1 of draft-ietf-plants-merkle-tree-certs.
+//
+// Note that:
+//   - If the provided [start, end) range is already a valid subtree, then it is returned as the only entry in the slice.
+//   - The 2nd subtree, if present, is adjacent to the first, and may not be a perfect subtree.
+//   - The returned subtree(s) fully cover the [start, end) range.
+//   - There are no "extra" entries covered past end, but there may be covered entries prior to start.
+//   - The number of entries covered before start is always less than half the size of the first returned subtree.
+func FindSubtrees(start, end uint64) ([]Subtree, error) {
+	if start >= end {
+		return nil, fmt.Errorf("start %d must be strictly less than end %d", start, end)
+	}
+	if end-start == 1 || isSubtreeValid(start, end) == nil {
+		return []Subtree{{Start: start, End: end}}, nil
+	}
+	last := end - 1
+	// Find where start and last's tree paths diverge.
+	split := bits.Len64(start^last) - 1
+	mask := (uint64(1) << split) - 1
+	mid := last & ^mask
+
+	// Maximize the left endpoint.
+	leftSplit := bits.Len64(^start & mask)
+	leftStart := start & ^((uint64(1) << leftSplit) - 1)
+	return []Subtree{
+		{Start: leftStart, End: mid},
+		{Start: mid, End: end},
+	}, nil
+}
+
 // isSubtreeValid returns whether a subtree covers a valid range.
 // A subtree is valid if there exist a parent tree node to:
 // - all the subtree nodes

proof/proof_test.go

@@ -820,3 +820,41 @@ func inclusion(t *testing.T, index, size uint64) Nodes {
 	}
 	return n
 }
+
+func TestFindSubtrees(t *testing.T) {
+	for _, tc := range []struct {
+		start, end uint64
+		want       []Subtree
+		wantErr    bool
+	}{
+		// Already-valid subtrees are returned as-is.
+		// Single entry subtrees:
+		{start: 0, end: 1, want: []Subtree{{Start: 0, End: 1}}},
+		{start: 3, end: 4, want: []Subtree{{Start: 3, End: 4}}},
+		// Perfectly aligned subtrees:
+		{start: 4, end: 6, want: []Subtree{{Start: 4, End: 6}}},
+		{start: 16, end: 32, want: []Subtree{{Start: 16, End: 32}}},
+		// Non-perfect trees are split into two:
+		{start: 5, end: 13, want: []Subtree{{Start: 4, End: 8}, {Start: 8, End: 13}}},
+		{start: 7, end: 9, want: []Subtree{{Start: 7, End: 8}, {Start: 8, End: 9}}},
+		// Invalid inputs:
+		{start: 5, end: 5, wantErr: true},
+		{start: 6, end: 5, wantErr: true},
+	} {
+		t.Run(fmt.Sprintf("%d:%d", tc.start, tc.end), func(t *testing.T) {
+			got, err := FindSubtrees(tc.start, tc.end)
+			if tc.wantErr {
+				if err == nil {
+					t.Fatal("expected error, got nil")
+				}
+				return
+			}
+			if err != nil {
+				t.Fatalf("FindSubtrees: %v", err)
+			}
+			if diff := cmp.Diff(tc.want, got); diff != "" {
+				t.Errorf("FindSubtrees mismatch (-want +got):\n%s", diff)
+			}
+		})
+	}
+}

testonly/vectors_test.go

@@ -19,6 +19,8 @@ import (
 	"fmt"
 	"io"
 	"testing"
+
+	"github.com/transparency-dev/merkle/proof"
 )
 
 // These tests reproduce the accumulated test vectors from the "Subtree Test
@@ -123,3 +125,31 @@ func TestSubtreeConsistencyProofVectors(t *testing.T) {
 		t.Errorf("subtree consistency proof vector = %s, want %s", got, want)
 	}
 }
+
+func TestSubtreeCoveringVectors(t *testing.T) {
+	h := sha256.New()
+	for end := uint64(1); end <= subtreeVectorMax; end++ {
+		for start := range end {
+			if err := isSubtreeValid(start, end); err == nil {
+				if _, err := fmt.Fprintf(h, "[%d, %d)\n", start, end); err != nil {
+					t.Fatalf("fmt.Fprintf: %v", err)
+				}
+			} else {
+				subtrees, err := proof.FindSubtrees(start, end)
+				if err != nil {
+					t.Fatalf("FindSubtrees(%d, %d): %v", start, end, err)
+				}
+				if l := len(subtrees); l != 2 {
+					t.Fatalf("FindSubtrees(%d, %d) returned unexpected number of subtrees: %d", start, end, l)
+				}
+				if _, err := fmt.Fprintf(h, "[%d, %d) [%d, %d)\n", subtrees[0].Start, subtrees[0].End, subtrees[1].Start, subtrees[1].End); err != nil {
+					t.Fatalf("fmt.Fprintf: %v", err)
+				}
+			}
+		}
+	}
+	const want = "e0aecb912a10c57d753b6ecc64db73217f9bc4ed10fcb4e9062be3b6fbe1ebfd"
+	if got := fmt.Sprintf("%x", h.Sum(nil)); got != want {
+		t.Errorf("subtree covering vector = %s, want %s", got, want)
+	}
+}