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