Produce trace on panic (#2788)

**Stack**:
- #2788 ⬅
- #2787


⚠️ *Part of a stack created by [spr](https://github.com/ejoffe/spr). Do
not merge manually using the UI - doing so may have unexpected results.*

Author: ksew1

CHANGELOG.md

@@ -7,6 +7,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Forge
+
+#### Fixed
+
+- `snforge` produces trace for contracts even if they fail or panic (assuming test passed)
+
 ## [0.35.0] - 2024-12-13
 
 ### Forge

crates/cheatnet/src/runtime_extensions/call_to_blockifier_runtime_extension/execution/cairo1_execution.rs

@@ -1,6 +1,5 @@
 use crate::runtime_extensions::call_to_blockifier_runtime_extension::execution::entry_point::{
-    ContractClassEntryPointExecutionResult, EntryPointExecutionErrorWithLastPc, GetLastPc,
-    OnErrorLastPc,
+    ContractClassEntryPointExecutionResult, EntryPointExecutionErrorWithTrace, OnErrorLastPc,
 };
 use crate::runtime_extensions::call_to_blockifier_runtime_extension::CheatnetState;
 use crate::runtime_extensions::cheatable_starknet_runtime_extension::CheatableStarknetRuntimeExtension;
@@ -72,23 +71,14 @@ pub fn execute_entry_point_call_cairo1(
     )
     .on_error_get_last_pc(&mut runner)?;
 
-    let vm_trace = if cheatable_runtime
-        .extension
-        .cheatnet_state
-        .trace_data
-        .is_vm_trace_needed
-    {
-        Some(get_relocated_vm_trace(&runner))
-    } else {
-        None
-    };
+    let trace = get_relocated_vm_trace(&mut runner);
+
     let syscall_counter = cheatable_runtime
         .extended_runtime
         .hint_handler
         .syscall_counter
         .clone();
 
-    let last_pc = runner.get_last_pc();
     let call_info = finalize_execution(
         runner,
         cheatable_runtime.extended_runtime.hint_handler,
@@ -97,15 +87,15 @@ pub fn execute_entry_point_call_cairo1(
         program_extra_data_length,
     )?;
     if call_info.execution.failed {
-        return Err(EntryPointExecutionErrorWithLastPc {
+        return Err(EntryPointExecutionErrorWithTrace {
             source: EntryPointExecutionError::ExecutionFailed {
                 error_data: call_info.execution.retdata.0,
             },
-            last_pc,
+            trace,
         });
     }
 
-    Ok((call_info, syscall_counter, vm_trace))
+    Ok((call_info, syscall_counter, trace))
     // endregion
 }
 

crates/cheatnet/src/runtime_extensions/call_to_blockifier_runtime_extension/execution/entry_point.rs

@@ -31,12 +31,12 @@ use cairo_vm::vm::trace::trace_entry::RelocatedTraceEntry;
 use thiserror::Error;
 use conversions::FromConv;
 use crate::runtime_extensions::call_to_blockifier_runtime_extension::rpc::{AddressOrClassHash, CallResult};
-use crate::runtime_extensions::common::sum_syscall_counters;
+use crate::runtime_extensions::common::{get_relocated_vm_trace, sum_syscall_counters};
 use conversions::string::TryFromHexStr;
 
 pub(crate) type ContractClassEntryPointExecutionResult = Result<
     (CallInfo, SyscallCounter, Option<Vec<RelocatedTraceEntry>>),
-    EntryPointExecutionErrorWithLastPc,
+    EntryPointExecutionErrorWithTrace,
 >;
 
 // blockifier/src/execution/entry_point.rs:180 (CallEntryPoint::execute)
@@ -157,16 +157,17 @@ pub fn execute_call_entry_point(
             );
             Ok(call_info)
         }
-        Err(EntryPointExecutionErrorWithLastPc {
-            source: err,
-            last_pc: pc,
-        }) => {
-            if let Some(pc) = pc {
+        Err(EntryPointExecutionErrorWithTrace { source: err, trace }) => {
+            if let Some(pc) = trace
+                .as_ref()
+                .and_then(|trace| trace.last())
+                .map(|entry| entry.pc)
+            {
                 cheatnet_state
                     .encountered_errors
                     .push(EncounteredError { pc, class_hash });
             }
-            exit_error_call(&err, cheatnet_state, resources, entry_point);
+            exit_error_call(&err, cheatnet_state, resources, entry_point, trace);
             Err(err)
         }
     }
@@ -203,6 +204,7 @@ fn exit_error_call(
     cheatnet_state: &mut CheatnetState,
     resources: &mut ExecutionResources,
     entry_point: &CallEntryPoint,
+    vm_trace: Option<Vec<RelocatedTraceEntry>>,
 ) {
     let identifier = match entry_point.call_type {
         CallType::Call => AddressOrClassHash::ContractAddress(entry_point.storage_address),
@@ -213,7 +215,7 @@ fn exit_error_call(
         Default::default(),
         CallResult::from_err(error, &identifier),
         &[],
-        None,
+        vm_trace,
     );
 }
 
@@ -314,19 +316,19 @@ fn aggregate_syscall_counters(trace: &Rc<RefCell<CallTrace>>) -> SyscallCounter
 
 #[derive(Debug, Error)]
 #[error("{}", source)]
-pub struct EntryPointExecutionErrorWithLastPc {
+pub struct EntryPointExecutionErrorWithTrace {
     pub source: EntryPointExecutionError,
-    pub last_pc: Option<usize>,
+    pub trace: Option<Vec<RelocatedTraceEntry>>,
 }
 
-impl<T> From<T> for EntryPointExecutionErrorWithLastPc
+impl<T> From<T> for EntryPointExecutionErrorWithTrace
 where
     T: Into<EntryPointExecutionError>,
 {
     fn from(value: T) -> Self {
         Self {
             source: value.into(),
-            last_pc: None,
+            trace: None,
         }
     }
 }
@@ -335,37 +337,21 @@ pub(crate) trait OnErrorLastPc<T>: Sized {
     fn on_error_get_last_pc(
         self,
         runner: &mut CairoRunner,
-    ) -> Result<T, EntryPointExecutionErrorWithLastPc>;
+    ) -> Result<T, EntryPointExecutionErrorWithTrace>;
 }
 
 impl<T> OnErrorLastPc<T> for Result<T, EntryPointExecutionError> {
     fn on_error_get_last_pc(
         self,
         runner: &mut CairoRunner,
-    ) -> Result<T, EntryPointExecutionErrorWithLastPc> {
+    ) -> Result<T, EntryPointExecutionErrorWithTrace> {
         match self {
             Err(source) => {
-                let last_pc = runner.get_last_pc();
+                let trace = get_relocated_vm_trace(runner);
 
-                Err(EntryPointExecutionErrorWithLastPc { source, last_pc })
+                Err(EntryPointExecutionErrorWithTrace { source, trace })
             }
             Ok(value) => Ok(value),
         }
     }
 }
-
-pub trait GetLastPc {
-    fn get_last_pc(&mut self) -> Option<usize>;
-}
-
-impl GetLastPc for CairoRunner {
-    fn get_last_pc(&mut self) -> Option<usize> {
-        if self.relocated_trace.is_none() {
-            self.relocate(true).ok()?;
-        }
-        self.relocated_trace
-            .as_ref()
-            .and_then(|trace| trace.last())
-            .map(|entry| entry.pc)
-    }
-}

crates/cheatnet/src/runtime_extensions/common.rs

@@ -18,6 +18,10 @@ pub fn sum_syscall_counters(mut a: SyscallCounter, b: &SyscallCounter) -> Syscal
 }
 
 #[must_use]
-pub fn get_relocated_vm_trace(cairo_runner: &CairoRunner) -> Vec<RelocatedTraceEntry> {
-    cairo_runner.relocated_trace.clone().unwrap()
+pub fn get_relocated_vm_trace(cairo_runner: &mut CairoRunner) -> Option<Vec<RelocatedTraceEntry>> {
+    // if vm execution failed, the trace is not relocated so we need to relocate it
+    if cairo_runner.relocated_trace.is_none() {
+        cairo_runner.relocate(true).ok()?;
+    }
+    cairo_runner.relocated_trace.clone()
 }

crates/cheatnet/src/runtime_extensions/forge_runtime_extension/mod.rs

@@ -591,7 +591,7 @@ pub fn update_top_call_l1_resources(runtime: &mut ForgeRuntime) {
     top_call.borrow_mut().used_l1_resources.l2_l1_message_sizes = all_l2_l1_message_sizes;
 }
 
-pub fn update_top_call_vm_trace(runtime: &mut ForgeRuntime, cairo_runner: &CairoRunner) {
+pub fn update_top_call_vm_trace(runtime: &mut ForgeRuntime, cairo_runner: &mut CairoRunner) {
     let trace_data = &mut runtime
         .extended_runtime
         .extended_runtime
@@ -601,7 +601,7 @@ pub fn update_top_call_vm_trace(runtime: &mut ForgeRuntime, cairo_runner: &Cairo
 
     if trace_data.is_vm_trace_needed {
         trace_data.current_call_stack.top().borrow_mut().vm_trace =
-            Some(get_relocated_vm_trace(cairo_runner));
+            get_relocated_vm_trace(cairo_runner);
     }
 }
 fn add_syscall_resources(

crates/forge-runner/src/running.rs

@@ -211,7 +211,7 @@ pub fn run_test_case(
 
     let run_result =
         match run_assembled_program(&assembled_program, builtins, hints_dict, &mut forge_runtime) {
-            Ok(runner) => {
+            Ok(mut runner) => {
                 let vm_resources_without_inner_calls = runner
                     .get_execution_resources()
                     .unwrap()
@@ -242,7 +242,7 @@ pub fn run_test_case(
                     &runner.relocated_memory,
                 );
 
-                update_top_call_vm_trace(&mut forge_runtime, &runner);
+                update_top_call_vm_trace(&mut forge_runtime, &mut runner);
 
                 Ok((gas_counter, runner.relocated_memory, value))
             }

crates/forge/tests/e2e/build_trace_data.rs

@@ -166,3 +166,32 @@ fn trace_has_deploy_with_no_constructor_phantom_nodes() {
         cairo_annotations::trace_data::CallTraceNode::DeployWithoutConstructor
     );
 }
+
+#[test]
+fn trace_is_produced_even_if_contract_panics() {
+    let temp = setup_package("backtrace_panic");
+    test_runner(&temp)
+        .arg("--save-trace-data")
+        .assert()
+        .success();
+
+    let trace_data = fs::read_to_string(
+        temp.join(TRACE_DIR)
+            .join("backtrace_panic::Test::test_contract_panics.json"),
+    )
+    .unwrap();
+
+    let call_trace: ProfilerCallTrace = serde_json::from_str(&trace_data).unwrap();
+
+    assert_all_execution_info_exists(&call_trace);
+}
+
+fn assert_all_execution_info_exists(trace: &ProfilerCallTrace) {
+    assert!(trace.cairo_execution_info.is_some());
+
+    for trace_node in &trace.nested_calls {
+        if let ProfilerCallTraceNode::EntryPointCall(trace) = trace_node {
+            assert_all_execution_info_exists(trace);
+        }
+    }
+}