Write storage change in trace

src/tracing/writer.rs

@@ -8,7 +8,10 @@ use super::{
 use alloy_primitives::{address, hex, Address};
 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");
 
@@ -161,6 +164,8 @@ impl<W: Write> TraceWriter<W> {
         self.indentation_level += 1;
         self.write_items(nodes, idx)?;
 
+        self.write_storage_changes(node)?;
+
         // Write return data.
         self.write_edge()?;
         self.write_trace_footer(&node.trace)?;
@@ -291,6 +296,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;
 
@@ -409,6 +415,47 @@ 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(ref change) = step.storage_change {
+                let entry =
+                    changes_map.entry((&step.contract, &change.key)).or_insert((None, None));
+                if entry.0.is_none() {
+                    entry.0 = Some(change);
+                }
+                entry.1 = Some(change);
+            }
+        }
+
+        for ((contract, key), (first_change, last_change)) in changes_map {
+            if let (Some(first), Some(last)) = (first_change, last_change) {
+                let had_value = first.had_value;
+                let value = last.value;
+
+                self.write_branch()?;
+                writeln!(
+                    self.writer,
+                    " storage write {} [0x{}]:",
+                    contract.to_checksum_buffer(None).as_str(),
+                    hex::encode(key.to_be_bytes_vec())
+                )?;
+                self.write_pipes()?;
+                writeln!(
+                    self.writer,
+                    "                0x{:>64} → 0x{:>64}",
+                    had_value.map_or("".into(), |v| hex::encode(v.to_be_bytes_vec())),
+                    hex::encode(value.to_be_bytes_vec())
+                )?;
+            }
+        }
+
+        Ok(())
+    }
 }
 
 fn use_colors(choice: ColorChoice) -> bool {