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