reth_revm/
cached.rs

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

reth_revm/
cached.rs
