anvil_zksync_core/node/
state_override.rs

1// Copy-pasted from the core API sandbox implementation.
2
3use zksync_multivm::interface::storage::{ReadStorage, StorageWithOverrides};
4use zksync_types::{
5    api::state_override::{BytecodeOverride, OverrideState, StateOverride},
6    bytecode::{pad_evm_bytecode, BytecodeHash, BytecodeMarker},
7    get_code_key, get_evm_code_hash_key, get_known_code_key, get_nonce_key, h256_to_u256,
8    u256_to_h256,
9    utils::{decompose_full_nonce, nonces_to_full_nonce, storage_key_for_eth_balance},
10    web3, AccountTreeId, StorageKey, H256,
11};
12
13pub(super) fn apply_state_override<S: ReadStorage>(
14    storage: S,
15    state_override: StateOverride,
16) -> StorageWithOverrides<S> {
17    let mut storage = StorageWithOverrides::new(storage);
18
19    for (account, overrides) in state_override {
20        if let Some(balance) = overrides.balance {
21            let balance_key = storage_key_for_eth_balance(&account);
22            storage.set_value(balance_key, u256_to_h256(balance));
23        }
24
25        if let Some(nonce) = overrides.nonce {
26            let nonce_key = get_nonce_key(&account);
27            let full_nonce = storage.read_value(&nonce_key);
28            let (_, deployment_nonce) = decompose_full_nonce(h256_to_u256(full_nonce));
29            let new_full_nonce = u256_to_h256(nonces_to_full_nonce(nonce, deployment_nonce));
30            storage.set_value(nonce_key, new_full_nonce);
31        }
32
33        if let Some(code) = overrides.code {
34            let (bytecode_kind, code) = match code {
35                BytecodeOverride::Unspecified(code) => (BytecodeMarker::detect(&code.0), code),
36                BytecodeOverride::EraVm(code) => (BytecodeMarker::EraVm, code),
37                BytecodeOverride::Evm(code) => (BytecodeMarker::Evm, code),
38            };
39            let code_key = get_code_key(&account);
40
41            let (code_hash, prepared_code) = match bytecode_kind {
42                BytecodeMarker::EraVm => (BytecodeHash::for_bytecode(&code.0).value(), code.0),
43                BytecodeMarker::Evm => {
44                    // For better usability, we allow overriding EVM bytecodes even if EVM contract deployment is not enabled for the chain.
45                    let versioned_hash = BytecodeHash::for_raw_evm_bytecode(&code.0).value();
46                    let evm_bytecode_hash_key = get_evm_code_hash_key(versioned_hash);
47                    storage.set_value(evm_bytecode_hash_key, H256(web3::keccak256(&code.0)));
48                    (versioned_hash, pad_evm_bytecode(&code.0))
49                }
50            };
51
52            storage.set_value(code_key, code_hash);
53            let known_code_key = get_known_code_key(&code_hash);
54            storage.set_value(known_code_key, H256::from_low_u64_be(1));
55            storage.store_factory_dep(code_hash, prepared_code);
56        }
57
58        match overrides.state {
59            Some(OverrideState::State(state)) => {
60                let account = AccountTreeId::new(account);
61                for (key, value) in state {
62                    storage.set_value(StorageKey::new(account, key), value);
63                }
64                storage.insert_erased_account(account);
65            }
66            Some(OverrideState::StateDiff(state_diff)) => {
67                let account = AccountTreeId::new(account);
68                for (key, value) in state_diff {
69                    storage.set_value(StorageKey::new(account, key), value);
70                }
71            }
72            None => { /* do nothing */ }
73        }
74    }
75    storage
76}