1use 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#[derive(Debug, Clone, Default)]
32pub struct CachedReads {
33 pub accounts: AddressMap<CachedAccount>,
35 pub contracts: B256Map<Bytecode>,
37 pub block_hashes: HashMap<u64, B256>,
39}
40
41impl CachedReads {
44 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 pub const fn as_db<DB>(&mut self, db: DB) -> CachedReadsDBRef<'_, DB> {
54 self.as_db_mut(db).into_db()
55 }
56
57 pub const fn as_db_mut<DB>(&mut self, db: DB) -> CachedReadsDbMut<'_, DB> {
59 CachedReadsDbMut { cached: self, db }
60 }
61
62 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 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#[derive(Debug)]
83pub struct CachedReadsDbMut<'a, DB> {
84 pub cached: &'a mut CachedReads,
86 pub db: DB,
88}
89
90impl<'a, DB> CachedReadsDbMut<'a, DB> {
91 pub const fn into_db(self) -> CachedReadsDBRef<'a, DB> {
94 CachedReadsDBRef { inner: RefCell::new(self) }
95 }
96
97 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 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#[derive(Debug)]
175pub struct CachedReadsDBRef<'a, DB> {
176 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#[derive(Debug, Clone)]
203pub struct CachedAccount {
204 pub info: Option<AccountInfo>,
206 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 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 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 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 primary.extend(additional);
248
249 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 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}