Skip to main content

trevm/db/
traits.rs

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
8/// Trait for types that can be used to connect to a database.
9///
10/// Connectors should contain configuration information like filesystem paths.
11/// They are intended to enable parallel instantiation of multiple EVMs in
12/// multiple threads sharing some database configuration
13///
14/// `DbConnect` is blanket implemented for clonable [`Database`] types by
15/// simply cloning the database instance. This allows already-instantiated DBs
16/// to be used as connectors, however, if the [`Database`] uses a shared
17/// resource like a file or network connection, care should be taken to ensure
18/// that the implementation does not share uintended state between EVM
19/// instances.
20pub trait DbConnect: Sync {
21    /// The database type returned when connecting.
22    type Database: Database;
23
24    /// The error type returned when connecting to the database.
25    type Error: core::error::Error;
26
27    /// Connect to the database.
28    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
44/// Abstraction trait covering types that accumulate state changes into a
45/// [`BundleState`]. The prime example of this is [`State`]. These types are
46/// use to accumulate state changes during the execution of a sequence of
47/// transactions, and then provide access to the net changes in the form of a
48/// [`BundleState`].
49pub trait StateAcc {
50    /// Set the state clear flag. See [`State::set_state_clear_flag`].
51    fn set_state_clear_flag(&mut self, flag: bool);
52
53    /// Merge transitions into the bundle. See [`State::merge_transitions`].
54    fn merge_transitions(&mut self, retention: BundleRetention);
55
56    /// Take the bundle. See [`State::take_bundle`].
57    fn take_bundle(&mut self) -> BundleState;
58
59    /// Set the block hashes, overriding any existing values, and inserting any
60    /// absent values.
61    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
82/// Fallible version of [`StateAcc`].
83///
84/// Abstraction trait covering types that accumulate state changes into a
85/// [`BundleState`]. The prime example of this is [`State`]. These types are
86/// use to accumulate state changes during the execution of a sequence of
87/// transactions, and then provide access to the net changes in the form of a
88/// [`BundleState`].
89///
90/// The primary motivator for this trait is to allow for the implementation of
91/// [`StateAcc`] for [`Arc`]-wrapped DBs, which may fail to mutate if the
92/// reference is not unique.
93pub trait TryStateAcc: Sync {
94    /// Error type to be thrown when state accumulation fails.
95    type Error: core::error::Error;
96
97    /// Attempt to set the state clear flag. See [`State::set_state_clear_flag`].
98    fn try_set_state_clear_flag(&mut self, flag: bool) -> Result<(), Self::Error>;
99
100    /// Attempt to merge transitions into the bundle. See
101    /// [`State::merge_transitions`].
102    fn try_merge_transitions(&mut self, retention: BundleRetention) -> Result<(), Self::Error>;
103
104    /// Attempt to take the bundle. See [`State::take_bundle`].
105    fn try_take_bundle(&mut self) -> Result<BundleState, Self::Error>;
106
107    /// Attempt to set the block hashes, overriding any existing values, and
108    /// inserting any absent values.
109    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/// Error type for implementation of [`TryStateAcc`] for [`Arc`]-wrapped
145/// DBs.
146#[derive(thiserror::Error, Debug, Clone, Copy, PartialEq, Eq)]
147pub enum ArcUpgradeError {
148    /// Arc reference is not unique. Ensure that all other references are
149    /// dropped before attempting to mutate the state.
150    #[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
183/// Trait for Databases that have a [`Cache`].
184pub trait CachingDb {
185    /// Get the cache.
186    fn cache(&self) -> &Cache;
187
188    /// Get the cache mutably.
189    fn cache_mut(&mut self) -> &mut Cache;
190
191    /// Deconstruct into the cache
192    fn into_cache(self) -> Cache;
193
194    /// Extend the cache with the given cache by copying data.
195    ///
196    /// The behavior is as follows:
197    /// - Accounts are overridden with outer accounts
198    /// - Contracts are overridden with outer contracts
199    /// - Logs are appended
200    /// - Block hashes are overridden with outer block hashes
201    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    /// Extend the cache with the given cache by moving data.
209    ///
210    /// The behavior is as follows:
211    /// - Accounts are overridden with outer accounts
212    /// - Contracts are overridden with outer contracts
213    /// - Logs are appended
214    /// - Block hashes are overridden with outer block hashes
215    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
237/// Trait for Databases that have a [`Cache`] and can fail to mutably access
238/// it. E.g. `Arc<CacheDB<Db>>`
239pub trait TryCachingDb {
240    /// Error type to be thrown when cache access fails.
241    type Error: core::error::Error;
242
243    /// Attempt to get the cache.
244    fn cache(&self) -> &Cache;
245
246    /// Attempt to get the cache mutably.
247    fn try_cache_mut(&mut self) -> Result<&mut Cache, Self::Error>;
248
249    /// Attempt to deconstruct into the cache
250    fn try_into_cache(self) -> Result<Cache, Self::Error>;
251
252    /// Attempt to fold a cache into the database.
253    ///
254    /// The behavior is as follows:
255    /// - Accounts are overridden with outer accounts
256    /// - Contracts are overridden with outer contracts
257    /// - Logs are appended
258    /// - Block hashes are overridden with outer block hashes
259    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    /// Attempt to extend the cache with the given cache by moving data.
272    ///
273    /// The behavior is as follows:
274    /// - Accounts are overridden with outer accounts
275    /// - Contracts are overridden with outer contracts
276    /// - Logs are appended
277    /// - Block hashes are overridden with outer block hashes
278    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}