feat: retain proofs of non-target nodes in certain edge-cases. (#109)

* feat: retain proofs of non-target nodes in certain edge-cases.

There are specific edge-cases where it's necessary to retain the proofs
for nodes which aren't given in the target set, namely leaf removals
which result in the deletion of a branch, and leaf additions which
result in the creation of a branch.

Documentation of each case is provided in the code at the point it is
handled.

This change will cause more proofs than are strictly necessary to be
retained, because the `target` set we are given does not tell us if
paths are added, updated or removed. This extra work is made up for by
us not needing to fetch the nodes which would otherwise be missing later
down the pipeline, and in benchmarking the overall change comes out very
slightly faster.

* Update src/proof/retainer.rs

Co-authored-by: Dan Cline <6798349+Rjected@users.noreply.github.com>

* Fully track branch target children, to properly handle collapse case

* Introduce AddedRemovedKeys, use that for trackinge extra proofs

* debug

* Don't track removed leaves as non-targets

* Make AddedRemovedKeys into a concrete type

* PR suggestion

* Remove assume_added

* Revert "Remove assume_added"

This reverts commit 1e4be6f.

* Fix unnecessary boolean being included

* Add clear_added method

* PR feedback

---------

Co-authored-by: Dan Cline <6798349+Rjected@users.noreply.github.com>

Author: mediocregopher

src/hash_builder/mod.rs

@@ -3,12 +3,11 @@
 use super::{
     BranchNodeCompact, EMPTY_ROOT_HASH, Nibbles, TrieMask,
     nodes::{BranchNodeRef, ExtensionNodeRef, LeafNodeRef},
-    proof::ProofRetainer,
+    proof::{ProofNodes, ProofRetainer},
 };
-use crate::{HashMap, nodes::RlpNode, proof::ProofNodes};
+use crate::{HashMap, nodes::RlpNode, proof::AddedRemovedKeys};
 use alloc::vec::Vec;
 use alloy_primitives::{B256, keccak256};
-use alloy_rlp::EMPTY_STRING_CODE;
 use core::cmp;
 use tracing::trace;
 
@@ -39,9 +38,9 @@ pub use value::{HashBuilderValue, HashBuilderValueRef};
 /// up, combining the hashes of child nodes and ultimately generating the root hash. The root hash
 /// can then be used to verify the integrity and authenticity of the trie's data by constructing and
 /// verifying Merkle proofs.
-#[derive(Debug, Clone, Default)]
+#[derive(Debug, Clone)]
 #[allow(missing_docs)]
-pub struct HashBuilder {
+pub struct HashBuilder<K = AddedRemovedKeys> {
     pub key: Nibbles,
     pub value: HashBuilderValue,
     pub stack: Vec<RlpNode>,
@@ -53,12 +52,29 @@ pub struct HashBuilder {
     pub stored_in_database: bool,
 
     pub updated_branch_nodes: Option<HashMap<Nibbles, BranchNodeCompact>>,
-    pub proof_retainer: Option<ProofRetainer>,
+    pub proof_retainer: Option<ProofRetainer<K>>,
 
     pub rlp_buf: Vec<u8>,
 }
 
-impl HashBuilder {
+impl Default for HashBuilder {
+    fn default() -> Self {
+        Self {
+            key: Default::default(),
+            value: Default::default(),
+            stack: Default::default(),
+            state_masks: Default::default(),
+            tree_masks: Default::default(),
+            hash_masks: Default::default(),
+            stored_in_database: Default::default(),
+            updated_branch_nodes: None,
+            proof_retainer: None,
+            rlp_buf: Default::default(),
+        }
+    }
+}
+
+impl<K> HashBuilder<K> {
     /// Enables the Hash Builder to store updated branch nodes.
     ///
     /// Call [HashBuilder::split] to get the updates to branch nodes.
@@ -68,9 +84,19 @@ impl HashBuilder {
     }
 
     /// Enable specified proof retainer.
-    pub fn with_proof_retainer(mut self, retainer: ProofRetainer) -> Self {
-        self.proof_retainer = Some(retainer);
-        self
+    pub fn with_proof_retainer<K2>(self, retainer: ProofRetainer<K2>) -> HashBuilder<K2> {
+        HashBuilder {
+            key: self.key,
+            value: self.value,
+            stack: self.stack,
+            state_masks: self.state_masks,
+            tree_masks: self.tree_masks,
+            hash_masks: self.hash_masks,
+            stored_in_database: self.stored_in_database,
+            updated_branch_nodes: self.updated_branch_nodes,
+            proof_retainer: Some(retainer),
+            rlp_buf: self.rlp_buf,
+        }
     }
 
     /// Enables the Hash Builder to store updated branch nodes.
@@ -81,7 +107,9 @@ impl HashBuilder {
             self.updated_branch_nodes = Some(HashMap::default());
         }
     }
+}
 
+impl<K: AsRef<AddedRemovedKeys>> HashBuilder<K> {
     /// Splits the [HashBuilder] into a [HashBuilder] and hash builder updates.
     pub fn split(mut self) -> (Self, HashMap<Nibbles, BranchNodeCompact>) {
         let updates = self.updated_branch_nodes.take();
@@ -159,7 +187,7 @@ impl HashBuilder {
         let root = self.current_root();
         if root == EMPTY_ROOT_HASH {
             if let Some(proof_retainer) = self.proof_retainer.as_mut() {
-                proof_retainer.retain(&Nibbles::default(), &[EMPTY_STRING_CODE])
+                proof_retainer.retain_empty_root_proof();
             }
         }
         root
@@ -267,7 +295,10 @@ impl HashBuilder {
                             "pushing leaf node",
                         );
                         self.stack.push(rlp);
-                        self.retain_proof_from_buf(&path);
+
+                        if let Some(proof_retainer) = self.proof_retainer.as_mut() {
+                            proof_retainer.retain_leaf_proof(&path, &self.rlp_buf)
+                        }
                     }
                     HashBuilderValueRef::Hash(hash) => {
                         trace!(target: "trie::hash_builder", ?hash, "pushing branch node hash");
@@ -302,7 +333,11 @@ impl HashBuilder {
                     "pushing extension node",
                 );
                 self.stack.push(rlp);
-                self.retain_proof_from_buf(&path);
+
+                if let Some(proof_retainer) = self.proof_retainer.as_mut() {
+                    proof_retainer.retain_extension_proof(&path, &short_node_key, &self.rlp_buf)
+                }
+
                 self.resize_masks(len_from);
             }
 
@@ -368,7 +403,10 @@ impl HashBuilder {
             ?rlp,
             "pushing branch node",
         );
-        self.retain_proof_from_buf(&path);
+
+        if let Some(proof_retainer) = self.proof_retainer.as_mut() {
+            proof_retainer.retain_branch_proof(&path, state_mask, &self.rlp_buf);
+        }
 
         // Clears the stack from the branch node elements
         let first_child_idx = self.stack.len() - state_mask.count_ones() as usize;
@@ -417,12 +455,6 @@ impl HashBuilder {
         }
     }
 
-    fn retain_proof_from_buf(&mut self, prefix: &Nibbles) {
-        if let Some(proof_retainer) = self.proof_retainer.as_mut() {
-            proof_retainer.retain(prefix, &self.rlp_buf)
-        }
-    }
-
     fn update_masks(&mut self, current: &Nibbles, len_from: usize) {
         if len_from > 0 {
             let flag = TrieMask::from_nibble(current.get_unchecked(len_from - 1));

src/proof/added_removed_keys.rs

@@ -0,0 +1,90 @@
+use crate::{Nibbles, TrieMask};
+use alloy_primitives::{B256, map::B256Set};
+
+/// Provides added and removed keys for an account or storage trie.
+///
+/// Used by the [`crate::proof::ProofRetainer`] to determine which nodes may be ancestors of newly
+/// added or removed leaves. This information allows for generation of more complete proofs which
+/// include the nodes necessary for adding and removing leaves from the trie.
+#[derive(Debug, Default, Clone)]
+pub struct AddedRemovedKeys {
+    /// Keys which are known to be removed from the trie.
+    removed_keys: B256Set,
+    /// Keys which are known to be added to the trie.
+    added_keys: B256Set,
+    /// Assume that all keys have been added.
+    assume_added: bool,
+}
+
+impl AsRef<Self> for AddedRemovedKeys {
+    fn as_ref(&self) -> &Self {
+        self
+    }
+}
+
+impl AddedRemovedKeys {
+    /// Sets the `assume_added` flag, which can be used instead of `insert_added` if exact
+    /// additions aren't known and you want to optimistically collect all proofs which _might_ be
+    /// necessary.
+    pub fn with_assume_added(mut self, assume_added: bool) -> Self {
+        self.assume_added = assume_added;
+        self
+    }
+
+    /// Sets the key as being a removed key. This removes the key from the `added_keys` set if it
+    /// was previously inserted into it.
+    pub fn insert_removed(&mut self, key: B256) {
+        self.added_keys.remove(&key);
+        self.removed_keys.insert(key);
+    }
+
+    /// Unsets the key as being a removed key.
+    pub fn remove_removed(&mut self, key: &B256) {
+        self.removed_keys.remove(key);
+    }
+
+    /// Sets the key as being an added key. This removes the key from the `removed_keys` set if it
+    /// was previously inserted into it.
+    pub fn insert_added(&mut self, key: B256) {
+        self.removed_keys.remove(&key);
+        self.added_keys.insert(key);
+    }
+
+    /// Clears all keys which have been added via `insert_added`.
+    pub fn clear_added(&mut self) {
+        self.added_keys.clear();
+    }
+
+    /// Returns true if the given key path is marked as removed.
+    pub fn is_removed(&self, path: &B256) -> bool {
+        self.removed_keys.contains(path)
+    }
+
+    /// Returns true if the given key path is marked as added.
+    pub fn is_added(&self, path: &B256) -> bool {
+        self.assume_added || self.added_keys.contains(path)
+    }
+
+    /// Returns true if the given path prefix is the prefix of an added key.
+    pub fn is_prefix_added(&self, prefix: &Nibbles) -> bool {
+        self.assume_added
+            || self.added_keys.iter().any(|key| {
+                let key_nibbles = Nibbles::unpack(key);
+                key_nibbles.starts_with(prefix)
+            })
+    }
+
+    /// Returns a mask containing a bit set for each child of the branch which is a prefix of a
+    /// removed leaf.
+    pub fn get_removed_mask(&self, branch_path: &Nibbles) -> TrieMask {
+        let mut mask = TrieMask::default();
+        for key in &self.removed_keys {
+            let key_nibbles = Nibbles::unpack(key);
+            if key_nibbles.starts_with(branch_path) {
+                let child_bit = key_nibbles.get_unchecked(branch_path.len());
+                mask.set_bit(child_bit);
+            }
+        }
+        mask
+    }
+}

src/proof/mod.rs

@@ -20,3 +20,6 @@ pub use proof_nodes::ProofNodes;
 
 mod retainer;
 pub use retainer::ProofRetainer;
+
+mod added_removed_keys;
+pub use added_removed_keys::AddedRemovedKeys;