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.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 #[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 #[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#[derive(Debug, Default, Clone, PartialEq, Eq)]
147#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
148pub struct AccountInfoBal {
149 pub nonce: BalWrites<u64>,
151 pub balance: BalWrites<U256>,
153 pub code: BalWrites<(B256, Bytecode)>,
155}
156
157impl AccountInfoBal {
158 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 #[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 #[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 #[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 #[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 #[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#[derive(Debug, Default, Clone, PartialEq, Eq)]
227#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
228pub struct StorageBal {
229 pub storage: BTreeMap<StorageKey, BalWrites<StorageValue>>,
231}
232
233impl StorageBal {
234 #[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 #[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 #[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 #[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 #[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 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 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}