tycho_simulation/evm/
account_storage.rs

1use std::collections::{hash_map::Entry::Vacant, HashMap};
2
3use alloy::primitives::{Address, U256};
4use revm::state::AccountInfo;
5use tracing::{debug, trace, warn};
6
7/// Represents an account in the account storage.
8///
9/// # Fields
10///
11/// * `info` - The account information of type `AccountInfo`.
12/// * `permanent_storage` - The permanent storage of the account.
13/// * `temp_storage` - The temporary storage of the account.
14/// * `mocked` - A boolean flag indicating whether the account is mocked.
15#[derive(Clone, Default, Debug)]
16pub struct Account {
17    pub info: AccountInfo,
18    pub permanent_storage: HashMap<U256, U256>,
19    pub temp_storage: HashMap<U256, U256>,
20    pub mocked: bool,
21}
22
23#[derive(Default, Clone, PartialEq, Eq, Debug)]
24pub struct StateUpdate {
25    pub storage: Option<HashMap<U256, U256>>,
26    pub balance: Option<U256>,
27}
28#[derive(Clone, Default, Debug)]
29/// A simpler implementation of CacheDB that can't query a node. It just stores data.
30pub struct AccountStorage {
31    accounts: HashMap<Address, Account>,
32}
33
34impl AccountStorage {
35    pub fn new() -> Self {
36        Self::default()
37    }
38
39    /// Clear all accounts from the storage.
40    pub fn clear(&mut self) {
41        self.accounts.clear();
42    }
43
44    /// Inserts account data into the current instance.
45    ///
46    /// # Arguments
47    ///
48    /// * `address` - The address of the account to insert.
49    /// * `info` - The account information to insert.
50    /// * `permanent_storage` - Optional storage information associated with the account.
51    /// * `mocked` - Whether this account should be considered mocked.
52    ///
53    /// # Notes
54    ///
55    /// This function checks if the `address` is already present in the `accounts`
56    /// collection. If so, it logs a trace message and returns without modifying the instance.
57    /// Otherwise, it stores a new `Account` instance with the provided data at the given address.
58    ///
59    /// Init-if-absent is only safe for seeding placeholder accounts during engine setup;
60    /// authoritative data that can arrive again for a known account (e.g. a re-fetched contract
61    /// snapshot) must use [`overwrite_account`](Self::overwrite_account) instead.
62    pub fn init_account(
63        &mut self,
64        address: Address,
65        info: AccountInfo,
66        permanent_storage: Option<HashMap<U256, U256>>,
67        mocked: bool,
68    ) {
69        if let Vacant(e) = self.accounts.entry(address) {
70            e.insert(Account {
71                info,
72                permanent_storage: permanent_storage.unwrap_or_default(),
73                temp_storage: HashMap::new(),
74                mocked,
75            });
76            debug!(
77                "Inserted a {} account {:x?}",
78                if mocked { "mocked" } else { "non-mocked" },
79                address
80            );
81        } else {
82            trace!("Skipped init for already-existing account {:x?}", address);
83        }
84    }
85
86    /// Inserts account data into the current instance, replacing any existing entry.
87    ///
88    /// # Arguments
89    ///
90    /// * `address` - The address of the account to insert.
91    /// * `info` - The account information to insert.
92    /// * `permanent_storage` - Optional storage information associated with the account.
93    /// * `mocked` - Whether this account should be considered mocked.
94    ///
95    /// # Notes
96    ///
97    /// Unlike `init_account`, this function always replaces the account at the given address,
98    /// even if one already exists. Use this for `ChangeType::Creation` updates where the latest
99    /// snapshot data must win over any previously inserted placeholder.
100    pub fn overwrite_account(
101        &mut self,
102        address: Address,
103        info: AccountInfo,
104        permanent_storage: Option<HashMap<U256, U256>>,
105        mocked: bool,
106    ) {
107        self.accounts.insert(
108            address,
109            Account {
110                info,
111                permanent_storage: permanent_storage.unwrap_or_default(),
112                temp_storage: HashMap::new(),
113                mocked,
114            },
115        );
116        debug!(
117            "Overwrote a {} account {:x?}",
118            if mocked { "mocked" } else { "non-mocked" },
119            address
120        );
121    }
122
123    /// Updates the account information and storage associated with the given address.
124    ///
125    /// # Arguments
126    ///
127    /// * `address` - The address of the account to update.
128    /// * `update` - The state update containing the new information to apply.
129    ///
130    /// # Notes
131    ///
132    /// This function looks for the account information and storage associated with the provided
133    /// `address`. If the `address` exists in the `accounts` collection, it updates the account
134    /// information based on the `balance` field in the `update` parameter. If the `address` exists
135    /// in the `storage` collection, it updates the storage information based on the `storage` field
136    /// in the `update` parameter.
137    ///
138    /// If the `address` is not found in either collection, a warning is logged and no changes are
139    /// made.
140    pub fn update_account(&mut self, address: &Address, update: &StateUpdate) {
141        if let Some(account) = self.accounts.get_mut(address) {
142            if let Some(new_balance) = update.balance {
143                account.info.balance = new_balance;
144            }
145            if let Some(new_storage) = &update.storage {
146                for (index, value) in new_storage {
147                    account
148                        .permanent_storage
149                        .insert(*index, *value);
150                }
151            }
152        } else {
153            warn!(?address, "Tried to update account {:x?} that was not initialized", address);
154        }
155    }
156
157    /// Retrieves the account information for a given address.
158    ///
159    /// This function retrieves the account information associated with the specified address from
160    /// the storage.
161    ///
162    /// # Arguments
163    ///
164    /// * `address`: The address of the account to retrieve the information for.
165    ///
166    /// # Returns
167    ///
168    /// Returns an `Option` that holds a reference to the `AccountInfo`. If the account is not
169    /// found, `None` is returned.
170    pub fn get_account_info(&self, address: &Address) -> Option<&AccountInfo> {
171        self.accounts
172            .get(address)
173            .map(|acc| &acc.info)
174    }
175
176    /// Checks if an account with the given address is present in the storage.
177    ///
178    /// # Arguments
179    ///
180    /// * `address`: A reference to the address of the account to check.
181    ///
182    /// # Returns
183    ///
184    /// Returns `true` if an account with the specified address is present in the storage,
185    /// otherwise returns `false`.
186    pub fn account_present(&self, address: &Address) -> bool {
187        self.accounts.contains_key(address)
188    }
189
190    /// Sets the storage value at the specified index for the given account.
191    ///
192    /// If the account exists in the storage, the storage value at the specified `index` is updated.
193    /// If the account does not exist, a warning message is logged indicating an attempt to set
194    /// storage on an uninitialized account.
195    ///
196    /// # Arguments
197    ///
198    /// * `address`: The address of the account to set the storage value for.
199    /// * `index`: The index of the storage value to set.
200    /// * `value`: The new value to set for the storage.
201    pub fn set_temp_storage(&mut self, address: Address, index: U256, value: U256) {
202        if let Some(acc) = self.accounts.get_mut(&address) {
203            acc.temp_storage.insert(index, value);
204        } else {
205            warn!("Trying to set storage on unitialized account {:x?}.", address);
206        }
207    }
208
209    /// Retrieves the storage value at the specified index for the given account, if it exists.
210    ///
211    /// If the account exists in the storage, the storage value at the specified `index` is returned
212    /// as a reference. Temp storage takes priority over permanent storage.
213    /// If the account does not exist, `None` is returned.
214    ///
215    /// # Arguments
216    ///
217    /// * `address`: A reference to the address of the account to retrieve the storage value from.
218    /// * `index`: A reference to the index of the storage value to retrieve.
219    ///
220    /// # Returns
221    ///
222    /// Returns an `Option` containing a reference to the storage value if it exists, otherwise
223    /// returns `None`.
224    pub fn get_storage(&self, address: &Address, index: &U256) -> Option<U256> {
225        if let Some(acc) = self.accounts.get(address) {
226            if let Some(s) = acc.temp_storage.get(index) {
227                Some(*s)
228            } else {
229                acc.permanent_storage
230                    .get(index)
231                    .copied()
232            }
233        } else {
234            None
235        }
236    }
237
238    /// Retrieves the permanent storage value for the given address and index.
239    ///
240    /// If an account with the specified address exists in the account storage, this function
241    /// retrieves the corresponding permanent storage value associated with the given index.
242    ///
243    /// # Arguments
244    ///
245    /// * `address` - The address of the account.
246    /// * `index` - The index of the desired storage value.
247    pub fn get_permanent_storage(&self, address: &Address, index: &U256) -> Option<U256> {
248        if let Some(acc) = self.accounts.get(address) {
249            acc.permanent_storage
250                .get(index)
251                .copied()
252        } else {
253            None
254        }
255    }
256
257    /// Removes all temp storage values.
258    ///
259    /// Iterates over the accounts in the storage and removes all temp storage values
260    pub fn clear_temp_storage(&mut self) {
261        self.accounts
262            .values_mut()
263            .for_each(|acc| acc.temp_storage.clear());
264    }
265
266    /// Checks if an account is mocked based on its address.
267    ///
268    /// # Arguments
269    ///
270    /// * `address` - A reference to the account address.
271    pub fn is_mocked_account(&self, address: &Address) -> Option<bool> {
272        self.accounts
273            .get(address)
274            .map(|acc| acc.mocked)
275    }
276}
277
278#[cfg(test)]
279mod tests {
280    use std::{error::Error, str::FromStr};
281
282    use revm::primitives::KECCAK_EMPTY;
283
284    use super::*;
285    use crate::evm::account_storage::{Account, AccountStorage};
286
287    #[test]
288    fn test_insert_account() -> Result<(), Box<dyn Error>> {
289        let mut account_storage = AccountStorage::default();
290        let expected_nonce = 100;
291        let expected_balance = U256::from(500);
292        let acc_address = Address::from_str("0xb4e16d0168e52d35cacd2c6185b44281ec28c9dc")?;
293        let info: AccountInfo = AccountInfo {
294            nonce: expected_nonce,
295            balance: expected_balance,
296            code: None,
297            code_hash: KECCAK_EMPTY,
298        };
299        let mut storage_new = HashMap::new();
300        let expected_storage_value = U256::from_str("5").unwrap();
301        storage_new.insert(U256::from_str("1").unwrap(), expected_storage_value);
302
303        account_storage.init_account(acc_address, info, Some(storage_new), false);
304
305        let acc = account_storage
306            .get_account_info(&acc_address)
307            .unwrap();
308        let storage_value = account_storage
309            .get_storage(&acc_address, &U256::from_str("1").unwrap())
310            .unwrap();
311        assert_eq!(acc.nonce, expected_nonce, "Nonce should match expected value");
312        assert_eq!(acc.balance, expected_balance, "Balance should match expected value");
313        assert_eq!(acc.code_hash, KECCAK_EMPTY, "Code hash should match expected value");
314        assert_eq!(
315            storage_value, expected_storage_value,
316            "Storage value should match expected value"
317        );
318        Ok(())
319    }
320
321    #[test]
322    fn test_update_account_info() -> Result<(), Box<dyn Error>> {
323        let mut account_storage = AccountStorage::default();
324        let acc_address = Address::from_str("0xb4e16d0168e52d35cacd2c6185b44281ec28c9dc")?;
325        let info: AccountInfo = AccountInfo {
326            nonce: 100,
327            balance: U256::from(500),
328            code: None,
329            code_hash: KECCAK_EMPTY,
330        };
331        let mut original_storage = HashMap::new();
332        let storage_index = U256::from_str("1").unwrap();
333        original_storage.insert(storage_index, U256::from_str("5").unwrap());
334        account_storage.accounts.insert(
335            acc_address,
336            Account {
337                info,
338                permanent_storage: original_storage,
339                temp_storage: HashMap::new(),
340                mocked: false,
341            },
342        );
343        let updated_balance = U256::from(100);
344        let updated_storage_value = U256::from_str("999").unwrap();
345        let mut updated_storage = HashMap::new();
346        updated_storage.insert(storage_index, updated_storage_value);
347        let state_update =
348            StateUpdate { balance: Some(updated_balance), storage: Some(updated_storage) };
349
350        account_storage.update_account(&acc_address, &state_update);
351
352        assert_eq!(
353            account_storage
354                .get_account_info(&acc_address)
355                .unwrap()
356                .balance,
357            updated_balance,
358            "Account balance should be updated"
359        );
360        assert_eq!(
361            account_storage
362                .get_storage(&acc_address, &storage_index)
363                .unwrap(),
364            updated_storage_value,
365            "Storage value should be updated"
366        );
367        Ok(())
368    }
369
370    #[test]
371    fn test_get_account_info() {
372        let mut account_storage = AccountStorage::default();
373        let address_1 = Address::from_str("0xb4e16d0168e52d35cacd2c6185b44281ec28c9dc").unwrap();
374        let address_2 = Address::from_str("0xb4e16d0168e52d35cacd2c6185b44281ec28c9dd").unwrap();
375        let account_info_1 = AccountInfo::default();
376        let account_info_2 = AccountInfo { nonce: 500, ..Default::default() };
377        account_storage.init_account(address_1, account_info_1, None, false);
378        account_storage.init_account(address_2, account_info_2, None, false);
379
380        let existing_account = account_storage.get_account_info(&address_1);
381        let address_3 = Address::from_str("0xb4e16d0168e52d35cacd2c6185b44281ec28c9de").unwrap();
382        let non_existing_account = account_storage.get_account_info(&address_3);
383
384        assert_eq!(
385            existing_account.unwrap().nonce,
386            AccountInfo::default().nonce,
387            "Existing account's nonce should match the expected value"
388        );
389        assert_eq!(non_existing_account, None, "Non-existing account should return None");
390    }
391
392    #[test]
393    fn test_account_present() {
394        let mut account_storage = AccountStorage::default();
395        let existing_account =
396            Address::from_str("0xb4e16d0168e52d35cacd2c6185b44281ec28c9dc").unwrap();
397        let address_2 = Address::from_str("0xb4e16d0168e52d35cacd2c6185b44281ec28c9dd").unwrap();
398        let non_existing_account =
399            Address::from_str("0xb4e16d0168e52d35cacd2c6185b44281ec28c9de").unwrap();
400        account_storage
401            .accounts
402            .insert(existing_account, Account::default());
403        account_storage
404            .accounts
405            .insert(address_2, Account::default());
406
407        assert!(
408            account_storage.account_present(&existing_account),
409            "Existing account should be present in the AccountStorage"
410        );
411        assert!(
412            !account_storage.account_present(&non_existing_account),
413            "Non-existing account should not be present in the AccountStorage"
414        );
415    }
416
417    #[test]
418    fn test_set_get_storage() {
419        // Create a new instance of the struct for testing
420        let mut account_storage = AccountStorage::default();
421        // Add a test account
422        let address = Address::from_str("0xb4e16d0168e52d35cacd2c6185b44281ec28c9dc").unwrap();
423        let non_existing_address =
424            Address::from_str("0xb4e16d0168e52d35cacd2c6185b44281ec28c9dd").unwrap();
425        let account = Account::default();
426        account_storage
427            .accounts
428            .insert(address, account);
429        let index = U256::from_str("1").unwrap();
430        let value = U256::from_str("1").unwrap();
431        let non_existing_index = U256::from_str("2").unwrap();
432        let non_existing_value = U256::from_str("2").unwrap();
433        account_storage.set_temp_storage(
434            non_existing_address,
435            non_existing_index,
436            non_existing_value,
437        );
438        account_storage.set_temp_storage(address, index, value);
439
440        let storage = account_storage.get_storage(&address, &index);
441        let empty_storage = account_storage.get_storage(&non_existing_address, &non_existing_index);
442
443        assert_eq!(storage, Some(value), "Storage value should match the value that was set");
444        assert_eq!(empty_storage, None, "Storage value should be None for a non-existing account");
445    }
446
447    #[test]
448    fn test_get_storage() {
449        let mut account_storage = AccountStorage::default();
450        let existing_address =
451            Address::from_str("0xb4e16d0168e52d35cacd2c6185b44281ec28c9dc").unwrap();
452        let non_existent_address =
453            Address::from_str("0xb4e16d0168e52d35cacd2c6185b44281ec28c9dd").unwrap();
454        let index = U256::from(42);
455        let value = U256::from(100);
456        let non_existent_index = U256::from(999);
457        let mut account = Account::default();
458        account
459            .temp_storage
460            .insert(index, value);
461        account_storage
462            .accounts
463            .insert(existing_address, account);
464
465        assert_eq!(
466            account_storage.get_storage(&existing_address, &index),
467            Some(value), "If the storage features the address and index the value at that position should be retunred."
468        );
469
470        // Test with non-existent address
471        assert_eq!(
472            account_storage.get_storage(&non_existent_address, &index),
473            None,
474            "If the storage does not feature the address None should be returned."
475        );
476
477        // Test with non-existent index
478        assert_eq!(
479            account_storage.get_storage(&existing_address, &non_existent_index),
480            None,
481            "If the storage does not feature the index None should be returned."
482        );
483    }
484
485    #[test]
486    fn test_get_storage_priority() {
487        let mut account_storage = AccountStorage::default();
488        let address = Address::from_str("0xb4e16d0168e52d35cacd2c6185b44281ec28c9dc").unwrap();
489        let index = U256::from(69);
490        let temp_value = U256::from(100);
491        let permanent_value = U256::from(200);
492        let mut account = Account::default();
493        account
494            .temp_storage
495            .insert(index, temp_value);
496        account
497            .permanent_storage
498            .insert(index, permanent_value);
499        account_storage
500            .accounts
501            .insert(address, account);
502
503        assert_eq!(
504            account_storage.get_storage(&address, &index),
505            Some(temp_value),
506            "Temp storage value should take priority over permanent storage value"
507        );
508    }
509
510    #[test]
511    fn test_is_mocked_account() {
512        let mut account_storage = AccountStorage::default();
513        let mocked_account_address =
514            Address::from_str("0xb4e16d0168e52d35cacd2c6185b44281ec28c9dc").unwrap();
515        let not_mocked_account_address =
516            Address::from_str("0xb4e16d0168e52d35cacd2c6185b44281ec28c9dd").unwrap();
517        let unknown_address =
518            Address::from_str("0xb4e16d0168e52d35cacd2c6185b44281ec28c9de").unwrap();
519        let mocked_account = Account { mocked: true, ..Default::default() };
520        let not_mocked_account = Account { mocked: false, ..Default::default() };
521        account_storage
522            .accounts
523            .insert(mocked_account_address, mocked_account);
524        account_storage
525            .accounts
526            .insert(not_mocked_account_address, not_mocked_account);
527
528        assert_eq!(account_storage.is_mocked_account(&mocked_account_address), Some(true));
529        assert_eq!(account_storage.is_mocked_account(&not_mocked_account_address), Some(false));
530        assert_eq!(account_storage.is_mocked_account(&unknown_address), None);
531    }
532
533    #[test]
534    fn test_clear_temp_storage() {
535        let mut account_storage = AccountStorage::default();
536        let address_1 = Address::from_str("0xb4e16d0168e52d35cacd2c6185b44281ec28c9dc").unwrap();
537        let address_2 = Address::from_str("0xb4e16d0168e52d35cacd2c6185b44281ec28c9dd").unwrap();
538        let mut account_1 = Account::default();
539        account_1
540            .temp_storage
541            .insert(U256::from(1), U256::from(10));
542        let mut account_2 = Account::default();
543        account_2
544            .temp_storage
545            .insert(U256::from(2), U256::from(20));
546        account_storage
547            .accounts
548            .insert(address_1, account_1);
549        account_storage
550            .accounts
551            .insert(address_2, account_2);
552
553        account_storage.clear_temp_storage();
554
555        let account_1_temp_storage = account_storage.accounts[&address_1]
556            .temp_storage
557            .len();
558        let account_2_temp_storage = account_storage.accounts[&address_2]
559            .temp_storage
560            .len();
561        assert_eq!(account_1_temp_storage, 0, "Temporary storage of account 1 should be cleared");
562        assert_eq!(account_2_temp_storage, 0, "Temporary storage of account 2 should be cleared");
563    }
564
565    #[test]
566    fn test_get_permanent_storage() {
567        let mut account_storage = AccountStorage::default();
568        let address = Address::from_str("0xb4e16d0168e52d35cacd2c6185b44281ec28c9dc").unwrap();
569        let non_existing_address =
570            Address::from_str("0xb4e16d0168e52d35cacd2c6185b44281ec28c9dd").unwrap();
571        let index = U256::from_str("123").unwrap();
572        let value = U256::from_str("456").unwrap();
573        let mut account = Account::default();
574        account
575            .permanent_storage
576            .insert(index, value);
577        account_storage
578            .accounts
579            .insert(address, account);
580
581        let result = account_storage.get_permanent_storage(&address, &index);
582        let not_existing_result =
583            account_storage.get_permanent_storage(&non_existing_address, &index);
584        let empty_index = U256::from_str("789").unwrap();
585        let no_storage = account_storage.get_permanent_storage(&address, &empty_index);
586
587        assert_eq!(
588            result,
589            Some(value),
590            "Expected value for existing account with permanent storage"
591        );
592        assert_eq!(not_existing_result, None, "Expected None for non-existing account");
593        assert_eq!(
594            no_storage, None,
595            "Expected None for existing account without permanent storage"
596        );
597    }
598}
tycho_simulation/evm/account_storage.rs

tycho_simulation/evm/
account_storage.rs
