anvil/eth/backend/mem/
state.rs
1use crate::eth::error::BlockchainError;
4use alloy_primitives::{keccak256, map::HashMap, Address, B256, U256};
5use alloy_rlp::Encodable;
6use alloy_rpc_types::{state::StateOverride, BlockOverrides};
7use alloy_trie::{HashBuilder, Nibbles};
8use foundry_evm::backend::DatabaseError;
9use revm::{
10 bytecode::Bytecode,
11 context::BlockEnv,
12 database::{CacheDB, DatabaseRef, DbAccount},
13 state::AccountInfo,
14};
15
16pub fn build_root(values: impl IntoIterator<Item = (Nibbles, Vec<u8>)>) -> B256 {
17 let mut builder = HashBuilder::default();
18 for (key, value) in values {
19 builder.add_leaf(key, value.as_ref());
20 }
21 builder.root()
22}
23
24pub fn state_root(accounts: &HashMap<Address, DbAccount>) -> B256 {
26 build_root(trie_accounts(accounts))
27}
28
29pub fn storage_root(storage: &HashMap<U256, U256>) -> B256 {
31 build_root(trie_storage(storage))
32}
33
34pub fn trie_storage(storage: &HashMap<U256, U256>) -> Vec<(Nibbles, Vec<u8>)> {
36 let mut storage = storage
37 .iter()
38 .map(|(key, value)| {
39 let data = alloy_rlp::encode(value);
40 (Nibbles::unpack(keccak256(key.to_be_bytes::<32>())), data)
41 })
42 .collect::<Vec<_>>();
43 storage.sort_by(|(key1, _), (key2, _)| key1.cmp(key2));
44
45 storage
46}
47
48pub fn trie_accounts(accounts: &HashMap<Address, DbAccount>) -> Vec<(Nibbles, Vec<u8>)> {
50 let mut accounts = accounts
51 .iter()
52 .map(|(address, account)| {
53 let data = trie_account_rlp(&account.info, &account.storage);
54 (Nibbles::unpack(keccak256(*address)), data)
55 })
56 .collect::<Vec<_>>();
57 accounts.sort_by(|(key1, _), (key2, _)| key1.cmp(key2));
58
59 accounts
60}
61
62pub fn trie_account_rlp(info: &AccountInfo, storage: &HashMap<U256, U256>) -> Vec<u8> {
64 let mut out: Vec<u8> = Vec::new();
65 let list: [&dyn Encodable; 4] =
66 [&info.nonce, &info.balance, &storage_root(storage), &info.code_hash];
67
68 alloy_rlp::encode_list::<_, dyn Encodable>(&list, &mut out);
69
70 out
71}
72
73pub fn apply_state_overrides<D>(
75 overrides: StateOverride,
76 cache_db: &mut CacheDB<D>,
77) -> Result<(), BlockchainError>
78where
79 D: DatabaseRef<Error = DatabaseError>,
80{
81 for (account, account_overrides) in &overrides {
82 let mut account_info = cache_db.basic_ref(*account)?.unwrap_or_default();
83
84 if let Some(nonce) = account_overrides.nonce {
85 account_info.nonce = nonce;
86 }
87 if let Some(code) = &account_overrides.code {
88 account_info.code = Some(Bytecode::new_raw(code.to_vec().into()));
89 }
90 if let Some(balance) = account_overrides.balance {
91 account_info.balance = balance;
92 }
93
94 cache_db.insert_account_info(*account, account_info);
95
96 match (&account_overrides.state, &account_overrides.state_diff) {
100 (Some(_), Some(_)) => {
101 return Err(BlockchainError::StateOverrideError(
102 "state and state_diff can't be used together".to_string(),
103 ))
104 }
105 (None, None) => (),
106 (Some(new_account_state), None) => {
107 cache_db.replace_account_storage(
108 *account,
109 new_account_state
110 .iter()
111 .map(|(key, value)| ((*key).into(), (*value).into()))
112 .collect(),
113 )?;
114 }
115 (None, Some(account_state_diff)) => {
116 for (key, value) in account_state_diff {
117 cache_db.insert_account_storage(*account, (*key).into(), (*value).into())?;
118 }
119 }
120 };
121 }
122 Ok(())
123}
124
125pub fn apply_block_overrides<DB>(
127 overrides: BlockOverrides,
128 cache_db: &mut CacheDB<DB>,
129 env: &mut BlockEnv,
130) {
131 let BlockOverrides {
132 number,
133 difficulty,
134 time,
135 gas_limit,
136 coinbase,
137 random,
138 base_fee,
139 block_hash,
140 } = overrides;
141
142 if let Some(block_hashes) = block_hash {
143 cache_db
145 .cache
146 .block_hashes
147 .extend(block_hashes.into_iter().map(|(num, hash)| (U256::from(num), hash)))
148 }
149
150 if let Some(number) = number {
151 env.number = number.saturating_to();
152 }
153 if let Some(difficulty) = difficulty {
154 env.difficulty = difficulty;
155 }
156 if let Some(time) = time {
157 env.timestamp = time;
158 }
159 if let Some(gas_limit) = gas_limit {
160 env.gas_limit = gas_limit;
161 }
162 if let Some(coinbase) = coinbase {
163 env.beneficiary = coinbase;
164 }
165 if let Some(random) = random {
166 env.prevrandao = Some(random);
167 }
168 if let Some(base_fee) = base_fee {
169 env.basefee = base_fee.saturating_to();
170 }
171}