Write storage change in trace (#213)

This PR writes the storage change in the trace. Example output when
using cast run from [the
PR](foundry-rs/foundry#9013):

```
❯ cargo run --bin cast run 0xefc789b63631b255aeb6f97d948c9eac14bae4b7f021122fe24c7c6e4f34667f   -r https://eth.llamarpc.com  -q --decode-internal --with-state-changes

Compiling: TransparentUpgradeableProxy 0x6bE457e04092B28865E0cBa84E3b2CFa0f871E67
Compiling: TransparentUpgradeableProxy 0x7a7f0b3c23C23a31cFcb0c44709be70d4D545c6e
Compiling: TransparentUpgradeableProxy 0xD523794C879D9eC028960a231F866758e405bE34
Compiling: Pool 0xDEbbf61098642C7c06fAd1E116C1a00e50405D0d
Traces:
  [247473] TransparentUpgradeableProxy::stake{value: 1949677189193480698}(1)
    ├─ [2232] TransparentUpgradeableProxy::_beforeFallback()
    │   ├─ [2150] ERC1967Upgrade::_getAdmin()
    │   │   └─ ← 0xD491302a6621128f4b5a6A49ce6657230732b0cb
    │   └─ ←
    ├─ [2196] ERC1967Proxy::_implementation()
    │   └─ ← 0xDEbbf61098642C7c06fAd1E116C1a00e50405D0d
    ├─ [240157] Pool::stake{value: 1949677189193480698}(1) [delegatecall]
    │   ├─ [213932] Pool::_stake(0x93386C72aa57082820Ad6Aa29998B820971a8d61, 1949677189193480698 [1.949e18])
    │   │   ├─ [228681] TransparentUpgradeableProxy::deposit(0x93386C72aa57082820Ad6Aa29998B820971a8d61, 1949677189193480698 [1.949e18])
    │   │   │   ├─ [2232] TransparentUpgradeableProxy::_beforeFallback()
    │   │   │   │   ├─ [2150] ERC1967Upgrade::_getAdmin()
    │   │   │   │   │   └─ ← 0xD491302a6621128f4b5a6A49ce6657230732b0cb
    │   │   │   │   └─ ←
    │   │   │   ├─ [2196] ERC1967Proxy::_implementation()
    │   │   │   │   └─ ← 0x1a5b89b2ef0028A059EAD7D9E648B533f87c8558
    │   │   │   ├─ [221356] 0x1a5b89b2ef0028A059EAD7D9E648B533f87c8558::deposit(0x93386C72aa57082820Ad6Aa29998B820971a8d61, 1949677189193480698 [1.949e18]) [delegatecall]
    │   │   │   │   ├─ emit Update(: 1305166715630598691 [1.305e18], : 17341785000000000 [1.734e16], : 1926865000000000 [1.926e15])
    │   │   │   │   ├─ [22215] TransparentUpgradeableProxy::reStake(17341785000000000 [1.734e16], 0)
    │   │   │   │   │   ├─ [2232] TransparentUpgradeableProxy::_beforeFallback()
    │   │   │   │   │   │   ├─ [2150] ERC1967Upgrade::_getAdmin()
    │   │   │   │   │   │   │   └─ ← 0xD491302a6621128f4b5a6A49ce6657230732b0cb
    │   │   │   │   │   │   └─ ←
    │   │   │   │   │   ├─ [2196] ERC1967Proxy::_implementation()
    │   │   │   │   │   │   └─ ← 0xbE3366a14D0c87094DEB6DFbA667299b4EaC489d
    │   │   │   │   │   ├─ [14896] 0xbE3366a14D0c87094DEB6DFbA667299b4EaC489d::reStake(17341785000000000 [1.734e16], 0) [delegatecall]
    │   │   │   │   │   │   ├─ [3266] TransparentUpgradeableProxy::restake{value: 17341785000000000}(0)
    │   │   │   │   │   │   │   ├─ [232] TransparentUpgradeableProxy::_beforeFallback()
    │   │   │   │   │   │   │   │   ├─ [150] ERC1967Upgrade::_getAdmin()
    │   │   │   │   │   │   │   │   │   └─ ← 0xD491302a6621128f4b5a6A49ce6657230732b0cb
    │   │   │   │   │   │   │   │   └─ ←
    │   │   │   │   │   │   │   ├─ [196] ERC1967Proxy::_implementation()
    │   │   │   │   │   │   │   │   └─ ← 0xDEbbf61098642C7c06fAd1E116C1a00e50405D0d
    │   │   │   │   │   │   │   ├─ [2450] Pool::restake{value: 17341785000000000}(0) [delegatecall]
    │   │   │   │   │   │   │   │   └─ ← [Stop]
    │   │   │   │   │   │   │   └─ ← [Return]
    │   │   │   │   │   │   └─ ← [Stop]
    │   │   │   │   │   └─ ← [Return]
    │   │   │   │   ├─  emit topic 0: 0xb0ec6c271a891e04f03d43bf454839acfbf574e269cb5599c1f70b9257cf9cd7
    │   │   │   │   │           data: 0x000000000000000000000000000000000000000000000000003d9c42f6b23a00
    │   │   │   │   ├─  emit topic 0: 0xc8724ec5e59eea00f3f35419c3139ead03ff07766e7e9cf00a62381692aac8c7
    │   │   │   │   │        topic 1: 0x00000000000000000000000093386c72aa57082820ad6aa29998b820971a8d61
    │   │   │   │   │           data: 0x0000000000000000000000000000000000000000000000001b0ea512e08f39fa
    │   │   │   │   ├─  storage write 0x7a7f0b3c23C23a31cFcb0c44709be70d4D545c6e [0x501d093c0316d001addb7dc8913312112ff29225e2343de5aed242436c83ecaf]:
    │   │   │   │   │                 0x0000000000000000000000000000000000000000000000000000000000000000 → 0x0000000000000000000000000000000000000000000000001b0ea512e08f39fa
...
    │   │   │   │   │                 0x0000000000000000000000000000000000000000000000000000000000000003 → 0x0000000000000000000000000000000000000000000000000000000000000004
    │   │   │   │   ├─  storage write 0x7a7f0b3c23C23a31cFcb0c44709be70d4D545c6e [0xb4c87350b3618bf3b0453372aae234908cc3b6cf61d458f195edd9160dac1f7f]:
    │   │   │   │   │                 0x000000000000000000000000000000000000000000000000ce8cbfb644313c26 → 0x000000000000000000000000000000000000000000000000e9d9010c1b72b020
    │   │   │   │   └─ ← [Return] 0x00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000
    │   │   │   └─ ← [Return] 0x00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000
    │   │   └─ ←
    │   ├─ emit StakeAdded(staker: 0x93386C72aa57082820Ad6Aa29998B820971a8d61, value: 1949677189193480698 [1.949e18], source: 1)
    │   └─ ← [Stop]
    └─ ← [Return]


Transaction successfully executed.
Gas used: 248777
```

---------

Co-authored-by: DaniPopes <[email protected]>

src/tracing/writer.rs

@@ -5,10 +5,13 @@ use super::{
     },
     CallTraceArena,
 };
-use alloy_primitives::{address, hex, Address};
+use alloy_primitives::{address, hex, Address, B256, U256};
 use anstyle::{AnsiColor, Color, Style};
 use colorchoice::ColorChoice;
-use std::io::{self, Write};
+use std::{
+    collections::HashMap,
+    io::{self, Write},
+};
 
 const CHEATCODE_ADDRESS: Address = address!("7109709ECfa91a80626fF3989D68f67F5b1DD12D");
 
@@ -28,6 +31,7 @@ pub struct TraceWriterConfig {
     use_colors: bool,
     color_cheatcodes: bool,
     write_bytecodes: bool,
+    write_storage_changes: bool,
 }
 
 impl Default for TraceWriterConfig {
@@ -43,6 +47,7 @@ impl TraceWriterConfig {
             use_colors: use_colors(ColorChoice::Auto),
             color_cheatcodes: false,
             write_bytecodes: false,
+            write_storage_changes: false,
         }
     }
 
@@ -79,6 +84,17 @@ impl TraceWriterConfig {
     pub fn get_write_bytecodes(&self) -> bool {
         self.write_bytecodes
     }
+
+    /// Sets whether to write storage changes.
+    pub fn write_storage_changes(mut self, yes: bool) -> Self {
+        self.write_storage_changes = yes;
+        self
+    }
+
+    /// Returns `true` if storage changes are written to the writer.
+    pub fn get_write_storage_changes(&self) -> bool {
+        self.write_storage_changes
+    }
 }
 
 /// Formats [call traces](CallTraceArena) to an [`Write`] writer.
@@ -131,6 +147,13 @@ impl<W: Write> TraceWriter<W> {
         self
     }
 
+    /// Sets whether to write storage changes.
+    #[inline]
+    pub fn with_storage_changes(mut self, yes: bool) -> Self {
+        self.config.write_storage_changes = yes;
+        self
+    }
+
     /// Returns a reference to the inner writer.
     #[inline]
     pub const fn writer(&self) -> &W {
@@ -218,6 +241,10 @@ impl<W: Write> TraceWriter<W> {
         self.indentation_level += 1;
         self.write_items(nodes, idx)?;
 
+        if self.config.write_storage_changes {
+            self.write_storage_changes(node)?;
+        }
+
         // Write return data.
         self.write_edge()?;
         self.write_trace_footer(&node.trace)?;
@@ -348,6 +375,7 @@ impl<W: Write> TraceWriter<W> {
         match decoded {
             DecodedTraceStep::InternalCall(call, end_idx) => {
                 let gas_used = node.trace.steps[*end_idx].gas_used.saturating_sub(step.gas_used);
+
                 self.write_branch()?;
                 self.indentation_level += 1;
 
@@ -463,6 +491,48 @@ impl<W: Write> TraceWriter<W> {
         }
         LOG_STYLE
     }
+
+    fn write_storage_changes(&mut self, node: &CallTraceNode) -> io::Result<()> {
+        let mut changes_map = HashMap::new();
+
+        // For each call trace, compact the results so we do not write the intermediate storage
+        // writes
+        for step in &node.trace.steps {
+            if let Some(change) = &step.storage_change {
+                let (_first, last) = changes_map.entry(&change.key).or_insert((change, change));
+                *last = change;
+            }
+        }
+
+        let changes = changes_map
+            .iter()
+            .filter_map(|(&&key, &(first, last))| {
+                let value_before = first.had_value.unwrap_or_default();
+                let value_after = last.value;
+                if value_before == value_after {
+                    return None;
+                }
+                Some((key, value_before, value_after))
+            })
+            .collect::<Vec<_>>();
+
+        if !changes.is_empty() {
+            self.write_branch()?;
+            writeln!(self.writer, " storage changes:")?;
+            for (key, value_before, value_after) in changes {
+                self.write_pipes()?;
+                writeln!(
+                    self.writer,
+                    "  @ {key}: {value_before} → {value_after}",
+                    key = num_or_hex(key),
+                    value_before = num_or_hex(value_before),
+                    value_after = num_or_hex(value_after),
+                )?;
+            }
+        }
+
+        Ok(())
+    }
 }
 
 fn use_colors(choice: ColorChoice) -> bool {
@@ -473,3 +543,13 @@ fn use_colors(choice: ColorChoice) -> bool {
         ColorChoice::Never => false,
     }
 }
+
+/// Formats the given U256 as a decimal number if it is short, otherwise as a hexadecimal
+/// byte-array.
+fn num_or_hex(x: U256) -> String {
+    if x < U256::from(1e6 as u128) {
+        x.to_string()
+    } else {
+        B256::from(x).to_string()
+    }
+}

tests/it/writer.rs

@@ -26,15 +26,15 @@ fn test_trace_printing() {
 
     let mut index = 0;
 
-    assert_traces(base_path, None, Some(index), true, &mut tracer);
+    assert_traces(base_path, None, Some(index), &mut tracer);
     index += 1;
 
     let mut call = |data: Vec<u8>| {
         let mut tracer = TracingInspector::new(TracingInspectorConfig::all());
         let r = evm.call(address, data.into(), &mut tracer).unwrap();
         assert!(r.is_success(), "evm.call reverted: {r:#?}");
 
-        assert_traces(base_path, None, Some(index), true, &mut tracer);
+        assert_traces(base_path, None, Some(index), &mut tracer);
 
         index += 1;
     };
@@ -72,13 +72,13 @@ fn deploy_fail() {
     let mut tracer = TracingInspector::new(TracingInspectorConfig::all());
     let _ = evm.try_deploy(bytes!("604260005260206000fd"), &mut tracer).unwrap();
 
-    assert_traces(base_path, Some("raw"), None, true, &mut tracer);
+    assert_traces(base_path, Some("raw"), None, &mut tracer);
 
     let node = &mut tracer.traces_mut().nodes_mut()[0];
     node.trace.decoded.label = Some("RevertingConstructor".to_string());
     node.trace.decoded.return_data = Some("42".to_string());
 
-    assert_traces(base_path, Some("decoded"), None, true, &mut tracer);
+    assert_traces(base_path, Some("decoded"), None, &mut tracer);
 }
 
 // (name, address)
@@ -205,33 +205,41 @@ fn assert_traces(
     base_path: &Path,
     name: Option<&str>,
     patch_index: Option<usize>,
-    write_bytecodes: bool,
     tracer: &mut TracingInspector,
 ) {
     let name = name.map_or_else(
         || patch_index.expect("at least one of name or patch_index must be provided").to_string(),
         ToString::to_string,
     );
-    let bytecodes = if write_bytecodes { &[false, true][..] } else { &[false][..] };
 
     let do_assert = |config: TraceWriterConfig, extra: &str, tracer: &TracingInspector| {
         let color = config.get_use_colors();
         let bytecodes = config.get_write_bytecodes();
+        let write_storage_changes = config.get_write_storage_changes();
 
         let file_kind = if color { DataFormat::TermSvg } else { DataFormat::Text };
         let extension = if color { "svg" } else { "txt" };
-        let bytecodes_extra = if bytecodes { ".write_bytecodes" } else { "" };
+        let bytecodes_extra = if bytecodes { ".bytecodes" } else { "" };
+        let storage_changes_extra = if write_storage_changes { ".storage" } else { "" };
 
         let s = write_traces_with(tracer, config);
-        let path = base_path.join(format!("{name}{bytecodes_extra}{extra}.{extension}"));
+        let path = base_path
+            .join(format!("{name}{bytecodes_extra}{storage_changes_extra}{extra}.{extension}"));
         let data = snapbox::Data::read_from(&path, Some(file_kind));
         assert_data_eq!(s, data);
     };
 
     let mut configs = vec![];
     for color in [ColorChoice::Never, ColorChoice::Always] {
-        for &bytecodes in bytecodes {
-            configs.push(TraceWriterConfig::new().color_choice(color).write_bytecodes(bytecodes));
+        for bytecodes in [false, true] {
+            for write_storage_changes in [false, true] {
+                configs.push(
+                    TraceWriterConfig::new()
+                        .color_choice(color)
+                        .write_bytecodes(bytecodes)
+                        .write_storage_changes(write_storage_changes),
+                );
+            }
         }
     }
 

tests/it/writer/deploy_fail/decoded.bytecodes.txt

@@ -0,0 +1,2 @@
+  [18] → new RevertingConstructor@0xBd770416a3345F91E4B34576cb804a576fa48EB1(0x604260005260206000fd)
+    └─ ← [Revert] 42

tests/it/writer/deploy_fail/decoded.storage.txt

@@ -0,0 +1,2 @@
+  [18] → new RevertingConstructor@0xBd770416a3345F91E4B34576cb804a576fa48EB1
+    └─ ← [Revert] 42

tests/it/writer/deploy_fail/raw.bytecodes.txt

@@ -0,0 +1,2 @@
+  [18] → new <unknown>@0xBd770416a3345F91E4B34576cb804a576fa48EB1(0x604260005260206000fd)
+    └─ ← [Revert] 0x0000000000000000000000000000000000000000000000000000000000000042

tests/it/writer/deploy_fail/raw.storage.txt

@@ -0,0 +1,2 @@
+  [18] → new <unknown>@0xBd770416a3345F91E4B34576cb804a576fa48EB1
+    └─ ← [Revert] 0x0000000000000000000000000000000000000000000000000000000000000042