tree_test

Author: phbnf

testonly/tree_test.go

@@ -17,6 +17,7 @@ package testonly
 import (
 	"bytes"
 	"fmt"
+	"math/bits"
 	"math/rand/v2"
 	"strconv"
 	"testing"
@@ -111,6 +112,36 @@ func TestTreeInclusionProof(t *testing.T) {
 	}
 }
 
+func TestTreeSubtreeInclusionProof(t *testing.T) {
+	test := func(desc string, entries [][]byte) {
+		t.Run(desc, func(t *testing.T) {
+			mt := newTree(entries)
+			end := uint64(len(entries))
+			for start := range end + 1 {
+				for index := start; index < end; index++ {
+					if bc := uint64(1) << bits.Len64(end-start-1); start%bc != 0 {
+						continue
+					}
+					got, err := mt.SubtreeInclusionProof(index, start, end)
+					if err != nil {
+						t.Fatalf("SubtreeInclusionProof(%d, %d, %d): %v", index, start, end, err)
+					}
+					want := refInclusionProof(entries[start:end], index-start, mt.hasher)
+					if diff := cmp.Diff(got, want, cmpopts.EquateEmpty()); diff != "" {
+						t.Fatalf("SubtreeInclusionProof(%d, %d, %d): diff (-got +want)\n%s", index, start, end, diff)
+					}
+				}
+			}
+		})
+	}
+
+	test("generated", genEntries(256))
+	entries := LeafInputs()
+	for end := range len(entries) {
+		test(fmt.Sprintf("golden:%d", end), entries[:end])
+	}
+}
+
 func TestTreeConsistencyProof(t *testing.T) {
 	entries := LeafInputs()
 	mt := newTree(entries)