perf(trie): reuse buf for rlp encoding on HashBuilder (#2374)

bin/reth/src/stage/mod.rs

@@ -16,7 +16,7 @@ use reth_staged_sync::{
     Config,
 };
 use reth_stages::{
-    stages::{BodyStage, ExecutionStage, SenderRecoveryStage, TransactionLookupStage},
+    stages::{BodyStage, ExecutionStage, MerkleStage, SenderRecoveryStage, TransactionLookupStage},
     ExecInput, Stage, StageId, UnwindInput,
 };
 use std::{net::SocketAddr, sync::Arc};
@@ -193,6 +193,16 @@ impl Command {
 
                 stage.execute(&mut tx, input).await?;
             }
+            StageEnum::Merkle => {
+                let mut stage = MerkleStage::default_execution();
+
+                // Unwind first
+                if !self.skip_unwind {
+                    stage.unwind(&mut tx, unwind).await?;
+                }
+
+                stage.execute(&mut tx, input).await?;
+            }
             _ => {}
         }
 

crates/trie/src/hash_builder.rs

@@ -46,6 +46,8 @@ pub struct HashBuilder {
     stored_in_database: bool,
 
     updated_branch_nodes: Option<HashMap<Nibbles, BranchNodeCompact>>,
+
+    rlp_buf: Vec<u8>,
 }
 
 impl From<HashBuilderState> for HashBuilder {
@@ -59,6 +61,7 @@ impl From<HashBuilderState> for HashBuilder {
             hash_masks: state.hash_masks,
             stored_in_database: state.stored_in_database,
             updated_branch_nodes: None,
+            rlp_buf: Vec::with_capacity(32),
         }
     }
 }
@@ -241,8 +244,13 @@ impl HashBuilder {
                     HashBuilderValue::Bytes(leaf_value) => {
                         let leaf_node = LeafNode::new(&short_node_key, leaf_value);
                         tracing::debug!(target: "trie::hash_builder", ?leaf_node, "pushing leaf node");
-                        tracing::trace!(target: "trie::hash_builder", rlp = hex::encode(&leaf_node.rlp()), "leaf node rlp");
-                        self.stack.push(leaf_node.rlp());
+                        tracing::trace!(target: "trie::hash_builder", rlp = {
+                            self.rlp_buf.clear();
+                            hex::encode(&leaf_node.rlp(&mut self.rlp_buf))
+                        }, "leaf node rlp");
+
+                        self.rlp_buf.clear();
+                        self.stack.push(leaf_node.rlp(&mut self.rlp_buf));
                     }
                     HashBuilderValue::Hash(hash) => {
                         tracing::debug!(target: "trie::hash_builder", ?hash, "pushing branch node hash");
@@ -266,8 +274,12 @@ impl HashBuilder {
                     self.stack.pop().expect("there should be at least one stack item; qed");
                 let extension_node = ExtensionNode::new(&short_node_key, &stack_last);
                 tracing::debug!(target: "trie::hash_builder", ?extension_node, "pushing extension node");
-                tracing::trace!(target: "trie::hash_builder", rlp = hex::encode(&extension_node.rlp()), "extension node rlp");
-                self.stack.push(extension_node.rlp());
+                tracing::trace!(target: "trie::hash_builder", rlp = {
+                    self.rlp_buf.clear();
+                    hex::encode(&extension_node.rlp(&mut self.rlp_buf))
+                }, "extension node rlp");
+                self.rlp_buf.clear();
+                self.stack.push(extension_node.rlp(&mut self.rlp_buf));
                 self.resize_masks(len_from);
             }
 
@@ -315,7 +327,9 @@ impl HashBuilder {
         let hash_mask = self.hash_masks[len];
         let branch_node = BranchNode::new(&self.stack);
         let children = branch_node.children(state_mask, hash_mask).collect();
-        let rlp = branch_node.rlp(state_mask);
+
+        self.rlp_buf.clear();
+        let rlp = branch_node.rlp(state_mask, &mut self.rlp_buf);
 
         // Clears the stack from the branch node elements
         let first_child_idx = self.stack.len() - state_mask.count_ones() as usize;

crates/trie/src/nodes/branch.rs

@@ -1,5 +1,5 @@
 use super::{rlp_node, CHILD_INDEX_RANGE};
-use reth_primitives::{bytes::BytesMut, trie::TrieMask, H256};
+use reth_primitives::{trie::TrieMask, H256};
 use reth_rlp::{BufMut, EMPTY_STRING_CODE};
 
 /// A Branch node is only a pointer to the stack of nodes and is used to
@@ -38,9 +38,8 @@ impl<'a> BranchNode<'a> {
     }
 
     /// Returns the RLP encoding of the branch node given the state mask of children present.
-    pub fn rlp(&self, state_mask: TrieMask) -> Vec<u8> {
+    pub fn rlp(&self, state_mask: TrieMask, buf: &mut Vec<u8>) -> Vec<u8> {
         let first_child_idx = self.stack.len() - state_mask.count_ones() as usize;
-        let mut buf = BytesMut::new();
 
         // Create the RLP header from the mask elements present.
         let mut i = first_child_idx;
@@ -56,7 +55,7 @@ impl<'a> BranchNode<'a> {
                 header
             },
         );
-        header.encode(&mut buf);
+        header.encode(buf);
 
         // Extend the RLP buffer with the present children
         let mut i = first_child_idx;
@@ -72,6 +71,6 @@ impl<'a> BranchNode<'a> {
         // Is this needed?
         buf.put_u8(EMPTY_STRING_CODE);
 
-        rlp_node(&buf)
+        rlp_node(buf)
     }
 }

crates/trie/src/nodes/extension.rs

@@ -1,6 +1,5 @@
 use super::rlp_node;
 use crate::Nibbles;
-use reth_primitives::bytes::BytesMut;
 use reth_rlp::{BufMut, Encodable};
 
 /// An intermediate node that exists solely to compress the trie's paths. It contains a path segment
@@ -24,10 +23,9 @@ impl<'a> ExtensionNode<'a> {
     }
 
     /// RLP encodes the node and returns either RLP(Node) or RLP(keccak(RLP(node))).
-    pub fn rlp(&self) -> Vec<u8> {
-        let mut buf = BytesMut::new();
-        self.encode(&mut buf);
-        rlp_node(&buf)
+    pub fn rlp(&self, buf: &mut Vec<u8>) -> Vec<u8> {
+        self.encode(buf);
+        rlp_node(buf)
     }
 }
 

crates/trie/src/nodes/leaf.rs

@@ -1,6 +1,5 @@
 use super::rlp_node;
 use crate::Nibbles;
-use reth_primitives::bytes::BytesMut;
 use reth_rlp::{BufMut, Encodable};
 
 /// A leaf node represents the endpoint or terminal node in the trie. In other words, a leaf node is
@@ -26,10 +25,9 @@ impl<'a> LeafNode<'a> {
 
     /// RLP encodes the node and returns either RLP(Node) or RLP(keccak(RLP(node)))
     /// depending on if the serialized node was longer than a keccak).
-    pub fn rlp(&self) -> Vec<u8> {
-        let mut out = BytesMut::new();
-        self.encode(&mut out);
-        rlp_node(&out)
+    pub fn rlp(&self, out: &mut Vec<u8>) -> Vec<u8> {
+        self.encode(out);
+        rlp_node(out)
     }
 }
 
@@ -73,7 +71,7 @@ mod tests {
         let nibble = Nibbles { hex_data: hex!("0604060f").to_vec() };
         let val = hex!("76657262").to_vec();
         let leaf = LeafNode::new(&nibble, &val);
-        let rlp = leaf.rlp();
+        let rlp = leaf.rlp(&mut vec![]);
 
         let expected = hex!("c98320646f8476657262").to_vec();
         assert_eq!(rlp, expected);

crates/trie/src/trie.rs

@@ -230,6 +230,8 @@ where
         walker.set_updates(retain_updates);
         hash_builder.set_updates(retain_updates);
 
+        let mut account_rlp = Vec::with_capacity(128);
+
         while let Some(key) = last_walker_key.take().or_else(|| walker.key()) {
             // Take the last account key to make sure we take it into consideration only once.
             let (next_key, mut next_account_entry) = match last_account_key.take() {
@@ -289,7 +291,8 @@ where
                 };
 
                 let account = EthAccount::from(account).with_storage_root(storage_root);
-                let mut account_rlp = Vec::with_capacity(account.length());
+
+                account_rlp.clear();
                 account.encode(&mut &mut account_rlp);
 
                 hash_builder.add_leaf(account_nibbles, &account_rlp);

crates/trie/src/trie_cursor/storage_cursor.rs

@@ -81,6 +81,6 @@ mod tests {
             .unwrap();
 
         let mut cursor = StorageTrieCursor::new(cursor, hashed_address);
-        assert_eq!(cursor.seek(key.clone().into()).unwrap().unwrap().1, value);
+        assert_eq!(cursor.seek(key.into()).unwrap().unwrap().1, value);
     }
 }