1use 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#[derive(Debug, Clone, Default)]
32pub struct CachedReads {
33 pub accounts: HashMap<Address, CachedAccount>,
35 pub contracts: HashMap<B256, Bytecode>,
37 pub block_hashes: HashMap<u64, B256>,
39}
40
41impl CachedReads {
44 pub const fn as_db<DB>(&mut self, db: DB) -> CachedReadsDBRef<'_, DB> {
46 self.as_db_mut(db).into_db()
47 }
48
49 pub const fn as_db_mut<DB>(&mut self, db: DB) -> CachedReadsDbMut<'_, DB> {
51 CachedReadsDbMut { cached: self, db }
52 }
53
54 pub fn insert_account(
56 &mut self,
57 address: Address,
58 info: AccountInfo,
59 storage: HashMap<U256, U256>,
60 ) {
61 self.accounts.insert(address, CachedAccount { info: Some(info), storage });
62 }
63
64 pub fn extend(&mut self, other: Self) {
68 self.accounts.extend(other.accounts);
69 self.contracts.extend(other.contracts);
70 self.block_hashes.extend(other.block_hashes);
71 }
72}
73
74#[derive(Debug)]
80pub struct CachedReadsDbMut<'a, DB> {
81 pub cached: &'a mut CachedReads,
83 pub db: DB,
85}
86
87impl<'a, DB> CachedReadsDbMut<'a, DB> {
88 pub const fn into_db(self) -> CachedReadsDBRef<'a, DB> {
91 CachedReadsDBRef { inner: RefCell::new(self) }
92 }
93
94 pub const fn inner(&self) -> &DB {
96 &self.db
97 }
98}
99
100impl<DB, T> AsRef<T> for CachedReadsDbMut<'_, DB>
101where
102 DB: AsRef<T>,
103{
104 fn as_ref(&self) -> &T {
105 self.inner().as_ref()
106 }
107}
108
109impl<DB: DatabaseRef> Database for CachedReadsDbMut<'_, DB> {
110 type Error = <DB as DatabaseRef>::Error;
111
112 fn basic(&mut self, address: Address) -> Result<Option<AccountInfo>, Self::Error> {
113 let basic = match self.cached.accounts.entry(address) {
114 Entry::Occupied(entry) => entry.get().info.clone(),
115 Entry::Vacant(entry) => {
116 entry.insert(CachedAccount::new(self.db.basic_ref(address)?)).info.clone()
117 }
118 };
119 Ok(basic)
120 }
121
122 fn code_by_hash(&mut self, code_hash: B256) -> Result<Bytecode, Self::Error> {
123 let code = match self.cached.contracts.entry(code_hash) {
124 Entry::Occupied(entry) => entry.get().clone(),
125 Entry::Vacant(entry) => entry.insert(self.db.code_by_hash_ref(code_hash)?).clone(),
126 };
127 Ok(code)
128 }
129
130 fn storage(&mut self, address: Address, index: U256) -> Result<U256, Self::Error> {
131 match self.cached.accounts.entry(address) {
132 Entry::Occupied(mut acc_entry) => match acc_entry.get_mut().storage.entry(index) {
133 Entry::Occupied(entry) => Ok(*entry.get()),
134 Entry::Vacant(entry) => Ok(*entry.insert(self.db.storage_ref(address, index)?)),
135 },
136 Entry::Vacant(acc_entry) => {
137 let info = self.db.basic_ref(address)?;
139 let (account, value) = if info.is_some() {
140 let value = self.db.storage_ref(address, index)?;
141 let mut account = CachedAccount::new(info);
142 account.storage.insert(index, value);
143 (account, value)
144 } else {
145 (CachedAccount::new(info), U256::ZERO)
146 };
147 acc_entry.insert(account);
148 Ok(value)
149 }
150 }
151 }
152
153 fn block_hash(&mut self, number: u64) -> Result<B256, Self::Error> {
154 let hash = match self.cached.block_hashes.entry(number) {
155 Entry::Occupied(entry) => *entry.get(),
156 Entry::Vacant(entry) => *entry.insert(self.db.block_hash_ref(number)?),
157 };
158 Ok(hash)
159 }
160}
161
162#[derive(Debug)]
172pub struct CachedReadsDBRef<'a, DB> {
173 pub inner: RefCell<CachedReadsDbMut<'a, DB>>,
175}
176
177impl<DB: DatabaseRef> DatabaseRef for CachedReadsDBRef<'_, DB> {
178 type Error = <DB as DatabaseRef>::Error;
179
180 fn basic_ref(&self, address: Address) -> Result<Option<AccountInfo>, Self::Error> {
181 self.inner.borrow_mut().basic(address)
182 }
183
184 fn code_by_hash_ref(&self, code_hash: B256) -> Result<Bytecode, Self::Error> {
185 self.inner.borrow_mut().code_by_hash(code_hash)
186 }
187
188 fn storage_ref(&self, address: Address, index: U256) -> Result<U256, Self::Error> {
189 self.inner.borrow_mut().storage(address, index)
190 }
191
192 fn block_hash_ref(&self, number: u64) -> Result<B256, Self::Error> {
193 self.inner.borrow_mut().block_hash(number)
194 }
195}
196
197#[derive(Debug, Clone)]
200pub struct CachedAccount {
201 pub info: Option<AccountInfo>,
203 pub storage: HashMap<U256, U256>,
205}
206
207impl CachedAccount {
208 fn new(info: Option<AccountInfo>) -> Self {
209 Self { info, storage: HashMap::default() }
210 }
211}
212
213#[cfg(test)]
214mod tests {
215 use super::*;
216
217 #[test]
218 fn test_extend_with_two_cached_reads() {
219 let hash1 = B256::from_slice(&[1u8; 32]);
221 let hash2 = B256::from_slice(&[2u8; 32]);
222 let address1 = Address::from_slice(&[1u8; 20]);
223 let address2 = Address::from_slice(&[2u8; 20]);
224
225 let mut primary = {
227 let mut cache = CachedReads::default();
228 cache.accounts.insert(address1, CachedAccount::new(Some(AccountInfo::default())));
229 cache.contracts.insert(hash1, Bytecode::default());
230 cache.block_hashes.insert(1, hash1);
231 cache
232 };
233
234 let additional = {
236 let mut cache = CachedReads::default();
237 cache.accounts.insert(address2, CachedAccount::new(Some(AccountInfo::default())));
238 cache.contracts.insert(hash2, Bytecode::default());
239 cache.block_hashes.insert(2, hash2);
240 cache
241 };
242
243 primary.extend(additional);
245
246 assert!(
248 primary.accounts.len() == 2 &&
249 primary.contracts.len() == 2 &&
250 primary.block_hashes.len() == 2,
251 "All maps should contain 2 entries"
252 );
253
254 assert!(
256 primary.accounts.contains_key(&address1) &&
257 primary.accounts.contains_key(&address2) &&
258 primary.contracts.contains_key(&hash1) &&
259 primary.contracts.contains_key(&hash2) &&
260 primary.block_hashes.get(&1) == Some(&hash1) &&
261 primary.block_hashes.get(&2) == Some(&hash2),
262 "All expected entries should be present"
263 );
264 }
265}