Write storage change in trace

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