1use 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#[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 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 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 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#[doc(hidden)] pub 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}