fix(chisel): final statement & fetch err with complex type fixes (fou…

…ndry-rs#9081) * fix(chisel): consider assembly block return as final statement * Fix 4938 * Start from first assembly block when checking for return statement * Fix 6618
rplusq · Nov 29, 2024 · e636d9b · e636d9b
1 parent 4feae7f
commit e636d9b
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 4 deletions.
diff --git 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 <https://github.com/foundry-rs/foundry/issues/6618>.
+                                                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::<Vec<_>>()
                                             .join(",")
                                     ));

diff --git 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 <https://github.com/foundry-rs/foundry/issues/4617>.
+            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 <https://github.com/foundry-rs/foundry/issues/4938>.
+                        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 = {