1use 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#[derive(Debug, Default, Clone, PartialEq, Eq)]
22#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
23pub struct AccountBal {
24 pub account_info: AccountInfoBal,
26 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 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 #[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
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 #[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 #[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#[derive(Debug, Default, Clone, PartialEq, Eq)]
150#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
151pub struct AccountInfoBal {
152 pub nonce: BalWrites<u64>,
154 pub balance: BalWrites<U256>,
156 pub code: BalWrites<(B256, Bytecode)>,
158}
159
160impl AccountInfoBal {
161 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 #[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 #[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 #[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 #[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 #[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#[derive(Debug, Default, Clone, PartialEq, Eq)]
232#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
233pub struct StorageBal {
234 pub storage: BTreeMap<StorageKey, BalWrites<StorageValue>>,
236}
237
238impl StorageBal {
239 #[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 #[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 #[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 #[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 #[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 #[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 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 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}