diff --git a/crates/chain-state/src/in_memory.rs b/crates/chain-state/src/in_memory.rs index 44a5fb8be66..8069046b159 100644 --- a/crates/chain-state/src/in_memory.rs +++ b/crates/chain-state/src/in_memory.rs @@ -958,7 +958,7 @@ impl> NewCanonicalChain { first.block_number(), first.execution_outcome().clone(), ), - first.trie_data_handle(), + first.trie_data_handle().to_lazy(), ); for exec in rest { chain.append_block( @@ -967,7 +967,7 @@ impl> NewCanonicalChain { exec.block_number(), exec.execution_outcome().clone(), ), - exec.trie_data_handle(), + exec.trie_data_handle().to_lazy(), ); } chain @@ -1577,14 +1577,14 @@ mod tests { // Compare execution outcome assert_eq!(*new.execution_outcome(), commit_execution_outcome); - // Compare trie data by waiting on deferred data + // Compare trie data for (block_num, expected_updates) in &expected_trie_updates { - let actual = new.trie_data_at(*block_num).unwrap().wait_cloned(); - assert_eq!(actual.trie_updates, *expected_updates); + let actual = new.trie_data_at(*block_num).unwrap(); + assert_eq!(actual.trie_updates(), *expected_updates); } for (block_num, expected_state) in &expected_hashed_state { - let actual = new.trie_data_at(*block_num).unwrap().wait_cloned(); - assert_eq!(actual.hashed_state, *expected_state); + let actual = new.trie_data_at(*block_num).unwrap(); + assert_eq!(actual.hashed_state(), *expected_state); } // Test reorg notification @@ -1639,12 +1639,12 @@ mod tests { // Compare old chain trie data for (block_num, expected_updates) in &old_trie_updates { - let actual = old.trie_data_at(*block_num).unwrap().wait_cloned(); - assert_eq!(actual.trie_updates, *expected_updates); + let actual = old.trie_data_at(*block_num).unwrap(); + assert_eq!(actual.trie_updates(), *expected_updates); } for (block_num, expected_state) in &old_hashed_state { - let actual = old.trie_data_at(*block_num).unwrap().wait_cloned(); - assert_eq!(actual.hashed_state, *expected_state); + let actual = old.trie_data_at(*block_num).unwrap(); + assert_eq!(actual.hashed_state(), *expected_state); } // Compare new chain blocks @@ -1658,12 +1658,12 @@ mod tests { // Compare new chain trie data for (block_num, expected_updates) in &new_trie_updates { - let actual = new.trie_data_at(*block_num).unwrap().wait_cloned(); - assert_eq!(actual.trie_updates, *expected_updates); + let actual = new.trie_data_at(*block_num).unwrap(); + assert_eq!(actual.trie_updates(), *expected_updates); } for (block_num, expected_state) in &new_hashed_state { - let actual = new.trie_data_at(*block_num).unwrap().wait_cloned(); - assert_eq!(actual.hashed_state, *expected_state); + let actual = new.trie_data_at(*block_num).unwrap(); + assert_eq!(actual.hashed_state(), *expected_state); } } } diff --git a/crates/evm/chain/src/chain.rs b/crates/evm/chain/src/chain.rs index 7cd3c4a88cd..915c24ff817 100644 --- a/crates/evm/chain/src/chain.rs +++ b/crates/evm/chain/src/chain.rs @@ -1,6 +1,5 @@ //! Contains [Chain], a chain of blocks and their final state. -use crate::DeferredTrieData; use alloy_consensus::{transaction::Recovered, BlockHeader}; use alloy_eips::{eip1898::ForkBlock, eip2718::Encodable2718, BlockNumHash}; use alloy_primitives::{Address, BlockHash, BlockNumber, TxHash}; @@ -9,7 +8,7 @@ use reth_primitives_traits::{ transaction::signed::SignedTransaction, Block, BlockBody, IndexedTx, NodePrimitives, RecoveredBlock, SealedHeader, }; -use reth_trie_common::{updates::TrieUpdatesSorted, HashedPostStateSorted}; +use reth_trie_common::{updates::TrieUpdatesSorted, HashedPostStateSorted, LazyTrieData}; use std::{borrow::Cow, collections::BTreeMap, fmt, ops::RangeInclusive, sync::Arc, vec::Vec}; /// A chain of blocks and their final state. @@ -33,12 +32,12 @@ pub struct Chain { /// /// Additionally, it includes the individual state changes that led to the current state. execution_outcome: ExecutionOutcome, - /// Deferred trie data for each block in the chain, keyed by block number. + /// Lazy trie data for each block in the chain, keyed by block number. /// - /// Contains handles to lazily-computed sorted trie updates and hashed state. + /// Contains handles to lazily-initialized sorted trie updates and hashed state. /// This allows Chain to be constructed without blocking on expensive trie /// computations - the data is only materialized when actually needed. - trie_data: BTreeMap, + trie_data: BTreeMap, } type ChainTxReceiptMeta<'a, N> = ( @@ -67,7 +66,7 @@ impl Chain { pub fn new( blocks: impl IntoIterator>, execution_outcome: ExecutionOutcome, - trie_data: BTreeMap, + trie_data: BTreeMap, ) -> Self { let blocks = blocks.into_iter().map(|b| (b.header().number(), b)).collect::>(); @@ -80,7 +79,7 @@ impl Chain { pub fn from_block( block: RecoveredBlock, execution_outcome: ExecutionOutcome, - trie_data: DeferredTrieData, + trie_data: LazyTrieData, ) -> Self { let block_number = block.header().number(); let trie_data_map = BTreeMap::from([(block_number, trie_data)]); @@ -102,42 +101,38 @@ impl Chain { self.blocks.values().map(|block| block.clone_sealed_header()) } - /// Get all deferred trie data for this chain. + /// Get all lazy trie data for this chain. /// - /// Returns handles to lazily-computed sorted trie updates and hashed state. - /// [`DeferredTrieData`] allows `Chain` to be constructed without blocking on - /// expensive trie computations - the data is only materialized when actually needed - /// via [`DeferredTrieData::wait_cloned`] or similar methods. - /// - /// This method does **not** block. To access the computed trie data, call - /// [`DeferredTrieData::wait_cloned`] on individual entries, which will block - /// if the background computation has not yet completed. - pub const fn trie_data(&self) -> &BTreeMap { + /// Returns handles to lazily-initialized sorted trie updates and hashed state. + /// [`LazyTrieData`] allows `Chain` to be constructed without blocking on + /// expensive trie computations - the data is only materialized when actually needed. + pub const fn trie_data(&self) -> &BTreeMap { &self.trie_data } - /// Get deferred trie data for a specific block number. + /// Get lazy trie data for a specific block number. /// - /// Returns a handle to the lazily-computed trie data. This method does **not** block. - /// Call [`DeferredTrieData::wait_cloned`] on the result to wait for and retrieve - /// the computed data, which will block if computation is still in progress. - pub fn trie_data_at(&self, block_number: BlockNumber) -> Option<&DeferredTrieData> { + /// Returns a handle to the lazily-initialized trie data. + pub fn trie_data_at(&self, block_number: BlockNumber) -> Option<&LazyTrieData> { self.trie_data.get(&block_number) } /// Get all trie updates for this chain. /// - /// Note: This blocks on deferred trie data for all blocks in the chain. - /// Prefer using [`trie_data`](Self::trie_data) when possible to avoid blocking. + /// # Panics + /// + /// Panics if any trie data has not been initialized. pub fn trie_updates(&self) -> BTreeMap> { - self.trie_data.iter().map(|(num, data)| (*num, data.wait_cloned().trie_updates)).collect() + self.trie_data.iter().map(|(num, data)| (*num, data.trie_updates())).collect() } /// Get trie updates for a specific block number. /// - /// Note: This waits for deferred trie data if not already computed. + /// # Panics + /// + /// Panics if the trie data for this block has not been initialized. pub fn trie_updates_at(&self, block_number: BlockNumber) -> Option> { - self.trie_data.get(&block_number).map(|data| data.wait_cloned().trie_updates) + self.trie_data.get(&block_number).map(|data| data.trie_updates()) } /// Remove all trie data for this chain. @@ -147,17 +142,20 @@ impl Chain { /// Get all hashed states for this chain. /// - /// Note: This blocks on deferred trie data for all blocks in the chain. - /// Prefer using [`trie_data`](Self::trie_data) when possible to avoid blocking. + /// # Panics + /// + /// Panics if any trie data has not been initialized. pub fn hashed_state(&self) -> BTreeMap> { - self.trie_data.iter().map(|(num, data)| (*num, data.wait_cloned().hashed_state)).collect() + self.trie_data.iter().map(|(num, data)| (*num, data.hashed_state())).collect() } /// Get hashed state for a specific block number. /// - /// Note: This waits for deferred trie data if not already computed. + /// # Panics + /// + /// Panics if the trie data for this block has not been initialized. pub fn hashed_state_at(&self, block_number: BlockNumber) -> Option> { - self.trie_data.get(&block_number).map(|data| data.wait_cloned().hashed_state) + self.trie_data.get(&block_number).map(|data| data.hashed_state()) } /// Get execution outcome of this chain @@ -205,14 +203,14 @@ impl Chain { /// Destructure the chain into its inner components: /// 1. The blocks contained in the chain. /// 2. The execution outcome representing the final state. - /// 3. The deferred trie data map. + /// 3. The lazy trie data map. #[allow(clippy::type_complexity)] pub fn into_inner( self, ) -> ( ChainBlocks<'static, N::Block>, ExecutionOutcome, - BTreeMap, + BTreeMap, ) { (ChainBlocks { blocks: Cow::Owned(self.blocks) }, self.execution_outcome, self.trie_data) } @@ -344,7 +342,7 @@ impl Chain { &mut self, block: RecoveredBlock, execution_outcome: ExecutionOutcome, - trie_data: DeferredTrieData, + trie_data: LazyTrieData, ) { let block_number = block.header().number(); self.blocks.insert(block_number, block); @@ -492,10 +490,9 @@ pub struct BlockReceipts { #[cfg(feature = "serde")] mod chain_serde { use super::*; - use crate::ComputedTrieData; use serde::{Deserialize, Deserializer, Serialize, Serializer}; - /// Serializable representation of Chain that waits for deferred trie data. + /// Serializable representation of Chain that waits for lazy trie data. #[derive(Serialize, Deserialize)] #[serde(bound = "")] struct ChainRepr { @@ -512,17 +509,11 @@ mod chain_serde { where S: Serializer, { - // Wait for deferred trie data for serialization - let trie_updates: BTreeMap<_, _> = self - .trie_data - .iter() - .map(|(num, data)| (*num, data.wait_cloned().trie_updates)) - .collect(); - let hashed_state: BTreeMap<_, _> = self - .trie_data - .iter() - .map(|(num, data)| (*num, data.wait_cloned().hashed_state)) - .collect(); + // Get trie data for serialization (panics if not initialized) + let trie_updates: BTreeMap<_, _> = + self.trie_data.iter().map(|(num, data)| (*num, data.trie_updates())).collect(); + let hashed_state: BTreeMap<_, _> = + self.trie_data.iter().map(|(num, data)| (*num, data.hashed_state())).collect(); let repr = ChainRepr:: { blocks: self.blocks.clone(), @@ -541,14 +532,13 @@ mod chain_serde { { let repr = ChainRepr::::deserialize(deserializer)?; - // Convert to ready DeferredTrieData handles + // Convert to ready LazyTrieData handles let trie_data = repr .trie_updates .into_iter() .map(|(num, trie_updates)| { let hashed_state = repr.hashed_state.get(&num).cloned().unwrap_or_default(); - let computed = ComputedTrieData::without_trie_input(hashed_state, trie_updates); - (num, DeferredTrieData::ready(computed)) + (num, LazyTrieData::ready(hashed_state, trie_updates)) }) .collect(); @@ -660,7 +650,7 @@ pub(super) mod serde_bincode_compat { >, { fn from(value: Chain<'a, N>) -> Self { - use crate::{ComputedTrieData, DeferredTrieData}; + use super::LazyTrieData; let trie_updates: BTreeMap<_, _> = value.trie_updates.into_iter().map(|(k, v)| (k, Arc::new(v.into()))).collect(); @@ -671,8 +661,7 @@ pub(super) mod serde_bincode_compat { .into_iter() .map(|(num, trie_updates)| { let hashed_state = hashed_state.get(&num).cloned().unwrap_or_default(); - let computed = ComputedTrieData::without_trie_input(hashed_state, trie_updates); - (num, DeferredTrieData::ready(computed)) + (num, LazyTrieData::ready(hashed_state, trie_updates)) }) .collect(); @@ -696,11 +685,11 @@ pub(super) mod serde_bincode_compat { { use reth_trie_common::serde_bincode_compat as trie_serde; - // Wait for deferred trie data and collect into maps we can borrow from + // Get trie data and collect into maps we can borrow from let trie_updates_data: BTreeMap = - source.trie_data.iter().map(|(k, v)| (*k, v.wait_cloned().trie_updates)).collect(); + source.trie_data.iter().map(|(k, v)| (*k, v.trie_updates())).collect(); let hashed_state_data: BTreeMap = - source.trie_data.iter().map(|(k, v)| (*k, v.wait_cloned().hashed_state)).collect(); + source.trie_data.iter().map(|(k, v)| (*k, v.hashed_state())).collect(); // Now create the serde-compatible struct borrowing from the collected data let chain: Chain<'_, N> = Chain { diff --git a/crates/evm/chain/src/deferred_trie.rs b/crates/evm/chain/src/deferred_trie.rs index 03d7184fd72..0ac198f8348 100644 --- a/crates/evm/chain/src/deferred_trie.rs +++ b/crates/evm/chain/src/deferred_trie.rs @@ -3,7 +3,7 @@ use parking_lot::Mutex; use reth_metrics::{metrics::Counter, Metrics}; use reth_trie::{ updates::{TrieUpdates, TrieUpdatesSorted}, - HashedPostState, HashedPostStateSorted, TrieInputSorted, + HashedPostState, HashedPostStateSorted, LazyTrieData, SortedTrieData, TrieInputSorted, }; use std::{ fmt, @@ -231,6 +231,24 @@ impl DeferredTrieData { } } } + + /// Converts this [`DeferredTrieData`] to a [`LazyTrieData`]. + /// + /// The computation is deferred - `sorted_trie_data()` will only be called when + /// the `LazyTrieData` is first accessed. This allows non-blocking conversion. + pub fn to_lazy(&self) -> LazyTrieData { + let this = self.clone(); + LazyTrieData::deferred(move || this.sorted_trie_data()) + } + + /// Returns the sorted trie data, waiting for computation if necessary. + /// + /// This is a convenience method that bundles both hashed state and trie updates + /// into a [`SortedTrieData`] container. + pub fn sorted_trie_data(&self) -> SortedTrieData { + let computed = self.wait_cloned(); + SortedTrieData::new(computed.hashed_state, computed.trie_updates) + } } /// Sorted trie data computed for an executed block. @@ -285,6 +303,11 @@ impl ComputedTrieData { pub fn trie_input(&self) -> Option<&Arc> { self.anchored_trie_input.as_ref().map(|anchored| &anchored.trie_input) } + + /// Returns a [`SortedTrieData`] containing the hashed state and trie updates. + pub fn sorted_trie_data(&self) -> SortedTrieData { + SortedTrieData::new(Arc::clone(&self.hashed_state), Arc::clone(&self.trie_updates)) + } } /// Trie input bundled with its anchor hash. diff --git a/crates/evm/chain/src/lib.rs b/crates/evm/chain/src/lib.rs index 38e7485de10..ab86be1ff85 100644 --- a/crates/evm/chain/src/lib.rs +++ b/crates/evm/chain/src/lib.rs @@ -1,7 +1,8 @@ -//! Chain and deferred trie data types for reth. +//! Chain and trie data types for reth. //! //! This crate contains the [`Chain`] type representing a chain of blocks and their final state, -//! as well as [`DeferredTrieData`] for handling asynchronously computed trie data. +//! as well as [`DeferredTrieData`] for handling asynchronously computed trie data and +//! re-exports [`LazyTrieData`] for lazy initialization. #![doc( html_logo_url = "https://raw.githubusercontent.com/paradigmxyz/reth/main/assets/reth-docs.png", @@ -17,6 +18,9 @@ pub use chain::*; mod deferred_trie; pub use deferred_trie::*; +// Re-export LazyTrieData from trie-common for convenience +pub use reth_trie_common::LazyTrieData; + /// Bincode-compatible serde implementations for chain types. /// /// `bincode` crate doesn't work with optionally serializable serde fields, but some of the diff --git a/crates/exex/exex/src/wal/storage.rs b/crates/exex/exex/src/wal/storage.rs index bb118c8a98a..bbf4a401767 100644 --- a/crates/exex/exex/src/wal/storage.rs +++ b/crates/exex/exex/src/wal/storage.rs @@ -357,7 +357,7 @@ mod tests { new: Arc::new(Chain::new( vec![block], Default::default(), - BTreeMap::from([(block_number, trie_data)]), + BTreeMap::from([(block_number, trie_data.to_lazy())]), )), }; Ok(notification) diff --git a/crates/trie/common/src/lazy.rs b/crates/trie/common/src/lazy.rs new file mode 100644 index 00000000000..bc72c43c1de --- /dev/null +++ b/crates/trie/common/src/lazy.rs @@ -0,0 +1,164 @@ +//! Lazy initialization wrapper for trie data. +//! +//! Provides a no-std compatible [`LazyTrieData`] type for lazily initialized +//! trie-related data containing [`HashedPostStateSorted`] and [`TrieUpdatesSorted`]. + +use crate::{updates::TrieUpdatesSorted, HashedPostStateSorted}; +use alloc::sync::Arc; +use core::fmt; +use reth_primitives_traits::sync::OnceLock; + +/// Container for sorted trie data: hashed state and trie updates. +/// +/// This bundles both [`HashedPostStateSorted`] and [`TrieUpdatesSorted`] together +/// for convenient passing and storage. +#[derive(Clone, Debug, Default)] +pub struct SortedTrieData { + /// Sorted hashed post-state produced by execution. + pub hashed_state: Arc, + /// Sorted trie updates produced by state root computation. + pub trie_updates: Arc, +} + +impl SortedTrieData { + /// Creates a new [`SortedTrieData`] with the given values. + pub const fn new( + hashed_state: Arc, + trie_updates: Arc, + ) -> Self { + Self { hashed_state, trie_updates } + } +} + +/// Lazily initialized trie data containing sorted hashed state and trie updates. +/// +/// This is a no-std compatible wrapper that supports two modes: +/// 1. **Ready mode**: Data is available immediately (created via `ready()`) +/// 2. **Deferred mode**: Data is computed on first access (created via `deferred()`) +/// +/// In deferred mode, the computation runs on the first call to `get()`, `hashed_state()`, +/// or `trie_updates()`, and results are cached for subsequent calls. +/// +/// Cloning is cheap (Arc clone) and clones share the cached state. +pub struct LazyTrieData { + /// Cached sorted trie data, computed on first access. + data: Arc>, + /// Optional deferred computation function. + compute: Option SortedTrieData + Send + Sync>>, +} + +impl Clone for LazyTrieData { + fn clone(&self) -> Self { + Self { data: Arc::clone(&self.data), compute: self.compute.clone() } + } +} + +impl fmt::Debug for LazyTrieData { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("LazyTrieData") + .field("data", &if self.data.get().is_some() { "initialized" } else { "pending" }) + .finish() + } +} + +impl LazyTrieData { + /// Creates a new [`LazyTrieData`] that is already initialized with the given values. + pub fn ready( + hashed_state: Arc, + trie_updates: Arc, + ) -> Self { + let data = OnceLock::new(); + let _ = data.set(SortedTrieData::new(hashed_state, trie_updates)); + Self { data: Arc::new(data), compute: None } + } + + /// Creates a new [`LazyTrieData`] from pre-computed [`SortedTrieData`]. + pub fn from_sorted(sorted: SortedTrieData) -> Self { + let data = OnceLock::new(); + let _ = data.set(sorted); + Self { data: Arc::new(data), compute: None } + } + + /// Creates a new [`LazyTrieData`] with a deferred computation function. + /// + /// The computation will run on the first call to `get()`, `hashed_state()`, + /// or `trie_updates()`. Results are cached for subsequent calls. + pub fn deferred(compute: impl Fn() -> SortedTrieData + Send + Sync + 'static) -> Self { + Self { data: Arc::new(OnceLock::new()), compute: Some(Arc::new(compute)) } + } + + /// Returns a reference to the sorted trie data, computing if necessary. + /// + /// # Panics + /// + /// Panics if created via `deferred()` and the computation function was not provided. + pub fn get(&self) -> &SortedTrieData { + self.data.get_or_init(|| { + self.compute.as_ref().expect("LazyTrieData::get called before initialization")() + }) + } + + /// Returns a clone of the hashed state Arc. + /// + /// If not initialized, computes from the deferred source or panics. + pub fn hashed_state(&self) -> Arc { + Arc::clone(&self.get().hashed_state) + } + + /// Returns a clone of the trie updates Arc. + /// + /// If not initialized, computes from the deferred source or panics. + pub fn trie_updates(&self) -> Arc { + Arc::clone(&self.get().trie_updates) + } + + /// Returns a clone of the [`SortedTrieData`]. + /// + /// If not initialized, computes from the deferred source or panics. + pub fn sorted_trie_data(&self) -> SortedTrieData { + self.get().clone() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_lazy_ready_is_initialized() { + let lazy = LazyTrieData::ready( + Arc::new(HashedPostStateSorted::default()), + Arc::new(TrieUpdatesSorted::default()), + ); + let _ = lazy.hashed_state(); + let _ = lazy.trie_updates(); + } + + #[test] + fn test_lazy_clone_shares_state() { + let lazy1 = LazyTrieData::ready( + Arc::new(HashedPostStateSorted::default()), + Arc::new(TrieUpdatesSorted::default()), + ); + let lazy2 = lazy1.clone(); + + // Both point to the same data + assert!(Arc::ptr_eq(&lazy1.hashed_state(), &lazy2.hashed_state())); + assert!(Arc::ptr_eq(&lazy1.trie_updates(), &lazy2.trie_updates())); + } + + #[test] + fn test_lazy_deferred() { + let lazy = LazyTrieData::deferred(SortedTrieData::default); + assert!(lazy.hashed_state().is_empty()); + assert!(lazy.trie_updates().is_empty()); + } + + #[test] + fn test_lazy_from_sorted() { + let sorted = SortedTrieData::default(); + let lazy = LazyTrieData::from_sorted(sorted); + assert!(lazy.hashed_state().is_empty()); + assert!(lazy.trie_updates().is_empty()); + } +} diff --git a/crates/trie/common/src/lib.rs b/crates/trie/common/src/lib.rs index 8faa44622fa..5cb346f2a2b 100644 --- a/crates/trie/common/src/lib.rs +++ b/crates/trie/common/src/lib.rs @@ -11,6 +11,9 @@ extern crate alloc; +mod lazy; +pub use lazy::{LazyTrieData, SortedTrieData}; + /// In-memory hashed state. mod hashed_state; pub use hashed_state::*;