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 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#[derive(Debug, Default, Clone, PartialEq, Eq)]
148#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
149pub struct AccountInfoBal {
150 pub nonce: BalWrites<u64>,
152 pub balance: BalWrites<U256>,
154 pub code: BalWrites<(B256, Bytecode)>,
156}
157
158impl AccountInfoBal {
159 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 #[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 #[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 #[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 #[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 #[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#[derive(Debug, Default, Clone, PartialEq, Eq)]
230#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
231pub struct StorageBal {
232 pub storage: BTreeMap<StorageKey, BalWrites<StorageValue>>,
234}
235
236impl StorageBal {
237 #[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 #[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 #[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 #[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 #[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 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 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}