1use crate::prelude::{EnvOverrides, Failure, HuffFunction, Outcome, RunnerError, SpecContext, TestResult, TestSpec, TestStatus};
2use crate::spec::skipped;
3use alloy_evm::EvmEnv;
4use alloy_primitives::{Address, Bytes, U256, hex};
5use foundry_evm::backend::{Backend, DatabaseExt};
6use foundry_evm::constants::CALLER;
7use foundry_evm::core::decode::SkipReason;
8use foundry_evm::core::evm::EthEvmNetwork;
9use foundry_evm::executors::{Executor, ExecutorBuilder, RawCallResult};
10use foundry_evm::fork::CreateFork;
11use foundry_evm::inspectors::CheatsConfig;
12use foundry_evm::traces::TraceRequirements;
13use foundry_evm_networks::NetworkConfigs;
14use huff_neo_codegen::Codegen;
15use huff_neo_utils::ast::huff::MacroDefinition;
16use huff_neo_utils::bytecode::LabelIndices;
17use huff_neo_utils::prelude::{Contract, EVMVersion, deploy_bootstrap};
18use huff_neo_utils::scope::ScopeManager;
19use revm::context::TxEnv;
20use revm::context::result::HaltReason;
21use revm::context_interface::block::BlobExcessGasAndPrice;
22use revm::context_interface::cfg::GasParams;
23use revm::context_interface::cfg::gas_params::Eip2780TxInfo;
24use revm::interpreter::SuccessOrHalt;
25use revm::primitives::hardfork::SpecId;
26use revm::state::Bytecode;
27use std::collections::HashMap;
28use std::path::Path;
29use std::sync::Arc;
30
31const DEFAULT_INITIAL_BALANCE: U256 = U256::from_limbs([u64::MAX, u32::MAX as u64, 0, 0]);
33
34#[derive(Clone, Debug, Default)]
36pub struct Env {
37 pub evm_env: EvmEnv,
38 pub tx: TxEnv,
39}
40
41#[derive(Debug)]
43pub struct TestRunner {
44 pub env: Env,
46 pub target_address: Option<Address>,
48 pub evm_version: EVMVersion,
50 pub cheats_config: Arc<CheatsConfig>,
52 pub trace_requirements: TraceRequirements,
54 pub collect_logs: bool,
56 pub live_logs: bool,
58 pub isolation: bool,
60 pub initial_balance: U256,
62 pub networks: NetworkConfigs,
64 pub fork: Option<CreateFork>,
66}
67
68impl Default for TestRunner {
69 fn default() -> Self {
70 let mut env = Env::default();
72 env.tx.caller = CALLER;
73 Self::new(env, None, EVMVersion::default(), Arc::new(cheats_config(CheatsConfig::default())))
74 }
75}
76
77pub fn cheats_config(config: CheatsConfig) -> CheatsConfig {
80 CheatsConfig { assertions_revert: true, ..config }
81}
82
83impl TestRunner {
84 pub fn new(env: Env, target_address: Option<Address>, evm_version: EVMVersion, cheats_config: Arc<CheatsConfig>) -> Self {
85 Self {
86 env,
87 target_address,
88 evm_version,
89 cheats_config,
90 trace_requirements: TraceRequirements::none(),
91 collect_logs: false,
92 live_logs: false,
93 isolation: false,
94 initial_balance: DEFAULT_INITIAL_BALANCE,
95 networks: NetworkConfigs::default(),
96 fork: None,
97 }
98 }
99
100 pub fn new_executor(&self) -> Result<Executor<EthEvmNetwork>, RunnerError> {
102 let backend = Backend::spawn(self.fork.clone()).map_err(|e| RunnerError::GenericError(format!("Failed to spawn backend: {e}")))?;
103 let executor = ExecutorBuilder::<EthEvmNetwork>::new()
104 .inspectors(|stack| {
105 let stack = stack
106 .cheatcodes(self.cheats_config.clone())
107 .trace_requirements(self.trace_requirements)
108 .enable_isolation(self.isolation);
109 if self.collect_logs || self.live_logs { stack.logs(self.live_logs) } else { stack }
110 })
111 .spec_id(self.env.evm_env.cfg_env.spec)
112 .gas_limit(self.env.tx.gas_limit)
113 .build(self.env.evm_env.clone(), self.env.tx.clone(), backend, self.networks);
114 Ok(executor)
115 }
116
117 pub fn set_balance(&self, executor: &mut Executor<EthEvmNetwork>, address: Address, amount: U256) -> Result<(), RunnerError> {
119 executor.set_balance(address, amount).map_err(db_error)
120 }
121
122 pub fn deploy_code(&self, executor: &mut Executor<EthEvmNetwork>, code: &str) -> Result<Address, RunnerError> {
124 let (bootstrap_code, _) = deploy_bootstrap(code.len() / 2, 0);
126 let bootstrap = decode_hex(&format!("{bootstrap_code}{code}"), "bytecode")?;
127
128 let sender = self.env.tx.caller;
129 self.set_balance(executor, sender, self.initial_balance)?;
130 use_sender_nonce(executor, sender)?;
131 let deployed = executor
132 .deploy(sender, bootstrap, U256::ZERO, None)
133 .map_err(|e| RunnerError::DeploymentError(format!("Deployment failed: {e}")))?;
134 Ok(deployed.address)
135 }
136
137 pub fn call(
139 &self,
140 name: String,
141 executor: &mut Executor<EthEvmNetwork>,
142 transact_to: Address,
143 value: U256,
144 data: String,
145 ) -> Result<TestResult, RunnerError> {
146 self.set_balance(executor, self.env.tx.caller, self.initial_balance)?;
147 let calldata = decode_hex(&data, "calldata")?;
148 self.call_from(name, executor, self.env.tx.caller, transact_to, value, calldata)
149 }
150
151 fn call_from(
154 &self,
155 name: String,
156 executor: &mut Executor<EthEvmNetwork>,
157 sender: Address,
158 transact_to: Address,
159 value: U256,
160 calldata: Bytes,
161 ) -> Result<TestResult, RunnerError> {
162 use_sender_nonce(executor, sender)?;
164 let result =
165 executor.transact_raw(sender, transact_to, calldata, value).map_err(|e| RunnerError::TransactError(format!("{e:?}")))?;
166
167 let outcome = match result.skip_reason() {
169 Some(SkipReason(reason)) => Outcome::Skip { reason },
170 None => outcome_of(&result),
171 };
172 let mut failures = Vec::new();
175 if executor.has_global_failure(&result.state_changeset) {
176 failures.push(Failure { expectation: String::from("failed"), message: String::from("the test set Foundry's failure flag") });
177 }
178 let passed = matches!(outcome, Outcome::Return { .. }) && failures.is_empty();
179 let (return_data, revert_reason) = match &outcome {
180 Outcome::Return { data } | Outcome::Revert { data } => ((!data.is_empty()).then(|| data.clone()), None),
181 Outcome::Halt { reason } => (None, Some(format!("Transaction halted with reason: {reason}"))),
182 Outcome::Skip { .. } => (None, None),
183 };
184
185 Ok(TestResult {
188 name,
189 return_data,
190 gas: result.gas_used.saturating_sub(call_base_gas(executor.spec_id(), value)),
191 status: TestStatus::from_outcome(&outcome, passed),
192 revert_reason,
193 passed,
194 outcome,
195 failures,
196 traces: result.traces,
197 logs: self.collect_logs.then_some(result.logs),
198 functions: Vec::new(),
199 source_map: None,
200 address: None,
201 source_code: None,
202 bytecode: None,
203 source_files: Vec::new(),
204 })
205 }
206
207 pub fn run_test(
209 &self,
210 executor: &mut Executor<EthEvmNetwork>,
211 m: &MacroDefinition,
212 contract: &Contract,
213 ) -> Result<TestResult, RunnerError> {
214 if let Some(reason) = skipped(m.decorator.as_ref())? {
215 return Ok(skipped_result(m.name.clone(), reason));
216 }
217 let evm_version = &self.evm_version;
218
219 let contract = Codegen::update_table_size(evm_version, contract)?;
221
222 let mut scope_mgr = ScopeManager::new();
224 scope_mgr.push_macro(m, 0);
225 let res = Codegen::macro_to_bytecode(evm_version, m, &contract, &mut scope_mgr, 0, false)?;
226 Codegen::reject_args_start(&res)?;
227
228 let functions = huff_functions(&contract, &res.label_indices);
229
230 let table_out = Codegen::gen_table_bytecode(&contract, res, 0, &HashMap::new())?;
233 let bytecode = format!("{}{}", table_out.body, table_out.tables);
234 let source_map = table_out.source_map;
235
236 let address = match self.target_address {
238 Some(address) => {
239 let code = Bytecode::new_raw(decode_hex(&bytecode, "bytecode")?);
240 executor.set_code(address, code).map_err(db_error)?;
241 address
242 }
243 None => {
244 let sender = self.env.tx.caller;
247 if executor.get_nonce(sender).map_err(db_error)? == 0 {
248 executor.set_account_nonce(sender, 1).map_err(db_error)?;
249 }
250 let address = self.deploy_code(executor, &bytecode)?;
251 self.set_balance(executor, address, self.initial_balance)?;
253 address
254 }
255 };
256 executor.backend_mut().allow_cheatcode_access(address);
258
259 let spec = TestSpec::from_decorator(m.decorator.as_ref(), &contract, SpecContext { this: address, caller: self.env.tx.caller })?;
263 let sender = spec.env.caller.unwrap_or(self.env.tx.caller);
264 apply_env_overrides(executor, &self.env, &spec.env);
265 self.set_balance(executor, sender, self.initial_balance)?;
266 for slot in &spec.storage {
267 executor.set_storage_slot(slot.address, slot.slot, slot.value).map_err(db_error)?;
268 }
269 for balance in &spec.balances {
270 self.set_balance(executor, balance.address, balance.value)?;
271 }
272 for code in &spec.code {
273 executor.set_code(code.address, Bytecode::new_raw(Bytes::copy_from_slice(&code.value))).map_err(db_error)?;
274 }
275 for nonce in &spec.nonces {
276 executor.set_account_nonce(nonce.address, nonce.value).map_err(db_error)?;
277 }
278
279 let mut res = self.call_from(m.name.clone(), executor, sender, address, spec.value, Bytes::from(spec.calldata))?;
281
282 res.failures.extend(spec.expectations.check(&res.outcome, res.gas, &contract));
285 if matches!(res.outcome, Outcome::Return { .. }) {
286 res.failures.extend(spec.expectations.check_post_state(executor.backend())?);
287 }
288 res.passed = res.failures.is_empty() && !matches!(res.outcome, Outcome::Skip { .. });
289 res.status = TestStatus::from_outcome(&res.outcome, res.passed);
290
291 res.source_map = Some(source_map);
293 res.functions = functions;
294 res.address = Some(address);
295
296 res.source_code = contract.flattened_source.clone();
299 res.bytecode = Some(bytecode);
300 res.source_files = contract.source_files.clone();
301
302 Ok(res)
303 }
304}
305
306fn huff_functions(contract: &Contract, label_indices: &LabelIndices) -> Vec<HuffFunction> {
308 contract
309 .macros
310 .values()
311 .filter(|f| f.outlined)
312 .filter_map(|f| {
313 let entry_pc = label_indices.get_any(&format!("goto_{}", f.name))?;
314 let file = f
316 .span
317 .0
318 .iter()
319 .find_map(|span| span.file.as_ref())
320 .and_then(|source| Path::new(&source.path).file_name().map(|name| name.to_string_lossy().into_owned()));
321 Some(HuffFunction { name: f.name.clone(), file, entry_pc, takes: f.takes, returns: f.returns })
322 })
323 .collect()
324}
325
326fn use_sender_nonce(executor: &mut Executor<EthEvmNetwork>, sender: Address) -> Result<(), RunnerError> {
328 let nonce = executor.get_nonce(sender).map_err(db_error)?;
329 executor.tx_env_mut().nonce = nonce;
330 Ok(())
331}
332
333fn outcome_of(result: &RawCallResult<EthEvmNetwork>) -> Outcome {
336 let data = hex::encode(&result.result);
337 match result.exit_reason.map(SuccessOrHalt::<HaltReason>::from) {
338 Some(SuccessOrHalt::Success(_)) => Outcome::Return { data },
339 Some(SuccessOrHalt::Revert) => Outcome::Revert { data },
340 Some(SuccessOrHalt::Halt(reason)) => Outcome::Halt { reason: format!("{reason:?}") },
341 other => Outcome::Halt { reason: format!("{other:?}") },
342 }
343}
344
345fn skipped_result(name: String, reason: Option<String>) -> TestResult {
347 TestResult {
348 name,
349 return_data: None,
350 gas: 0,
351 status: TestStatus::Skipped,
352 revert_reason: None,
353 passed: false,
354 outcome: Outcome::Skip { reason },
355 failures: Vec::new(),
356 traces: None,
357 logs: None,
358 functions: Vec::new(),
359 source_map: None,
360 address: None,
361 source_code: None,
362 bytecode: None,
363 source_files: Vec::new(),
364 }
365}
366
367fn decode_hex(hex: &str, what: &'static str) -> Result<Bytes, RunnerError> {
369 hex::decode(hex).map(Bytes::from).map_err(|e| RunnerError::InvalidHex(what, e.to_string()))
370}
371
372pub(crate) fn db_error(err: impl std::fmt::Debug) -> RunnerError {
374 RunnerError::GenericError(format!("{err:?}"))
375}
376
377fn apply_env_overrides(executor: &mut Executor<EthEvmNetwork>, env: &Env, overrides: &EnvOverrides) {
382 let EnvOverrides {
383 caller: _,
384 gas_limit,
385 gas_price,
386 block_number,
387 timestamp,
388 basefee,
389 blob_basefee,
390 coinbase,
391 chain_id,
392 prevrandao,
393 slot_number,
394 } = *overrides;
395 let mut block = env.evm_env.block_env.clone();
396
397 if let Some(gas_limit) = gas_limit {
398 executor.set_gas_limit(gas_limit);
399 block.gas_limit = block.gas_limit.max(gas_limit);
401 }
402 if let Some(number) = block_number {
403 block.number = number;
404 }
405 if let Some(timestamp) = timestamp {
406 block.timestamp = timestamp;
407 }
408 if let Some(basefee) = basefee {
409 block.basefee = basefee;
410 }
411 if let Some(blob_gasprice) = blob_basefee {
412 block.blob_excess_gas_and_price = Some(BlobExcessGasAndPrice { excess_blob_gas: 0, blob_gasprice });
413 }
414 if let Some(coinbase) = coinbase {
415 block.beneficiary = coinbase;
416 }
417 if let Some(chain_id) = chain_id {
418 executor.evm_env_mut().cfg_env.chain_id = chain_id;
419 }
420 if let Some(prevrandao) = prevrandao {
421 block.prevrandao = Some(prevrandao);
422 }
423 if let Some(slot_number) = slot_number {
424 block.slot_num = slot_number;
425 }
426
427 executor.evm_env_mut().block_env = block.clone();
428 executor.inspector_mut().set_block(block);
429 executor.inspector_mut().set_gas_price(gas_price.unwrap_or(env.tx.gas_price));
430}
431
432fn call_base_gas(spec: SpecId, value: U256) -> u64 {
437 let eip2780 = spec.is_enabled_in(SpecId::AMSTERDAM).then_some(Eip2780TxInfo { value, is_self_transfer: false });
438 GasParams::new_spec(spec).initial_tx_gas(&[], false, 0, 0, 0, eip2780).initial_regular_gas
439}