prefix: add membership proofs

Author: FiloSottile

prefix/storage.go

@@ -3,19 +3,12 @@ package prefix
 import (
 	"context"
 	"errors"
-	"slices"
 )
 
 type Storage interface {
 	// Load retrieves the node with the given label.
 	Load(ctx context.Context, label Label) (*Node, error)
 
-	// LoadPath loads the siblings of the path to reach the given node
-	// (intuitively, the inclusion proof). If the node is not present, the
-	// sequence stops with what would be its sibling if it were present. The
-	// returned nodes are ordered from the node sibling up to the root's child.
-	LoadPath(ctx context.Context, label Label) ([]*Node, error)
-
 	// Store stores the given nodes. If a node with the same label already
 	// exists, it is replaced. The nodes can be in any order.
 	Store(ctx context.Context, nodes ...*Node) error
@@ -40,29 +33,6 @@ func (s *memoryStorage) Load(ctx context.Context, label Label) (*Node, error) {
 	return nil, ErrNodeNotFound
 }
 
-func (s *memoryStorage) LoadPath(ctx context.Context, label Label) ([]*Node, error) {
-	var nodes []*Node
-	node := s.nodes[RootLabel]
-	for node.Label != label {
-		if !label.HasPrefix(node.Label) {
-			if node.Label != EmptyNodeLabel {
-				nodes = append(nodes, node)
-			}
-			break
-		}
-		switch label.SideOf(node.Label) {
-		case Left:
-			nodes = append(nodes, s.nodes[node.Right])
-			node = s.nodes[node.Left]
-		case Right:
-			nodes = append(nodes, s.nodes[node.Left])
-			node = s.nodes[node.Right]
-		}
-	}
-	slices.Reverse(nodes)
-	return nodes, nil
-}
-
 func (s *memoryStorage) Store(ctx context.Context, nodes ...*Node) error {
 	for _, node := range nodes {
 		s.nodes[node.Label] = node

prefix/tree.go

@@ -18,6 +18,7 @@ import (
 	"encoding/binary"
 	"errors"
 	"fmt"
+	"slices"
 	"strings"
 )
 
@@ -127,10 +128,86 @@ func InitStorage(ctx context.Context, h HashFunc, s Storage) error {
 	return s.Store(ctx, newEmptyNode(h), newRootNode(h))
 }
 
+func (t *Tree) RootHash(ctx context.Context) ([32]byte, error) {
+	root, err := t.s.Load(ctx, RootLabel)
+	if err != nil {
+		return [32]byte{}, fmt.Errorf("failed to load root: %w", err)
+	}
+	return root.Hash, nil
+}
+
+type ProofNode struct {
+	Label Label
+	Hash  [32]byte
+}
+
+func (t *Tree) Lookup(ctx context.Context, label [32]byte) ([]ProofNode, error) {
+	l := Label{256, label}
+	if _, err := t.s.Load(ctx, l); err != nil {
+		return nil, fmt.Errorf("failed to load node %s: %w", l, err)
+	}
+	path, err := loadPath(ctx, t.s, l)
+	if err != nil {
+		return nil, fmt.Errorf("failed to load path for node %s: %w", l, err)
+	}
+	proof := make([]ProofNode, 0, len(path))
+	for _, sibling := range path {
+		proof = append(proof, ProofNode{
+			Label: sibling.Label,
+			Hash:  sibling.Hash,
+		})
+	}
+	return proof, nil
+}
+
+// loadPath loads the siblings of the path to reach the given node
+// (intuitively, the inclusion proof). If the node is not present, the
+// sequence stops with what would be its sibling if it were present. The
+// returned nodes are ordered from the node sibling up to the root's child.
+func loadPath(ctx context.Context, s Storage, label Label) ([]*Node, error) {
+	// If the Storage has a custom implementation of LoadPath, use it.
+	if s, ok := s.(interface {
+		LoadPath(context.Context, Label) ([]*Node, error)
+	}); ok {
+		return s.LoadPath(ctx, label)
+	}
+	var nodes []*Node
+	node, err := s.Load(ctx, RootLabel)
+	if err != nil {
+		return nil, fmt.Errorf("failed to load root: %w", err)
+	}
+	for node.Label != label {
+		if !label.HasPrefix(node.Label) {
+			if node.Label != EmptyNodeLabel {
+				nodes = append(nodes, node)
+			}
+			break
+		}
+		left, err := s.Load(ctx, node.Left)
+		if err != nil {
+			return nil, fmt.Errorf("failed to load left node %s: %w", node.Left, err)
+		}
+		right, err := s.Load(ctx, node.Right)
+		if err != nil {
+			return nil, fmt.Errorf("failed to load left node %s: %w", node.Right, err)
+		}
+		switch label.SideOf(node.Label) {
+		case Left:
+			nodes = append(nodes, right)
+			node = left
+		case Right:
+			nodes = append(nodes, left)
+			node = right
+		}
+	}
+	slices.Reverse(nodes)
+	return nodes, nil
+}
+
 func (t *Tree) Insert(ctx context.Context, label, value [32]byte) error {
 	leaf := newLeaf(t.h, label, value)
 
-	path, err := t.s.LoadPath(ctx, leaf.Label)
+	path, err := loadPath(ctx, t.s, leaf.Label)
 	if err != nil {
 		return err
 	}
@@ -148,3 +225,24 @@ func (t *Tree) Insert(ctx context.Context, label, value [32]byte) error {
 
 	return t.s.Store(ctx, changed...)
 }
+
+func VerifyMembershipProof(h HashFunc, label, value [32]byte, proof []ProofNode, root [32]byte) error {
+	node := newLeaf(h, label, value)
+	for _, sibling := range proof {
+		var err error
+		node, err = newParentNode(h, node, &Node{
+			Label: sibling.Label,
+			Hash:  sibling.Hash,
+		})
+		if err != nil {
+			return fmt.Errorf("failed to compute parent node: %w", err)
+		}
+	}
+	if node.Label != RootLabel {
+		return fmt.Errorf("proof does not lead to root, got %s", node.Label)
+	}
+	if node.Hash != root {
+		return fmt.Errorf("proof does not match root hash, got %x, want %x", node.Hash, root)
+	}
+	return nil
+}

prefix/tree_test.go

@@ -6,6 +6,7 @@ import (
 	"encoding/binary"
 	"encoding/hex"
 	"math/rand/v2"
+	"runtime"
 	"testing"
 
 	"lukechampine.com/blake3"
@@ -48,9 +49,8 @@ func testFullTree(t *testing.T, newStorage func(t *testing.T) Storage) {
 		fatalIfErr(t, tree.Insert(t.Context(), label, value))
 	}
 
-	root, err := store.Load(t.Context(), RootLabel)
+	rootHash, err := tree.RootHash(t.Context())
 	fatalIfErr(t, err)
-	rootHash := root.Hash
 
 	store = newStorage(t)
 	fatalIfErr(t, InitStorage(t.Context(), blake3.Sum256, store))
@@ -63,10 +63,10 @@ func testFullTree(t *testing.T, newStorage func(t *testing.T) Storage) {
 		fatalIfErr(t, tree.Insert(t.Context(), label, value))
 	}
 
-	root, err = store.Load(t.Context(), RootLabel)
+	rootHash1, err := tree.RootHash(t.Context())
 	fatalIfErr(t, err)
-	if root.Hash != rootHash {
-		t.Fatalf("after inserting in reverse order: got %x, want %x", root.Hash, rootHash)
+	if rootHash1 != rootHash {
+		t.Fatalf("after inserting in reverse order: got %x, want %x", rootHash1, rootHash)
 	}
 
 	store = newStorage(t)
@@ -80,35 +80,39 @@ func testFullTree(t *testing.T, newStorage func(t *testing.T) Storage) {
 		fatalIfErr(t, tree.Insert(t.Context(), label, value))
 	}
 
-	root, err = store.Load(t.Context(), RootLabel)
+	rootHash1, err = tree.RootHash(t.Context())
 	fatalIfErr(t, err)
-	if root.Hash != rootHash {
-		t.Fatalf("after inserting in random order: got %x, want %x", root.Hash, rootHash)
+	if rootHash1 != rootHash {
+		t.Fatalf("after inserting in random order: got %x, want %x", rootHash1, rootHash)
 	}
 }
 
 func TestAccumulated(t *testing.T) {
 	testAllStorage(t, testAccumulated)
 }
 func testAccumulated(t *testing.T, newStorage func(t *testing.T) Storage) {
+	if _, ok := newStorage(t).(*prefixsqlite.Storage); ok && testing.Short() {
+		t.Skip("skipping accumulated test for sqlite storage in short mode")
+	}
+
 	source := blake3.New(0, nil).XOF()
 	sink := blake3.New(32, nil)
 
 	for range 100 {
 		store := newStorage(t)
 		fatalIfErr(t, InitStorage(t.Context(), blake3.Sum256, store))
 		tree := NewTree(blake3.Sum256, store)
-		root, err := store.Load(t.Context(), RootLabel)
+		rootHash, err := tree.RootHash(t.Context())
 		fatalIfErr(t, err)
-		sink.Write(root.Hash[:])
+		sink.Write(rootHash[:])
 		for range 1000 {
 			var label, value [32]byte
 			source.Read(label[:])
 			source.Read(value[:])
 			fatalIfErr(t, tree.Insert(t.Context(), label, value))
-			root, err := store.Load(t.Context(), RootLabel)
+			rootHash, err := tree.RootHash(t.Context())
 			fatalIfErr(t, err)
-			sink.Write(root.Hash[:])
+			sink.Write(rootHash[:])
 		}
 	}
 
@@ -119,6 +123,66 @@ func testAccumulated(t *testing.T, newStorage func(t *testing.T) Storage) {
 	}
 }
 
+func TestMemoryUsage(t *testing.T) {
+	if testing.Short() {
+		t.Skip("skipping memory usage test in short mode")
+	}
+
+	store := NewMemoryStorage()
+	fatalIfErr(t, InitStorage(t.Context(), blake3.Sum256, store))
+	tree := NewTree(blake3.Sum256, store)
+
+	runtime.GC()
+	var start runtime.MemStats
+	runtime.ReadMemStats(&start)
+
+	source := blake3.New(0, nil).XOF()
+	for n := range 1000000 {
+		var label, value [32]byte
+		source.Read(label[:])
+		source.Read(value[:])
+		fatalIfErr(t, tree.Insert(t.Context(), label, value))
+
+		switch n + 1 {
+		case 1000, 10000, 100000, 1000000:
+			runtime.GC()
+			var m runtime.MemStats
+			runtime.ReadMemStats(&m)
+			t.Logf("Memory usage after inserting % 8d nodes: % 10d bytes", n+1, int64(m.Alloc)-int64(start.Alloc))
+		}
+	}
+}
+
+func TestMembershipProof(t *testing.T) {
+	testAllStorage(t, testMembershipProof)
+}
+func testMembershipProof(t *testing.T, newStorage func(t *testing.T) Storage) {
+	store := newStorage(t)
+	fatalIfErr(t, InitStorage(t.Context(), blake3.Sum256, store))
+	tree := NewTree(blake3.Sum256, store)
+
+	var entries [][32]byte
+	for _, n := range rand.Perm(100) {
+		var label [32]byte
+		binary.LittleEndian.PutUint16(label[:], uint16(n))
+		value := blake3.Sum256(label[:])
+		fatalIfErr(t, tree.Insert(t.Context(), label, value))
+		entries = append(entries, label)
+
+		rootHash, err := tree.RootHash(t.Context())
+		fatalIfErr(t, err)
+
+		for _, label := range entries {
+			value := blake3.Sum256(label[:])
+			proof, err := tree.Lookup(t.Context(), label)
+			fatalIfErr(t, err)
+			if err := VerifyMembershipProof(blake3.Sum256, label, value, proof, rootHash); err != nil {
+				t.Fatalf("membership proof for %x with %d entries failed: %v", label, len(entries), err)
+			}
+		}
+	}
+}
+
 func fatalIfErr(t *testing.T, err error) {
 	if err != nil {
 		t.Helper()