Skip to main content

zksync_vm2/
testonly.rs

1//! Test-only tools for EraVM.
2
3use std::{
4    collections::{hash_map::DefaultHasher, BTreeMap},
5    hash::{Hash, Hasher},
6};
7
8use primitive_types::{H160, U256};
9use zkevm_opcode_defs::{
10    ethereum_types::Address, system_params::DEPLOYER_SYSTEM_CONTRACT_ADDRESS_LOW,
11};
12use zksync_vm2_interface::Tracer;
13
14use crate::{
15    instruction_handlers::address_into_u256, Program, StorageInterface, StorageSlot, World,
16};
17
18/// Test [`World`] implementation.
19#[derive(Debug)]
20pub struct TestWorld<T> {
21    pub(crate) address_to_hash: BTreeMap<U256, U256>,
22    pub(crate) hash_to_contract: BTreeMap<U256, Program<T, Self>>,
23    decommit_code_calls: usize,
24}
25
26impl<T: Tracer> TestWorld<T> {
27    /// Creates a test world with the provided programs.
28    ///
29    /// # Panics
30    ///
31    /// Panics if the provided `Program`s are malformed.
32    pub fn new(contracts: &[(Address, Program<T, Self>)]) -> Self {
33        let mut address_to_hash = BTreeMap::new();
34        let mut hash_to_contract = BTreeMap::new();
35        for (i, (address, code)) in contracts.iter().enumerate() {
36            // We add the index to the hash because tests may leave the code page blank.
37            let mut hasher = DefaultHasher::new();
38            i.hash(&mut hasher);
39            code.code_page().hash(&mut hasher);
40
41            let mut code_info_bytes = [0; 32];
42            code_info_bytes[24..].copy_from_slice(&hasher.finish().to_be_bytes());
43            let code_len = u16::try_from(code.code_page().len())
44                .expect("code length must not exceed u16::MAX");
45            code_info_bytes[2..=3].copy_from_slice(&code_len.to_be_bytes());
46            code_info_bytes[0] = 1;
47            let hash = U256::from_big_endian(&code_info_bytes);
48
49            address_to_hash.insert(address_into_u256(*address), hash);
50            hash_to_contract.insert(hash, code.clone());
51        }
52        Self {
53            address_to_hash,
54            hash_to_contract,
55            decommit_code_calls: 0,
56        }
57    }
58
59    /// Number of times [`World::decommit_code()`] was called on this world.
60    pub fn decommit_code_calls(&self) -> usize {
61        self.decommit_code_calls
62    }
63}
64
65impl<T: Tracer> World<T> for TestWorld<T> {
66    fn decommit(&mut self, hash: U256) -> Program<T, Self> {
67        if let Some(program) = self.hash_to_contract.get(&hash) {
68            program.clone()
69        } else {
70            panic!("unexpected decommit")
71        }
72    }
73
74    fn decommit_code(&mut self, hash: U256) -> Vec<u8> {
75        self.decommit_code_calls += 1;
76        self.decommit(hash)
77            .code_page()
78            .iter()
79            .flat_map(|u256| {
80                let mut buffer = [0u8; 32];
81                u256.to_big_endian(&mut buffer);
82                buffer
83            })
84            .collect()
85    }
86}
87
88impl<T> StorageInterface for TestWorld<T> {
89    fn read_storage(&mut self, contract: H160, key: U256) -> StorageSlot {
90        let deployer_system_contract_address =
91            Address::from_low_u64_be(DEPLOYER_SYSTEM_CONTRACT_ADDRESS_LOW.into());
92
93        if contract == deployer_system_contract_address {
94            let value = self
95                .address_to_hash
96                .get(&key)
97                .copied()
98                .unwrap_or_else(U256::zero);
99            StorageSlot {
100                value,
101                is_write_initial: false,
102            }
103        } else {
104            StorageSlot::EMPTY
105        }
106    }
107
108    fn cost_of_writing_storage(&mut self, _initial_slot: StorageSlot, _new_value: U256) -> u32 {
109        50
110    }
111
112    fn is_free_storage_slot(&self, _contract: &H160, _key: &U256) -> bool {
113        false
114    }
115}
116
117/// May be used to load code when the VM first starts up.
118/// Doesn't check for any errors.
119/// Doesn't cost anything but also doesn't make the code free in future decommits.
120#[doc(hidden)] // should be used only in low-level testing / benches
121pub fn initial_decommit<T: Tracer, W: World<T>>(world: &mut W, address: H160) -> Program<T, W> {
122    let deployer_system_contract_address =
123        Address::from_low_u64_be(DEPLOYER_SYSTEM_CONTRACT_ADDRESS_LOW.into());
124    let code_info =
125        world.read_storage_value(deployer_system_contract_address, address_into_u256(address));
126
127    let mut code_info_bytes = [0; 32];
128    code_info.to_big_endian(&mut code_info_bytes);
129
130    code_info_bytes[1] = 0;
131    let code_key: U256 = U256::from_big_endian(&code_info_bytes);
132
133    world.decommit(code_key)
134}