From a51bbe4f03aa04c8efceab46624b3b391140eada Mon Sep 17 00:00:00 2001
From: cl <chencassc@gmail.com>
Date: Thu, 28 Nov 2024 03:36:30 +0800
Subject: [PATCH] feat(`traces`): show state changes in `cast run` and `forge
 test` on `-vvvvv` (#9013)

* Add options for state changes output and json output in cast run command

* fix test

* add back serde_json in Cargo.lock

* format using nightly

* rename parameter

* update revm-inspectors

* supress clippy warning and merge master

* add serde_json

* disable some stdout print when --json option is used

* remove unnecessary check

* replace with sh_println

* replace with shell::is_json

* Show storage for verbosity > 1, add test

* Change verbosity to > 4 for both cast and forge test, add test, fix ci

---------

Co-authored-by: grandizzy <grandizzy.the.egg@gmail.com>
---
 Cargo.lock                            |  1 +
 crates/cast/bin/cmd/call.rs           | 13 ++++--
 crates/cast/bin/cmd/run.rs            |  7 ++-
 crates/cast/tests/cli/main.rs         | 65 ++++++++++++++++++++++++++-
 crates/cli/src/opts/global.rs         |  2 +-
 crates/cli/src/utils/cmd.rs           | 21 ++++++---
 crates/evm/evm/src/executors/trace.rs |  9 ++--
 crates/evm/traces/Cargo.toml          |  1 +
 crates/evm/traces/src/lib.rs          | 37 ++++++++++++---
 crates/forge/bin/cmd/test/mod.rs      |  8 ++--
 crates/forge/src/multi_runner.rs      |  5 ++-
 crates/forge/tests/cli/cmd.rs         |  2 +-
 crates/forge/tests/cli/test_cmd.rs    | 34 ++++++++++++++
 crates/verify/src/utils.rs            |  1 +
 14 files changed, 176 insertions(+), 30 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 1c065f741ccd9..f8af0b8b1ac08 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -4086,6 +4086,7 @@ dependencies = [
  "revm",
  "revm-inspectors",
  "serde",
+ "serde_json",
  "solar-parse",
  "tempfile",
  "tokio",
diff --git a/crates/cast/bin/cmd/call.rs b/crates/cast/bin/cmd/call.rs
index cdc3bd4bc6e2f..2d5692efe4c1a 100644
--- a/crates/cast/bin/cmd/call.rs
+++ b/crates/cast/bin/cmd/call.rs
@@ -8,7 +8,7 @@ use foundry_cli::{
     opts::{EthereumOpts, TransactionOpts},
     utils::{self, handle_traces, parse_ether_value, TraceResult},
 };
-use foundry_common::ens::NameOrAddress;
+use foundry_common::{ens::NameOrAddress, shell};
 use foundry_compilers::artifacts::EvmVersion;
 use foundry_config::{
     figment::{
@@ -182,8 +182,15 @@ impl CallArgs {
             env.cfg.disable_block_gas_limit = true;
             env.block.gas_limit = U256::MAX;
 
-            let mut executor =
-                TracingExecutor::new(env, fork, evm_version, debug, decode_internal, alphanet);
+            let mut executor = TracingExecutor::new(
+                env,
+                fork,
+                evm_version,
+                debug,
+                decode_internal,
+                shell::verbosity() > 4,
+                alphanet,
+            );
 
             let value = tx.value.unwrap_or_default();
             let input = tx.inner.input.into_input().unwrap_or_default();
diff --git a/crates/cast/bin/cmd/run.rs b/crates/cast/bin/cmd/run.rs
index 0b85d14feb987..cfad7263a1e4e 100644
--- a/crates/cast/bin/cmd/run.rs
+++ b/crates/cast/bin/cmd/run.rs
@@ -10,7 +10,7 @@ use foundry_cli::{
     opts::{EtherscanOpts, RpcOpts},
     utils::{handle_traces, init_progress, TraceResult},
 };
-use foundry_common::{is_known_system_sender, SYSTEM_TRANSACTION_TYPE};
+use foundry_common::{is_known_system_sender, shell, SYSTEM_TRANSACTION_TYPE};
 use foundry_compilers::artifacts::EvmVersion;
 use foundry_config::{
     figment::{
@@ -169,6 +169,7 @@ impl RunArgs {
             evm_version,
             self.debug,
             self.decode_internal,
+            shell::verbosity() > 4,
             alphanet,
         );
         let mut env =
@@ -176,7 +177,9 @@ impl RunArgs {
 
         // Set the state to the moment right before the transaction
         if !self.quick {
-            sh_println!("Executing previous transactions from the block.")?;
+            if !shell::is_json() {
+                sh_println!("Executing previous transactions from the block.")?;
+            }
 
             if let Some(block) = block {
                 let pb = init_progress(block.transactions.len() as u64, "tx");
diff --git a/crates/cast/tests/cli/main.rs b/crates/cast/tests/cli/main.rs
index a036d49ae0742..d142cbd599b22 100644
--- a/crates/cast/tests/cli/main.rs
+++ b/crates/cast/tests/cli/main.rs
@@ -64,7 +64,7 @@ Display options:
           - 2 (-vv): Print logs for all tests.
           - 3 (-vvv): Print execution traces for failing tests.
           - 4 (-vvvv): Print execution traces for all tests, and setup traces for failing tests.
-          - 5 (-vvvvv): Print execution and setup traces for all tests.
+          - 5 (-vvvvv): Print execution and setup traces for all tests, including storage changes.
 
 Find more information in the book: http://book.getfoundry.sh/reference/cast/cast.html
 
@@ -1778,3 +1778,66 @@ Transaction successfully executed.
 
 "#]]);
 });
+
+// tests cast can decode traces when running with verbosity level > 4
+forgetest_async!(show_state_changes_in_traces, |prj, cmd| {
+    let (api, handle) = anvil::spawn(NodeConfig::test()).await;
+
+    foundry_test_utils::util::initialize(prj.root());
+    // Deploy counter contract.
+    cmd.args([
+        "script",
+        "--private-key",
+        "0xac0974bec39a17e36ba4a6b4d238ff944bacb478cbed5efcae784d7bf4f2ff80",
+        "--rpc-url",
+        &handle.http_endpoint(),
+        "--broadcast",
+        "CounterScript",
+    ])
+    .assert_success();
+
+    // Send tx to change counter storage value.
+    cmd.cast_fuse()
+        .args([
+            "send",
+            "0x5FbDB2315678afecb367f032d93F642f64180aa3",
+            "setNumber(uint256)",
+            "111",
+            "--private-key",
+            "0xac0974bec39a17e36ba4a6b4d238ff944bacb478cbed5efcae784d7bf4f2ff80",
+            "--rpc-url",
+            &handle.http_endpoint(),
+        ])
+        .assert_success();
+
+    let tx_hash = api
+        .transaction_by_block_number_and_index(BlockNumberOrTag::Latest, Index::from(0))
+        .await
+        .unwrap()
+        .unwrap()
+        .tx_hash();
+
+    // Assert cast with verbosity displays storage changes.
+    cmd.cast_fuse()
+        .args([
+            "run",
+            format!("{tx_hash}").as_str(),
+            "-vvvvv",
+            "--rpc-url",
+            &handle.http_endpoint(),
+        ])
+        .assert_success()
+        .stdout_eq(str![[r#"
+Executing previous transactions from the block.
+Traces:
+  [22287] 0x5FbDB2315678afecb367f032d93F642f64180aa3::setNumber(111)
+    ├─  storage changes:
+    │   @ 0: 0 → 111
+    └─ ← [Stop] 
+
+
+Transaction successfully executed.
+[GAS]
+
+"#]]);
+});
diff --git a/crates/cli/src/opts/global.rs b/crates/cli/src/opts/global.rs
index ad715f24180a1..c820ca2cff7e7 100644
--- a/crates/cli/src/opts/global.rs
+++ b/crates/cli/src/opts/global.rs
@@ -15,7 +15,7 @@ pub struct GlobalOpts {
     /// - 2 (-vv): Print logs for all tests.
     /// - 3 (-vvv): Print execution traces for failing tests.
     /// - 4 (-vvvv): Print execution traces for all tests, and setup traces for failing tests.
-    /// - 5 (-vvvvv): Print execution and setup traces for all tests.
+    /// - 5 (-vvvvv): Print execution and setup traces for all tests, including storage changes.
     #[arg(help_heading = "Display options", global = true, short, long, verbatim_doc_comment, conflicts_with = "quiet", action = ArgAction::Count)]
     verbosity: Verbosity,
 
diff --git a/crates/cli/src/utils/cmd.rs b/crates/cli/src/utils/cmd.rs
index 523c10478fb64..505634996de8f 100644
--- a/crates/cli/src/utils/cmd.rs
+++ b/crates/cli/src/utils/cmd.rs
@@ -17,8 +17,7 @@ use foundry_evm::{
         debug::{ContractSources, DebugTraceIdentifier},
         decode_trace_arena,
         identifier::{CachedSignatures, SignaturesIdentifier, TraceIdentifiers},
-        render_trace_arena_with_bytecodes, CallTraceDecoder, CallTraceDecoderBuilder, TraceKind,
-        Traces,
+        render_trace_arena_inner, CallTraceDecoder, CallTraceDecoderBuilder, TraceKind, Traces,
     },
 };
 use std::{
@@ -450,7 +449,7 @@ pub async fn handle_traces(
         decoder.debug_identifier = Some(DebugTraceIdentifier::new(sources));
     }
 
-    print_traces(&mut result, &decoder, shell::verbosity() > 0).await?;
+    print_traces(&mut result, &decoder, shell::verbosity() > 0, shell::verbosity() > 4).await?;
 
     Ok(())
 }
@@ -459,23 +458,31 @@ pub async fn print_traces(
     result: &mut TraceResult,
     decoder: &CallTraceDecoder,
     verbose: bool,
+    state_changes: bool,
 ) -> Result<()> {
     let traces = result.traces.as_mut().expect("No traces found");
 
-    sh_println!("Traces:")?;
+    if !shell::is_json() {
+        sh_println!("Traces:")?;
+    }
+
     for (_, arena) in traces {
         decode_trace_arena(arena, decoder).await?;
-        sh_println!("{}", render_trace_arena_with_bytecodes(arena, verbose))?;
+        sh_println!("{}", render_trace_arena_inner(arena, verbose, state_changes))?;
+    }
+
+    if shell::is_json() {
+        return Ok(());
     }
-    sh_println!()?;
 
+    sh_println!()?;
     if result.success {
         sh_println!("{}", "Transaction successfully executed.".green())?;
     } else {
         sh_err!("Transaction failed.")?;
     }
-
     sh_println!("Gas used: {}", result.gas_used)?;
+
     Ok(())
 }
 
diff --git a/crates/evm/evm/src/executors/trace.rs b/crates/evm/evm/src/executors/trace.rs
index 69c68442b65cd..ceea6e67248a4 100644
--- a/crates/evm/evm/src/executors/trace.rs
+++ b/crates/evm/evm/src/executors/trace.rs
@@ -18,15 +18,18 @@ impl TracingExecutor {
         version: Option<EvmVersion>,
         debug: bool,
         decode_internal: bool,
+        with_state_changes: bool,
         alphanet: bool,
     ) -> Self {
         let db = Backend::spawn(fork);
-        let trace_mode =
-            TraceMode::Call.with_debug(debug).with_decode_internal(if decode_internal {
+        let trace_mode = TraceMode::Call
+            .with_debug(debug)
+            .with_decode_internal(if decode_internal {
                 InternalTraceMode::Full
             } else {
                 InternalTraceMode::None
-            });
+            })
+            .with_state_changes(with_state_changes);
         Self {
             // configures a bare version of the evm executor: no cheatcode inspector is enabled,
             // tracing will be enabled only for the targeted transaction
diff --git a/crates/evm/traces/Cargo.toml b/crates/evm/traces/Cargo.toml
index 53bf8b3bb2cb7..f555d619fa228 100644
--- a/crates/evm/traces/Cargo.toml
+++ b/crates/evm/traces/Cargo.toml
@@ -36,6 +36,7 @@ eyre.workspace = true
 futures.workspace = true
 itertools.workspace = true
 serde.workspace = true
+serde_json.workspace = true
 tokio = { workspace = true, features = ["time", "macros"] }
 tracing.workspace = true
 tempfile.workspace = true
diff --git a/crates/evm/traces/src/lib.rs b/crates/evm/traces/src/lib.rs
index 18136c481e7c9..a0fa7e1fca498 100644
--- a/crates/evm/traces/src/lib.rs
+++ b/crates/evm/traces/src/lib.rs
@@ -11,7 +11,10 @@ extern crate foundry_common;
 #[macro_use]
 extern crate tracing;
 
-use foundry_common::contracts::{ContractsByAddress, ContractsByArtifact};
+use foundry_common::{
+    contracts::{ContractsByAddress, ContractsByArtifact},
+    shell,
+};
 use revm::interpreter::OpCode;
 use revm_inspectors::tracing::{
     types::{DecodedTraceStep, TraceMemberOrder},
@@ -183,15 +186,23 @@ pub async fn decode_trace_arena(
 
 /// Render a collection of call traces to a string.
 pub fn render_trace_arena(arena: &SparsedTraceArena) -> String {
-    render_trace_arena_with_bytecodes(arena, false)
+    render_trace_arena_inner(arena, false, false)
 }
 
-/// Render a collection of call traces to a string optionally including contract creation bytecodes.
-pub fn render_trace_arena_with_bytecodes(
+/// Render a collection of call traces to a string optionally including contract creation bytecodes
+/// and in JSON format.
+pub fn render_trace_arena_inner(
     arena: &SparsedTraceArena,
     with_bytecodes: bool,
+    with_storage_changes: bool,
 ) -> String {
-    let mut w = TraceWriter::new(Vec::<u8>::new()).write_bytecodes(with_bytecodes);
+    if shell::is_json() {
+        return serde_json::to_string(&arena.resolve_arena()).expect("Failed to write traces");
+    }
+
+    let mut w = TraceWriter::new(Vec::<u8>::new())
+        .write_bytecodes(with_bytecodes)
+        .with_storage_changes(with_storage_changes);
     w.write_arena(&arena.resolve_arena()).expect("Failed to write traces");
     String::from_utf8(w.into_writer()).expect("trace writer wrote invalid UTF-8")
 }
@@ -289,6 +300,8 @@ pub enum TraceMode {
     ///
     /// Used by debugger.
     Debug,
+    /// Debug trace with storage changes.
+    RecordStateDiff,
 }
 
 impl TraceMode {
@@ -308,6 +321,10 @@ impl TraceMode {
         matches!(self, Self::Jump)
     }
 
+    pub const fn record_state_diff(self) -> bool {
+        matches!(self, Self::RecordStateDiff)
+    }
+
     pub const fn is_debug(self) -> bool {
         matches!(self, Self::Debug)
     }
@@ -324,6 +341,14 @@ impl TraceMode {
         std::cmp::max(self, mode.into())
     }
 
+    pub fn with_state_changes(self, yes: bool) -> Self {
+        if yes {
+            std::cmp::max(self, Self::RecordStateDiff)
+        } else {
+            self
+        }
+    }
+
     pub fn with_verbosity(self, verbosiy: u8) -> Self {
         if verbosiy >= 3 {
             std::cmp::max(self, Self::Call)
@@ -345,7 +370,7 @@ impl TraceMode {
                     StackSnapshotType::None
                 },
                 record_logs: true,
-                record_state_diff: false,
+                record_state_diff: self.record_state_diff(),
                 record_returndata_snapshots: self.is_debug(),
                 record_opcodes_filter: (self.is_jump() || self.is_jump_simple())
                     .then(|| OpcodeFilter::new().enabled(OpCode::JUMP).enabled(OpCode::JUMPDEST)),
diff --git a/crates/forge/bin/cmd/test/mod.rs b/crates/forge/bin/cmd/test/mod.rs
index 1a409b33a822f..ad2e52df6725d 100644
--- a/crates/forge/bin/cmd/test/mod.rs
+++ b/crates/forge/bin/cmd/test/mod.rs
@@ -12,7 +12,7 @@ use forge::{
         debug::{ContractSources, DebugTraceIdentifier},
         decode_trace_arena, folded_stack_trace,
         identifier::SignaturesIdentifier,
-        render_trace_arena, CallTraceDecoderBuilder, InternalTraceMode, TraceKind,
+        CallTraceDecoderBuilder, InternalTraceMode, TraceKind,
     },
     MultiContractRunner, MultiContractRunnerBuilder, TestFilter, TestOptions, TestOptionsBuilder,
 };
@@ -56,7 +56,7 @@ use summary::TestSummaryReporter;
 
 use crate::cmd::test::summary::print_invariant_metrics;
 pub use filter::FilterArgs;
-use forge::result::TestKind;
+use forge::{result::TestKind, traces::render_trace_arena_inner};
 
 // Loads project's figment and merges the build cli arguments into it
 foundry_config::merge_impl_figment_convert!(TestArgs, opts, evm_opts);
@@ -652,7 +652,7 @@ impl TestArgs {
                     // - 0..3: nothing
                     // - 3: only display traces for failed tests
                     // - 4: also display the setup trace for failed tests
-                    // - 5..: display all traces for all tests
+                    // - 5..: display all traces for all tests, including storage changes
                     let should_include = match kind {
                         TraceKind::Execution => {
                             (verbosity == 3 && result.status.is_failure()) || verbosity >= 4
@@ -665,7 +665,7 @@ impl TestArgs {
 
                     if should_include {
                         decode_trace_arena(arena, &decoder).await?;
-                        decoded_traces.push(render_trace_arena(arena));
+                        decoded_traces.push(render_trace_arena_inner(arena, false, verbosity > 4));
                     }
                 }
 
diff --git a/crates/forge/src/multi_runner.rs b/crates/forge/src/multi_runner.rs
index a9f2a93eb37a6..dfb498c060e66 100644
--- a/crates/forge/src/multi_runner.rs
+++ b/crates/forge/src/multi_runner.rs
@@ -7,7 +7,7 @@ use crate::{
 use alloy_json_abi::{Function, JsonAbi};
 use alloy_primitives::{Address, Bytes, U256};
 use eyre::Result;
-use foundry_common::{get_contract_name, ContractsByArtifact, TestFunctionExt};
+use foundry_common::{get_contract_name, shell::verbosity, ContractsByArtifact, TestFunctionExt};
 use foundry_compilers::{
     artifacts::Libraries, compilers::Compiler, Artifact, ArtifactId, ProjectCompileOutput,
 };
@@ -249,7 +249,8 @@ impl MultiContractRunner {
         let trace_mode = TraceMode::default()
             .with_debug(self.debug)
             .with_decode_internal(self.decode_internal)
-            .with_verbosity(self.evm_opts.verbosity);
+            .with_verbosity(self.evm_opts.verbosity)
+            .with_state_changes(verbosity() > 4);
 
         let executor = ExecutorBuilder::new()
             .inspectors(|stack| {
diff --git a/crates/forge/tests/cli/cmd.rs b/crates/forge/tests/cli/cmd.rs
index e0000e01bee02..35f5c2314c16e 100644
--- a/crates/forge/tests/cli/cmd.rs
+++ b/crates/forge/tests/cli/cmd.rs
@@ -67,7 +67,7 @@ Display options:
           - 2 (-vv): Print logs for all tests.
           - 3 (-vvv): Print execution traces for failing tests.
           - 4 (-vvvv): Print execution traces for all tests, and setup traces for failing tests.
-          - 5 (-vvvvv): Print execution and setup traces for all tests.
+          - 5 (-vvvvv): Print execution and setup traces for all tests, including storage changes.
 
 Find more information in the book: http://book.getfoundry.sh/reference/forge/forge.html
 
diff --git a/crates/forge/tests/cli/test_cmd.rs b/crates/forge/tests/cli/test_cmd.rs
index 8e064c63c9fdd..819c3e9407e0b 100644
--- a/crates/forge/tests/cli/test_cmd.rs
+++ b/crates/forge/tests/cli/test_cmd.rs
@@ -2665,3 +2665,37 @@ Suite result: FAILED. 0 passed; 1 failed; 0 skipped; [ELAPSED]
 
 "#]]);
 });
+
+// Tests that test traces display state changes when running with verbosity.
+forgetest_init!(should_show_state_changes, |prj, cmd| {
+    cmd.args(["test", "--mt", "test_Increment", "-vvvvv"]).assert_success().stdout_eq(str![[r#"
+...
+Ran 1 test for test/Counter.t.sol:CounterTest
+[PASS] test_Increment() ([GAS])
+Traces:
+  [87464] CounterTest::setUp()
+    ├─ [47297] → new Counter@0x5615dEB798BB3E4dFa0139dFa1b3D433Cc23b72f
+    │   └─ ← [Return] 236 bytes of code
+    ├─ [2387] Counter::setNumber(0)
+    │   └─ ← [Stop] 
+    └─ ← [Stop] 
+
+  [31293] CounterTest::test_Increment()
+    ├─ [22337] Counter::increment()
+    │   ├─  storage changes:
+    │   │   @ 0: 0 → 1
+    │   └─ ← [Stop] 
+    ├─ [281] Counter::number() [staticcall]
+    │   └─ ← [Return] 1
+    ├─ [0] VM::assertEq(1, 1) [staticcall]
+    │   └─ ← [Return] 
+    ├─  storage changes:
+    │   @ 0: 0 → 1
+    └─ ← [Stop] 
+
+Suite result: ok. 1 passed; 0 failed; 0 skipped; [ELAPSED]
+
+Ran 1 test suite [ELAPSED]: 1 tests passed, 0 failed, 0 skipped (1 total tests)
+
+"#]]);
+});
diff --git a/crates/verify/src/utils.rs b/crates/verify/src/utils.rs
index a14d6af6df964..56ec035abcead 100644
--- a/crates/verify/src/utils.rs
+++ b/crates/verify/src/utils.rs
@@ -334,6 +334,7 @@ pub async fn get_tracing_executor(
         Some(fork_config.evm_version),
         false,
         false,
+        false,
         is_alphanet,
     );