feat: add TraceWriterConfig (paradigmxyz#216)

src/tracing/mod.rs

@@ -46,7 +46,7 @@ use types::{CallLog, CallTrace, CallTraceStep};
 mod utils;
 
 mod writer;
-pub use writer::TraceWriter;
+pub use writer::{TraceWriter, TraceWriterConfig};
 
 #[cfg(feature = "js-tracer")]
 pub mod js;

src/tracing/writer.rs

@@ -21,42 +21,99 @@ const RETURN: &str = "← ";
 const TRACE_KIND_STYLE: Style = AnsiColor::Yellow.on_default();
 const LOG_STYLE: Style = AnsiColor::Cyan.on_default();
 
+/// Configuration for a [`TraceWriter`].
+#[derive(Clone, Debug)]
+#[allow(missing_copy_implementations)]
+pub struct TraceWriterConfig {
+    use_colors: bool,
+    color_cheatcodes: bool,
+    write_bytecodes: bool,
+}
+
+impl Default for TraceWriterConfig {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl TraceWriterConfig {
+    /// Create a new `TraceWriterConfig` with default settings.
+    pub fn new() -> Self {
+        Self {
+            use_colors: use_colors(ColorChoice::Auto),
+            color_cheatcodes: false,
+            write_bytecodes: false,
+        }
+    }
+
+    /// Use colors in the output. Default: [`Auto`](ColorChoice::Auto).
+    pub fn color_choice(mut self, choice: ColorChoice) -> Self {
+        self.use_colors = use_colors(choice);
+        self
+    }
+
+    /// Get the current color choice. `Auto` is lost, so this returns `true` if colors are enabled.
+    pub fn get_use_colors(&self) -> bool {
+        self.use_colors
+    }
+
+    /// Color calls to the cheatcode address differently. Default: false.
+    pub fn color_cheatcodes(mut self, yes: bool) -> Self {
+        self.color_cheatcodes = yes;
+        self
+    }
+
+    /// Returns `true` if calls to the cheatcode address are colored differently.
+    pub fn get_color_cheatcodes(&self) -> bool {
+        self.color_cheatcodes
+    }
+
+    /// Write contract creation codes and deployed codes when writing "create" traces.
+    /// Default: false.
+    pub fn write_bytecodes(mut self, yes: bool) -> Self {
+        self.write_bytecodes = yes;
+        self
+    }
+
+    /// Returns `true` if contract creation codes and deployed codes are written.
+    pub fn get_write_bytecodes(&self) -> bool {
+        self.write_bytecodes
+    }
+}
+
 /// Formats [call traces](CallTraceArena) to an [`Write`] writer.
 ///
 /// Will never write invalid UTF-8.
 #[derive(Clone, Debug)]
 pub struct TraceWriter<W> {
     writer: W,
-    use_colors: bool,
-    color_cheatcodes: bool,
     indentation_level: u16,
-    write_bytecodes: bool,
+    config: TraceWriterConfig,
 }
 
 impl<W: Write> TraceWriter<W> {
     /// Create a new `TraceWriter` with the given writer.
     #[inline]
     pub fn new(writer: W) -> Self {
-        Self {
-            writer,
-            use_colors: use_colors(ColorChoice::global()),
-            color_cheatcodes: false,
-            indentation_level: 0,
-            write_bytecodes: false,
-        }
+        Self::with_config(writer, TraceWriterConfig::new())
+    }
+
+    /// Create a new `TraceWriter` with the given writer and configuration.
+    pub fn with_config(writer: W, config: TraceWriterConfig) -> Self {
+        Self { writer, indentation_level: 0, config }
     }
 
     /// Sets the color choice.
     #[inline]
     pub fn use_colors(mut self, color_choice: ColorChoice) -> Self {
-        self.use_colors = use_colors(color_choice);
+        self.config.use_colors = use_colors(color_choice);
         self
     }
 
     /// Sets whether to color calls to the cheatcode address differently.
     #[inline]
     pub fn color_cheatcodes(mut self, yes: bool) -> Self {
-        self.color_cheatcodes = yes;
+        self.config.color_cheatcodes = yes;
         self
     }
 
@@ -70,7 +127,7 @@ impl<W: Write> TraceWriter<W> {
     /// Sets whether contract creation codes and deployed codes should be written.
     #[inline]
     pub fn write_bytecodes(mut self, yes: bool) -> Self {
-        self.write_bytecodes = yes;
+        self.config.write_bytecodes = yes;
         self
     }
 
@@ -184,7 +241,7 @@ 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_bytecodes {
+            if self.config.write_bytecodes {
                 write!(self.writer, "({})", hex::encode(&trace.data))?;
             }
         } else {
@@ -344,7 +401,7 @@ impl<W: Write> TraceWriter<W> {
 
         if trace.kind.is_any_create() && trace.status.is_ok() {
             write!(self.writer, "{} bytes of code", trace.output.len())?;
-            if self.write_bytecodes {
+            if self.config.write_bytecodes {
                 write!(self.writer, " ({})", hex::encode(&trace.output))?;
             }
         } else if !trace.output.is_empty() {
@@ -383,10 +440,10 @@ impl<W: Write> TraceWriter<W> {
     }
 
     fn trace_style(&self, trace: &CallTrace) -> Style {
-        if !self.use_colors {
+        if !self.config.use_colors {
             return Style::default();
         }
-        let color = if self.color_cheatcodes && trace.address == CHEATCODE_ADDRESS {
+        let color = if self.config.color_cheatcodes && trace.address == CHEATCODE_ADDRESS {
             AnsiColor::Blue
         } else if trace.success {
             AnsiColor::Green
@@ -397,14 +454,14 @@ impl<W: Write> TraceWriter<W> {
     }
 
     fn trace_kind_style(&self) -> Style {
-        if !self.use_colors {
+        if !self.config.use_colors {
             return Style::default();
         }
         TRACE_KIND_STYLE
     }
 
     fn log_style(&self) -> Style {
-        if !self.use_colors {
+        if !self.config.use_colors {
             return Style::default();
         }
         LOG_STYLE

tests/it/utils.rs

@@ -9,7 +9,7 @@ use revm::{
     },
     Database, DatabaseCommit, GetInspector,
 };
-use revm_inspectors::tracing::TracingInspector;
+use revm_inspectors::tracing::{TraceWriter, TraceWriterConfig, TracingInspector};
 use std::convert::Infallible;
 
 type TestDb = CacheDB<EmptyDB>;
@@ -113,26 +113,16 @@ where
 }
 
 pub fn write_traces(tracer: &TracingInspector) -> String {
-    write_traces_with(tracer, ColorChoice::Never, false)
+    write_traces_with(tracer, TraceWriterConfig::new().color_choice(ColorChoice::Never))
 }
 
-pub fn write_traces_with_bytecodes(tracer: &TracingInspector) -> String {
-    write_traces_with(tracer, ColorChoice::Never, true)
-}
-
-pub fn write_traces_with(
-    tracer: &TracingInspector,
-    color: ColorChoice,
-    write_bytecodes: bool,
-) -> String {
-    let mut w = revm_inspectors::tracing::TraceWriter::new(Vec::<u8>::new())
-        .use_colors(color)
-        .write_bytecodes(write_bytecodes);
+pub fn write_traces_with(tracer: &TracingInspector, config: TraceWriterConfig) -> String {
+    let mut w = TraceWriter::with_config(Vec::<u8>::new(), config);
     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, false));
+    println!("{}", write_traces_with(tracer, TraceWriterConfig::new()));
 }

tests/it/writer.rs

@@ -4,7 +4,7 @@ use alloy_sol_types::{sol, SolCall};
 use colorchoice::ColorChoice;
 use revm_inspectors::tracing::{
     types::{DecodedCallData, DecodedInternalCall, DecodedTraceStep},
-    TracingInspector, TracingInspectorConfig,
+    TraceWriterConfig, TracingInspector, TracingInspectorConfig,
 };
 use snapbox::{assert_data_eq, data::DataFormat};
 use std::path::Path;
@@ -214,31 +214,36 @@ fn assert_traces(
     );
     let bytecodes = if write_bytecodes { &[false, true][..] } else { &[false][..] };
 
-    let do_assert = |color: bool, bytecodes: bool, extra: &str, tracer: &TracingInspector| {
+    let do_assert = |config: TraceWriterConfig, extra: &str, tracer: &TracingInspector| {
+        let color = config.get_use_colors();
+        let bytecodes = config.get_write_bytecodes();
+
         let file_kind = if color { DataFormat::TermSvg } else { DataFormat::Text };
         let extension = if color { "svg" } else { "txt" };
-        let color_choice = if color { ColorChoice::Always } else { ColorChoice::Never };
         let bytecodes_extra = if bytecodes { ".write_bytecodes" } else { "" };
 
-        let s = write_traces_with(tracer, color_choice, bytecodes);
+        let s = write_traces_with(tracer, config);
         let path = base_path.join(format!("{name}{bytecodes_extra}{extra}.{extension}"));
         let data = snapbox::Data::read_from(&path, Some(file_kind));
         assert_data_eq!(s, data);
     };
 
-    for color in [false, true] {
+    let mut configs = vec![];
+    for color in [ColorChoice::Never, ColorChoice::Always] {
         for &bytecodes in bytecodes {
-            do_assert(color, bytecodes, "", tracer);
+            configs.push(TraceWriterConfig::new().color_choice(color).write_bytecodes(bytecodes));
         }
     }
 
+    for config in &configs {
+        do_assert(config.clone(), "", tracer);
+    }
+
     if let Some(patch) = patch_index {
         patch_traces(patch, tracer);
 
-        for color in [false, true] {
-            for &bytecodes in bytecodes {
-                do_assert(color, bytecodes, ".decoded", tracer);
-            }
+        for config in &configs {
+            do_assert(config.clone(), ".decoded", tracer);
         }
     }
 }