Skip to main content

huff_neo_test_runner/
runner.rs

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
31/// Foundry's default `--initial-balance`: 2^96 - 1 wei
32const DEFAULT_INITIAL_BALANCE: U256 = U256::from_limbs([u64::MAX, u32::MAX as u64, 0, 0]);
33
34/// Container for EVM environment (cfg + block) and transaction.
35#[derive(Clone, Debug, Default)]
36pub struct Env {
37    pub evm_env: EvmEnv,
38    pub tx: TxEnv,
39}
40
41/// Runs test macros on Foundry's executor, with its cheatcodes, tracer and log collector.
42#[derive(Debug)]
43pub struct TestRunner {
44    /// Environment every test starts from, including the sender of its call
45    pub env: Env,
46    /// Address the test macros run at instead of being deployed
47    pub target_address: Option<Address>,
48    /// EVM version the test macros are compiled for
49    pub evm_version: EVMVersion,
50    /// Settings of the Foundry cheatcodes the tests can call
51    pub cheats_config: Arc<CheatsConfig>,
52    /// What the tracer records of each test
53    pub trace_requirements: TraceRequirements,
54    /// Whether to collect the logs of each test, for console logs in the report
55    pub collect_logs: bool,
56    /// Whether to print console logs while a test runs
57    pub live_logs: bool,
58    /// Whether to run every call of a test as its own transaction, as Foundry's `--isolate` does
59    pub isolation: bool,
60    /// The balance a deployed test and its sender start with
61    pub initial_balance: U256,
62    /// The network families the executor runs for
63    pub networks: NetworkConfigs,
64    /// The chain every test forks, or none to start each test on an empty chain
65    pub fork: Option<CreateFork>,
66}
67
68impl Default for TestRunner {
69    fn default() -> Self {
70        // Send from Foundry's default sender, as `hnc test` does
71        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
77/// Settings for Foundry cheatcodes in `hnc test`. A failed `vm.assert*` reverts, as in `forge test`,
78/// so the test stops there and the failure survives a `vm.revertToState`.
79pub 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    /// A fresh executor for one test, on an empty chain or the runner's fork
101    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    /// Set the balance of an account.
118    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    /// Deploy arbitrary bytecode and return the contract address.
123    pub fn deploy_code(&self, executor: &mut Executor<EthEvmNetwork>, code: &str) -> Result<Address, RunnerError> {
124        // Wrap code in a bootstrap constructor
125        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    /// Perform a call to a deployed contract
138    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    /// Perform a call to a deployed contract from the given sender, in the executor's current
152    /// environment; the sender must already hold the funds for the call
153    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        // The sender may be the deployer of the test or an address set by a `caller` decorator
163        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        // `vm.skip(true)` ends the test with a payload Foundry records, so it is told apart from a revert
168        let outcome = match result.skip_reason() {
169            Some(SkipReason(reason)) => Outcome::Skip { reason },
170            None => outcome_of(&result),
171        };
172        // Foundry's failure flag, which a DSTest-style `fail()` sets by writing the `failed` slot of the
173        // cheatcode address, fails the test as in forge
174        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        // NOTE: We subtract the intrinsic base cost of the CALL transaction, which depends on the
186        // active hardfork (EIP-2780 reprices it from Amsterdam).
187        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    /// Compile a test macro and run it on the executor, which holds the chain state of this test
208    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        // Update the table size in the contract
220        let contract = Codegen::update_table_size(evm_version, contract)?;
221
222        // Compile the passed test macro
223        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        // Generate table bytecode for compiled test macro. Test macros run standalone (no
231        // bootstrap, no upstream dedup target), so base_offset = 0 and the external map is empty.
232        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        // Deploy compiled test macro
237        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                // Like forge, a fresh sender deploys the test with nonce 1, so the test runs at
245                // Foundry's test contract address
246                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                // Like forge, the test contract starts with the initial balance
252                self.set_balance(executor, address, self.initial_balance)?;
253                address
254            }
255        };
256        // On a fork, only allowed accounts may call cheatcodes
257        executor.backend_mut().allow_cheatcode_access(address);
258
259        // Apply the setup requested by the test decorator. The sender is funded first, so a `deal`
260        // to the sender sets its balance instead. The environment is applied after the deployment,
261        // which resets it.
262        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        // Call the deployed test
280        let mut res = self.call_from(m.name.clone(), executor, sender, address, spec.value, Bytes::from(spec.calldata))?;
281
282        // Check how the test ended against what its decorators expect. A test that did not return
283        // kept none of its state, and the outcome failure already says why.
284        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        // Add source map and address for debugging
292        res.source_map = Some(source_map);
293        res.functions = functions;
294        res.address = Some(address);
295
296        // Get the full source code for debugging
297        // Use the flattened source from the contract if available
298        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
306/// Where each outlined function of the compiled test starts, so traces can show its calls
307fn 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            // The file name without its directory, like forge shows the contract of an internal call
315            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
326/// Sends the next transaction of the sender with its current nonce
327fn 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
333/// How a call ended, with Foundry's exit reason turned back into the EVM's halt reason, so
334/// `OutOfGas` covers every kind of out-of-gas halt
335fn 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
345/// The result of a test a `skip` flag skips, which does not run
346fn 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
367/// Decodes hex the runner was given, such as compiled code or calldata, naming it in the error
368fn 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
372/// Wraps a database failure as a runner error
373pub(crate) fn db_error(err: impl std::fmt::Debug) -> RunnerError {
374    RunnerError::GenericError(format!("{err:?}"))
375}
376
377/// Sets the transaction and block fields a test overrides, on top of the runner's environment
378///
379/// The executor sends every call with a zero base fee and gas price so it can never be rejected for
380/// its fees, and its cheatcode inspector shows the test the block and gas price set here.
381fn 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        // A transaction may not use more gas than its block allows
400        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
432/// Intrinsic gas of a call transaction to another account, excluding calldata and access lists
433///
434/// This is 21000 before Amsterdam. From Amsterdam (EIP-2780) it is the sender base cost plus a cold
435/// access of the recipient, and an extra charge when value is transferred.
436fn 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}