diff --git a/crates/chisel/src/dispatcher.rs b/crates/chisel/src/dispatcher.rs index 0f2b7afe38b6..22ad242d091f 100644 --- a/crates/chisel/src/dispatcher.rs +++ b/crates/chisel/src/dispatcher.rs @@ -10,7 +10,7 @@ use crate::{ }, session_source::SessionSource, }; -use alloy_json_abi::JsonAbi; +use alloy_json_abi::{InternalType, JsonAbi}; use alloy_primitives::{hex, Address}; use forge_fmt::FormatterConfig; use foundry_config::{Config, RpcEndpoint}; @@ -528,7 +528,22 @@ impl ChiselDispatcher { err.name, err.inputs .iter() - .map(|input| format_param!(input)) + .map(|input| { + let mut param_type = &input.ty; + // If complex type then add the name of custom type. + // see . + if input.is_complex_type() { + if let Some( + InternalType::Enum { contract: _, ty } | + InternalType::Struct { contract: _, ty } | + InternalType::Other { contract: _, ty }, + ) = &input.internal_type + { + param_type = ty; + } + } + format!("{} {}", param_type, input.name) + }) .collect::>() .join(",") )); diff --git a/crates/chisel/src/executor.rs b/crates/chisel/src/executor.rs index 16387a69b589..b63ebe3e4113 100644 --- a/crates/chisel/src/executor.rs +++ b/crates/chisel/src/executor.rs @@ -49,6 +49,22 @@ impl SessionSource { // Fetch the run function's body statement let run_func_statements = compiled.intermediate.run_func_body()?; + // Record loc of first yul block return statement (if any). + // This is used to decide which is the final statement within the `run()` method. + // see . + let last_yul_return = run_func_statements.iter().find_map(|statement| { + if let pt::Statement::Assembly { loc: _, dialect: _, flags: _, block } = statement { + if let Some(statement) = block.statements.last() { + if let pt::YulStatement::FunctionCall(yul_call) = statement { + if yul_call.id.name == "return" { + return Some(statement.loc()) + } + } + } + } + None + }); + // Find the last statement within the "run()" method and get the program // counter via the source map. if let Some(final_statement) = run_func_statements.last() { @@ -58,9 +74,13 @@ impl SessionSource { // // There is some code duplication within the arms due to the difference between // the [pt::Statement] type and the [pt::YulStatement] types. - let source_loc = match final_statement { + let mut source_loc = match final_statement { pt::Statement::Assembly { loc: _, dialect: _, flags: _, block } => { - if let Some(statement) = block.statements.last() { + // Select last non variable declaration statement, see . + let last_statement = block.statements.iter().rev().find(|statement| { + !matches!(statement, pt::YulStatement::VariableDeclaration(_, _, _)) + }); + if let Some(statement) = last_statement { statement.loc() } else { // In the case where the block is empty, attempt to grab the statement @@ -88,6 +108,13 @@ impl SessionSource { _ => final_statement.loc(), }; + // Consider yul return statement as final statement (if it's loc is lower) . + if let Some(yul_return) = last_yul_return { + if yul_return.end() < source_loc.start() { + source_loc = yul_return; + } + } + // Map the source location of the final statement of the `run()` function to its // corresponding runtime program counter let final_pc = {