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