From af5ae5a792762d2b02abc31a99807ed185c3b617 Mon Sep 17 00:00:00 2001 From: Alexey Shekhirin Date: Mon, 28 Oct 2024 13:55:31 +0000 Subject: [PATCH] perf(trie): reduce allocations in sparse trie rlp node calculation (#12092) --- crates/trie/sparse/src/trie.rs | 86 ++++++++++++++++++++++------------ 1 file changed, 57 insertions(+), 29 deletions(-) diff --git a/crates/trie/sparse/src/trie.rs b/crates/trie/sparse/src/trie.rs index 035eeaf73e81..11ca100791ec 100644 --- a/crates/trie/sparse/src/trie.rs +++ b/crates/trie/sparse/src/trie.rs @@ -558,7 +558,7 @@ impl RevealedSparseTrie { pub fn root(&mut self) -> B256 { // take the current prefix set. let mut prefix_set = std::mem::take(&mut self.prefix_set).freeze(); - let root_rlp = self.rlp_node(Nibbles::default(), &mut prefix_set); + let root_rlp = self.rlp_node_allocate(Nibbles::default(), &mut prefix_set); if let Some(root_hash) = root_rlp.as_hash() { root_hash } else { @@ -570,10 +570,12 @@ impl RevealedSparseTrie { /// depth. Root node has a level of 0. pub fn update_rlp_node_level(&mut self, depth: usize) { let mut prefix_set = self.prefix_set.clone().freeze(); + let mut buffers = RlpNodeBuffers::default(); let targets = self.get_changed_nodes_at_depth(&mut prefix_set, depth); for target in targets { - self.rlp_node(target, &mut prefix_set); + buffers.path_stack.push((target, Some(true))); + self.rlp_node(&mut prefix_set, &mut buffers); } } @@ -629,17 +631,13 @@ impl RevealedSparseTrie { targets } - fn rlp_node(&mut self, path: Nibbles, prefix_set: &mut PrefixSet) -> RlpNode { - // stack of paths we need rlp nodes for - let mut path_stack = Vec::from([(path, None)]); - // stack of rlp nodes - let mut rlp_node_stack = Vec::<(Nibbles, RlpNode)>::new(); - // reusable branch child path - let mut branch_child_buf = SmallVec::<[Nibbles; 16]>::new_const(); - // reusable branch value stack - let mut branch_value_stack_buf = SmallVec::<[RlpNode; 16]>::new_const(); - - 'main: while let Some((path, mut is_in_prefix_set)) = path_stack.pop() { + fn rlp_node_allocate(&mut self, path: Nibbles, prefix_set: &mut PrefixSet) -> RlpNode { + let mut buffers = RlpNodeBuffers::new_with_path(path); + self.rlp_node(prefix_set, &mut buffers) + } + + fn rlp_node(&mut self, prefix_set: &mut PrefixSet, buffers: &mut RlpNodeBuffers) -> RlpNode { + 'main: while let Some((path, mut is_in_prefix_set)) = buffers.path_stack.pop() { // Check if the path is in the prefix set. // First, check the cached value. If it's `None`, then check the prefix set, and update // the cached value. @@ -667,63 +665,68 @@ impl RevealedSparseTrie { child_path.extend_from_slice_unchecked(key); if let Some(hash) = hash.filter(|_| !prefix_set_contains(&path)) { RlpNode::word_rlp(&hash) - } else if rlp_node_stack.last().map_or(false, |e| e.0 == child_path) { - let (_, child) = rlp_node_stack.pop().unwrap(); + } else if buffers.rlp_node_stack.last().map_or(false, |e| e.0 == child_path) { + let (_, child) = buffers.rlp_node_stack.pop().unwrap(); self.rlp_buf.clear(); let rlp_node = ExtensionNodeRef::new(key, &child).rlp(&mut self.rlp_buf); *hash = rlp_node.as_hash(); rlp_node } else { // need to get rlp node for child first - path_stack.extend([(path, is_in_prefix_set), (child_path, None)]); + buffers.path_stack.extend([(path, is_in_prefix_set), (child_path, None)]); continue } } SparseNode::Branch { state_mask, hash } => { if let Some(hash) = hash.filter(|_| !prefix_set_contains(&path)) { - rlp_node_stack.push((path, RlpNode::word_rlp(&hash))); + buffers.rlp_node_stack.push((path, RlpNode::word_rlp(&hash))); continue } - branch_child_buf.clear(); + buffers.branch_child_buf.clear(); // Walk children in a reverse order from `f` to `0`, so we pop the `0` first // from the stack. for bit in CHILD_INDEX_RANGE.rev() { if state_mask.is_bit_set(bit) { let mut child = path.clone(); child.push_unchecked(bit); - branch_child_buf.push(child); + buffers.branch_child_buf.push(child); } } - branch_value_stack_buf.resize(branch_child_buf.len(), Default::default()); + buffers + .branch_value_stack_buf + .resize(buffers.branch_child_buf.len(), Default::default()); let mut added_children = false; - for (i, child_path) in branch_child_buf.iter().enumerate() { - if rlp_node_stack.last().map_or(false, |e| &e.0 == child_path) { - let (_, child) = rlp_node_stack.pop().unwrap(); + for (i, child_path) in buffers.branch_child_buf.iter().enumerate() { + if buffers.rlp_node_stack.last().map_or(false, |e| &e.0 == child_path) { + let (_, child) = buffers.rlp_node_stack.pop().unwrap(); // Insert children in the resulting buffer in a normal order, because // initially we iterated in reverse. - branch_value_stack_buf[branch_child_buf.len() - i - 1] = child; + buffers.branch_value_stack_buf + [buffers.branch_child_buf.len() - i - 1] = child; added_children = true; } else { debug_assert!(!added_children); - path_stack.push((path, is_in_prefix_set)); - path_stack.extend(branch_child_buf.drain(..).map(|p| (p, None))); + buffers.path_stack.push((path, is_in_prefix_set)); + buffers + .path_stack + .extend(buffers.branch_child_buf.drain(..).map(|p| (p, None))); continue 'main } } self.rlp_buf.clear(); - let rlp_node = BranchNodeRef::new(&branch_value_stack_buf, *state_mask) + let rlp_node = BranchNodeRef::new(&buffers.branch_value_stack_buf, *state_mask) .rlp(&mut self.rlp_buf); *hash = rlp_node.as_hash(); rlp_node } }; - rlp_node_stack.push((path, rlp_node)); + buffers.rlp_node_stack.push((path, rlp_node)); } - rlp_node_stack.pop().unwrap().1 + buffers.rlp_node_stack.pop().unwrap().1 } } @@ -803,6 +806,31 @@ struct RemovedSparseNode { unset_branch_nibble: Option, } +/// Collection of reusable buffers for [`RevealedSparseTrie::rlp_node`]. +#[derive(Debug, Default)] +struct RlpNodeBuffers { + /// Stack of paths we need rlp nodes for and whether the path is in the prefix set. + path_stack: Vec<(Nibbles, Option)>, + /// Stack of rlp nodes + rlp_node_stack: Vec<(Nibbles, RlpNode)>, + /// Reusable branch child path + branch_child_buf: SmallVec<[Nibbles; 16]>, + /// Reusable branch value stack + branch_value_stack_buf: SmallVec<[RlpNode; 16]>, +} + +impl RlpNodeBuffers { + /// Creates a new instance of buffers with the given path on the stack. + fn new_with_path(path: Nibbles) -> Self { + Self { + path_stack: vec![(path, None)], + rlp_node_stack: Vec::new(), + branch_child_buf: SmallVec::<[Nibbles; 16]>::new_const(), + branch_value_stack_buf: SmallVec::<[RlpNode; 16]>::new_const(), + } + } +} + #[cfg(test)] mod tests { use std::collections::BTreeMap;