feat: add decoding for individual trace steps (#157)

ref foundry-rs/foundry#8222
ref foundry-rs/foundry#8198

Adds structs and extends `TraceWriter` to support formatting of decoded
trace steps. Currently two decoding formats are supported:
- Internal calls. Similar to a decoded call trace, decoded internal
function invocation which spans over multiple steps. Kept as decoded
function name, inputs, outputs and index of the last step.
- Arbitrary strings. This might be useful for formatting decoded opcodes
(e.g. adding `├─ [sload] <slot>` to trace. It might make sense to extend
it to something more configurable once we start implementing this

src/tracing/mod.rs

@@ -416,6 +416,7 @@ impl TracingInspector {
             gas_remaining: interp.gas.remaining(),
             gas_refund_counter: interp.gas.refunded() as u64,
             gas_used,
+            decoded: None,
 
             // fields will be populated end of call
             gas_cost: 0,

src/tracing/types.rs

@@ -541,6 +541,31 @@ pub enum TraceMemberOrder {
     Step(usize),
 }
 
+/// Represents a decoded internal function call.
+#[derive(Clone, Debug, PartialEq, Eq)]
+#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
+pub struct DecodedInternalCall {
+    /// Name of the internal function.
+    pub func_name: String,
+    /// Input arguments of the internal function.
+    pub args: Option<Vec<String>>,
+    /// Optional decoded return data.
+    pub return_data: Option<Vec<String>>,
+}
+
+/// Represents a decoded trace step. Currently two formats are supported.
+#[derive(Clone, Debug, PartialEq, Eq)]
+#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
+pub enum DecodedTraceStep {
+    /// Decoded internal function call. Displayed similarly to external calls.
+    ///
+    /// Keeps decoded internal call data and an index of the step where the internal call execution
+    /// ends.
+    InternalCall(DecodedInternalCall, usize),
+    /// Arbitrary line reperesenting the step. Might be used for displaying individual opcodes.
+    Line(String),
+}
+
 /// Represents a tracked call step during execution
 #[derive(Clone, Debug, PartialEq, Eq)]
 #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
@@ -580,6 +605,8 @@ pub struct CallTraceStep {
     ///
     /// This is set after the step was executed.
     pub status: InstructionResult,
+    /// Optional complementary decoded step data.
+    pub decoded: Option<DecodedTraceStep>,
 }
 
 // === impl CallTraceStep ===

src/tracing/writer.rs

@@ -1,5 +1,8 @@
 use super::{
-    types::{CallKind, CallLog, CallTrace, CallTraceNode, DecodedCallData, TraceMemberOrder},
+    types::{
+        CallKind, CallLog, CallTrace, CallTraceNode, DecodedCallData, DecodedTraceStep,
+        TraceMemberOrder,
+    },
     CallTraceArena,
 };
 use alloy_primitives::{address, hex, Address};
@@ -86,6 +89,56 @@ impl<W: Write> TraceWriter<W> {
         self.writer.flush()
     }
 
+    /// Writes a single item of a single node to the writer. Returns the index of the next item to
+    /// be written.
+    ///
+    /// Note: this will return length of [CallTraceNode::ordering] when last item will get
+    /// processed.
+    fn write_item(
+        &mut self,
+        nodes: &[CallTraceNode],
+        node_idx: usize,
+        item_idx: usize,
+    ) -> io::Result<usize> {
+        let node = &nodes[node_idx];
+        match &node.ordering[item_idx] {
+            TraceMemberOrder::Log(index) => {
+                self.write_log(&node.logs[*index])?;
+                Ok(item_idx + 1)
+            }
+            TraceMemberOrder::Call(index) => {
+                self.write_node(nodes, node.children[*index])?;
+                Ok(item_idx + 1)
+            }
+            TraceMemberOrder::Step(index) => self.write_step(nodes, node_idx, item_idx, *index),
+        }
+    }
+
+    /// Writes items of a single node to the writer, starting from the given index, and until the
+    /// given predicate is false.
+    ///
+    /// Returns the index of the next item to be written.
+    fn write_items_until(
+        &mut self,
+        nodes: &[CallTraceNode],
+        node_idx: usize,
+        first_item_idx: usize,
+        f: impl Fn(usize) -> bool,
+    ) -> io::Result<usize> {
+        let mut item_idx = first_item_idx;
+        while !f(item_idx) {
+            item_idx = self.write_item(nodes, node_idx, item_idx)?;
+        }
+        Ok(item_idx)
+    }
+
+    /// Writes all items of a single node to the writer.
+    fn write_items(&mut self, nodes: &[CallTraceNode], node_idx: usize) -> io::Result<()> {
+        let items_cnt = nodes[node_idx].ordering.len();
+        self.write_items_until(nodes, node_idx, 0, |idx| idx == items_cnt)?;
+        Ok(())
+    }
+
     /// Writes a single node and its children to the writer.
     fn write_node(&mut self, nodes: &[CallTraceNode], idx: usize) -> io::Result<()> {
         let node = &nodes[idx];
@@ -97,13 +150,7 @@ impl<W: Write> TraceWriter<W> {
 
         // Write logs and subcalls.
         self.indentation_level += 1;
-        for child in &node.ordering {
-            match *child {
-                TraceMemberOrder::Log(index) => self.write_log(&node.logs[index]),
-                TraceMemberOrder::Call(index) => self.write_node(nodes, node.children[index]),
-                TraceMemberOrder::Step(_) => Ok(()),
-            }?;
-        }
+        self.write_items(nodes, idx)?;
 
         // Write return data.
         self.write_edge()?;
@@ -212,6 +259,64 @@ impl<W: Write> TraceWriter<W> {
         Ok(())
     }
 
+    /// Writes a single step of a single node to the writer. Returns the index of the next item to
+    /// be written.
+    fn write_step(
+        &mut self,
+        nodes: &[CallTraceNode],
+        node_idx: usize,
+        item_idx: usize,
+        step_idx: usize,
+    ) -> io::Result<usize> {
+        let node = &nodes[node_idx];
+        let step = &node.trace.steps[step_idx];
+
+        let Some(decoded) = &step.decoded else {
+            // We only write explicitly decoded steps to avoid bloating the output.
+            return Ok(item_idx + 1);
+        };
+
+        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;
+
+                writeln!(
+                    self.writer,
+                    "[{}] {}{}",
+                    gas_used,
+                    call.func_name,
+                    call.args.as_ref().map(|v| format!("({})", v.join(", "))).unwrap_or_default()
+                )?;
+
+                let end_item_idx =
+                    self.write_items_until(nodes, node_idx, item_idx + 1, |item_idx: usize| {
+                        matches!(&node.ordering[item_idx], TraceMemberOrder::Step(idx) if *idx == *end_idx)
+                    })?;
+
+                self.write_edge()?;
+                write!(self.writer, "{RETURN}")?;
+
+                if let Some(outputs) = &call.return_data {
+                    write!(self.writer, "{}", outputs.join(", "))?;
+                }
+
+                writeln!(self.writer)?;
+
+                self.indentation_level -= 1;
+
+                Ok(end_item_idx + 1)
+            }
+            DecodedTraceStep::Line(line) => {
+                self.write_branch()?;
+                writeln!(self.writer, "{line}")?;
+
+                Ok(item_idx + 1)
+            }
+        }
+    }
+
     /// Writes the footer of a call trace.
     fn write_trace_footer(&mut self, trace: &CallTrace) -> io::Result<()> {
         write!(

tests/it/writer.rs

@@ -2,7 +2,10 @@ use crate::utils::{write_traces, TestEvm};
 use alloy_primitives::{bytes, Bytes, U256};
 use alloy_sol_types::{sol, SolCall};
 use expect_test::expect;
-use revm_inspectors::tracing::{types::DecodedCallData, TracingInspector, TracingInspectorConfig};
+use revm_inspectors::tracing::{
+    types::{DecodedCallData, DecodedInternalCall, DecodedTraceStep},
+    TracingInspector, TracingInspectorConfig,
+};
 
 #[test]
 fn test_basic_trace_printing() {
@@ -194,7 +197,7 @@ fn test_decoded_trace_printing() {
     let mut index = 0;
 
     let mut call = |data: Vec<u8>| -> String {
-        let mut tracer = TracingInspector::new(TracingInspectorConfig::all().disable_steps());
+        let mut tracer = TracingInspector::new(TracingInspectorConfig::all());
         let r = evm.call(address, data.into(), &mut tracer).unwrap();
         assert!(r.is_success());
 
@@ -244,17 +247,33 @@ fn test_decoded_trace_printing() {
     "#]]
     .assert_eq(&s);
 
+    let mut s = call(Counter::log2Call {}.abi_encode());
+    patch_output(&mut s);
+    expect![[r#"
+        . [4642] Counter::log2()
+            ├─ [sload] 0x0000000000000000000000000000000000000000000000000000000000000045
+            ├─ emit Log2(foo: 0x0000000000000000000000000000000000000000000000000000000000000045, bar: 0x000000000000000000000000000000000000000000000000000000000000007b)
+            └─ ← [Stop] 
+    "#]]
+    .assert_eq(&s);
+
     let mut s = call(Counter::nest1Call {}.abi_encode());
     patch_output(&mut s);
     expect![[r#"
         . [13556] Counter::nest1()
-            ├─ emit Log1(foo: 0x0000000000000000000000000000000000000000000000000000000000000045)
-            ├─ [8575] Counter::nest2()
-            │   ├─ [2337] Counter::nest3()
-            │   │   ├─ emit Log1(foo: 0x0000000000000000000000000000000000000000000000000000000000000046)
+            ├─ [13264] Counter::_nest1()
+            │   ├─ emit Log1(foo: 0x0000000000000000000000000000000000000000000000000000000000000045)
+            │   ├─ [8575] Counter::nest2()
+            │   │   ├─ [2337] Counter::nest3()
+            │   │   │   ├─ [2220] Counter::_nest3Internal(arg1, arg2, 3)
+            │   │   │   │   ├─ [mstore]
+            │   │   │   │   ├─ emit Log1(foo: 0x0000000000000000000000000000000000000000000000000000000000000046)
+            │   │   │   │   └─ ← ret1
+            │   │   │   ├─ [before_return]
+            │   │   │   └─ ← [Return] 
+            │   │   ├─ emit Log2(foo: 0x0000000000000000000000000000000000000000000000000000000000000046, bar: 0x000000000000000000000000000000000000000000000000000000000000007b)
             │   │   └─ ← [Return] 
-            │   ├─ emit Log2(foo: 0x0000000000000000000000000000000000000000000000000000000000000046, bar: 0x000000000000000000000000000000000000000000000000000000000000007b)
-            │   └─ ← [Return] 
+            │   └─ ← ret1, ret2
             └─ ← [Return] 
     "#]]
     .assert_eq(&s);
@@ -368,6 +387,37 @@ fn patch_traces(patch: usize, t: &mut TracingInspector) {
                 node.trace.decoded.return_data = Some("69".to_string())
             }
             4 => node.trace.decoded.return_data = Some("69".to_string()),
+            5 => {
+                node.trace.steps[0].decoded = Some(DecodedTraceStep::Line(
+                    "[sload] 0x0000000000000000000000000000000000000000000000000000000000000045"
+                        .to_string(),
+                ));
+            }
+            6 if node.trace.depth == 2 => {
+                node.trace.steps[30].decoded = Some(DecodedTraceStep::InternalCall(
+                    DecodedInternalCall {
+                        func_name: "Counter::_nest3Internal".to_string(),
+                        args: Some(vec!["arg1".to_string(), "arg2".to_string(), "3".to_string()]),
+                        return_data: Some(vec!["ret1".to_string()]),
+                    },
+                    89,
+                ));
+                node.trace.steps[87].decoded = Some(DecodedTraceStep::Line("[mstore]".to_string()));
+                node.trace.steps[90].decoded =
+                    Some(DecodedTraceStep::Line("[before_return]".to_string()));
+                println!("{:?}", node.ordering);
+            }
+            6 if node.trace.depth == 0 => {
+                node.trace.steps[10].decoded = Some(DecodedTraceStep::InternalCall(
+                    DecodedInternalCall {
+                        func_name: "Counter::_nest1".to_string(),
+                        args: Some(vec![]),
+                        return_data: Some(vec!["ret1".to_string(), "ret2".to_string()]),
+                    },
+                    150,
+                ));
+                println!("{:?}", node.ordering);
+            }
             _ => continue,
         }
     }