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            // Selfdestruct wipes all storage to zero, record writes accordingly.
58            self.storage
59                .update_selfdestruct(bal_index, &account.storage);
60            return;
61        }
62
63        self.account_info
64            .update(bal_index, &account.original_info, &account.info);
65
66        self.storage.update(bal_index, &account.storage);
67    }
68
69    /// Create account from alloy account changes.
70    #[inline]
71    pub fn try_from_alloy(
72        alloy_account: AlloyAccountChanges,
73    ) -> Result<(Address, Self), BytecodeDecodeError> {
74        Ok((
75            alloy_account.address,
76            AccountBal {
77                account_info: AccountInfoBal {
78                    nonce: BalWrites::from(alloy_account.nonce_changes),
79                    balance: BalWrites::from(alloy_account.balance_changes),
80                    code: BalWrites::try_from(alloy_account.code_changes)?,
81                },
82                storage: StorageBal::from_iter(
83                    alloy_account
84                        .storage_changes
85                        .into_iter()
86                        .chain(
87                            alloy_account
88                                .storage_reads
89                                .into_iter()
90                                .map(|key| AlloySlotChanges::new(key, Default::default())),
91                        )
92                        .map(|slot| (slot.slot, BalWrites::from(slot.changes))),
93                ),
94            },
95        ))
96    }
97
98    /// Consumes AccountBal and converts it into [`AlloyAccountChanges`].
99    #[inline]
100    pub fn into_alloy_account(self, address: Address) -> AlloyAccountChanges {
101        let storage_len = self.storage.storage.len();
102        let mut storage_reads = Vec::with_capacity(storage_len);
103        let mut storage_changes = Vec::with_capacity(storage_len);
104        for (key, value) in self.storage.storage {
105            if value.writes.is_empty() {
106                storage_reads.push(key);
107            } else {
108                storage_changes.push(AlloySlotChanges::new(
109                    key,
110                    value
111                        .writes
112                        .into_iter()
113                        .map(|(index, value)| AlloyStorageChange::new(index, value))
114                        .collect(),
115                ));
116            }
117        }
118
119        AlloyAccountChanges {
120            address,
121            storage_changes,
122            storage_reads,
123            balance_changes: self
124                .account_info
125                .balance
126                .writes
127                .into_iter()
128                .map(|(index, value)| AlloyBalanceChange::new(index, value))
129                .collect(),
130            nonce_changes: self
131                .account_info
132                .nonce
133                .writes
134                .into_iter()
135                .map(|(index, value)| AlloyNonceChange::new(index, value))
136                .collect(),
137            code_changes: self
138                .account_info
139                .code
140                .writes
141                .into_iter()
142                .map(|(index, (_, value))| AlloyCodeChange::new(index, value.original_bytes()))
143                .collect(),
144        }
145    }
146}
147
148/// Account info bal structure.
149#[derive(Debug, Default, Clone, PartialEq, Eq)]
150#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
151pub struct AccountInfoBal {
152    /// Nonce builder.
153    pub nonce: BalWrites<u64>,
154    /// Balance builder.
155    pub balance: BalWrites<U256>,
156    /// Code builder.
157    pub code: BalWrites<(B256, Bytecode)>,
158}
159
160impl AccountInfoBal {
161    /// Populate account info from BAL. Return true if account info got changed
162    pub fn populate_account_info(&self, bal_index: BalIndex, account: &mut AccountInfo) -> bool {
163        let mut changed = false;
164        if let Some(nonce) = self.nonce.get(bal_index) {
165            account.nonce = nonce;
166            changed = true;
167        }
168        if let Some(balance) = self.balance.get(bal_index) {
169            account.balance = balance;
170            changed = true;
171        }
172        if let Some(code) = self.code.get(bal_index) {
173            account.code_hash = code.0;
174            account.code = Some(code.1);
175            changed = true;
176        }
177        changed
178    }
179
180    /// Extend account info from another account info.
181    #[inline]
182    pub fn update(&mut self, index: BalIndex, original: &AccountInfo, present: &AccountInfo) {
183        self.nonce.update(index, &original.nonce, present.nonce);
184        self.balance
185            .update(index, &original.balance, present.balance);
186        if original.code_hash != present.code_hash {
187            self.code.update_with_key(
188                index,
189                &original.code_hash,
190                (present.code_hash, present.code.clone().unwrap_or_default()),
191                |i| &i.0,
192            );
193        }
194    }
195
196    /// Extend account info from another account info.
197    #[inline]
198    pub fn extend(&mut self, bal_account: AccountInfoBal) {
199        self.nonce.extend(bal_account.nonce);
200        self.balance.extend(bal_account.balance);
201        self.code.extend(bal_account.code);
202    }
203
204    /// Update account balance in BAL.
205    #[inline]
206    pub fn balance_update(&mut self, bal_index: BalIndex, original_balance: &U256, balance: U256) {
207        self.balance.update(bal_index, original_balance, balance);
208    }
209
210    /// Update account nonce in BAL.
211    #[inline]
212    pub fn nonce_update(&mut self, bal_index: BalIndex, original_nonce: &u64, nonce: u64) {
213        self.nonce.update(bal_index, original_nonce, nonce);
214    }
215
216    /// Update account code in BAL.
217    #[inline]
218    pub fn code_update(
219        &mut self,
220        bal_index: BalIndex,
221        original_code_hash: &B256,
222        code_hash: B256,
223        code: Bytecode,
224    ) {
225        self.code
226            .update_with_key(bal_index, original_code_hash, (code_hash, code), |i| &i.0);
227    }
228}
229
230/// Storage BAL
231#[derive(Debug, Default, Clone, PartialEq, Eq)]
232#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
233pub struct StorageBal {
234    /// Storage with writes and reads.
235    pub storage: BTreeMap<StorageKey, BalWrites<StorageValue>>,
236}
237
238impl StorageBal {
239    /// Get storage from the builder.
240    #[inline]
241    pub fn get(
242        &self,
243        key: StorageKey,
244        bal_index: BalIndex,
245    ) -> Result<Option<StorageValue>, BalError> {
246        Ok(self.get_bal_writes(key)?.get(bal_index))
247    }
248
249    /// Get storage writes from the builder.
250    #[inline]
251    pub fn get_bal_writes(&self, key: StorageKey) -> Result<&BalWrites<StorageValue>, BalError> {
252        self.storage.get(&key).ok_or(BalError::SlotNotFound)
253    }
254
255    /// Extend storage from another storage.
256    #[inline]
257    pub fn extend(&mut self, storage: StorageBal) {
258        for (key, value) in storage.storage {
259            match self.storage.entry(key) {
260                Entry::Occupied(mut entry) => {
261                    entry.get_mut().extend(value);
262                }
263                Entry::Vacant(entry) => {
264                    entry.insert(value);
265                }
266            }
267        }
268    }
269
270    /// Update storage from [`EvmStorage`].
271    #[inline]
272    pub fn update(&mut self, bal_index: BalIndex, storage: &EvmStorage) {
273        for (key, value) in storage {
274            self.storage.entry(*key).or_default().update(
275                bal_index,
276                &value.original_value,
277                value.present_value,
278            );
279        }
280    }
281
282    /// Update storage for a selfdestructed account.
283    ///
284    /// All accessed slots are recorded as written to zero since selfdestruct wipes storage.
285    #[inline]
286    pub fn update_selfdestruct(&mut self, bal_index: BalIndex, storage: &EvmStorage) {
287        for (key, value) in storage {
288            self.storage.entry(*key).or_default().update(
289                bal_index,
290                &value.original_value,
291                StorageValue::ZERO,
292            );
293        }
294    }
295
296    /// Update reads from [`EvmStorage`].
297    ///
298    /// It will expend inner map with new reads.
299    #[inline]
300    pub fn update_reads(&mut self, storage: impl Iterator<Item = StorageKey>) {
301        for key in storage {
302            self.storage.entry(key).or_default();
303        }
304    }
305
306    /// Insert storage into the builder.
307    pub fn extend_iter(
308        &mut self,
309        storage: impl Iterator<Item = (StorageKey, BalWrites<StorageValue>)>,
310    ) {
311        for (key, value) in storage {
312            self.storage.insert(key, value);
313        }
314    }
315
316    /// Convert the storage into a vector of reads and writes
317    pub fn into_vecs(self) -> (Vec<StorageKey>, Vec<(StorageKey, BalWrites<StorageValue>)>) {
318        let mut reads = Vec::new();
319        let mut writes = Vec::new();
320
321        for (key, value) in self.storage {
322            if value.writes.is_empty() {
323                reads.push(key);
324            } else {
325                writes.push((key, value));
326            }
327        }
328
329        (reads, writes)
330    }
331}
332
333impl FromIterator<(StorageKey, BalWrites<StorageValue>)> for StorageBal {
334    fn from_iter<I: IntoIterator<Item = (StorageKey, BalWrites<StorageValue>)>>(iter: I) -> Self {
335        Self {
336            storage: iter.into_iter().collect(),
337        }
338    }
339}