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]>
paradigmxyz · Oct 16, 2024 · ad6455f · ad6455f
1 parent d410e81
commit ad6455f
Show file tree

Hide file tree

Showing 170 changed files with 2,302 additions and 12 deletions.
diff --git a/src/tracing/writer.rs b/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()
+    }
+}
diff --git a/tests/it/writer.rs b/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),
+                );
+            }
         }
     }
 

diff --git a/...r/deploy_fail/decoded.write_bytecodes.svg → ...deploy_fail/decoded.bytecodes.storage.svg b/...r/deploy_fail/decoded.write_bytecodes.svg → ...deploy_fail/decoded.bytecodes.storage.svg
diff --git a/...r/deploy_fail/decoded.write_bytecodes.txt → ...deploy_fail/decoded.bytecodes.storage.txt b/...r/deploy_fail/decoded.write_bytecodes.txt → ...deploy_fail/decoded.bytecodes.storage.txt
diff --git a/tests/it/writer/deploy_fail/decoded.bytecodes.svg b/tests/it/writer/deploy_fail/decoded.bytecodes.svg
diff --git a/tests/it/writer/deploy_fail/decoded.bytecodes.txt b/tests/it/writer/deploy_fail/decoded.bytecodes.txt
@@ -0,0 +1,2 @@
+  [18] → new RevertingConstructor@0xBd770416a3345F91E4B34576cb804a576fa48EB1(0x604260005260206000fd)
+    └─ ← [Revert] 42
diff --git a/tests/it/writer/deploy_fail/decoded.storage.svg b/tests/it/writer/deploy_fail/decoded.storage.svg
diff --git a/tests/it/writer/deploy_fail/decoded.storage.txt b/tests/it/writer/deploy_fail/decoded.storage.txt
@@ -0,0 +1,2 @@
+  [18] → new RevertingConstructor@0xBd770416a3345F91E4B34576cb804a576fa48EB1
+    └─ ← [Revert] 42
diff --git a/...riter/deploy_fail/raw.write_bytecodes.svg → ...ter/deploy_fail/raw.bytecodes.storage.svg b/...riter/deploy_fail/raw.write_bytecodes.svg → ...ter/deploy_fail/raw.bytecodes.storage.svg
diff --git a/...riter/deploy_fail/raw.write_bytecodes.txt → ...ter/deploy_fail/raw.bytecodes.storage.txt b/...riter/deploy_fail/raw.write_bytecodes.txt → ...ter/deploy_fail/raw.bytecodes.storage.txt
diff --git a/tests/it/writer/deploy_fail/raw.bytecodes.svg b/tests/it/writer/deploy_fail/raw.bytecodes.svg
diff --git a/tests/it/writer/deploy_fail/raw.bytecodes.txt b/tests/it/writer/deploy_fail/raw.bytecodes.txt
@@ -0,0 +1,2 @@
+  [18] → new <unknown>@0xBd770416a3345F91E4B34576cb804a576fa48EB1(0x604260005260206000fd)
+    └─ ← [Revert] 0x0000000000000000000000000000000000000000000000000000000000000042
diff --git a/tests/it/writer/deploy_fail/raw.storage.svg b/tests/it/writer/deploy_fail/raw.storage.svg
diff --git a/tests/it/writer/deploy_fail/raw.storage.txt b/tests/it/writer/deploy_fail/raw.storage.txt
@@ -0,0 +1,2 @@
+  [18] → new <unknown>@0xBd770416a3345F91E4B34576cb804a576fa48EB1
+    └─ ← [Revert] 0x0000000000000000000000000000000000000000000000000000000000000042
diff --git a/...ce_printing/0.write_bytecodes.decoded.svg → ...st_trace_printing/0.bytecodes.decoded.svg b/...ce_printing/0.write_bytecodes.decoded.svg → ...st_trace_printing/0.bytecodes.decoded.svg
diff --git a/...ce_printing/0.write_bytecodes.decoded.txt → ...st_trace_printing/0.bytecodes.decoded.txt b/...ce_printing/0.write_bytecodes.decoded.txt → ...st_trace_printing/0.bytecodes.decoded.txt
diff --git a/tests/it/writer/test_trace_printing/0.bytecodes.storage.decoded.svg b/tests/it/writer/test_trace_printing/0.bytecodes.storage.decoded.svg
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,2 @@
		[18] → new RevertingConstructor@0xBd770416a3345F91E4B34576cb804a576fa48EB1(0x604260005260206000fd)
		└─ ← [Revert] 42
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,2 @@
		[18] → new <unknown>@0xBd770416a3345F91E4B34576cb804a576fa48EB1(0x604260005260206000fd)
		└─ ← [Revert] 0x0000000000000000000000000000000000000000000000000000000000000042