1use 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#[derive(Debug, Clone, Default)]
35pub struct CachedReads {
36 pub accounts: HashMap<Address, CachedAccount>,
38 pub contracts: HashMap<B256, Bytecode>,
40 pub block_hashes: HashMap<u64, B256>,
42}
43
44impl CachedReads {
47 pub const fn as_db<DB>(&mut self, db: DB) -> CachedReadsDBRef<'_, DB> {
49 self.as_db_mut(db).into_db()
50 }
51
52 pub const fn as_db_mut<DB>(&mut self, db: DB) -> CachedReadsDbMut<'_, DB> {
54 CachedReadsDbMut { cached: self, db }
55 }
56
57 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 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)]
79pub struct CachedReadsDbMut<'a, DB> {
80 pub cached: &'a mut CachedReads,
82 pub db: DB,
84}
85
86impl<'a, DB> CachedReadsDbMut<'a, DB> {
87 pub const fn into_db(self) -> CachedReadsDBRef<'a, DB> {
90 CachedReadsDBRef { inner: RefCell::new(self) }
91 }
92
93 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 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#[derive(Debug)]
166pub struct CachedReadsDBRef<'a, DB> {
167 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#[derive(Debug, Clone)]
194pub struct CachedAccount {
195 pub info: Option<AccountInfo>,
197 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 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 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 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 primary.extend(additional);
239
240 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 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}