Skip to main content

revm_database_interface/
bal.rs

1//! Database implementation for BAL.
2use core::{
3    error::Error,
4    fmt::Display,
5    ops::{Deref, DerefMut},
6};
7use primitives::{Address, StorageKey, StorageValue, B256};
8use state::{
9    bal::{alloy::AlloyBal, Bal, BalError, BlockAccessIndex},
10    Account, AccountId, AccountInfo, Bytecode, EvmState,
11};
12use std::sync::Arc;
13
14use crate::{DBErrorMarker, Database, DatabaseCommit};
15
16/// Contains both the BAL for reads and BAL builders.
17#[derive(Clone, Default, Debug, PartialEq, Eq)]
18#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
19pub struct BalState {
20    /// BAL used to execute transactions.
21    pub bal: Option<Arc<Bal>>,
22    /// BAL builder that is used to build BAL.
23    /// It is create from State output of transaction execution.
24    pub bal_builder: Option<Bal>,
25    /// BAL index, used by bal to fetch appropriate values and used by bal_builder on commit
26    /// to submit changes.
27    pub bal_index: BlockAccessIndex,
28}
29
30impl BalState {
31    /// Create a new BAL manager.
32    #[inline]
33    pub fn new() -> Self {
34        Self::default()
35    }
36
37    /// Reset BAL index to pre-execution.
38    #[inline]
39    pub const fn reset_bal_index(&mut self) {
40        self.bal_index = BlockAccessIndex::PRE_EXECUTION;
41    }
42
43    /// Bump BAL index.
44    #[inline]
45    pub const fn bump_bal_index(&mut self) {
46        self.bal_index.increment();
47    }
48
49    /// Get BAL index.
50    #[inline]
51    pub const fn bal_index(&self) -> BlockAccessIndex {
52        self.bal_index
53    }
54
55    /// Get BAL.
56    #[inline]
57    pub fn bal(&self) -> Option<Arc<Bal>> {
58        self.bal.clone()
59    }
60
61    /// Get BAL builder.
62    #[inline]
63    pub fn bal_builder(&self) -> Option<Bal> {
64        self.bal_builder.clone()
65    }
66
67    /// Set BAL.
68    #[inline]
69    pub fn with_bal(mut self, bal: Arc<Bal>) -> Self {
70        self.bal = Some(bal);
71        self
72    }
73
74    /// Set BAL builder.
75    #[inline]
76    pub fn with_bal_builder(mut self) -> Self {
77        self.bal_builder = Some(Bal::new());
78        self
79    }
80
81    /// Take BAL builder.
82    #[inline]
83    pub const fn take_built_bal(&mut self) -> Option<Bal> {
84        self.reset_bal_index();
85        self.bal_builder.take()
86    }
87
88    /// Take built BAL as AlloyBAL.
89    #[inline]
90    pub fn take_built_alloy_bal(&mut self) -> Option<AlloyBal> {
91        self.take_built_bal().map(|bal| bal.into_alloy_bal())
92    }
93
94    /// Get account id from BAL.
95    ///
96    /// Return Error if BAL is not found and Account is not
97    #[inline]
98    pub fn get_account_id(&self, address: &Address) -> Result<Option<AccountId>, BalError> {
99        self.bal
100            .as_ref()
101            .map(|bal| {
102                bal.accounts
103                    .get_full(address)
104                    .map(|i| AccountId::new(i.0).expect("too many bals"))
105                    .ok_or(BalError::AccountNotFound { address: *address })
106            })
107            .transpose()
108    }
109
110    /// Fetch account from database and apply bal changes to it.
111    ///
112    /// Return Some if BAL is existing, None if not.
113    /// Return Err if Accounts is not found inside BAL.
114    /// And return true
115    #[inline]
116    pub fn basic(
117        &self,
118        address: Address,
119        basic: &mut Option<AccountInfo>,
120    ) -> Result<bool, BalError> {
121        let Some(account_id) = self.get_account_id(&address)? else {
122            return Ok(false);
123        };
124        self.basic_by_account_id(account_id, basic)
125    }
126
127    /// Fetch account from database and apply bal changes to it by account id.
128    #[inline]
129    pub fn basic_by_account_id(
130        &self,
131        account_id: AccountId,
132        basic: &mut Option<AccountInfo>,
133    ) -> Result<bool, BalError> {
134        let Some(bal) = &self.bal else {
135            return Ok(false);
136        };
137        let is_none = basic.is_none();
138        let mut bal_basic = core::mem::take(basic).unwrap_or_default();
139        let changed = bal.populate_account_info(account_id, self.bal_index, &mut bal_basic)?;
140
141        // If account was not in DB and BAL has no changes, keep it as None.
142        if !changed && is_none {
143            return Ok(true);
144        }
145
146        *basic = Some(bal_basic);
147        Ok(true)
148    }
149
150    /// Get storage value from BAL.
151    ///
152    /// Return Err if bal is present but account or storage is not found inside BAL.
153    #[inline]
154    pub fn storage(
155        &self,
156        account: &Address,
157        storage_key: StorageKey,
158    ) -> Result<Option<StorageValue>, BalError> {
159        let Some(bal) = &self.bal else {
160            return Ok(None);
161        };
162
163        let Some(bal_account) = bal.accounts.get(account) else {
164            return Err(BalError::AccountNotFound { address: *account });
165        };
166
167        Ok(bal_account
168            .storage
169            .get_bal_writes(account, storage_key)?
170            .get(self.bal_index))
171    }
172
173    /// Get the storage value by account id.
174    ///
175    /// Return Err if bal is present but account or storage is not found inside BAL.
176    #[inline]
177    pub fn storage_by_account_id(
178        &self,
179        account_id: AccountId,
180        storage_key: StorageKey,
181    ) -> Result<Option<StorageValue>, BalError> {
182        let Some(bal) = &self.bal else {
183            return Ok(None);
184        };
185
186        let Some((address, bal_account)) = bal.accounts.get_index(account_id.get()) else {
187            return Err(BalError::InvalidAccountId { account_id });
188        };
189
190        Ok(bal_account
191            .storage
192            .get_bal_writes(address, storage_key)?
193            .get(self.bal_index))
194    }
195
196    /// Apply changed from EvmState to the bal_builder
197    #[inline]
198    pub fn commit(&mut self, changes: &EvmState) {
199        if let Some(bal_builder) = &mut self.bal_builder {
200            for (address, account) in changes.iter() {
201                bal_builder.update_account(self.bal_index, *address, account);
202            }
203        }
204    }
205
206    /// Commit one account to the BAL builder.
207    #[inline]
208    pub fn commit_one(&mut self, address: Address, account: &Account) {
209        if let Some(bal_builder) = &mut self.bal_builder {
210            bal_builder.update_account(self.bal_index, address, account);
211        }
212    }
213}
214
215/// Database implementation for BAL.
216#[derive(Clone, Debug, PartialEq, Eq)]
217#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
218pub struct BalDatabase<DB> {
219    /// BAL manager.
220    pub bal_state: BalState,
221    /// Database.
222    pub db: DB,
223}
224
225impl<DB> Deref for BalDatabase<DB> {
226    type Target = DB;
227
228    fn deref(&self) -> &Self::Target {
229        &self.db
230    }
231}
232
233impl<DB> DerefMut for BalDatabase<DB> {
234    fn deref_mut(&mut self) -> &mut Self::Target {
235        &mut self.db
236    }
237}
238
239impl<DB> BalDatabase<DB> {
240    /// Create a new BAL database.
241    #[inline]
242    pub fn new(db: DB) -> Self {
243        Self {
244            bal_state: BalState::default(),
245            db,
246        }
247    }
248
249    /// With BAL.
250    #[inline]
251    pub fn with_bal_option(self, bal: Option<Arc<Bal>>) -> Self {
252        Self {
253            bal_state: BalState {
254                bal,
255                ..self.bal_state
256            },
257            ..self
258        }
259    }
260
261    /// With BAL builder.
262    #[inline]
263    pub fn with_bal_builder(self) -> Self {
264        Self {
265            bal_state: self.bal_state.with_bal_builder(),
266            ..self
267        }
268    }
269
270    /// Reset BAL index.
271    #[inline]
272    pub const fn reset_bal_index(mut self) -> Self {
273        self.bal_state.reset_bal_index();
274        self
275    }
276
277    /// Bump BAL index.
278    #[inline]
279    pub const fn bump_bal_index(&mut self) {
280        self.bal_state.bump_bal_index();
281    }
282}
283
284/// Error type from database.
285#[derive(Clone, Debug, PartialEq, Eq)]
286#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
287pub enum EvmDatabaseError<ERROR> {
288    /// BAL error.
289    Bal(BalError),
290    /// External database error.
291    Database(ERROR),
292}
293
294impl<ERROR> From<BalError> for EvmDatabaseError<ERROR> {
295    fn from(error: BalError) -> Self {
296        Self::Bal(error)
297    }
298}
299
300impl<ERROR: core::error::Error + Send + Sync + 'static> DBErrorMarker for EvmDatabaseError<ERROR> {}
301
302impl<ERROR: Display> Display for EvmDatabaseError<ERROR> {
303    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
304        match self {
305            Self::Bal(error) => write!(f, "Bal error: {error}"),
306            Self::Database(error) => write!(f, "Database error: {error}"),
307        }
308    }
309}
310
311impl<ERROR: Error> Error for EvmDatabaseError<ERROR> {}
312
313impl<ERROR> EvmDatabaseError<ERROR> {
314    /// Convert BAL database error to database error.
315    ///
316    /// Panics if BAL error is present.
317    pub fn into_external_error(self) -> ERROR {
318        match self {
319            Self::Bal(_) => panic!("Expected database error, got BAL error"),
320            Self::Database(error) => error,
321        }
322    }
323}
324
325impl<DB: Database> Database for BalDatabase<DB> {
326    type Error = EvmDatabaseError<DB::Error>;
327
328    #[inline]
329    fn basic(&mut self, address: Address) -> Result<Option<AccountInfo>, Self::Error> {
330        let account_id = self.bal_state.get_account_id(&address)?;
331
332        let mut account = self.db.basic(address).map_err(EvmDatabaseError::Database)?;
333
334        if let Some(account_id) = account_id {
335            self.bal_state
336                .basic_by_account_id(account_id, &mut account)?;
337        }
338
339        Ok(account)
340    }
341
342    #[inline]
343    fn code_by_hash(&mut self, code_hash: B256) -> Result<Bytecode, Self::Error> {
344        self.db
345            .code_by_hash(code_hash)
346            .map_err(EvmDatabaseError::Database)
347    }
348
349    #[inline]
350    fn storage(&mut self, address: Address, key: StorageKey) -> Result<StorageValue, Self::Error> {
351        if let Some(storage) = self.bal_state.storage(&address, key)? {
352            return Ok(storage);
353        }
354
355        self.db
356            .storage(address, key)
357            .map_err(EvmDatabaseError::Database)
358    }
359
360    #[inline]
361    fn storage_by_account_id(
362        &mut self,
363        address: Address,
364        account_id: AccountId,
365        storage_key: StorageKey,
366    ) -> Result<StorageValue, Self::Error> {
367        if let Some(value) = self
368            .bal_state
369            .storage_by_account_id(account_id, storage_key)?
370        {
371            return Ok(value);
372        }
373
374        self.db
375            .storage(address, storage_key)
376            .map_err(EvmDatabaseError::Database)
377    }
378
379    fn block_hash(&mut self, number: u64) -> Result<B256, Self::Error> {
380        self.db
381            .block_hash(number)
382            .map_err(EvmDatabaseError::Database)
383    }
384}
385
386impl<DB: DatabaseCommit> DatabaseCommit for BalDatabase<DB> {
387    fn commit(&mut self, changes: EvmState) {
388        self.bal_state.commit(&changes);
389        self.db.commit(changes);
390    }
391}