feat: small updates for steps tracing (#152)

- Adds `record_returndata_snapshots` flag to config which enables
snapshots of `interpreter.return_data_buffer`
- Adds `record_opcodes_filter` parameter which allows to only record
specific opcodes. ref
foundry-rs/foundry#8222 (comment)
- Adds `gas_used` field for `CallTraceStep`

This should be enough to migrate foundry's debugger to using
`TracingInspector` from here, I will open PR for this later today.

src/tracing/config.rs

@@ -1,9 +1,44 @@
+use alloy_primitives::U256;
 use alloy_rpc_types::trace::{
     geth::{CallConfig, GethDefaultTracingOptions, PreStateConfig},
     parity::TraceType,
 };
+use revm::interpreter::OpCode;
 use std::collections::HashSet;
 
+/// 256 bits each marking whether an opcode should be included into steps trace or not.
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub struct OpcodeFilter(U256);
+
+impl Default for OpcodeFilter {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl OpcodeFilter {
+    /// Returns a new [OpcodeFilter] that does not trace any opcodes.
+    pub const fn new() -> Self {
+        Self(U256::ZERO)
+    }
+
+    /// Returns whether steps with given [OpCode] should be traced.
+    pub fn is_enabled(&self, op: &OpCode) -> bool {
+        self.0.bit(op.get() as usize)
+    }
+
+    /// Enables tracing of given [OpCode].
+    #[must_use]
+    pub const fn enable(self, op: OpCode) -> Self {
+        let bit = op.get() as usize;
+
+        let mut bytes = self.0.to_le_bytes::<32>();
+        bytes[bit / 8] |= 1 << (bit % 8);
+
+        Self(U256::from_le_bytes(bytes))
+    }
+}
+
 /// Gives guidance to the [TracingInspector](crate::tracing::TracingInspector).
 ///
 /// Use [TracingInspectorConfig::default_parity] or [TracingInspectorConfig::default_geth] to get
@@ -18,6 +53,11 @@ pub struct TracingInspectorConfig {
     pub record_stack_snapshots: StackSnapshotType,
     /// Whether to record state diffs.
     pub record_state_diff: bool,
+    /// Whether to record returndata buffer snapshots.
+    pub record_returndata_snapshots: bool,
+    /// Optional filter for opcodes to record. If provided, only steps with opcode in this set will
+    /// be recorded.
+    pub record_opcodes_filter: Option<OpcodeFilter>,
     /// Whether to ignore precompile calls.
     pub exclude_precompile_calls: bool,
     /// Whether to record logs
@@ -32,6 +72,8 @@ impl TracingInspectorConfig {
             record_memory_snapshots: true,
             record_stack_snapshots: StackSnapshotType::Full,
             record_state_diff: false,
+            record_returndata_snapshots: true,
+            record_opcodes_filter: None,
             exclude_precompile_calls: false,
             record_logs: true,
         }
@@ -44,8 +86,10 @@ impl TracingInspectorConfig {
             record_memory_snapshots: false,
             record_stack_snapshots: StackSnapshotType::None,
             record_state_diff: false,
+            record_returndata_snapshots: false,
             exclude_precompile_calls: false,
             record_logs: false,
+            record_opcodes_filter: None,
         }
     }
 
@@ -58,8 +102,10 @@ impl TracingInspectorConfig {
             record_memory_snapshots: false,
             record_stack_snapshots: StackSnapshotType::None,
             record_state_diff: false,
+            record_returndata_snapshots: false,
             exclude_precompile_calls: true,
             record_logs: false,
+            record_opcodes_filter: None,
         }
     }
 
@@ -75,8 +121,10 @@ impl TracingInspectorConfig {
             record_memory_snapshots: false,
             record_stack_snapshots: StackSnapshotType::Full,
             record_state_diff: true,
+            record_returndata_snapshots: false,
             exclude_precompile_calls: false,
             record_logs: false,
+            record_opcodes_filter: None,
         }
     }
 
@@ -233,6 +281,12 @@ impl TracingInspectorConfig {
         self.record_logs = record_logs;
         self
     }
+
+    /// If [OpcodeFilter] is configured, returns whether the given opcode should be recorded.
+    /// Otherwise, always returns true.
+    pub fn should_record_opcode(&self, op: &OpCode) -> bool {
+        self.record_opcodes_filter.map_or(true, |filter| filter.is_enabled(op))
+    }
 }
 
 /// How much of the stack to record. Nothing, just the items pushed, or the full stack

src/tracing/mod.rs

@@ -27,7 +27,7 @@ pub use builder::{
 };
 
 mod config;
-pub use config::{StackSnapshotType, TracingInspectorConfig};
+pub use config::{OpcodeFilter, StackSnapshotType, TracingInspectorConfig};
 
 mod fourbyte;
 pub use fourbyte::FourByteInspector;
@@ -334,7 +334,17 @@ impl TracingInspector {
         let trace = &mut self.traces.arena[trace_idx];
 
         let step_idx = trace.trace.steps.len();
-        self.step_stack.push(StackStep { trace_idx, step_idx });
+        // we always want an OpCode, even it is unknown because it could be an additional opcode
+        // that not a known constant
+        let op = unsafe { OpCode::new_unchecked(interp.current_opcode()) };
+
+        let record = self.config.should_record_opcode(&op);
+
+        self.step_stack.push(StackStep { trace_idx, step_idx, record });
+
+        if !record {
+            return;
+        }
 
         let memory = self
             .config
@@ -346,10 +356,17 @@ impl TracingInspector {
         } else {
             None
         };
+        let returndata = self
+            .config
+            .record_returndata_snapshots
+            .then(|| interp.return_data_buffer.clone())
+            .unwrap_or_default();
 
-        // we always want an OpCode, even it is unknown because it could be an additional opcode
-        // that not a known constant
-        let op = unsafe { OpCode::new_unchecked(interp.current_opcode()) };
+        let gas_used = gas_used(
+            context.spec_id(),
+            interp.gas.limit().saturating_sub(interp.gas.remaining()),
+            interp.gas.refunded() as u64,
+        );
 
         trace.trace.steps.push(CallTraceStep {
             depth: context.journaled_state.depth(),
@@ -360,8 +377,10 @@ impl TracingInspector {
             push_stack: None,
             memory_size: memory.len(),
             memory,
+            returndata,
             gas_remaining: interp.gas.remaining(),
             gas_refund_counter: interp.gas.refunded() as u64,
+            gas_used,
 
             // fields will be populated end of call
             gas_cost: 0,
@@ -381,8 +400,13 @@ impl TracingInspector {
         interp: &mut Interpreter,
         context: &mut EvmContext<DB>,
     ) {
-        let StackStep { trace_idx, step_idx } =
+        let StackStep { trace_idx, step_idx, record } =
             self.step_stack.pop().expect("can't fill step without starting a step first");
+
+        if !record {
+            return;
+        }
+
         let step = &mut self.traces.arena[trace_idx].trace.steps[step_idx];
 
         if self.config.record_stack_snapshots.is_pushes() {
@@ -558,8 +582,19 @@ where
     }
 }
 
+/// Struct keeping track of internal inspector steps stack.
 #[derive(Clone, Copy, Debug)]
 struct StackStep {
+    /// Whether this step should be recorded.
+    ///
+    /// This is set to `false` if [OpcodeFilter] is configured and this step's opcode is not
+    /// enabled for tracking
+    record: bool,
+    /// Idx of the trace node this step belongs.
     trace_idx: usize,
+    /// Idx of this step in the [CallTrace::steps].
+    ///
+    /// Please note that if `record` is `false`, this will still contain a value, but the step will
+    /// not appear in the steps list.
     step_idx: usize,
 }

src/tracing/types.rs

@@ -506,10 +506,14 @@ pub struct CallTraceStep {
     pub memory: RecordedMemory,
     /// Size of memory at the beginning of the step
     pub memory_size: usize,
+    /// Returndata before step execution
+    pub returndata: Bytes,
     /// Remaining gas before step execution
     pub gas_remaining: u64,
     /// Gas refund counter before step execution
     pub gas_refund_counter: u64,
+    /// Total gas used before step execution
+    pub gas_used: u64,
     // Fields filled in `step_end`
     /// Gas cost of step execution
     pub gas_cost: u64,