1use revm::{
2 database::{states::bundle_state::BundleRetention, BundleState, Cache, CacheDB, State},
3 primitives::B256,
4 Database,
5};
6use std::{collections::BTreeMap, convert::Infallible, sync::Arc};
7
8pub trait DbConnect: Sync {
21 type Database: Database;
23
24 type Error: core::error::Error;
26
27 fn connect(&self) -> Result<Self::Database, Self::Error>;
29}
30
31impl<Db> DbConnect for Db
32where
33 Db: Database + Clone + Sync,
34{
35 type Database = Self;
36
37 type Error = Infallible;
38
39 fn connect(&self) -> Result<Self::Database, Self::Error> {
40 Ok(self.clone())
41 }
42}
43
44pub trait StateAcc {
50 fn set_state_clear_flag(&mut self, flag: bool);
52
53 fn merge_transitions(&mut self, retention: BundleRetention);
55
56 fn take_bundle(&mut self) -> BundleState;
58
59 fn set_block_hashes(&mut self, block_hashes: &BTreeMap<u64, B256>);
62}
63
64impl<Db: Database> StateAcc for State<Db> {
65 fn set_state_clear_flag(&mut self, flag: bool) {
66 Self::set_state_clear_flag(self, flag)
67 }
68
69 fn merge_transitions(&mut self, retention: BundleRetention) {
70 Self::merge_transitions(self, retention)
71 }
72
73 fn take_bundle(&mut self) -> BundleState {
74 Self::take_bundle(self)
75 }
76
77 fn set_block_hashes(&mut self, block_hashes: &BTreeMap<u64, B256>) {
78 self.block_hashes.extend(block_hashes)
79 }
80}
81
82pub trait TryStateAcc: Sync {
94 type Error: core::error::Error;
96
97 fn try_set_state_clear_flag(&mut self, flag: bool) -> Result<(), Self::Error>;
99
100 fn try_merge_transitions(&mut self, retention: BundleRetention) -> Result<(), Self::Error>;
103
104 fn try_take_bundle(&mut self) -> Result<BundleState, Self::Error>;
106
107 fn try_set_block_hashes(
110 &mut self,
111 block_hashes: &BTreeMap<u64, B256>,
112 ) -> Result<(), Self::Error>;
113}
114
115impl<Db> TryStateAcc for Db
116where
117 Db: StateAcc + Sync,
118{
119 type Error = Infallible;
120
121 fn try_set_state_clear_flag(&mut self, flag: bool) -> Result<(), Infallible> {
122 self.set_state_clear_flag(flag);
123 Ok(())
124 }
125
126 fn try_merge_transitions(&mut self, retention: BundleRetention) -> Result<(), Infallible> {
127 self.merge_transitions(retention);
128 Ok(())
129 }
130
131 fn try_take_bundle(&mut self) -> Result<BundleState, Infallible> {
132 Ok(self.take_bundle())
133 }
134
135 fn try_set_block_hashes(
136 &mut self,
137 block_hashes: &BTreeMap<u64, B256>,
138 ) -> Result<(), Infallible> {
139 self.set_block_hashes(block_hashes);
140 Ok(())
141 }
142}
143
144#[derive(thiserror::Error, Debug, Clone, Copy, PartialEq, Eq)]
147pub enum ArcUpgradeError {
148 #[error("Arc reference is not unique, cannot mutate")]
151 NotUnique,
152}
153
154impl<Db> TryStateAcc for Arc<Db>
155where
156 Db: StateAcc + Sync + Send,
157{
158 type Error = ArcUpgradeError;
159
160 fn try_set_state_clear_flag(&mut self, flag: bool) -> Result<(), ArcUpgradeError> {
161 Self::get_mut(self).ok_or(ArcUpgradeError::NotUnique)?.set_state_clear_flag(flag);
162 Ok(())
163 }
164
165 fn try_merge_transitions(&mut self, retention: BundleRetention) -> Result<(), ArcUpgradeError> {
166 Self::get_mut(self).ok_or(ArcUpgradeError::NotUnique)?.merge_transitions(retention);
167 Ok(())
168 }
169
170 fn try_take_bundle(&mut self) -> Result<BundleState, ArcUpgradeError> {
171 Ok(Self::get_mut(self).ok_or(ArcUpgradeError::NotUnique)?.take_bundle())
172 }
173
174 fn try_set_block_hashes(
175 &mut self,
176 block_hashes: &BTreeMap<u64, B256>,
177 ) -> Result<(), ArcUpgradeError> {
178 Self::get_mut(self).ok_or(ArcUpgradeError::NotUnique)?.set_block_hashes(block_hashes);
179 Ok(())
180 }
181}
182
183pub trait CachingDb {
185 fn cache(&self) -> &Cache;
187
188 fn cache_mut(&mut self) -> &mut Cache;
190
191 fn into_cache(self) -> Cache;
193
194 fn extend_ref(&mut self, cache: &Cache) {
202 self.cache_mut().accounts.extend(cache.accounts.iter().map(|(k, v)| (*k, v.clone())));
203 self.cache_mut().contracts.extend(cache.contracts.iter().map(|(k, v)| (*k, v.clone())));
204 self.cache_mut().logs.extend(cache.logs.iter().cloned());
205 self.cache_mut().block_hashes.extend(cache.block_hashes.iter().map(|(k, v)| (*k, *v)));
206 }
207
208 fn extend(&mut self, cache: Cache) {
216 self.cache_mut().accounts.extend(cache.accounts);
217 self.cache_mut().contracts.extend(cache.contracts);
218 self.cache_mut().logs.extend(cache.logs);
219 self.cache_mut().block_hashes.extend(cache.block_hashes);
220 }
221}
222
223impl<Db> CachingDb for CacheDB<Db> {
224 fn cache(&self) -> &Cache {
225 &self.cache
226 }
227
228 fn cache_mut(&mut self) -> &mut Cache {
229 &mut self.cache
230 }
231
232 fn into_cache(self) -> Cache {
233 self.cache
234 }
235}
236
237pub trait TryCachingDb {
240 type Error: core::error::Error;
242
243 fn cache(&self) -> &Cache;
245
246 fn try_cache_mut(&mut self) -> Result<&mut Cache, Self::Error>;
248
249 fn try_into_cache(self) -> Result<Cache, Self::Error>;
251
252 fn try_extend_ref(&mut self, cache: &Cache) -> Result<(), Self::Error>
260 where
261 Self: Sized,
262 {
263 let inner_cache = self.try_cache_mut()?;
264 inner_cache.accounts.extend(cache.accounts.iter().map(|(k, v)| (*k, v.clone())));
265 inner_cache.contracts.extend(cache.contracts.iter().map(|(k, v)| (*k, v.clone())));
266 inner_cache.logs.extend(cache.logs.iter().cloned());
267 inner_cache.block_hashes.extend(cache.block_hashes.iter().map(|(k, v)| (*k, *v)));
268 Ok(())
269 }
270
271 fn try_extend(&mut self, cache: Cache) -> Result<(), Self::Error>
279 where
280 Self: Sized,
281 {
282 let inner_cache = self.try_cache_mut()?;
283 inner_cache.accounts.extend(cache.accounts);
284 inner_cache.contracts.extend(cache.contracts);
285 inner_cache.logs.extend(cache.logs);
286 inner_cache.block_hashes.extend(cache.block_hashes);
287 Ok(())
288 }
289}
290
291impl<Db> TryCachingDb for Arc<Db>
292where
293 Db: CachingDb,
294{
295 type Error = ArcUpgradeError;
296
297 fn cache(&self) -> &Cache {
298 self.as_ref().cache()
299 }
300
301 fn try_cache_mut(&mut self) -> Result<&mut Cache, Self::Error> {
302 Self::get_mut(self).ok_or(ArcUpgradeError::NotUnique).map(|db| db.cache_mut())
303 }
304
305 fn try_into_cache(self) -> Result<Cache, Self::Error> {
306 Self::into_inner(self).ok_or(ArcUpgradeError::NotUnique).map(|db| db.into_cache())
307 }
308}