From f5471bee1e76dee439e7847019977c8c17fd6a18 Mon Sep 17 00:00:00 2001 From: Alexey Shekhirin Date: Sat, 26 Oct 2024 13:17:43 +0100 Subject: [PATCH 1/2] perf(trie): reduce allocations in sparse trie rlp node calculation --- crates/trie/sparse/src/trie.rs | 44 ++++++++++++++++++++++++++++++---- 1 file changed, 39 insertions(+), 5 deletions(-) diff --git a/crates/trie/sparse/src/trie.rs b/crates/trie/sparse/src/trie.rs index 8d65378f614e..aa1a31b9f6ba 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 { @@ -571,8 +571,25 @@ impl RevealedSparseTrie { 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(); + + // stack of paths we need rlp nodes for + let mut path_stack = Vec::new(); + // 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(); + for target in targets { - self.rlp_node(target, &mut prefix_set); + path_stack.push(target); + self.rlp_node( + &mut prefix_set, + &mut path_stack, + &mut rlp_node_stack, + &mut branch_child_buf, + &mut branch_value_stack_buf, + ); } } @@ -616,7 +633,7 @@ impl RevealedSparseTrie { targets } - fn rlp_node(&mut self, path: Nibbles, prefix_set: &mut PrefixSet) -> RlpNode { + fn rlp_node_allocate(&mut self, path: Nibbles, prefix_set: &mut PrefixSet) -> RlpNode { // stack of paths we need rlp nodes for let mut path_stack = Vec::from([path]); // stack of rlp nodes @@ -626,6 +643,23 @@ impl RevealedSparseTrie { // reusable branch value stack let mut branch_value_stack_buf = SmallVec::<[RlpNode; 16]>::new_const(); + self.rlp_node( + prefix_set, + &mut path_stack, + &mut rlp_node_stack, + &mut branch_child_buf, + &mut branch_value_stack_buf, + ) + } + + fn rlp_node( + &mut self, + prefix_set: &mut PrefixSet, + path_stack: &mut Vec, + rlp_node_stack: &mut Vec<(Nibbles, RlpNode)>, + branch_child_buf: &mut SmallVec<[Nibbles; 16]>, + branch_value_stack_buf: &mut SmallVec<[RlpNode; 16]>, + ) -> RlpNode { 'main: while let Some(path) = path_stack.pop() { let rlp_node = match self.nodes.get_mut(&path).unwrap() { SparseNode::Empty => RlpNode::word_rlp(&EMPTY_ROOT_HASH), @@ -675,7 +709,7 @@ impl RevealedSparseTrie { } branch_value_stack_buf.clear(); - for child_path in &branch_child_buf { + for child_path in branch_child_buf.iter() { if rlp_node_stack.last().map_or(false, |e| &e.0 == child_path) { let (_, child) = rlp_node_stack.pop().unwrap(); branch_value_stack_buf.push(child); @@ -688,7 +722,7 @@ impl RevealedSparseTrie { } self.rlp_buf.clear(); - let rlp_node = BranchNodeRef::new(&branch_value_stack_buf, *state_mask) + let rlp_node = BranchNodeRef::new(branch_value_stack_buf, *state_mask) .rlp(&mut self.rlp_buf); *hash = rlp_node.as_hash(); rlp_node From 737d86740fb876691d298148cce2cc62039dab66 Mon Sep 17 00:00:00 2001 From: Alexey Shekhirin Date: Mon, 28 Oct 2024 10:21:02 +0000 Subject: [PATCH 2/2] create a struct with buffers --- crates/trie/sparse/src/trie.rs | 108 +++++++++++++++------------------ 1 file changed, 50 insertions(+), 58 deletions(-) diff --git a/crates/trie/sparse/src/trie.rs b/crates/trie/sparse/src/trie.rs index aa1a31b9f6ba..f7827163cec7 100644 --- a/crates/trie/sparse/src/trie.rs +++ b/crates/trie/sparse/src/trie.rs @@ -572,24 +572,11 @@ impl RevealedSparseTrie { let targets = self.get_nodes_at_depth(depth); let mut prefix_set = self.prefix_set.clone().freeze(); - // stack of paths we need rlp nodes for - let mut path_stack = Vec::new(); - // 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(); + let mut buffers = RlpNodeBuffers::default(); for target in targets { - path_stack.push(target); - self.rlp_node( - &mut prefix_set, - &mut path_stack, - &mut rlp_node_stack, - &mut branch_child_buf, - &mut branch_value_stack_buf, - ); + buffers.path_stack.push(target); + self.rlp_node(&mut prefix_set, &mut buffers); } } @@ -634,33 +621,13 @@ impl RevealedSparseTrie { } fn rlp_node_allocate(&mut self, path: Nibbles, prefix_set: &mut PrefixSet) -> RlpNode { - // stack of paths we need rlp nodes for - let mut path_stack = Vec::from([path]); - // 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(); - - self.rlp_node( - prefix_set, - &mut path_stack, - &mut rlp_node_stack, - &mut branch_child_buf, - &mut branch_value_stack_buf, - ) + let mut buffers = RlpNodeBuffers::new_with_path(path); + + self.rlp_node(prefix_set, &mut buffers) } - fn rlp_node( - &mut self, - prefix_set: &mut PrefixSet, - path_stack: &mut Vec, - rlp_node_stack: &mut Vec<(Nibbles, RlpNode)>, - branch_child_buf: &mut SmallVec<[Nibbles; 16]>, - branch_value_stack_buf: &mut SmallVec<[RlpNode; 16]>, - ) -> RlpNode { - 'main: while let Some(path) = path_stack.pop() { + fn rlp_node(&mut self, prefix_set: &mut PrefixSet, buffers: &mut RlpNodeBuffers) -> RlpNode { + 'main: while let Some(path) = buffers.path_stack.pop() { let rlp_node = match self.nodes.get_mut(&path).unwrap() { SparseNode::Empty => RlpNode::word_rlp(&EMPTY_ROOT_HASH), SparseNode::Hash(hash) => RlpNode::word_rlp(hash), @@ -682,56 +649,56 @@ 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 { - path_stack.extend([path, child_path]); // need to get rlp node for child first + buffers.path_stack.extend([path, child_path]); // need to get rlp node for child first 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(); for bit in CHILD_INDEX_RANGE { 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.clear(); - for child_path in branch_child_buf.iter() { - if rlp_node_stack.last().map_or(false, |e| &e.0 == child_path) { - let (_, child) = rlp_node_stack.pop().unwrap(); - branch_value_stack_buf.push(child); + buffers.branch_value_stack_buf.clear(); + for child_path in &buffers.branch_child_buf { + if buffers.rlp_node_stack.last().map_or(false, |e| &e.0 == child_path) { + let (_, child) = buffers.rlp_node_stack.pop().unwrap(); + buffers.branch_value_stack_buf.push(child); } else { - debug_assert!(branch_value_stack_buf.is_empty()); - path_stack.push(path); - path_stack.extend(branch_child_buf.drain(..)); + debug_assert!(buffers.branch_value_stack_buf.is_empty()); + buffers.path_stack.push(path); + buffers.path_stack.extend(buffers.branch_child_buf.drain(..)); 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 } } @@ -811,6 +778,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 + path_stack: Vec, + /// 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], + 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;