Skip to main content

reth_revm/
cached.rs

1//! Database adapters for payload building.
2use alloy_primitives::{
3    map::{AddressMap, B256Map, Entry, HashMap, U256Map},
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: AddressMap<CachedAccount>,
35    /// Created contracts.
36    pub contracts: B256Map<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    /// Creates a new [`CachedReads`] with capacity for account entries.
45    pub fn with_account_capacity(capacity: usize) -> Self {
46        Self {
47            accounts: AddressMap::with_capacity_and_hasher(capacity, Default::default()),
48            ..Default::default()
49        }
50    }
51
52    /// Gets a [`DatabaseRef`] that will cache reads from the given database.
53    pub const fn as_db<DB>(&mut self, db: DB) -> CachedReadsDBRef<'_, DB> {
54        self.as_db_mut(db).into_db()
55    }
56
57    /// Gets a mutable [`Database`] that will cache reads from the underlying database.
58    pub const fn as_db_mut<DB>(&mut self, db: DB) -> CachedReadsDbMut<'_, DB> {
59        CachedReadsDbMut { cached: self, db }
60    }
61
62    /// Inserts an account info into the cache.
63    pub fn insert_account(&mut self, address: Address, info: AccountInfo, storage: U256Map<U256>) {
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///
79/// The lifetime parameter `'a` is tied to the lifetime of the underlying [`CachedReads`] instance.
80/// This ensures that the cache remains valid for the entire duration this wrapper is used.
81/// The original [`CachedReads`] must outlive this wrapper to prevent use-after-free.
82#[derive(Debug)]
83pub struct CachedReadsDbMut<'a, DB> {
84    /// The cache of reads.
85    pub cached: &'a mut CachedReads,
86    /// The underlying database.
87    pub db: DB,
88}
89
90impl<'a, DB> CachedReadsDbMut<'a, DB> {
91    /// Converts this [`Database`] implementation into a [`DatabaseRef`] that will still cache
92    /// reads.
93    pub const fn into_db(self) -> CachedReadsDBRef<'a, DB> {
94        CachedReadsDBRef { inner: RefCell::new(self) }
95    }
96
97    /// Returns access to wrapped [`DatabaseRef`].
98    pub const fn inner(&self) -> &DB {
99        &self.db
100    }
101}
102
103impl<DB, T> AsRef<T> for CachedReadsDbMut<'_, DB>
104where
105    DB: AsRef<T>,
106{
107    fn as_ref(&self) -> &T {
108        self.inner().as_ref()
109    }
110}
111
112impl<DB: DatabaseRef> Database for CachedReadsDbMut<'_, DB> {
113    type Error = <DB as DatabaseRef>::Error;
114
115    fn basic(&mut self, address: Address) -> Result<Option<AccountInfo>, Self::Error> {
116        let basic = match self.cached.accounts.entry(address) {
117            Entry::Occupied(entry) => entry.get().info.clone(),
118            Entry::Vacant(entry) => {
119                entry.insert(CachedAccount::new(self.db.basic_ref(address)?)).info.clone()
120            }
121        };
122        Ok(basic)
123    }
124
125    fn code_by_hash(&mut self, code_hash: B256) -> Result<Bytecode, Self::Error> {
126        let code = match self.cached.contracts.entry(code_hash) {
127            Entry::Occupied(entry) => entry.get().clone(),
128            Entry::Vacant(entry) => entry.insert(self.db.code_by_hash_ref(code_hash)?).clone(),
129        };
130        Ok(code)
131    }
132
133    fn storage(&mut self, address: Address, index: U256) -> Result<U256, Self::Error> {
134        match self.cached.accounts.entry(address) {
135            Entry::Occupied(mut acc_entry) => match acc_entry.get_mut().storage.entry(index) {
136                Entry::Occupied(entry) => Ok(*entry.get()),
137                Entry::Vacant(entry) => Ok(*entry.insert(self.db.storage_ref(address, index)?)),
138            },
139            Entry::Vacant(acc_entry) => {
140                // acc needs to be loaded for us to access slots.
141                let info = self.db.basic_ref(address)?;
142                let (account, value) = if info.is_some() {
143                    let value = self.db.storage_ref(address, index)?;
144                    let mut account = CachedAccount::new(info);
145                    account.storage.insert(index, value);
146                    (account, value)
147                } else {
148                    (CachedAccount::new(info), U256::ZERO)
149                };
150                acc_entry.insert(account);
151                Ok(value)
152            }
153        }
154    }
155
156    fn block_hash(&mut self, number: u64) -> Result<B256, Self::Error> {
157        let hash = match self.cached.block_hashes.entry(number) {
158            Entry::Occupied(entry) => *entry.get(),
159            Entry::Vacant(entry) => *entry.insert(self.db.block_hash_ref(number)?),
160        };
161        Ok(hash)
162    }
163}
164
165/// A [`DatabaseRef`] that caches reads inside [`CachedReads`].
166///
167/// This is intended to be used as the [`DatabaseRef`] for
168/// `revm::db::State` for repeated payload build jobs.
169///
170/// The lifetime parameter `'a` matches the lifetime of the underlying [`CachedReadsDbMut`],
171/// which in turn is tied to the [`CachedReads`] cache. [`RefCell`] is used here to provide
172/// interior mutability for the [`DatabaseRef`] trait (which requires `&self`), while the
173/// lifetime ensures the cache remains valid throughout the wrapper's usage.
174#[derive(Debug)]
175pub struct CachedReadsDBRef<'a, DB> {
176    /// The inner cache reads db mut.
177    pub inner: RefCell<CachedReadsDbMut<'a, DB>>,
178}
179
180impl<DB: DatabaseRef> DatabaseRef for CachedReadsDBRef<'_, DB> {
181    type Error = <DB as DatabaseRef>::Error;
182
183    fn basic_ref(&self, address: Address) -> Result<Option<AccountInfo>, Self::Error> {
184        self.inner.borrow_mut().basic(address)
185    }
186
187    fn code_by_hash_ref(&self, code_hash: B256) -> Result<Bytecode, Self::Error> {
188        self.inner.borrow_mut().code_by_hash(code_hash)
189    }
190
191    fn storage_ref(&self, address: Address, index: U256) -> Result<U256, Self::Error> {
192        self.inner.borrow_mut().storage(address, index)
193    }
194
195    fn block_hash_ref(&self, number: u64) -> Result<B256, Self::Error> {
196        self.inner.borrow_mut().block_hash(number)
197    }
198}
199
200/// Cached account contains the account state with storage
201/// but lacks the account status.
202#[derive(Debug, Clone)]
203pub struct CachedAccount {
204    /// Account state.
205    pub info: Option<AccountInfo>,
206    /// Account's storage.
207    pub storage: U256Map<U256>,
208}
209
210impl CachedAccount {
211    fn new(info: Option<AccountInfo>) -> Self {
212        Self { info, storage: U256Map::default() }
213    }
214}
215
216#[cfg(test)]
217mod tests {
218    use super::*;
219
220    #[test]
221    fn test_extend_with_two_cached_reads() {
222        // Setup test data
223        let hash1 = B256::from_slice(&[1u8; 32]);
224        let hash2 = B256::from_slice(&[2u8; 32]);
225        let address1 = Address::from_slice(&[1u8; 20]);
226        let address2 = Address::from_slice(&[2u8; 20]);
227
228        // Create primary cache
229        let mut primary = {
230            let mut cache = CachedReads::default();
231            cache.accounts.insert(address1, CachedAccount::new(Some(AccountInfo::default())));
232            cache.contracts.insert(hash1, Bytecode::default());
233            cache.block_hashes.insert(1, hash1);
234            cache
235        };
236
237        // Create additional cache
238        let additional = {
239            let mut cache = CachedReads::default();
240            cache.accounts.insert(address2, CachedAccount::new(Some(AccountInfo::default())));
241            cache.contracts.insert(hash2, Bytecode::default());
242            cache.block_hashes.insert(2, hash2);
243            cache
244        };
245
246        // Extending primary with additional cache
247        primary.extend(additional);
248
249        // Verify the combined state
250        assert!(
251            primary.accounts.len() == 2 &&
252                primary.contracts.len() == 2 &&
253                primary.block_hashes.len() == 2,
254            "All maps should contain 2 entries"
255        );
256
257        // Verify specific entries
258        assert!(
259            primary.accounts.contains_key(&address1) &&
260                primary.accounts.contains_key(&address2) &&
261                primary.contracts.contains_key(&hash1) &&
262                primary.contracts.contains_key(&hash2) &&
263                primary.block_hashes.get(&1) == Some(&hash1) &&
264                primary.block_hashes.get(&2) == Some(&hash2),
265            "All expected entries should be present"
266        );
267    }
268}