From 6c057df21bf69ff9d2abdc14c6199d4771845f5c Mon Sep 17 00:00:00 2001
From: David Pearce <david.pearce@consensys.net>
Date: Wed, 16 Oct 2024 04:23:47 +1300
Subject: [PATCH] feat: Support custom class name for trace file. (#272)

This allows a custom class name to be used in place of `Trace.java` when
desired.
---
 src/exporters/besu.rs                 | 11 +++++++++--
 src/exporters/besu_trace_columns.java |  2 +-
 src/main.rs                           | 10 ++++++++++
 3 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/src/exporters/besu.rs b/src/exporters/besu.rs
index 5043773..d24d924 100644
--- a/src/exporters/besu.rs
+++ b/src/exporters/besu.rs
@@ -46,6 +46,7 @@ struct BesuConstant {
 }
 #[derive(Serialize)]
 struct TemplateData {
+    class: String,
     module: String,
     module_prefix: String,
     columns: Vec<BesuColumn>,
@@ -201,7 +202,12 @@ fn perspectivize_name(h: &Handle, p: &str) -> String {
     )
 }
 
-pub fn render(cs: &ConstraintSet, package: &str, output_path: Option<&String>) -> Result<()> {
+pub fn render(
+    cs: &ConstraintSet,
+    package: &str,
+    class: &str,
+    output_path: Option<&String>,
+) -> Result<()> {
     let registers = cs
         .columns
         .registers
@@ -277,6 +283,7 @@ pub fn render(cs: &ConstraintSet, package: &str, output_path: Option<&String>) -
     handlebars.register_escape_fn(handlebars::no_escape);
 
     let template_data = TemplateData {
+        class: class.to_owned(),
         module: package.to_owned(),
         module_prefix: package.to_case(Case::Pascal),
         constants,
@@ -294,7 +301,7 @@ pub fn render(cs: &ConstraintSet, package: &str, output_path: Option<&String>) -
                 bail!("{} is not a directory", f.bold().yellow());
             }
 
-            let trace_columns_java_filepath = Path::new(f).join("Trace.java");
+            let trace_columns_java_filepath = Path::new(f).join(format!("{class}.java"));
 
             File::create(&trace_columns_java_filepath)?
                 .write_all(trace_columns_render.as_bytes())
diff --git a/src/exporters/besu_trace_columns.java b/src/exporters/besu_trace_columns.java
index b12cf03..79fd9b0 100644
--- a/src/exporters/besu_trace_columns.java
+++ b/src/exporters/besu_trace_columns.java
@@ -31,7 +31,7 @@
  * <p>Any modifications to this code may be overwritten and could lead to unexpected behavior.
  * Please DO NOT ATTEMPT TO MODIFY this code directly.
  */
-public class Trace {
+public class {{ class }} {
   {{#each constants}}
   public static final {{ this.tupe }} {{ this.name }} = {{ this.value }};
   {{/each}}
diff --git a/src/main.rs b/src/main.rs
index 69107b0..6f4d98f 100755
--- a/src/main.rs
+++ b/src/main.rs
@@ -129,6 +129,14 @@ enum Commands {
         )]
         package: String,
 
+        #[arg(
+            short = 'c',
+            long = "class",
+            default_value = "Trace",
+            help = "class name for generated trace file"
+        )]
+        class: String,
+
         #[arg(short = 'o', long = "out", help = "where to render the columns")]
         output_file_path: Option<String>,
     },
@@ -674,11 +682,13 @@ fn main() -> Result<()> {
         #[cfg(feature = "exporters")]
         Commands::Besu {
             package,
+            class,
             output_file_path: output_path,
         } => {
             exporters::besu::render(
                 &builder.into_constraint_set()?,
                 &package,
+                &class,
                 output_path.as_ref(),
             )?;
         }