mpt: add support for prefix scans and proofs

mpt/disk.go

@@ -201,7 +201,7 @@ func memOpen(file1, file2, disk File, op string) (_ Tree, err error) {
 		if err != nil {
 			return nil, err
 		}
-		h := emptyTreeHash()
+		h := emptyTreeHash
 		if err := t.mutate(mem[hdrHash:], h[:]); err != nil {
 			return nil, err
 		}

mpt/dmem.go

@@ -8,6 +8,7 @@ import (
 	"bytes"
 	"encoding/binary"
 	"fmt"
+	"iter"
 )
 
 // hash returns the hash for the given tree node.
@@ -338,44 +339,43 @@ func (t *diskTree) predict(s []node, a addr, pbit int, list []KeyVal) ([]node, [
 	return s, list, nil
 }
 
-// Prove returns a proof of the presence or absence of key in t.
-func (t *diskTree) Prove(key Key) (val Val, ok bool, proof Proof, err error) {
+// Path returns the path proof for key in t.
+func (t *diskTree) Path(key Key) (Proof, error) {
 	t.mmu.RLock()
 	defer t.mmu.RUnlock()
 
 	if t.err != nil {
-		return Val{}, false, nil, t.err
+		return nil, t.err
 	}
 	if t.hdr().dirty() {
-		return Val{}, false, nil, ErrModifiedTree
+		return nil, ErrModifiedTree
 	}
 	root, err := t.node(t.hdr().root())
 	if err != nil {
-		return Val{}, false, nil, err
+		return nil, err
 	}
 	if root == nil {
-		return Val{}, false, Proof{}, nil
+		return Proof{}, nil
 	}
-	return root.prove(t, -1, key)
+	return root.path(t, -1, key)
 }
 
-func (n *diskNode) prove(t *diskTree, pbit int, key Key) (val Val, ok bool, proof Proof, err error) {
+// path returns the path proof for key in the subtree rooted at n.
+// pbit is the parent bit depth, controlling whether n is viewed as a leaf.
+func (n *diskNode) path(t *diskTree, pbit int, key Key) (Proof, error) {
 	nbit := n.bit()
 	if nbit <= pbit {
 		// view n as leaf
 		nkey, nval, err := n.keyVal(t)
 		if err != nil {
-			return Val{}, false, nil, err
-		}
-		if bytes.Equal(nkey, key) {
-			return nval, true, Proof{}, nil
+			return nil, err
 		}
 		var p Proof
 		p = binary.AppendUvarint(p, uint64(len(nkey)))
-		p = append(p, nkey[:]...)
+		p = append(p, nkey...)
 		p = binary.AppendUvarint(p, uint64(len(nval)))
-		p = append(p, nval[:]...)
-		return Val{}, false, p, nil
+		p = append(p, nval...)
+		return p, nil
 	}
 
 	childAddr, sibAddr := n.left(), n.right()
@@ -384,24 +384,141 @@ func (n *diskNode) prove(t *diskTree, pbit int, key Key) (val Val, ok bool, proo
 	}
 	child, err := t.node(childAddr)
 	if err != nil {
-		return Val{}, false, nil, err
+		return nil, err
 	}
 	sib, err := t.node(sibAddr)
 	if err != nil {
-		return Val{}, false, nil, err
+		return nil, err
 	}
 	sibHash, err := sib.hash(t, nbit)
 	if err != nil {
-		return Val{}, false, nil, err
+		return nil, err
 	}
 
-	val, ok, proof, err = child.prove(t, nbit, key)
+	proof, err := child.path(t, nbit, key)
 	if err != nil {
-		return
+		return nil, err
 	}
 	proof = binary.AppendUvarint(proof, uint64(nbit))
 	proof = append(proof, sibHash[:]...)
-	return
+	return proof, nil
+}
+
+// Scan returns an iterator over key-val pairs whose keys start with prefix.
+func (t *diskTree) Scan(prefix []byte) iter.Seq2[KeyVal, error] {
+	return func(yield func(KeyVal, error) bool) {
+		t.mmu.RLock()
+		defer t.mmu.RUnlock()
+
+		if t.err != nil {
+			yield(KeyVal{}, t.err)
+			return
+		}
+		n, pbit, err := t.subtree(prefix)
+		if err != nil {
+			yield(KeyVal{}, err)
+			return
+		}
+		if n == nil {
+			return
+		}
+		n.scan(t, pbit, yield)
+	}
+}
+
+// subtree returns the root of the subtree holding every key that starts
+// with prefix, along with its parent's bit depth, or nil if the tree
+// holds no such key.
+func (t *diskTree) subtree(prefix []byte) (*diskNode, int, error) {
+	n, err := t.node(t.hdr().root())
+	if err != nil || n == nil {
+		return nil, 0, err
+	}
+
+	// Look up prefix as if it were a key, stopping at the first node
+	// that splits at a bit index at or past the end of the prefix:
+	// every key below that node agrees with the others there,
+	// so either all of them start with prefix or none do.
+	pbits := 8 * len(prefix)
+	pbit := -1
+	for n.bit() > pbit && n.bit() < pbits {
+		nbit := n.bit()
+		a := n.left()
+		if bit(prefix, nbit) != 0 {
+			a = n.right()
+		}
+		if n, err = t.child(a); err != nil {
+			return nil, 0, err
+		}
+		pbit = nbit
+	}
+
+	// Check one key to decide for all of them.
+	left, err := n.leftmost(t, pbit)
+	if err != nil {
+		return nil, 0, err
+	}
+	key, err := left.key(t)
+	if err != nil {
+		return nil, 0, err
+	}
+	if !key.HasPrefix(prefix) {
+		return nil, 0, nil
+	}
+	return n, pbit, nil
+}
+
+// child returns the node at address a, which must not be a nil address.
+func (t *diskTree) child(a addr) (*diskNode, error) {
+	n, err := t.node(a)
+	if err != nil {
+		return nil, err
+	}
+	if n == nil {
+		return nil, t.broken(errCorrupt)
+	}
+	return n, nil
+}
+
+// leftmost returns the leaf holding the smallest key
+// in the subtree rooted at n.
+func (n *diskNode) leftmost(t *diskTree, pbit int) (*diskNode, error) {
+	for n.bit() > pbit {
+		pbit = n.bit()
+		next, err := t.child(n.left())
+		if err != nil {
+			return nil, err
+		}
+		n = next
+	}
+	return n, nil
+}
+
+// scan yields the key-val pairs in the subtree rooted at n, in key order,
+// reporting whether iteration should continue.
+func (n *diskNode) scan(t *diskTree, pbit int, yield func(KeyVal, error) bool) bool {
+	nbit := n.bit()
+	if nbit <= pbit {
+		// view n as leaf
+		key, val, err := n.keyVal(t)
+		if err != nil {
+			yield(KeyVal{}, err)
+			return false
+		}
+		return yield(KeyVal{key, val}, nil)
+	}
+
+	left, err := t.child(n.left())
+	if err != nil {
+		yield(KeyVal{}, err)
+		return false
+	}
+	right, err := t.child(n.right())
+	if err != nil {
+		yield(KeyVal{}, err)
+		return false
+	}
+	return left.scan(t, nbit, yield) && right.scan(t, nbit, yield)
 }
 
 func (t *diskTree) check() {

mpt/mem.go

@@ -9,6 +9,7 @@ import (
 	"encoding/binary"
 	"errors"
 	"fmt"
+	"iter"
 )
 
 // A memTree is an in-memory [Tree].
@@ -42,7 +43,7 @@ func (n *memNode) bit() int {
 // NewMemTree returns a new in-memory [Tree].
 func NewMemTree() Tree {
 	t := &memTree{
-		hash:  emptyTreeHash(),
+		hash:  emptyTreeHash,
 		exact: true,
 	}
 	return t
@@ -245,33 +246,32 @@ func (t *memTree) predict(s []node, n *memNode, pbit int, list []KeyVal) ([]node
 	return s, list
 }
 
-// Prove returns a proof of the presence or absence of key in t.
-func (t *memTree) Prove(key Key) (val Val, ok bool, proof Proof, err error) {
+// Path returns the path proof for key in t.
+func (t *memTree) Path(key Key) (Proof, error) {
 	if t.err != nil {
-		return Val{}, false, nil, t.err
+		return nil, t.err
 	}
 	if t.dirty {
-		return Val{}, false, nil, ErrModifiedTree
+		return nil, ErrModifiedTree
 	}
 	if t.root == nil {
-		return Val{}, false, Proof{}, nil
+		return Proof{}, nil
 	}
-	return t.root.prove(-1, key)
+	return t.root.path(-1, key), nil
 }
 
-func (n *memNode) prove(pbit int, key Key) (val Val, ok bool, proof Proof, err error) {
+// path returns the path proof for key in the subtree rooted at n.
+// pbit is the parent bit depth, controlling whether n is viewed as a leaf.
+func (n *memNode) path(pbit int, key Key) Proof {
 	nbit := n.bit()
 	if nbit <= pbit {
 		// view n as leaf
-		if bytes.Equal(n.key, key) {
-			return n.val, true, Proof{}, nil
-		}
 		var p Proof
 		p = binary.AppendUvarint(p, uint64(len(n.key)))
-		p = append(p, n.key[:]...)
+		p = append(p, n.key...)
 		p = binary.AppendUvarint(p, uint64(len(n.val)))
-		p = append(p, n.val[:]...)
-		return Val{}, false, p, nil
+		p = append(p, n.val...)
+		return p
 	}
 
 	var sib Hash
@@ -284,10 +284,78 @@ func (n *memNode) prove(pbit int, key Key) (val Val, ok bool, proof Proof, err e
 		sib = n.left.hash(nbit)
 	}
 
-	val, ok, proof, _ = child.prove(nbit, key)
+	proof := child.path(nbit, key)
 	proof = binary.AppendUvarint(proof, uint64(nbit))
 	proof = append(proof, sib[:]...)
-	return
+	return proof
+}
+
+// Scan returns an iterator over key-val pairs whose keys start with prefix.
+func (t *memTree) Scan(prefix []byte) iter.Seq2[KeyVal, error] {
+	return func(yield func(KeyVal, error) bool) {
+		if t.err != nil {
+			yield(KeyVal{}, t.err)
+			return
+		}
+		n, pbit := t.subtree(prefix)
+		if n == nil {
+			return
+		}
+		n.scan(pbit, yield)
+	}
+}
+
+// subtree returns the root of the subtree holding every key that starts
+// with prefix, along with its parent's bit depth, or nil if the tree
+// holds no such key.
+func (t *memTree) subtree(prefix []byte) (*memNode, int) {
+	n := t.root
+	if n == nil {
+		return nil, 0
+	}
+
+	// Look up prefix as if it were a key, stopping at the first node
+	// that splits at a bit index at or past the end of the prefix:
+	// every key below that node agrees with the others there,
+	// so either all of them start with prefix or none do.
+	pbits := 8 * len(prefix)
+	pbit := -1
+	for n.bit() > pbit && n.bit() < pbits {
+		nbit := n.bit()
+		if bit(prefix, nbit) == 0 {
+			n = n.left
+		} else {
+			n = n.right
+		}
+		pbit = nbit
+	}
+
+	// Check one key to decide for all of them.
+	if !n.leftmost(pbit).key.HasPrefix(prefix) {
+		return nil, 0
+	}
+	return n, pbit
+}
+
+// leftmost returns the leaf holding the smallest key
+// in the subtree rooted at n.
+func (n *memNode) leftmost(pbit int) *memNode {
+	for n.bit() > pbit {
+		pbit = n.bit()
+		n = n.left
+	}
+	return n
+}
+
+// scan yields the key-val pairs in the subtree rooted at n, in key order,
+// reporting whether iteration should continue.
+func (n *memNode) scan(pbit int, yield func(KeyVal, error) bool) bool {
+	nbit := n.bit()
+	if nbit <= pbit {
+		// view n as leaf
+		return yield(KeyVal{n.key, n.val}, nil)
+	}
+	return n.left.scan(nbit, yield) && n.right.scan(nbit, yield)
 }
 
 // check checks all the tree invariants, walking the entire tree.