From dbb0a357087c0c8b89bd15209b4ceb2a32bbc59b Mon Sep 17 00:00:00 2001 From: joshieDo <93316087+joshieDo@users.noreply.github.com> Date: Tue, 25 Apr 2023 05:31:09 +0800 Subject: [PATCH] perf(trie): reuse `buf` for rlp encoding on `HashBuilder` (#2374) --- bin/reth/src/stage/mod.rs | 12 +++++++++- crates/trie/src/hash_builder.rs | 24 +++++++++++++++---- crates/trie/src/nodes/branch.rs | 9 ++++--- crates/trie/src/nodes/extension.rs | 8 +++---- crates/trie/src/nodes/leaf.rs | 10 ++++---- crates/trie/src/trie.rs | 5 +++- crates/trie/src/trie_cursor/storage_cursor.rs | 2 +- 7 files changed, 46 insertions(+), 24 deletions(-) diff --git a/bin/reth/src/stage/mod.rs b/bin/reth/src/stage/mod.rs index 2f1f122bee..8fd40c9528 100644 --- a/bin/reth/src/stage/mod.rs +++ b/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?; + } _ => {} } diff --git a/crates/trie/src/hash_builder.rs b/crates/trie/src/hash_builder.rs index d99fcf0796..0e071323b3 100644 --- a/crates/trie/src/hash_builder.rs +++ b/crates/trie/src/hash_builder.rs @@ -46,6 +46,8 @@ pub struct HashBuilder { stored_in_database: bool, updated_branch_nodes: Option>, + + rlp_buf: Vec, } impl From for HashBuilder { @@ -59,6 +61,7 @@ impl From 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; diff --git a/crates/trie/src/nodes/branch.rs b/crates/trie/src/nodes/branch.rs index c6b58dfeb1..899b165f45 100644 --- a/crates/trie/src/nodes/branch.rs +++ b/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 { + pub fn rlp(&self, state_mask: TrieMask, buf: &mut Vec) -> Vec { 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) } } diff --git a/crates/trie/src/nodes/extension.rs b/crates/trie/src/nodes/extension.rs index a998d4d8a7..73f8467617 100644 --- a/crates/trie/src/nodes/extension.rs +++ b/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 { - let mut buf = BytesMut::new(); - self.encode(&mut buf); - rlp_node(&buf) + pub fn rlp(&self, buf: &mut Vec) -> Vec { + self.encode(buf); + rlp_node(buf) } } diff --git a/crates/trie/src/nodes/leaf.rs b/crates/trie/src/nodes/leaf.rs index 516be22e64..9a7559f513 100644 --- a/crates/trie/src/nodes/leaf.rs +++ b/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 { - let mut out = BytesMut::new(); - self.encode(&mut out); - rlp_node(&out) + pub fn rlp(&self, out: &mut Vec) -> Vec { + 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); diff --git a/crates/trie/src/trie.rs b/crates/trie/src/trie.rs index 9596a46fef..bf82df800f 100644 --- a/crates/trie/src/trie.rs +++ b/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); diff --git a/crates/trie/src/trie_cursor/storage_cursor.rs b/crates/trie/src/trie_cursor/storage_cursor.rs index 4cabb8d9d3..ae0362b4c8 100644 --- a/crates/trie/src/trie_cursor/storage_cursor.rs +++ b/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); } }