Skip to main content

revm_state/bal/
account.rs

1//! BAL builder module
2
3use crate::{
4    bal::{writes::BalWrites, BalError, BalIndex},
5    Account, AccountInfo, EvmStorage,
6};
7use alloy_eip7928::{
8    AccountChanges as AlloyAccountChanges, BalanceChange as AlloyBalanceChange,
9    CodeChange as AlloyCodeChange, NonceChange as AlloyNonceChange,
10    SlotChanges as AlloySlotChanges, StorageChange as AlloyStorageChange,
11};
12use bytecode::{Bytecode, BytecodeDecodeError};
13use core::ops::{Deref, DerefMut};
14use primitives::{Address, StorageKey, StorageValue, B256, U256};
15use std::{
16    collections::{btree_map::Entry, BTreeMap},
17    vec::Vec,
18};
19
20/// Account BAL structure.
21#[derive(Debug, Default, Clone, PartialEq, Eq)]
22#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
23pub struct AccountBal {
24    /// Account info bal.
25    pub account_info: AccountInfoBal,
26    /// Storage bal.
27    pub storage: StorageBal,
28}
29
30impl Deref for AccountBal {
31    type Target = AccountInfoBal;
32
33    fn deref(&self) -> &Self::Target {
34        &self.account_info
35    }
36}
37
38impl DerefMut for AccountBal {
39    fn deref_mut(&mut self) -> &mut Self::Target {
40        &mut self.account_info
41    }
42}
43
44impl AccountBal {
45    /// Populate account from BAL. Return true if account info got changed
46    pub fn populate_account_info(&self, bal_index: BalIndex, account: &mut AccountInfo) -> bool {
47        self.account_info.populate_account_info(bal_index, account)
48    }
49
50    /// Extend account from another account.
51    #[inline]
52    pub fn update(&mut self, bal_index: BalIndex, account: &Account) {
53        if account.is_selfdestructed_locally() {
54            let empty_info = AccountInfo::default();
55            self.account_info
56                .update(bal_index, &account.original_info, &empty_info);
57            self.storage.update_reads(account.storage.keys().copied());
58            return;
59        }
60
61        self.account_info
62            .update(bal_index, &account.original_info, &account.info);
63
64        self.storage.update(bal_index, &account.storage);
65    }
66
67    /// Create account from alloy account changes.
68    #[inline]
69    pub fn try_from_alloy(
70        alloy_account: AlloyAccountChanges,
71    ) -> Result<(Address, Self), BytecodeDecodeError> {
72        Ok((
73            alloy_account.address,
74            AccountBal {
75                account_info: AccountInfoBal {
76                    nonce: BalWrites::from(alloy_account.nonce_changes),
77                    balance: BalWrites::from(alloy_account.balance_changes),
78                    code: BalWrites::try_from(alloy_account.code_changes)?,
79                },
80                storage: StorageBal::from_iter(
81                    alloy_account
82                        .storage_changes
83                        .into_iter()
84                        .chain(
85                            alloy_account
86                                .storage_reads
87                                .into_iter()
88                                .map(|key| AlloySlotChanges::new(key, Default::default())),
89                        )
90                        .map(|slot| (slot.slot, BalWrites::from(slot.changes))),
91                ),
92            },
93        ))
94    }
95
96    /// Consumes AccountBal and converts it into [`AlloyAccountChanges`].
97    #[inline]
98    pub fn into_alloy_account(self, address: Address) -> AlloyAccountChanges {
99        let storage_len = self.storage.storage.len();
100        let mut storage_reads = Vec::with_capacity(storage_len);
101        let mut storage_changes = Vec::with_capacity(storage_len);
102        for (key, value) in self.storage.storage {
103            if value.writes.is_empty() {
104                storage_reads.push(key);
105            } else {
106                storage_changes.push(AlloySlotChanges::new(
107                    key,
108                    value
109                        .writes
110                        .into_iter()
111                        .map(|(index, value)| AlloyStorageChange::new(index, value))
112                        .collect(),
113                ));
114            }
115        }
116
117        AlloyAccountChanges {
118            address,
119            storage_changes,
120            storage_reads,
121            balance_changes: self
122                .account_info
123                .balance
124                .writes
125                .into_iter()
126                .map(|(index, value)| AlloyBalanceChange::new(index, value))
127                .collect(),
128            nonce_changes: self
129                .account_info
130                .nonce
131                .writes
132                .into_iter()
133                .map(|(index, value)| AlloyNonceChange::new(index, value))
134                .collect(),
135            code_changes: self
136                .account_info
137                .code
138                .writes
139                .into_iter()
140                .map(|(index, (_, value))| AlloyCodeChange::new(index, value.original_bytes()))
141                .collect(),
142        }
143    }
144}
145
146/// Account info bal structure.
147#[derive(Debug, Default, Clone, PartialEq, Eq)]
148#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
149pub struct AccountInfoBal {
150    /// Nonce builder.
151    pub nonce: BalWrites<u64>,
152    /// Balance builder.
153    pub balance: BalWrites<U256>,
154    /// Code builder.
155    pub code: BalWrites<(B256, Bytecode)>,
156}
157
158impl AccountInfoBal {
159    /// Populate account info from BAL. Return true if account info got changed
160    pub fn populate_account_info(&self, bal_index: BalIndex, account: &mut AccountInfo) -> bool {
161        let mut changed = false;
162        if let Some(nonce) = self.nonce.get(bal_index) {
163            account.nonce = nonce;
164            changed = true;
165        }
166        if let Some(balance) = self.balance.get(bal_index) {
167            account.balance = balance;
168            changed = true;
169        }
170        if let Some(code) = self.code.get(bal_index) {
171            account.code_hash = code.0;
172            account.code = Some(code.1);
173            changed = true;
174        }
175        changed
176    }
177
178    /// Extend account info from another account info.
179    #[inline]
180    pub fn update(&mut self, index: BalIndex, original: &AccountInfo, present: &AccountInfo) {
181        self.nonce.update(index, &original.nonce, present.nonce);
182        self.balance
183            .update(index, &original.balance, present.balance);
184        if original.code_hash != present.code_hash {
185            self.code.update_with_key(
186                index,
187                &original.code_hash,
188                (present.code_hash, present.code.clone().unwrap_or_default()),
189                |i| &i.0,
190            );
191        }
192    }
193
194    /// Extend account info from another account info.
195    #[inline]
196    pub fn extend(&mut self, bal_account: AccountInfoBal) {
197        self.nonce.extend(bal_account.nonce);
198        self.balance.extend(bal_account.balance);
199        self.code.extend(bal_account.code);
200    }
201
202    /// Update account balance in BAL.
203    #[inline]
204    pub fn balance_update(&mut self, bal_index: BalIndex, original_balance: &U256, balance: U256) {
205        self.balance.update(bal_index, original_balance, balance);
206    }
207
208    /// Update account nonce in BAL.
209    #[inline]
210    pub fn nonce_update(&mut self, bal_index: BalIndex, original_nonce: &u64, nonce: u64) {
211        self.nonce.update(bal_index, original_nonce, nonce);
212    }
213
214    /// Update account code in BAL.
215    #[inline]
216    pub fn code_update(
217        &mut self,
218        bal_index: BalIndex,
219        original_code_hash: &B256,
220        code_hash: B256,
221        code: Bytecode,
222    ) {
223        self.code
224            .update_with_key(bal_index, original_code_hash, (code_hash, code), |i| &i.0);
225    }
226}
227
228/// Storage BAL
229#[derive(Debug, Default, Clone, PartialEq, Eq)]
230#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
231pub struct StorageBal {
232    /// Storage with writes and reads.
233    pub storage: BTreeMap<StorageKey, BalWrites<StorageValue>>,
234}
235
236impl StorageBal {
237    /// Get storage from the builder.
238    #[inline]
239    pub fn get(
240        &self,
241        key: StorageKey,
242        bal_index: BalIndex,
243    ) -> Result<Option<StorageValue>, BalError> {
244        Ok(self.get_bal_writes(key)?.get(bal_index))
245    }
246
247    /// Get storage writes from the builder.
248    #[inline]
249    pub fn get_bal_writes(&self, key: StorageKey) -> Result<&BalWrites<StorageValue>, BalError> {
250        self.storage.get(&key).ok_or(BalError::SlotNotFound)
251    }
252
253    /// Extend storage from another storage.
254    #[inline]
255    pub fn extend(&mut self, storage: StorageBal) {
256        for (key, value) in storage.storage {
257            match self.storage.entry(key) {
258                Entry::Occupied(mut entry) => {
259                    entry.get_mut().extend(value);
260                }
261                Entry::Vacant(entry) => {
262                    entry.insert(value);
263                }
264            }
265        }
266    }
267
268    /// Update storage from [`EvmStorage`].
269    #[inline]
270    pub fn update(&mut self, bal_index: BalIndex, storage: &EvmStorage) {
271        for (key, value) in storage {
272            self.storage.entry(*key).or_default().update(
273                bal_index,
274                &value.original_value,
275                value.present_value,
276            );
277        }
278    }
279
280    /// Update reads from [`EvmStorage`].
281    ///
282    /// It will expend inner map with new reads.
283    #[inline]
284    pub fn update_reads(&mut self, storage: impl Iterator<Item = StorageKey>) {
285        for key in storage {
286            self.storage.entry(key).or_default();
287        }
288    }
289
290    /// Insert storage into the builder.
291    pub fn extend_iter(
292        &mut self,
293        storage: impl Iterator<Item = (StorageKey, BalWrites<StorageValue>)>,
294    ) {
295        for (key, value) in storage {
296            self.storage.insert(key, value);
297        }
298    }
299
300    /// Convert the storage into a vector of reads and writes
301    pub fn into_vecs(self) -> (Vec<StorageKey>, Vec<(StorageKey, BalWrites<StorageValue>)>) {
302        let mut reads = Vec::new();
303        let mut writes = Vec::new();
304
305        for (key, value) in self.storage {
306            if value.writes.is_empty() {
307                reads.push(key);
308            } else {
309                writes.push((key, value));
310            }
311        }
312
313        (reads, writes)
314    }
315}
316
317impl FromIterator<(StorageKey, BalWrites<StorageValue>)> for StorageBal {
318    fn from_iter<I: IntoIterator<Item = (StorageKey, BalWrites<StorageValue>)>>(iter: I) -> Self {
319        Self {
320            storage: iter.into_iter().collect(),
321        }
322    }
323}