[DNM] Trace Dump Feature

Cargo.toml

@@ -55,6 +55,7 @@ with-cheatcode = []
 with-debug-utils = []
 with-mem-tracing = []
 with-segfault-catcher = []
+with-trace-dump = ["dep:sierra-emu"]
 
 # the aquamarine dep is only used in docs and cannot be detected as used by cargo udeps
 [package.metadata.cargo-udeps.ignore]
@@ -127,6 +128,9 @@ num-integer = "0.1.46"
 rand = "0.9.0"
 starknet-curve = "0.5.1"
 
+# trace dump
+sierra-emu = { git = "https://github.com/lambdaclass/sierra-emu", branch = "main", optional = true}
+
 [dev-dependencies]
 cairo-vm = { git = "https://github.com/lambdaclass/cairo-vm", rev = "368e3fb311601a33ff8945e784eaa332f1fd499a", features = ["cairo-1-hints"] }
 cairo-lang-semantic = { version = "=2.12.0-dev.0", features = ["testing"] }

src/bin/cairo-native-run.rs

@@ -11,7 +11,7 @@ use cairo_native::{
     starknet_stub::StubSyscallHandler,
 };
 use clap::{Parser, ValueEnum};
-use std::path::PathBuf;
+use std::{path::PathBuf, fs::File};
 use tracing_subscriber::{EnvFilter, FmtSubscriber};
 use utils::{find_function, result_to_runresult};
 
@@ -45,6 +45,14 @@ struct Args {
     /// Optimization level, Valid: 0, 1, 2, 3. Values higher than 3 are considered as 3.
     #[arg(short = 'O', long, default_value_t = 0)]
     opt_level: u8,
+
+    #[cfg(feature = "with-trace-dump")]
+    #[arg(long)]
+    trace_output: Option<PathBuf>,
+
+    #[cfg(feature = "with-trace-dump")]
+    #[arg(long)]
+    sierra_output: Option<PathBuf>,
 }
 
 fn main() -> anyhow::Result<()> {
@@ -70,6 +78,13 @@ fn main() -> anyhow::Result<()> {
     )?
     .program;
 
+    #[cfg(feature = "with-trace-dump")]
+    if let Some(sierra_output) = args.sierra_output {
+        use std::io::Write;
+        let mut file = File::create(sierra_output).unwrap();
+        write!(file, "{}", &sierra_program).unwrap();
+    }
+
     let native_context = NativeContext::new();
 
     // Compile the sierra program into a MLIR module.
@@ -81,6 +96,16 @@ fn main() -> anyhow::Result<()> {
         RunMode::Aot => {
             let executor =
                 AotNativeExecutor::from_native_module(native_module, args.opt_level.into())?;
+
+            #[cfg(feature = "with-trace-dump")]
+            {
+                use cairo_native::metadata::trace_dump::TraceBinding;
+                if let Some(trace_id) = executor.find_symbol_ptr(TraceBinding::TraceId.symbol()) {
+                    let trace_id = trace_id.cast::<u64>();
+                    unsafe { *trace_id = 0 };
+                }
+            }
+
             Box::new(move |function_id, args, gas, syscall_handler| {
                 executor.invoke_dynamic_with_syscall_handler(
                     function_id,
@@ -93,6 +118,16 @@ fn main() -> anyhow::Result<()> {
         RunMode::Jit => {
             let executor =
                 JitNativeExecutor::from_native_module(native_module, args.opt_level.into())?;
+
+            #[cfg(feature = "with-trace-dump")]
+            {
+                use cairo_native::metadata::trace_dump::TraceBinding;
+                if let Some(trace_id) = executor.find_symbol_ptr(TraceBinding::TraceId.symbol()) {
+                    let trace_id = trace_id.cast::<u64>();
+                    unsafe { *trace_id = 0 };
+                }
+            }
+
             Box::new(move |function_id, args, gas, syscall_handler| {
                 executor.invoke_dynamic_with_syscall_handler(
                     function_id,
@@ -104,6 +139,21 @@ fn main() -> anyhow::Result<()> {
         }
     };
 
+    #[cfg(feature = "with-trace-dump")]
+    {
+        use cairo_lang_sierra::program_registry::ProgramRegistry;
+        use cairo_native::runtime::trace_dump::{TraceDump, TRACE_DUMP};
+        use cairo_native::types::TypeBuilder;
+
+        TRACE_DUMP.lock().unwrap().insert(
+            0,
+            TraceDump::new(
+                ProgramRegistry::new(&sierra_program).unwrap(),
+                |ty, registry| ty.layout(registry).unwrap(),
+            ),
+        );
+    }
+
     let gas_metadata =
         GasMetadata::new(&sierra_program, Some(MetadataComputationConfig::default())).unwrap();
 
@@ -139,5 +189,20 @@ fn main() -> anyhow::Result<()> {
         println!("Remaining gas: {gas}");
     }
 
+    #[cfg(feature = "with-trace-dump")]
+    if let Some(trace_output) = args.trace_output {
+        let traces = cairo_native::runtime::trace_dump::TRACE_DUMP
+            .lock()
+            .unwrap();
+        assert_eq!(traces.len(), 1);
+
+        let trace_dump = traces.values().next().unwrap();
+        serde_json::to_writer_pretty(
+            std::fs::File::create(trace_output).unwrap(),
+            &trace_dump.trace,
+        )
+        .unwrap();
+    }
+
     Ok(())
 }