reth_revm/
cached.rs

1//! Database adapters for payload building.
2use alloy_primitives::{
3    map::{Entry, HashMap},
4    Address, B256, U256,
5};
6use core::cell::RefCell;
7use revm::{bytecode::Bytecode, state::AccountInfo, Database, DatabaseRef};
8
9/// A container type that caches reads from an underlying [`DatabaseRef`].
10///
11/// This is intended to be used in conjunction with `revm::db::State`
12/// during payload building which repeatedly accesses the same data.
13///
14/// [`CachedReads::as_db_mut`] transforms this type into a [`Database`] implementation that uses
15/// [`CachedReads`] as a caching layer for operations, and records any cache misses.
16///
17/// # Example
18///
19/// ```
20/// use reth_revm::{cached::CachedReads, DatabaseRef, db::State};
21///
22/// fn build_payload<DB: DatabaseRef>(db: DB) {
23///     let mut cached_reads = CachedReads::default();
24///     let db = cached_reads.as_db_mut(db);
25///     // this is `Database` and can be used to build a payload, it never commits to `CachedReads` or the underlying database, but all reads from the underlying database are cached in `CachedReads`.
26///     // Subsequent payload build attempts can use cached reads and avoid hitting the underlying database.
27///     // Note: `cached_reads` must outlive `db` to satisfy lifetime requirements.
28///     let state = State::builder().with_database(db).build();
29/// }
30/// ```
31#[derive(Debug, Clone, Default)]
32pub struct CachedReads {
33    /// Block state account with storage.
34    pub accounts: HashMap<Address, CachedAccount>,
35    /// Created contracts.
36    pub contracts: HashMap<B256, Bytecode>,
37    /// Block hash mapped to the block number.
38    pub block_hashes: HashMap<u64, B256>,
39}
40
41// === impl CachedReads ===
42
43impl CachedReads {
44    /// Gets a [`DatabaseRef`] that will cache reads from the given database.
45    pub const fn as_db<DB>(&mut self, db: DB) -> CachedReadsDBRef<'_, DB> {
46        self.as_db_mut(db).into_db()
47    }
48
49    /// Gets a mutable [`Database`] that will cache reads from the underlying database.
50    pub const fn as_db_mut<DB>(&mut self, db: DB) -> CachedReadsDbMut<'_, DB> {
51        CachedReadsDbMut { cached: self, db }
52    }
53
54    /// Inserts an account info into the cache.
55    pub fn insert_account(
56        &mut self,
57        address: Address,
58        info: AccountInfo,
59        storage: HashMap<U256, U256>,
60    ) {
61        self.accounts.insert(address, CachedAccount { info: Some(info), storage });
62    }
63
64    /// Extends current cache with entries from another [`CachedReads`] instance.
65    ///
66    /// Note: It is expected that both instances are based on the exact same state.
67    pub fn extend(&mut self, other: Self) {
68        self.accounts.extend(other.accounts);
69        self.contracts.extend(other.contracts);
70        self.block_hashes.extend(other.block_hashes);
71    }
72}
73
74/// A [Database] that caches reads inside [`CachedReads`].
75///
76/// The lifetime parameter `'a` is tied to the lifetime of the underlying [`CachedReads`] instance.
77/// This ensures that the cache remains valid for the entire duration this wrapper is used.
78/// The original [`CachedReads`] must outlive this wrapper to prevent use-after-free.
79#[derive(Debug)]
80pub struct CachedReadsDbMut<'a, DB> {
81    /// The cache of reads.
82    pub cached: &'a mut CachedReads,
83    /// The underlying database.
84    pub db: DB,
85}
86
87impl<'a, DB> CachedReadsDbMut<'a, DB> {
88    /// Converts this [`Database`] implementation into a [`DatabaseRef`] that will still cache
89    /// reads.
90    pub const fn into_db(self) -> CachedReadsDBRef<'a, DB> {
91        CachedReadsDBRef { inner: RefCell::new(self) }
92    }
93
94    /// Returns access to wrapped [`DatabaseRef`].
95    pub const fn inner(&self) -> &DB {
96        &self.db
97    }
98}
99
100impl<DB, T> AsRef<T> for CachedReadsDbMut<'_, DB>
101where
102    DB: AsRef<T>,
103{
104    fn as_ref(&self) -> &T {
105        self.inner().as_ref()
106    }
107}
108
109impl<DB: DatabaseRef> Database for CachedReadsDbMut<'_, DB> {
110    type Error = <DB as DatabaseRef>::Error;
111
112    fn basic(&mut self, address: Address) -> Result<Option<AccountInfo>, Self::Error> {
113        let basic = match self.cached.accounts.entry(address) {
114            Entry::Occupied(entry) => entry.get().info.clone(),
115            Entry::Vacant(entry) => {
116                entry.insert(CachedAccount::new(self.db.basic_ref(address)?)).info.clone()
117            }
118        };
119        Ok(basic)
120    }
121
122    fn code_by_hash(&mut self, code_hash: B256) -> Result<Bytecode, Self::Error> {
123        let code = match self.cached.contracts.entry(code_hash) {
124            Entry::Occupied(entry) => entry.get().clone(),
125            Entry::Vacant(entry) => entry.insert(self.db.code_by_hash_ref(code_hash)?).clone(),
126        };
127        Ok(code)
128    }
129
130    fn storage(&mut self, address: Address, index: U256) -> Result<U256, Self::Error> {
131        match self.cached.accounts.entry(address) {
132            Entry::Occupied(mut acc_entry) => match acc_entry.get_mut().storage.entry(index) {
133                Entry::Occupied(entry) => Ok(*entry.get()),
134                Entry::Vacant(entry) => Ok(*entry.insert(self.db.storage_ref(address, index)?)),
135            },
136            Entry::Vacant(acc_entry) => {
137                // acc needs to be loaded for us to access slots.
138                let info = self.db.basic_ref(address)?;
139                let (account, value) = if info.is_some() {
140                    let value = self.db.storage_ref(address, index)?;
141                    let mut account = CachedAccount::new(info);
142                    account.storage.insert(index, value);
143                    (account, value)
144                } else {
145                    (CachedAccount::new(info), U256::ZERO)
146                };
147                acc_entry.insert(account);
148                Ok(value)
149            }
150        }
151    }
152
153    fn block_hash(&mut self, number: u64) -> Result<B256, Self::Error> {
154        let hash = match self.cached.block_hashes.entry(number) {
155            Entry::Occupied(entry) => *entry.get(),
156            Entry::Vacant(entry) => *entry.insert(self.db.block_hash_ref(number)?),
157        };
158        Ok(hash)
159    }
160}
161
162/// A [`DatabaseRef`] that caches reads inside [`CachedReads`].
163///
164/// This is intended to be used as the [`DatabaseRef`] for
165/// `revm::db::State` for repeated payload build jobs.
166///
167/// The lifetime parameter `'a` matches the lifetime of the underlying [`CachedReadsDbMut`],
168/// which in turn is tied to the [`CachedReads`] cache. [`RefCell`] is used here to provide
169/// interior mutability for the [`DatabaseRef`] trait (which requires `&self`), while the
170/// lifetime ensures the cache remains valid throughout the wrapper's usage.
171#[derive(Debug)]
172pub struct CachedReadsDBRef<'a, DB> {
173    /// The inner cache reads db mut.
174    pub inner: RefCell<CachedReadsDbMut<'a, DB>>,
175}
176
177impl<DB: DatabaseRef> DatabaseRef for CachedReadsDBRef<'_, DB> {
178    type Error = <DB as DatabaseRef>::Error;
179
180    fn basic_ref(&self, address: Address) -> Result<Option<AccountInfo>, Self::Error> {
181        self.inner.borrow_mut().basic(address)
182    }
183
184    fn code_by_hash_ref(&self, code_hash: B256) -> Result<Bytecode, Self::Error> {
185        self.inner.borrow_mut().code_by_hash(code_hash)
186    }
187
188    fn storage_ref(&self, address: Address, index: U256) -> Result<U256, Self::Error> {
189        self.inner.borrow_mut().storage(address, index)
190    }
191
192    fn block_hash_ref(&self, number: u64) -> Result<B256, Self::Error> {
193        self.inner.borrow_mut().block_hash(number)
194    }
195}
196
197/// Cached account contains the account state with storage
198/// but lacks the account status.
199#[derive(Debug, Clone)]
200pub struct CachedAccount {
201    /// Account state.
202    pub info: Option<AccountInfo>,
203    /// Account's storage.
204    pub storage: HashMap<U256, U256>,
205}
206
207impl CachedAccount {
208    fn new(info: Option<AccountInfo>) -> Self {
209        Self { info, storage: HashMap::default() }
210    }
211}
212
213#[cfg(test)]
214mod tests {
215    use super::*;
216
217    #[test]
218    fn test_extend_with_two_cached_reads() {
219        // Setup test data
220        let hash1 = B256::from_slice(&[1u8; 32]);
221        let hash2 = B256::from_slice(&[2u8; 32]);
222        let address1 = Address::from_slice(&[1u8; 20]);
223        let address2 = Address::from_slice(&[2u8; 20]);
224
225        // Create primary cache
226        let mut primary = {
227            let mut cache = CachedReads::default();
228            cache.accounts.insert(address1, CachedAccount::new(Some(AccountInfo::default())));
229            cache.contracts.insert(hash1, Bytecode::default());
230            cache.block_hashes.insert(1, hash1);
231            cache
232        };
233
234        // Create additional cache
235        let additional = {
236            let mut cache = CachedReads::default();
237            cache.accounts.insert(address2, CachedAccount::new(Some(AccountInfo::default())));
238            cache.contracts.insert(hash2, Bytecode::default());
239            cache.block_hashes.insert(2, hash2);
240            cache
241        };
242
243        // Extending primary with additional cache
244        primary.extend(additional);
245
246        // Verify the combined state
247        assert!(
248            primary.accounts.len() == 2 &&
249                primary.contracts.len() == 2 &&
250                primary.block_hashes.len() == 2,
251            "All maps should contain 2 entries"
252        );
253
254        // Verify specific entries
255        assert!(
256            primary.accounts.contains_key(&address1) &&
257                primary.accounts.contains_key(&address2) &&
258                primary.contracts.contains_key(&hash1) &&
259                primary.contracts.contains_key(&hash2) &&
260                primary.block_hashes.get(&1) == Some(&hash1) &&
261                primary.block_hashes.get(&2) == Some(&hash2),
262            "All expected entries should be present"
263        );
264    }
265}