perf(trie): reuse stored hashes and traversed nodes in Prove

core/trie/proof.go

@@ -67,7 +67,13 @@
 // The result contains the proof nodes on the path from the root to the leaf.
 // The value is included in the proof if the key is present in the trie.
 // If the key is not present, the proof will contain the nodes on the path to the closest ancestor.
+// Proof hashes are taken from stored node values, so the trie's hashes must be
+// current: call Hash after the last write and before Prove.
 func (t *Trie) Prove(key *felt.Felt, proof *ProofNodeSet) error {
+	if t.rootKeyIsDirty || len(t.dirtyNodes) > 0 {
+		return errors.New("cannot prove a trie with unhashed writes")
+	}
+
 	k := t.FeltToKey(key)
 
 	nodesFromRoot, err := t.nodesFromRoot(&k)
@@ -76,29 +82,67 @@
 	}
 
 	var parentKey *BitArray
+	// Parent-facing hash of the current node, computed by the previous
+	// iteration; nil for the root.
+	var carriedHash *felt.Felt
 
+	// Proof nodes alias felts of every node read below, so none of these nodes
+	// can go back to nodePool: reuse would overwrite hashes the caller holds.
 	for i, sNode := range nodesFromRoot {
-		sNodeEdge, sNodeBinary, err := storageNodeToProofNode(t, parentKey, sNode)
+		isLeaf := sNode.key.len == t.height
+
+		var sNodeEdge *Edge
+		if isEdge(parentKey, sNode) {
+			edgePath := path(sNode.key, parentKey)
+			sNodeEdge = &Edge{Path: &edgePath, Child: sNode.node.Value}
+		}
+
+		if isLeaf {
+			if sNodeEdge != nil { // Leaf Edge
+				proof.Put(edgeHash(sNodeEdge, carriedHash, t.hash), sNodeEdge)
+			}
+			break // sNode is a binary leaf otherwise; nothing to add
+		}
+
+		var onPathChild *StorageNode
+		if i+1 < len(nodesFromRoot) {
+			onPathChild = &nodesFromRoot[i+1]
+		}
+		sNodeBinary, err := binaryProofNode(t, sNode, onPathChild)
 		if err != nil {
 			return err
 		}
-		isLeaf := sNode.key.len == t.height
 
-		if sNodeEdge != nil && !isLeaf { // Internal Edge
-			proof.Put(sNodeEdge.Hash(t.hash), sNodeEdge)
-			proof.Put(sNodeBinary.Hash(t.hash), sNodeBinary)
-		} else if sNodeEdge == nil && !isLeaf { // Internal Binary
-			proof.Put(sNodeBinary.Hash(t.hash), sNodeBinary)
-		} else if sNodeEdge != nil && isLeaf { // Leaf Edge
-			proof.Put(sNodeEdge.Hash(t.hash), sNodeEdge)
-		} else if sNodeEdge == nil && sNodeBinary == nil { // sNode is a binary leaf
-			break
+		if sNodeEdge != nil { // Internal Edge
+			proof.Put(edgeHash(sNodeEdge, carriedHash, t.hash), sNodeEdge)
+		}
+		// A hashed internal node stores hash(leftHash, rightHash) as its value.
+		proof.Put(*sNode.node.Value, sNodeBinary)
+
+		// Carry the on-path child's parent-facing hash from the Binary; a nil
+		// carry only costs a recomputation, never a wrong hash.
+		carriedHash = nil
+		switch {
+		case onPathChild == nil:
+		case onPathChild.key.Equal(sNode.node.Left):
+			carriedHash = sNodeBinary.LeftHash
+		case onPathChild.key.Equal(sNode.node.Right):
+			carriedHash = sNodeBinary.RightHash
 		}
-		parentKey = nodesFromRoot[i].key
+		parentKey = sNode.key
 	}
 	return nil
 }
 
+// edgeHash returns the parent-facing hash of an edge node, reusing the value
+// the parent iteration already computed when there is one.
+func edgeHash(edge *Edge, carried *felt.Felt, hash crypto.HashFn) felt.Felt {
+	if carried != nil {
+		return *carried
+	}
+	return edge.Hash(hash)
+}
+
 // GetRangeProof generates a range proof for the given range of keys.
 // The proof contains the proof nodes on the path from the root to the closest ancestor of the left and right keys.
 func (t *Trie) GetRangeProof(leftKey, rightKey *felt.Felt, proofSet *ProofNodeSet) error {
@@ -339,63 +383,57 @@
 
 // isEdge checks if the storage node is an edge node.
 func isEdge(parentKey *BitArray, sNode StorageNode) bool {
-	sNodeLen := sNode.key.len
+	return isEdgeKey(parentKey, sNode.key)
+}
+
+// isEdgeKey reports whether childKey hangs off parentKey via an edge, i.e. the
+// path between them is longer than the single branching bit.
+func isEdgeKey(parentKey, childKey *BitArray) bool {
 	if parentKey == nil { // Root
-		return sNodeLen != 0
+		return childKey.len != 0
 	}
-	return sNodeLen-parentKey.len > 1
+	return childKey.len-parentKey.len > 1
 }
 
-// storageNodeToProofNode converts a StorageNode to the ProofNode(s).
-// Juno's Trie has nodes that are Binary AND Edge, whereas the protocol requires nodes that are Binary XOR Edge.
-// We need to convert the former to the latter for proof generation.
-func storageNodeToProofNode(tri *Trie, parentKey *BitArray, sNode StorageNode) (*Edge, *Binary, error) {
-	var edge *Edge
-	if isEdge(parentKey, sNode) {
-		edgePath := path(sNode.key, parentKey)
-		edge = &Edge{
-			Path:  &edgePath,
-			Child: sNode.node.Value,
+// binaryProofNode builds the Binary proof node of an internal StorageNode.
+// Juno's Trie has nodes that are Binary AND Edge, whereas the protocol requires
+// nodes that are Binary XOR Edge. We need to convert the former to the latter for
+// proof generation. onPathChild is the next node the traversal already read, so
+// it is not read again; only the off-path sibling costs a database lookup.
+func binaryProofNode(
+	tri *Trie, sNode StorageNode, onPathChild *StorageNode,
+) (*Binary, error) {
+	childHash := func(childKey *BitArray) (*felt.Felt, error) {
+		var child *Node
+		if onPathChild != nil && childKey.Equal(onPathChild.key) {
+			child = onPathChild.node
+		} else {
+			var err error
+			if child, err = tri.GetNodeFromKey(childKey); err != nil {
+				return nil, err
+			}
 		}
-	}
-	if sNode.key.len == tri.height { // Leaf
-		return edge, nil, nil
-	}
-	lNode, err := tri.GetNodeFromKey(sNode.node.Left)
-	if err != nil {
-		return nil, nil, err
-	}
-	rNode, err := tri.GetNodeFromKey(sNode.node.Right)
-	if err != nil {
-		return nil, nil, err
-	}
 
-	rightHash := rNode.Value
-	if isEdge(sNode.key, StorageNode{node: rNode, key: sNode.node.Right}) {
-		edgePath := path(sNode.node.Right, sNode.key)
-		rEdge := &Edge{
-			Path:  &edgePath,
-			Child: rNode.Value,
+		// Not Node.HashFromParent: taking the address of its returned value
+		// costs an allocation per non-edge child.
+		if isEdgeKey(sNode.key, childKey) {
+			edgePath := path(childKey, sNode.key)
+			wrapped := child.Hash(&edgePath, tri.hash)
+			return &wrapped, nil
 		}
-		hash := rEdge.Hash(tri.hash)
-		rightHash = &hash
+		return child.Value, nil
 	}
-	leftHash := lNode.Value
-	if isEdge(sNode.key, StorageNode{node: lNode, key: sNode.node.Left}) {
-		edgePath := path(sNode.node.Left, sNode.key)
-		lEdge := &Edge{
-			Path:  &edgePath,
-			Child: lNode.Value,
-		}
-		hash := lEdge.Hash(tri.hash)
-		leftHash = &hash
+
+	leftHash, err := childHash(sNode.node.Left)
+	if err != nil {
+		return nil, err
 	}
-	binary := &Binary{
-		LeftHash:  leftHash,
-		RightHash: rightHash,
+	rightHash, err := childHash(sNode.node.Right)
+	if err != nil {
+		return nil, err
 	}
 
-	return edge, binary, nil
+	return &Binary{LeftHash: leftHash, RightHash: rightHash}, nil
 }
 
 // proofToPath converts a Merkle proof to trie node path. All necessary nodes will be resolved and leave the remaining

core/trie/proof_test.go

@@ -855,3 +855,51 @@ type keyValue struct {
 	key   *felt.Felt
 	value *felt.Felt
 }
+
+func TestProveErrsOnStaleHashes(t *testing.T) {
+	t.Parallel()
+
+	memdb := memory.New()
+	txn := memdb.NewIndexedBatch()
+	tempTrie, err := trie.NewTriePedersen(txn, []byte{0}, 251)
+	require.NoError(t, err)
+
+	key := new(felt.Felt).SetUint64(1)
+	_, err = tempTrie.Put(key, new(felt.Felt).SetUint64(2))
+	require.NoError(t, err)
+
+	err = tempTrie.Prove(key, trie.NewProofNodeSet())
+	require.EqualError(t, err, "cannot prove a trie with unhashed writes")
+
+	require.NoError(t, tempTrie.Commit())
+	require.NoError(t, tempTrie.Prove(key, trie.NewProofNodeSet()))
+}
+
+// TestProveSetInvariant checks every proof entry is keyed by its own hash.
+// Prove reuses stored and carried hashes instead of recomputing them, so a
+// drift between the reused values and the node contents would mis-key entries
+// and only surface at the verifier.
+func TestProveSetInvariant(t *testing.T) {
+	t.Parallel()
+
+	tempTrie, records := randomTrie(t, 100)
+
+	keys := make([]*felt.Felt, 0, len(records)+1)
+	for _, record := range records {
+		keys = append(keys, record.key)
+	}
+	// Non-membership proofs must hold the invariant too.
+	keys = append(keys, new(felt.Felt).SetUint64(0xdead))
+
+	for _, key := range keys {
+		proofSet := trie.NewProofNodeSet()
+		require.NoError(t, tempTrie.Prove(key, proofSet))
+
+		hashes := proofSet.Keys()
+		nodes := proofSet.List()
+		for i, node := range nodes {
+			hash := node.Hash(crypto.Pedersen)
+			require.True(t, hash.Equal(&hashes[i]), "entry %d for key %s", i, key)
+		}
+	}
+}