From 48a1704ecfe8679bc378a80603e68d8ff06323f1 Mon Sep 17 00:00:00 2001
From: Gabriel de Quadros Ligneul <gligneul@offchainlabs.com>
Date: Tue, 20 Aug 2024 10:42:47 -0300
Subject: [PATCH] Add option to use native trace in stylus replay

---
 replay/src/main.rs  | 24 ++++++++++--------------
 replay/src/trace.rs | 12 ++++++++++--
 2 files changed, 20 insertions(+), 16 deletions(-)

diff --git a/replay/src/main.rs b/replay/src/main.rs
index 79fd1da..9303553 100644
--- a/replay/src/main.rs
+++ b/replay/src/main.rs
@@ -48,15 +48,8 @@ enum Subcommands {
 
 #[derive(Args, Clone, Debug)]
 struct ReplayArgs {
-    /// RPC endpoint.
-    #[arg(short, long, default_value = "http://localhost:8547")]
-    endpoint: String,
-    /// Tx to replay.
-    #[arg(short, long)]
-    tx: TxHash,
-    /// Project path.
-    #[arg(short, long, default_value = ".")]
-    project: PathBuf,
+    #[command(flatten)]
+    trace: TraceArgs,
     /// Whether to use stable Rust. Note that nightly is needed to expand macros.
     #[arg(short, long)]
     stable_rust: bool,
@@ -76,6 +69,9 @@ struct TraceArgs {
     /// Project path.
     #[arg(short, long, default_value = ".")]
     project: PathBuf,
+    /// If set, use the native tracer instead of the JavaScript one. Notice the native tracer might not be available in the node.
+    #[arg(short, long, default_value_t = false)]
+    use_native_tracer: bool,
 }
 
 fn main() -> Result<()> {
@@ -106,7 +102,7 @@ async fn main_impl(args: Opts) -> Result<()> {
 
 async fn trace(args: TraceArgs) -> Result<()> {
     let provider = sys::new_provider(&args.endpoint)?;
-    let trace = Trace::new(provider, args.tx).await?;
+    let trace = Trace::new(provider, args.tx, args.use_native_tracer).await?;
     println!("{}", trace.json);
     Ok(())
 }
@@ -144,11 +140,11 @@ async fn replay(args: ReplayArgs) -> Result<()> {
         bail!("failed to exec gdb {:?}", err);
     }
 
-    let provider = sys::new_provider(&args.endpoint)?;
-    let trace = Trace::new(provider, args.tx).await?;
+    let provider = sys::new_provider(&args.trace.endpoint)?;
+    let trace = Trace::new(provider, args.trace.tx, args.trace.use_native_tracer).await?;
 
-    build_so(&args.project)?;
-    let so = find_so(&args.project)?;
+    build_so(&args.trace.project)?;
+    let so = find_so(&args.trace.project)?;
 
     // TODO: don't assume the contract is top-level
     let args_len = trace.tx.input.len();
diff --git a/replay/src/trace.rs b/replay/src/trace.rs
index 13be439..04f2a5c 100644
--- a/replay/src/trace.rs
+++ b/replay/src/trace.rs
@@ -23,7 +23,11 @@ pub struct Trace {
 }
 
 impl Trace {
-    pub async fn new<T: JsonRpcClient>(provider: Provider<T>, tx: TxHash) -> Result<Self> {
+    pub async fn new<T: JsonRpcClient>(
+        provider: Provider<T>,
+        tx: TxHash,
+        use_native_tracer: bool,
+    ) -> Result<Self> {
         let hash = tx.0.into();
 
         let Some(receipt) = provider.get_transaction_receipt(hash).await? else {
@@ -33,7 +37,11 @@ impl Trace {
             bail!("failed to get tx data: {}", hash)
         };
 
-        let query = include_str!("query.js");
+        let query = if use_native_tracer {
+            "stylusTracer"
+        } else {
+            include_str!("query.js")
+        };
         let tracer = GethDebugTracingOptions {
             tracer: Some(GethDebugTracerType::JsTracer(query.to_owned())),
             ..GethDebugTracingOptions::default()