alloy_trie/proof/
added_removed_keys.rs

1use crate::{Nibbles, TrieMask};
2use alloy_primitives::{B256, map::B256Set};
3
4/// Provides added and removed keys for an account or storage trie.
5///
6/// Used by the [`crate::proof::ProofRetainer`] to determine which nodes may be ancestors of newly
7/// added or removed leaves. This information allows for generation of more complete proofs which
8/// include the nodes necessary for adding and removing leaves from the trie.
9#[derive(Debug, Default, Clone)]
10pub struct AddedRemovedKeys {
11    /// Keys which are known to be removed from the trie.
12    removed_keys: B256Set,
13    /// Keys which are known to be added to the trie.
14    added_keys: B256Set,
15    /// Assume that all keys have been added.
16    assume_added: bool,
17}
18
19impl AsRef<Self> for AddedRemovedKeys {
20    fn as_ref(&self) -> &Self {
21        self
22    }
23}
24
25impl AddedRemovedKeys {
26    /// Sets the `assume_added` flag, which can be used instead of `insert_added` if exact
27    /// additions aren't known and you want to optimistically collect all proofs which _might_ be
28    /// necessary.
29    pub fn with_assume_added(mut self, assume_added: bool) -> Self {
30        self.assume_added = assume_added;
31        self
32    }
33
34    /// Sets the key as being a removed key. This removes the key from the `added_keys` set if it
35    /// was previously inserted into it.
36    pub fn insert_removed(&mut self, key: B256) {
37        self.added_keys.remove(&key);
38        self.removed_keys.insert(key);
39    }
40
41    /// Unsets the key as being a removed key.
42    pub fn remove_removed(&mut self, key: &B256) {
43        self.removed_keys.remove(key);
44    }
45
46    /// Sets the key as being an added key. This removes the key from the `removed_keys` set if it
47    /// was previously inserted into it.
48    pub fn insert_added(&mut self, key: B256) {
49        self.removed_keys.remove(&key);
50        self.added_keys.insert(key);
51    }
52
53    /// Clears all keys which have been added via `insert_added`.
54    pub fn clear_added(&mut self) {
55        self.added_keys.clear();
56    }
57
58    /// Returns true if the given key path is marked as removed.
59    pub fn is_removed(&self, path: &B256) -> bool {
60        self.removed_keys.contains(path)
61    }
62
63    /// Returns true if the given key path is marked as added.
64    pub fn is_added(&self, path: &B256) -> bool {
65        self.assume_added || self.added_keys.contains(path)
66    }
67
68    /// Returns true if the given path prefix is the prefix of an added key.
69    pub fn is_prefix_added(&self, prefix: &Nibbles) -> bool {
70        self.assume_added
71            || self.added_keys.iter().any(|key| {
72                let key_nibbles = Nibbles::unpack(key);
73                key_nibbles.starts_with(prefix)
74            })
75    }
76
77    /// Returns a mask containing a bit set for each child of the branch which is a prefix of a
78    /// removed leaf.
79    pub fn get_removed_mask(&self, branch_path: &Nibbles) -> TrieMask {
80        let mut mask = TrieMask::default();
81        for key in &self.removed_keys {
82            let key_nibbles = Nibbles::unpack(key);
83            if key_nibbles.starts_with(branch_path) {
84                let child_bit = key_nibbles.get_unchecked(branch_path.len());
85                mask.set_bit(child_bit);
86            }
87        }
88        mask
89    }
90}
alloy_trie/proof/added_removed_keys.rs

alloy_trie/proof/
added_removed_keys.rs
