perf(trie): reduce allocations in sparse trie rlp node calculation (#…

…12092)

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<u8>,
 }
 
+/// 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<bool>)>,
+    /// 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;