From 09afce4ccdbdd98d4efc4a0cca3784a028f11efb Mon Sep 17 00:00:00 2001
From: CodeSandwich <igor.zuk@protonmail.com>
Date: Fri, 20 Sep 2024 21:32:25 +0200
Subject: [PATCH] feat: Add creation code printing in traces

---
 src/tracing/writer.rs | 12 ++++++++++++
 tests/it/utils.rs     | 18 ++++++++++++++----
 tests/it/writer.rs    | 12 +++++++++++-
 3 files changed, 37 insertions(+), 5 deletions(-)

diff --git a/src/tracing/writer.rs b/src/tracing/writer.rs
index 5ae00fbb..ea284d71 100644
--- a/src/tracing/writer.rs
+++ b/src/tracing/writer.rs
@@ -30,6 +30,7 @@ pub struct TraceWriter<W> {
     use_colors: bool,
     color_cheatcodes: bool,
     indentation_level: u16,
+    write_creation_codes: bool,
 }
 
 impl<W: Write> TraceWriter<W> {
@@ -41,6 +42,7 @@ impl<W: Write> TraceWriter<W> {
             use_colors: use_colors(ColorChoice::global()),
             color_cheatcodes: false,
             indentation_level: 0,
+            write_creation_codes: false,
         }
     }
 
@@ -65,6 +67,13 @@ impl<W: Write> TraceWriter<W> {
         self
     }
 
+    /// Sets the starting indentation level.
+    #[inline]
+    pub fn write_creation_codes(mut self, yes: bool) -> Self {
+        self.write_creation_codes = yes;
+        self
+    }
+
     /// Returns a reference to the inner writer.
     #[inline]
     pub const fn writer(&self) -> &W {
@@ -175,6 +184,9 @@ impl<W: Write> TraceWriter<W> {
                 "{trace_kind_style}{CALL}new{trace_kind_style:#} {label}@{address}",
                 label = trace.decoded.label.as_deref().unwrap_or("<unknown>")
             )?;
+            if self.write_creation_codes {
+                write!(self.writer, "({})", hex::encode(&trace.data))?;
+            }
         } else {
             let (func_name, inputs) = match &trace.decoded.call_data {
                 Some(DecodedCallData { signature, args }) => {
diff --git a/tests/it/utils.rs b/tests/it/utils.rs
index 9a6c1b90..004f3c0c 100644
--- a/tests/it/utils.rs
+++ b/tests/it/utils.rs
@@ -113,16 +113,26 @@ where
 }
 
 pub fn write_traces(tracer: &TracingInspector) -> String {
-    write_traces_with(tracer, ColorChoice::Never)
+    write_traces_with(tracer, ColorChoice::Never, false)
 }
 
-pub fn write_traces_with(tracer: &TracingInspector, color: ColorChoice) -> String {
-    let mut w = revm_inspectors::tracing::TraceWriter::new(Vec::<u8>::new()).use_colors(color);
+pub fn write_traces_with_creation_codes(tracer: &TracingInspector) -> String {
+    write_traces_with(tracer, ColorChoice::Never, true)
+}
+
+pub fn write_traces_with(
+    tracer: &TracingInspector,
+    color: ColorChoice,
+    write_creation_codes: bool,
+) -> String {
+    let mut w = revm_inspectors::tracing::TraceWriter::new(Vec::<u8>::new())
+        .use_colors(color)
+        .write_creation_codes(write_creation_codes);
     w.write_arena(tracer.traces()).expect("failed to write traces to Vec<u8>");
     String::from_utf8(w.into_writer()).expect("trace writer wrote invalid UTF-8")
 }
 
 pub fn print_traces(tracer: &TracingInspector) {
     // Use `println!` so that the output is captured by the test runner.
-    println!("{}", write_traces_with(tracer, ColorChoice::Auto));
+    println!("{}", write_traces_with(tracer, ColorChoice::Auto, false));
 }
diff --git a/tests/it/writer.rs b/tests/it/writer.rs
index 2b0859cd..8f55589a 100644
--- a/tests/it/writer.rs
+++ b/tests/it/writer.rs
@@ -1,4 +1,4 @@
-use crate::utils::{write_traces, TestEvm};
+use crate::utils::{write_traces, write_traces_with_creation_codes, TestEvm};
 use alloy_primitives::{address, b256, bytes, hex, Address, Bytes, B256, U256};
 use alloy_sol_types::{sol, SolCall};
 use revm_inspectors::tracing::{
@@ -28,6 +28,16 @@ fn test_trace_printing() {
 "#]]
     );
 
+    let s = write_traces_with_creation_codes(&tracer);
+    let raw = r#"
+  [209257] → new <unknown>@0xBd770416a3345F91E4B34576cb804a576fa48EB1(<BYTECODE>)
+    └─ ← [Return] 1045 bytes of code
+"#
+    .strip_prefix("\n")
+    .unwrap()
+    .replace("<BYTECODE>", BYTECODE.to_string().strip_prefix("0x").unwrap());
+    assert_data_eq!(s, raw);
+
     let mut index = 0;
 
     let mut call = |data: Vec<u8>, raw: Inline, decoded: Inline| {