perf(trie): collect only changed sparse nodes at a depth (#12093)

This commit is contained in:
Alexey Shekhirin
2024-10-28 11:06:26 +00:00
committed by GitHub
parent 268090e879
commit 1b0f625f1d

View File

@@ -569,27 +569,36 @@ impl RevealedSparseTrie {
/// Update hashes of the nodes that are located at a level deeper than or equal to the provided
/// depth. Root node has a level of 0.
pub fn update_rlp_node_level(&mut self, depth: usize) {
let targets = self.get_nodes_at_depth(depth);
let mut prefix_set = self.prefix_set.clone().freeze();
let targets = self.get_changed_nodes_at_depth(&mut prefix_set, depth);
for target in targets {
self.rlp_node(target, &mut prefix_set);
}
}
/// Returns a list of paths to the nodes that are located at the provided depth when counting
/// from the root node. If there's a leaf at a depth less than the provided depth, it will be
/// included in the result.
fn get_nodes_at_depth(&self, depth: usize) -> Vec<Nibbles> {
/// Returns a list of paths to the nodes that were changed according to the prefix set and are
/// located at the provided depth when counting from the root node. If there's a leaf at a
/// depth less than the provided depth, it will be included in the result.
fn get_changed_nodes_at_depth(&self, prefix_set: &mut PrefixSet, depth: usize) -> Vec<Nibbles> {
let mut paths = Vec::from([(Nibbles::default(), 0)]);
let mut targets = Vec::new();
while let Some((mut path, level)) = paths.pop() {
match self.nodes.get(&path).unwrap() {
SparseNode::Empty | SparseNode::Hash(_) => {}
SparseNode::Leaf { .. } => {
SparseNode::Leaf { hash, .. } => {
if hash.is_some() && !prefix_set.contains(&path) {
continue
}
targets.push(path);
}
SparseNode::Extension { key, .. } => {
SparseNode::Extension { key, hash } => {
if hash.is_some() && !prefix_set.contains(&path) {
continue
}
if level >= depth {
targets.push(path);
} else {
@@ -597,7 +606,11 @@ impl RevealedSparseTrie {
paths.push((path, level + 1));
}
}
SparseNode::Branch { state_mask, .. } => {
SparseNode::Branch { state_mask, hash } => {
if hash.is_some() && !prefix_set.contains(&path) {
continue
}
if level >= depth {
targets.push(path);
} else {
@@ -1540,7 +1553,7 @@ mod tests {
}
#[test]
fn sparse_trie_get_nodes_at_depth() {
fn sparse_trie_get_changed_nodes_at_depth() {
let mut sparse = RevealedSparseTrie::default();
let value = alloy_rlp::encode_fixed_size(&U256::ZERO).to_vec();
@@ -1574,10 +1587,16 @@ mod tests {
.unwrap();
sparse.update_leaf(Nibbles::from_nibbles([0x5, 0x3, 0x3, 0x2, 0x0]), value).unwrap();
assert_eq!(sparse.get_nodes_at_depth(0), vec![Nibbles::default()]);
assert_eq!(sparse.get_nodes_at_depth(1), vec![Nibbles::from_nibbles_unchecked([0x5])]);
assert_eq!(
sparse.get_nodes_at_depth(2),
sparse.get_changed_nodes_at_depth(&mut PrefixSet::default(), 0),
vec![Nibbles::default()]
);
assert_eq!(
sparse.get_changed_nodes_at_depth(&mut PrefixSet::default(), 1),
vec![Nibbles::from_nibbles_unchecked([0x5])]
);
assert_eq!(
sparse.get_changed_nodes_at_depth(&mut PrefixSet::default(), 2),
vec![
Nibbles::from_nibbles_unchecked([0x5, 0x0]),
Nibbles::from_nibbles_unchecked([0x5, 0x2]),
@@ -1585,7 +1604,7 @@ mod tests {
]
);
assert_eq!(
sparse.get_nodes_at_depth(3),
sparse.get_changed_nodes_at_depth(&mut PrefixSet::default(), 3),
vec![
Nibbles::from_nibbles_unchecked([0x5, 0x0, 0x2, 0x3]),
Nibbles::from_nibbles_unchecked([0x5, 0x2]),
@@ -1594,7 +1613,7 @@ mod tests {
]
);
assert_eq!(
sparse.get_nodes_at_depth(4),
sparse.get_changed_nodes_at_depth(&mut PrefixSet::default(), 4),
vec![
Nibbles::from_nibbles_unchecked([0x5, 0x0, 0x2, 0x3, 0x1]),
Nibbles::from_nibbles_unchecked([0x5, 0x0, 0x2, 0x3, 0x3]),