feat: complete vm and statediff tracers (#3529)

Co-authored-by: N <mail@nuhhtyy.xyz>
Co-authored-by: Matthias Seitz <matthias.seitz@outlook.de>

crates/revm/revm-inspectors/src/tracing/builder/mod.rs

@@ -5,3 +5,6 @@ pub mod geth;
 
 /// Parity style trace builders for `trace_` namespace
 pub mod parity;
+
+/// Walker types used for traversing various callgraphs
+mod walker;

crates/revm/revm-inspectors/src/tracing/builder/parity.rs

@@ -1,11 +1,16 @@
-use crate::tracing::{types::CallTraceNode, TracingInspectorConfig};
+use super::walker::CallTraceNodeWalkerBF;
+use crate::tracing::{
+    types::{CallTraceNode, CallTraceStep},
+    TracingInspectorConfig,
+};
 use reth_primitives::{Address, U64};
 use reth_rpc_types::{trace::parity::*, TransactionInfo};
 use revm::{
     db::DatabaseRef,
-    primitives::{AccountInfo, ExecutionResult, ResultAndState},
+    interpreter::opcode,
+    primitives::{AccountInfo, ExecutionResult, ResultAndState, KECCAK_EMPTY},
 };
-use std::collections::HashSet;
+use std::collections::{HashSet, VecDeque};
 
 /// A type for creating parity style traces
 ///
@@ -14,6 +19,7 @@ use std::collections::HashSet;
 pub struct ParityTraceBuilder {
     /// Recorded trace nodes
     nodes: Vec<CallTraceNode>,
+
     /// How the traces were recorded
     _config: TracingInspectorConfig,
 }
@@ -154,7 +160,18 @@ impl ParityTraceBuilder {
         DB: DatabaseRef,
     {
         let ResultAndState { result, state } = res;
+
+        let breadth_first_addresses = if trace_types.contains(&TraceType::VmTrace) {
+            CallTraceNodeWalkerBF::new(&self.nodes)
+                .map(|node| node.trace.address)
+                .collect::<Vec<_>>()
+        } else {
+            vec![]
+        };
+
         let mut trace_res = self.into_trace_results(result, trace_types);
+
+        // check the state diff case
         if let Some(ref mut state_diff) = trace_res.state_diff {
             populate_account_balance_nonce_diffs(
                 state_diff,
@@ -162,6 +179,12 @@ impl ParityTraceBuilder {
                 state.into_iter().map(|(addr, acc)| (addr, acc.info)),
             )?;
         }
+
+        // check the vm trace case
+        if let Some(ref mut vm_trace) = trace_res.vm_trace {
+            populate_vm_trace_bytecodes(&db, vm_trace, breadth_first_addresses)?;
+        }
+
         Ok(trace_res)
     }
 
@@ -177,11 +200,8 @@ impl ParityTraceBuilder {
         let with_traces = trace_types.contains(&TraceType::Trace);
         let with_diff = trace_types.contains(&TraceType::StateDiff);
 
-        let vm_trace = if trace_types.contains(&TraceType::VmTrace) {
-            Some(vm_trace(&self.nodes))
-        } else {
-            None
-        };
+        let vm_trace =
+            if trace_types.contains(&TraceType::VmTrace) { Some(self.vm_trace()) } else { None };
 
         let mut traces = Vec::with_capacity(if with_traces { self.nodes.len() } else { 0 });
         let mut diff = StateDiff::default();
@@ -218,13 +238,142 @@ impl ParityTraceBuilder {
     pub fn into_transaction_traces(self) -> Vec<TransactionTrace> {
         self.into_transaction_traces_iter().collect()
     }
+
+    /// Creates a VM trace by walking over `CallTraceNode`s
+    ///
+    /// does not have the code fields filled in
+    pub fn vm_trace(&self) -> VmTrace {
+        match self.nodes.get(0) {
+            Some(current) => self.make_vm_trace(current),
+            None => VmTrace { code: Default::default(), ops: Vec::new() },
+        }
+    }
+
+    /// returns a VM trace without the code filled in
+    ///
+    /// iteratively creaters a VM trace by traversing an arena
+    fn make_vm_trace(&self, start: &CallTraceNode) -> VmTrace {
+        let mut child_idx_stack: Vec<usize> = Vec::with_capacity(self.nodes.len());
+        let mut sub_stack: VecDeque<Option<VmTrace>> = VecDeque::with_capacity(self.nodes.len());
+
+        let mut current = start;
+        let mut child_idx: usize = 0;
+
+        // finds the deepest nested calls of each call frame and fills them up bottom to top
+        let instructions = loop {
+            match current.children.get(child_idx) {
+                Some(child) => {
+                    child_idx_stack.push(child_idx + 1);
+
+                    child_idx = 0;
+                    current = self.nodes.get(*child).expect("there should be a child");
+                }
+                None => {
+                    let mut instructions: Vec<VmInstruction> =
+                        Vec::with_capacity(current.trace.steps.len());
+
+                    for step in &current.trace.steps {
+                        let maybe_sub = match step.op.u8() {
+                            opcode::CALL |
+                            opcode::CALLCODE |
+                            opcode::DELEGATECALL |
+                            opcode::STATICCALL |
+                            opcode::CREATE |
+                            opcode::CREATE2 => {
+                                sub_stack.pop_front().expect("there should be a sub trace")
+                            }
+                            _ => None,
+                        };
+
+                        instructions.push(Self::make_instruction(step, maybe_sub));
+                    }
+
+                    match current.parent {
+                        Some(parent) => {
+                            sub_stack.push_back(Some(VmTrace {
+                                code: Default::default(),
+                                ops: instructions,
+                            }));
+
+                            child_idx = child_idx_stack.pop().expect("there should be a child idx");
+
+                            current = self.nodes.get(parent).expect("there should be a parent");
+                        }
+                        None => break instructions,
+                    }
+                }
+            }
+        };
+
+        VmTrace { code: Default::default(), ops: instructions }
+    }
+
+    /// Creates a VM instruction from a [CallTraceStep] and a [VmTrace] for the subcall if there is
+    /// one
+    fn make_instruction(step: &CallTraceStep, maybe_sub: Option<VmTrace>) -> VmInstruction {
+        let maybe_storage = step.storage_change.map(|storage_change| StorageDelta {
+            key: storage_change.key,
+            val: storage_change.value,
+        });
+
+        let maybe_memory = match step.memory.len() {
+            0 => None,
+            _ => {
+                Some(MemoryDelta { off: step.memory_size, data: step.memory.data().clone().into() })
+            }
+        };
+
+        let maybe_execution = Some(VmExecutedOperation {
+            used: step.gas_cost,
+            push: step.new_stack.map(|new_stack| new_stack.into()),
+            mem: maybe_memory,
+            store: maybe_storage,
+        });
+
+        VmInstruction {
+            pc: step.pc,
+            cost: 0, // TODO: use op gas cost
+            ex: maybe_execution,
+            sub: maybe_sub,
+        }
+    }
 }
 
-/// Construct the vmtrace for the entire callgraph
-fn vm_trace(nodes: &[CallTraceNode]) -> VmTrace {
-    // TODO: populate vm trace
+/// addresses are presorted via breadth first walk thru [CallTraceNode]s, this  can be done by a
+/// walker in [crate::tracing::builder::walker]
+///
+/// iteratively fill the [VmTrace] code fields
+pub(crate) fn populate_vm_trace_bytecodes<DB, I>(
+    db: &DB,
+    trace: &mut VmTrace,
+    breadth_first_addresses: I,
+) -> Result<(), DB::Error>
+where
+    DB: DatabaseRef,
+    I: IntoIterator<Item = Address>,
+{
+    let mut stack: VecDeque<&mut VmTrace> = VecDeque::new();
+    stack.push_back(trace);
 
-    VmTrace { code: nodes[0].trace.data.clone().into(), ops: vec![] }
+    let mut addrs = breadth_first_addresses.into_iter();
+
+    while let Some(curr_ref) = stack.pop_front() {
+        for op in curr_ref.ops.iter_mut() {
+            if let Some(sub) = op.sub.as_mut() {
+                stack.push_back(sub);
+            }
+        }
+
+        let addr = addrs.next().expect("there should be an address");
+
+        let db_acc = db.basic(addr)?.unwrap_or_default();
+
+        let code_hash = if db_acc.code_hash != KECCAK_EMPTY { db_acc.code_hash } else { continue };
+
+        curr_ref.code = db.code_by_hash(code_hash)?.bytecode.into();
+    }
+
+    Ok(())
 }
 
 /// Loops over all state accounts in the accounts diff that contains all accounts that are included

crates/revm/revm-inspectors/src/tracing/builder/walker.rs

@@ -0,0 +1,39 @@
+use crate::tracing::types::CallTraceNode;
+use std::collections::VecDeque;
+
+/// Traverses Reths internal tracing structure breadth-first
+///
+/// This is a lazy iterator
+pub(crate) struct CallTraceNodeWalkerBF<'trace> {
+    /// the entire arena
+    nodes: &'trace Vec<CallTraceNode>,
+
+    /// holds indexes of nodes to visit as we traverse
+    queue: VecDeque<usize>,
+}
+
+impl<'trace> CallTraceNodeWalkerBF<'trace> {
+    pub(crate) fn new(nodes: &'trace Vec<CallTraceNode>) -> Self {
+        let mut queue = VecDeque::with_capacity(nodes.len());
+        queue.push_back(0);
+
+        Self { nodes, queue }
+    }
+}
+
+impl<'trace> Iterator for CallTraceNodeWalkerBF<'trace> {
+    type Item = &'trace CallTraceNode;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        match self.queue.pop_front() {
+            Some(idx) => {
+                let curr = self.nodes.get(idx).expect("there should be a node");
+
+                self.queue.extend(curr.children.iter());
+
+                Some(curr)
+            }
+            None => None,
+        }
+    }
+}

crates/revm/revm-inspectors/src/tracing/mod.rs

@@ -265,6 +265,7 @@ impl TracingInspector {
             op,
             contract: interp.contract.address,
             stack,
+            new_stack: None,
             memory,
             memory_size: interp.memory.len(),
             gas_remaining: self.gas_inspector.gas_remaining(),
@@ -290,6 +291,10 @@ impl TracingInspector {
             self.step_stack.pop().expect("can't fill step without starting a step first");
         let step = &mut self.traces.arena[trace_idx].trace.steps[step_idx];
 
+        if interp.stack.len() > step.stack.len() {
+            step.new_stack = interp.stack.data().last().copied();
+        }
+
         if self.config.record_memory_snapshots {
             // resize memory so opcodes that allocated memory is correctly displayed
             if interp.memory.len() > step.memory.len() {

crates/revm/revm-inspectors/src/tracing/types.rs

@@ -463,11 +463,13 @@ pub(crate) struct CallTraceStep {
     pub(crate) contract: Address,
     /// Stack before step execution
     pub(crate) stack: Stack,
+    /// The new stack item placed by this step if any
+    pub(crate) new_stack: Option<U256>,
     /// All allocated memory in a step
     ///
     /// This will be empty if memory capture is disabled
     pub(crate) memory: Memory,
-    /// Size of memory
+    /// Size of memory at the beginning of the step
     pub(crate) memory_size: usize,
     /// Remaining gas before step execution
     pub(crate) gas_remaining: u64,