From 9e685163deb57b2a2e1661195c093071d1e85647 Mon Sep 17 00:00:00 2001 From: Federico Gimenez Date: Fri, 19 Jul 2024 11:40:33 +0200 Subject: [PATCH] chore: remove InMemoryState trait (#9642) --- crates/engine/tree/src/tree/mod.rs | 150 +-------------------- crates/engine/tree/src/tree/state.rs | 187 +++++++++++++++++++++++---- 2 files changed, 163 insertions(+), 174 deletions(-) diff --git a/crates/engine/tree/src/tree/mod.rs b/crates/engine/tree/src/tree/mod.rs index a3e4179318..7d2ab8286a 100644 --- a/crates/engine/tree/src/tree/mod.rs +++ b/crates/engine/tree/src/tree/mod.rs @@ -34,7 +34,7 @@ use reth_rpc_types::{ ExecutionPayload, }; use reth_trie::{updates::TrieUpdates, HashedPostState}; -pub use state::{BlockState, CanonicalInMemoryState, InMemoryState, InMemoryStateImpl}; +pub use state::{BlockState, CanonicalInMemoryState, InMemoryState}; use std::{ collections::{BTreeMap, HashMap}, marker::PhantomData, @@ -913,18 +913,11 @@ impl PersistenceState { #[cfg(test)] mod tests { use super::*; - use crate::{ - static_files::StaticFileAction, - test_utils::{ - get_executed_block_with_number, get_executed_block_with_receipts, get_executed_blocks, - }, - }; - use rand::Rng; + use crate::{static_files::StaticFileAction, test_utils::get_executed_blocks}; use reth_beacon_consensus::EthBeaconConsensus; use reth_chainspec::{ChainSpecBuilder, MAINNET}; use reth_ethereum_engine_primitives::EthEngineTypes; use reth_evm::test_utils::MockExecutorProvider; - use reth_primitives::Receipt; use reth_provider::test_utils::MockEthProvider; use std::sync::mpsc::{channel, Sender}; use tokio::sync::mpsc::unbounded_channel; @@ -1003,10 +996,6 @@ mod tests { TestHarness { tree, to_tree_tx, blocks, sf_action_rx } } - fn create_mock_state(block_number: u64) -> BlockState { - BlockState::new(get_executed_block_with_number(block_number)) - } - #[tokio::test] async fn test_tree_persist_blocks() { // we need more than PERSISTENCE_THRESHOLD blocks to trigger the @@ -1058,139 +1047,4 @@ mod tests { assert_eq!(expected_state, *actual_state_by_number); } } - - #[tokio::test] - async fn test_in_memory_state_impl_state_by_hash() { - let mut state_by_hash = HashMap::new(); - let number = rand::thread_rng().gen::(); - let state = Arc::new(create_mock_state(number)); - state_by_hash.insert(state.hash(), state.clone()); - - let in_memory_state = InMemoryStateImpl::new(state_by_hash, HashMap::new(), None); - - assert_eq!(in_memory_state.state_by_hash(state.hash()), Some(state)); - assert_eq!(in_memory_state.state_by_hash(B256::random()), None); - } - - #[tokio::test] - async fn test_in_memory_state_impl_state_by_number() { - let mut state_by_hash = HashMap::new(); - let mut hash_by_number = HashMap::new(); - - let number = rand::thread_rng().gen::(); - let state = Arc::new(create_mock_state(number)); - let hash = state.hash(); - - state_by_hash.insert(hash, state.clone()); - hash_by_number.insert(number, hash); - - let in_memory_state = InMemoryStateImpl::new(state_by_hash, hash_by_number, None); - - assert_eq!(in_memory_state.state_by_number(number), Some(state)); - assert_eq!(in_memory_state.state_by_number(number + 1), None); - } - - #[tokio::test] - async fn test_in_memory_state_impl_head_state() { - let mut state_by_hash = HashMap::new(); - let mut hash_by_number = HashMap::new(); - let state1 = Arc::new(create_mock_state(1)); - let state2 = Arc::new(create_mock_state(2)); - let hash1 = state1.hash(); - let hash2 = state2.hash(); - hash_by_number.insert(1, hash1); - hash_by_number.insert(2, hash2); - state_by_hash.insert(hash1, state1); - state_by_hash.insert(hash2, state2); - - let in_memory_state = InMemoryStateImpl::new(state_by_hash, hash_by_number, None); - let head_state = in_memory_state.head_state().unwrap(); - - assert_eq!(head_state.hash(), hash2); - assert_eq!(head_state.number(), 2); - } - - #[tokio::test] - async fn test_in_memory_state_impl_pending_state() { - let pending_number = rand::thread_rng().gen::(); - let pending_state = create_mock_state(pending_number); - let pending_hash = pending_state.hash(); - - let in_memory_state = - InMemoryStateImpl::new(HashMap::new(), HashMap::new(), Some(pending_state)); - - let result = in_memory_state.pending_state(); - assert!(result.is_some()); - let actual_pending_state = result.unwrap(); - assert_eq!(actual_pending_state.0.block().hash(), pending_hash); - assert_eq!(actual_pending_state.0.block().number, pending_number); - } - - #[tokio::test] - async fn test_in_memory_state_impl_no_pending_state() { - let in_memory_state = InMemoryStateImpl::new(HashMap::new(), HashMap::new(), None); - - assert_eq!(in_memory_state.pending_state(), None); - } - - #[tokio::test] - async fn test_state_new() { - let number = rand::thread_rng().gen::(); - let block = get_executed_block_with_number(number); - - let state = BlockState::new(block.clone()); - - assert_eq!(state.block(), block); - } - - #[tokio::test] - async fn test_state_block() { - let number = rand::thread_rng().gen::(); - let block = get_executed_block_with_number(number); - - let state = BlockState::new(block.clone()); - - assert_eq!(state.block(), block); - } - - #[tokio::test] - async fn test_state_hash() { - let number = rand::thread_rng().gen::(); - let block = get_executed_block_with_number(number); - - let state = BlockState::new(block.clone()); - - assert_eq!(state.hash(), block.block().hash()); - } - - #[tokio::test] - async fn test_state_number() { - let number = rand::thread_rng().gen::(); - let block = get_executed_block_with_number(number); - - let state = BlockState::new(block); - - assert_eq!(state.number(), number); - } - - #[tokio::test] - async fn test_state_state_root() { - let number = rand::thread_rng().gen::(); - let block = get_executed_block_with_number(number); - - let state = BlockState::new(block.clone()); - - assert_eq!(state.state_root(), block.block().state_root); - } - - #[tokio::test] - async fn test_state_receipts() { - let receipts = Receipts { receipt_vec: vec![vec![Some(Receipt::default())]] }; - - let block = get_executed_block_with_receipts(receipts.clone()); - - let state = BlockState::new(block); - - assert_eq!(state.receipts(), &receipts); - } } diff --git a/crates/engine/tree/src/tree/state.rs b/crates/engine/tree/src/tree/state.rs index 9e5056f345..1ec67b8032 100644 --- a/crates/engine/tree/src/tree/state.rs +++ b/crates/engine/tree/src/tree/state.rs @@ -8,13 +8,13 @@ use std::{collections::HashMap, sync::Arc}; /// Container type for in memory state data. #[derive(Debug, Default)] -pub struct InMemoryStateImpl { +pub struct InMemoryState { blocks: RwLock>>, numbers: RwLock>, pending: RwLock>, } -impl InMemoryStateImpl { +impl InMemoryState { pub(crate) const fn new( blocks: HashMap>, numbers: HashMap, @@ -26,18 +26,19 @@ impl InMemoryStateImpl { pending: RwLock::new(pending), } } -} -impl InMemoryState for InMemoryStateImpl { - fn state_by_hash(&self, hash: B256) -> Option> { + /// Returns the state for a given block hash. + pub(crate) fn state_by_hash(&self, hash: B256) -> Option> { self.blocks.read().get(&hash).cloned() } - fn state_by_number(&self, number: u64) -> Option> { + /// Returns the state for a given block number. + pub(crate) fn state_by_number(&self, number: u64) -> Option> { self.numbers.read().get(&number).and_then(|hash| self.blocks.read().get(hash).cloned()) } - fn head_state(&self) -> Option> { + /// Returns the current chain head state. + pub(crate) fn head_state(&self) -> Option> { self.numbers .read() .iter() @@ -45,7 +46,9 @@ impl InMemoryState for InMemoryStateImpl { .and_then(|(_, hash)| self.blocks.read().get(hash).cloned()) } - fn pending_state(&self) -> Option> { + /// Returns the pending state corresponding to the current head plus one, + /// from the payload received in newPayload that does not have a FCU yet. + pub(crate) fn pending_state(&self) -> Option> { self.pending.read().as_ref().map(|state| Arc::new(BlockState(state.0.clone()))) } } @@ -55,7 +58,7 @@ impl InMemoryState for InMemoryStateImpl { #[derive(Debug)] pub(crate) struct CanonicalInMemoryStateInner { pub(crate) chain_info_tracker: ChainInfoTracker, - pub(crate) in_memory_state: InMemoryStateImpl, + pub(crate) in_memory_state: InMemoryState, } /// This type is responsible for providing the blocks, receipts, and state for @@ -73,7 +76,7 @@ impl CanonicalInMemoryState { numbers: HashMap, pending: Option, ) -> Self { - let in_memory_state = InMemoryStateImpl::new(blocks, numbers, pending); + let in_memory_state = InMemoryState::new(blocks, numbers, pending); let head_state = in_memory_state.head_state(); let header = match head_state { Some(state) => state.block().block().header.clone(), @@ -88,14 +91,12 @@ impl CanonicalInMemoryState { /// Create a new in memory state with the given local head. pub fn with_head(head: SealedHeader) -> Self { let chain_info_tracker = ChainInfoTracker::new(head); - let in_memory_state = InMemoryStateImpl::default(); + let in_memory_state = InMemoryState::default(); let inner = CanonicalInMemoryStateInner { chain_info_tracker, in_memory_state }; Self { inner: Arc::new(inner) } } -} -impl InMemoryState for CanonicalInMemoryState { fn state_by_hash(&self, hash: B256) -> Option> { self.inner.in_memory_state.state_by_hash(hash) } @@ -113,19 +114,6 @@ impl InMemoryState for CanonicalInMemoryState { } } -/// Represents the tree state kept in memory. -pub trait InMemoryState: Send + Sync { - /// Returns the state for a given block hash. - fn state_by_hash(&self, hash: B256) -> Option>; - /// Returns the state for a given block number. - fn state_by_number(&self, number: u64) -> Option>; - /// Returns the current chain head state. - fn head_state(&self) -> Option>; - /// Returns the pending state corresponding to the current head plus one, - /// from the payload received in newPayload that does not have a FCU yet. - fn pending_state(&self) -> Option>; -} - /// State after applying the given block. #[derive(Debug, PartialEq, Eq, Clone)] pub struct BlockState(pub(crate) ExecutedBlock); @@ -155,3 +143,150 @@ impl BlockState { &self.0.execution_outcome().receipts } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::test_utils::{get_executed_block_with_number, get_executed_block_with_receipts}; + use rand::Rng; + use reth_primitives::Receipt; + + fn create_mock_state(block_number: u64) -> BlockState { + BlockState::new(get_executed_block_with_number(block_number)) + } + + #[tokio::test] + async fn test_in_memory_state_impl_state_by_hash() { + let mut state_by_hash = HashMap::new(); + let number = rand::thread_rng().gen::(); + let state = Arc::new(create_mock_state(number)); + state_by_hash.insert(state.hash(), state.clone()); + + let in_memory_state = InMemoryState::new(state_by_hash, HashMap::new(), None); + + assert_eq!(in_memory_state.state_by_hash(state.hash()), Some(state)); + assert_eq!(in_memory_state.state_by_hash(B256::random()), None); + } + + #[tokio::test] + async fn test_in_memory_state_impl_state_by_number() { + let mut state_by_hash = HashMap::new(); + let mut hash_by_number = HashMap::new(); + + let number = rand::thread_rng().gen::(); + let state = Arc::new(create_mock_state(number)); + let hash = state.hash(); + + state_by_hash.insert(hash, state.clone()); + hash_by_number.insert(number, hash); + + let in_memory_state = InMemoryState::new(state_by_hash, hash_by_number, None); + + assert_eq!(in_memory_state.state_by_number(number), Some(state)); + assert_eq!(in_memory_state.state_by_number(number + 1), None); + } + + #[tokio::test] + async fn test_in_memory_state_impl_head_state() { + let mut state_by_hash = HashMap::new(); + let mut hash_by_number = HashMap::new(); + let state1 = Arc::new(create_mock_state(1)); + let state2 = Arc::new(create_mock_state(2)); + let hash1 = state1.hash(); + let hash2 = state2.hash(); + hash_by_number.insert(1, hash1); + hash_by_number.insert(2, hash2); + state_by_hash.insert(hash1, state1); + state_by_hash.insert(hash2, state2); + + let in_memory_state = InMemoryState::new(state_by_hash, hash_by_number, None); + let head_state = in_memory_state.head_state().unwrap(); + + assert_eq!(head_state.hash(), hash2); + assert_eq!(head_state.number(), 2); + } + + #[tokio::test] + async fn test_in_memory_state_impl_pending_state() { + let pending_number = rand::thread_rng().gen::(); + let pending_state = create_mock_state(pending_number); + let pending_hash = pending_state.hash(); + + let in_memory_state = + InMemoryState::new(HashMap::new(), HashMap::new(), Some(pending_state)); + + let result = in_memory_state.pending_state(); + assert!(result.is_some()); + let actual_pending_state = result.unwrap(); + assert_eq!(actual_pending_state.0.block().hash(), pending_hash); + assert_eq!(actual_pending_state.0.block().number, pending_number); + } + + #[tokio::test] + async fn test_in_memory_state_impl_no_pending_state() { + let in_memory_state = InMemoryState::new(HashMap::new(), HashMap::new(), None); + + assert_eq!(in_memory_state.pending_state(), None); + } + + #[tokio::test] + async fn test_state_new() { + let number = rand::thread_rng().gen::(); + let block = get_executed_block_with_number(number); + + let state = BlockState::new(block.clone()); + + assert_eq!(state.block(), block); + } + + #[tokio::test] + async fn test_state_block() { + let number = rand::thread_rng().gen::(); + let block = get_executed_block_with_number(number); + + let state = BlockState::new(block.clone()); + + assert_eq!(state.block(), block); + } + + #[tokio::test] + async fn test_state_hash() { + let number = rand::thread_rng().gen::(); + let block = get_executed_block_with_number(number); + + let state = BlockState::new(block.clone()); + + assert_eq!(state.hash(), block.block().hash()); + } + + #[tokio::test] + async fn test_state_number() { + let number = rand::thread_rng().gen::(); + let block = get_executed_block_with_number(number); + + let state = BlockState::new(block); + + assert_eq!(state.number(), number); + } + + #[tokio::test] + async fn test_state_state_root() { + let number = rand::thread_rng().gen::(); + let block = get_executed_block_with_number(number); + + let state = BlockState::new(block.clone()); + + assert_eq!(state.state_root(), block.block().state_root); + } + + #[tokio::test] + async fn test_state_receipts() { + let receipts = Receipts { receipt_vec: vec![vec![Some(Receipt::default())]] }; + + let block = get_executed_block_with_receipts(receipts.clone()); + + let state = BlockState::new(block); + + assert_eq!(state.receipts(), &receipts); + } +}