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 mut storage_reads = Vec::new();
100        let mut storage_changes = Vec::new();
101        for (key, value) in self.storage.storage {
102            if value.writes.is_empty() {
103                storage_reads.push(key);
104            } else {
105                storage_changes.push(AlloySlotChanges::new(
106                    key,
107                    value
108                        .writes
109                        .into_iter()
110                        .map(|(index, value)| AlloyStorageChange::new(index, value))
111                        .collect(),
112                ));
113            }
114        }
115
116        AlloyAccountChanges {
117            address,
118            storage_changes,
119            storage_reads,
120            balance_changes: self
121                .account_info
122                .balance
123                .writes
124                .into_iter()
125                .map(|(index, value)| AlloyBalanceChange::new(index, value))
126                .collect(),
127            nonce_changes: self
128                .account_info
129                .nonce
130                .writes
131                .into_iter()
132                .map(|(index, value)| AlloyNonceChange::new(index, value))
133                .collect(),
134            code_changes: self
135                .account_info
136                .code
137                .writes
138                .into_iter()
139                .map(|(index, (_, value))| AlloyCodeChange::new(index, value.original_bytes()))
140                .collect(),
141        }
142    }
143}
144
145/// Account info bal structure.
146#[derive(Debug, Default, Clone, PartialEq, Eq)]
147#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
148pub struct AccountInfoBal {
149    /// Nonce builder.
150    pub nonce: BalWrites<u64>,
151    /// Balance builder.
152    pub balance: BalWrites<U256>,
153    /// Code builder.
154    pub code: BalWrites<(B256, Bytecode)>,
155}
156
157impl AccountInfoBal {
158    /// Populate account info from BAL. Return true if account info got changed
159    pub fn populate_account_info(&self, bal_index: BalIndex, account: &mut AccountInfo) -> bool {
160        let mut changed = false;
161        if let Some(nonce) = self.nonce.get(bal_index) {
162            account.nonce = nonce;
163            changed = true;
164        }
165        if let Some(balance) = self.balance.get(bal_index) {
166            account.balance = balance;
167            changed = true;
168        }
169        if let Some(code) = self.code.get(bal_index) {
170            account.code_hash = code.0;
171            account.code = Some(code.1);
172            changed = true;
173        }
174        changed
175    }
176
177    /// Extend account info from another account info.
178    #[inline]
179    pub fn update(&mut self, index: BalIndex, original: &AccountInfo, present: &AccountInfo) {
180        self.nonce.update(index, &original.nonce, present.nonce);
181        self.balance
182            .update(index, &original.balance, present.balance);
183        self.code.update_with_key(
184            index,
185            &original.code_hash,
186            (present.code_hash, present.code.clone().unwrap_or_default()),
187            |i| &i.0,
188        );
189    }
190
191    /// Extend account info from another account info.
192    #[inline]
193    pub fn extend(&mut self, bal_account: AccountInfoBal) {
194        self.nonce.extend(bal_account.nonce);
195        self.balance.extend(bal_account.balance);
196        self.code.extend(bal_account.code);
197    }
198
199    /// Update account balance in BAL.
200    #[inline]
201    pub fn balance_update(&mut self, bal_index: BalIndex, original_balance: &U256, balance: U256) {
202        self.balance.update(bal_index, original_balance, balance);
203    }
204
205    /// Update account nonce in BAL.
206    #[inline]
207    pub fn nonce_update(&mut self, bal_index: BalIndex, original_nonce: &u64, nonce: u64) {
208        self.nonce.update(bal_index, original_nonce, nonce);
209    }
210
211    /// Update account code in BAL.
212    #[inline]
213    pub fn code_update(
214        &mut self,
215        bal_index: BalIndex,
216        original_code_hash: &B256,
217        code_hash: B256,
218        code: Bytecode,
219    ) {
220        self.code
221            .update_with_key(bal_index, original_code_hash, (code_hash, code), |i| &i.0);
222    }
223}
224
225/// Storage BAL
226#[derive(Debug, Default, Clone, PartialEq, Eq)]
227#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
228pub struct StorageBal {
229    /// Storage with writes and reads.
230    pub storage: BTreeMap<StorageKey, BalWrites<StorageValue>>,
231}
232
233impl StorageBal {
234    /// Get storage from the builder.
235    #[inline]
236    pub fn get(
237        &self,
238        key: StorageKey,
239        bal_index: BalIndex,
240    ) -> Result<Option<StorageValue>, BalError> {
241        Ok(self.get_bal_writes(key)?.get(bal_index))
242    }
243
244    /// Get storage writes from the builder.
245    #[inline]
246    pub fn get_bal_writes(&self, key: StorageKey) -> Result<&BalWrites<StorageValue>, BalError> {
247        self.storage.get(&key).ok_or(BalError::SlotNotFound)
248    }
249
250    /// Extend storage from another storage.
251    #[inline]
252    pub fn extend(&mut self, storage: StorageBal) {
253        for (key, value) in storage.storage {
254            match self.storage.entry(key) {
255                Entry::Occupied(mut entry) => {
256                    entry.get_mut().extend(value);
257                }
258                Entry::Vacant(entry) => {
259                    entry.insert(value);
260                }
261            }
262        }
263    }
264
265    /// Update storage from [`EvmStorage`].
266    #[inline]
267    pub fn update(&mut self, bal_index: BalIndex, storage: &EvmStorage) {
268        for (key, value) in storage {
269            self.storage.entry(*key).or_default().update(
270                bal_index,
271                &value.original_value,
272                value.present_value,
273            );
274        }
275    }
276
277    /// Update reads from [`EvmStorage`].
278    ///
279    /// It will expend inner map with new reads.
280    #[inline]
281    pub fn update_reads(&mut self, storage: impl Iterator<Item = StorageKey>) {
282        for key in storage {
283            self.storage.entry(key).or_default();
284        }
285    }
286
287    /// Insert storage into the builder.
288    pub fn extend_iter(
289        &mut self,
290        storage: impl Iterator<Item = (StorageKey, BalWrites<StorageValue>)>,
291    ) {
292        for (key, value) in storage {
293            self.storage.insert(key, value);
294        }
295    }
296
297    /// Convert the storage into a vector of reads and writes
298    pub fn into_vecs(self) -> (Vec<StorageKey>, Vec<(StorageKey, BalWrites<StorageValue>)>) {
299        let mut reads = Vec::new();
300        let mut writes = Vec::new();
301
302        for (key, value) in self.storage {
303            if value.writes.is_empty() {
304                reads.push(key);
305            } else {
306                writes.push((key, value));
307            }
308        }
309
310        (reads, writes)
311    }
312}
313
314impl FromIterator<(StorageKey, BalWrites<StorageValue>)> for StorageBal {
315    fn from_iter<I: IntoIterator<Item = (StorageKey, BalWrites<StorageValue>)>>(iter: I) -> Self {
316        Self {
317            storage: iter.into_iter().collect(),
318        }
319    }
320}