Partially Revert Optimize trace relocation #906

CHANGELOG.md

@@ -2,11 +2,17 @@
 
 #### Upcoming Changes
 
+* BREAKING: Partially Revert `Optimize trace relocation #906` [#1492](https://github.com/lambdaclass/cairo-vm/pull/1492)
+
+  * Remove methods `VirtualMachine::get_relocated_trace`& `VirtualMachine::relocate_trace`.
+  * Add `relocated_trace` field  & `relocate_trace` method to `CairoRunner`.
+  * Swap `TraceEntry` for `RelocatedTraceEntry` type in `write_encoded_trace` & `PublicInput::new` signatures.
+  * Now takes into account the program counter's segment index when building the execution trace instead of assuming it to be 0.
+
 * feat: add debugging capabilities behind `print` feature flag. [#1476](https://github.com/lambdaclass/cairo-vm/pull/1476)
 
 #### [0.9.1] - 2023-11-16
 
-
 * chore: bump `cairo-lang-` dependencies to 2.3.1 [#1482](https://github.com/lambdaclass/cairo-vm/pull/1482), [#1483](https://github.com/lambdaclass/cairo-vm/pull/1483)
 
 * feat: Make PublicInput fields public [#1474](https://github.com/lambdaclass/cairo-vm/pull/1474)

cairo-vm-cli/src/main.rs

@@ -149,7 +149,10 @@ fn run(args: impl Iterator<Item = String>) -> Result<(), Error> {
     }
 
     if let Some(trace_path) = args.trace_file {
-        let relocated_trace = vm.get_relocated_trace()?;
+        let relocated_trace = cairo_runner
+            .relocated_trace
+            .as_ref()
+            .ok_or(Error::Trace(TraceError::TraceNotRelocated))?;
 
         let trace_file = std::fs::File::create(trace_path)?;
         let mut trace_writer =

cairo1-run/src/main.rs

@@ -294,12 +294,14 @@ fn run(args: impl Iterator<Item = String>) -> Result<Vec<MaybeRelocatable>, Erro
     runner.relocate(&mut vm, true)?;
 
     if let Some(trace_path) = args.trace_file {
-        let relocated_trace = vm.get_relocated_trace()?;
+        let relocated_trace = runner
+            .relocated_trace
+            .ok_or(Error::Trace(TraceError::TraceNotRelocated))?;
         let trace_file = std::fs::File::create(trace_path)?;
         let mut trace_writer =
             FileWriter::new(io::BufWriter::with_capacity(3 * 1024 * 1024, trace_file));
 
-        cairo_run::write_encoded_trace(relocated_trace, &mut trace_writer)?;
+        cairo_run::write_encoded_trace(&relocated_trace, &mut trace_writer)?;
         trace_writer.flush()?;
     }
     if let Some(memory_path) = args.memory_file {

vm/src/air_public_input.rs

@@ -10,7 +10,7 @@ use crate::{
     types::layout::CairoLayout,
     vm::{
         errors::{trace_errors::TraceError, vm_errors::VirtualMachineError},
-        trace::trace_entry::TraceEntry,
+        trace::trace_entry::RelocatedTraceEntry,
     },
 };
 
@@ -73,7 +73,7 @@ impl<'a> PublicInput<'a> {
         dyn_layout_params: Option<&'a CairoLayout>,
         public_memory_addresses: &[(usize, usize)],
         memory_segment_addresses: HashMap<&'static str, (usize, usize)>,
-        trace: &[TraceEntry],
+        trace: &[RelocatedTraceEntry],
         rc_limits: (isize, isize),
     ) -> Result<Self, PublicInputError> {
         let memory_entry =

vm/src/cairo_run.rs

@@ -158,7 +158,7 @@ pub struct EncodeTraceError(usize, bincode::error::EncodeError);
 /// Bincode encodes to little endian by default and each trace entry is composed of
 /// 3 usize values that are padded to always reach 64 bit size.
 pub fn write_encoded_trace(
-    relocated_trace: &[crate::vm::trace::trace_entry::TraceEntry],
+    relocated_trace: &[crate::vm::trace::trace_entry::RelocatedTraceEntry],
     dest: &mut impl Writer,
 ) -> Result<(), EncodeTraceError> {
     for (i, entry) in relocated_trace.iter().enumerate() {
@@ -318,11 +318,11 @@ mod tests {
         // relocate memory so we can dump it to file
         assert!(cairo_runner.relocate(&mut vm, false).is_ok());
 
-        let trace_entries = vm.get_relocated_trace().unwrap();
+        let trace_entries = cairo_runner.relocated_trace.unwrap();
         let mut buffer = [0; 24];
         let mut buff_writer = SliceWriter::new(&mut buffer);
         // write cairo_rs vm trace file
-        write_encoded_trace(trace_entries, &mut buff_writer).unwrap();
+        write_encoded_trace(&trace_entries, &mut buff_writer).unwrap();
 
         // compare that the original cairo vm trace file and cairo_rs vm trace files are equal
         assert_eq!(buffer, *expected_encoded_trace);
@@ -369,6 +369,6 @@ mod tests {
             .run_until_pc(end, &mut vm, &mut hint_processor)
             .is_ok());
         assert!(cairo_runner.relocate(&mut vm, false).is_ok());
-        assert!(vm.get_relocated_trace().is_err());
+        assert!(cairo_runner.relocated_trace.is_none());
     }
 }

vm/src/tests/mod.rs

@@ -2,6 +2,7 @@
 use crate::vm::errors::cairo_run_errors::CairoRunError;
 #[cfg(feature = "cairo-1-hints")]
 use crate::vm::runners::cairo_runner::RunResources;
+use crate::vm::trace::trace_entry::RelocatedTraceEntry;
 #[cfg(feature = "cairo-1-hints")]
 use crate::{
     hint_processor::cairo_1_hint_processor::hint_processor::Cairo1HintProcessor,
@@ -22,7 +23,6 @@ use crate::stdlib::prelude::*;
 use crate::{
     cairo_run::{cairo_run, CairoRunConfig},
     hint_processor::builtin_hint_processor::builtin_hint_processor_definition::BuiltinHintProcessor,
-    vm::trace::trace_entry::TraceEntry,
 };
 
 #[cfg(target_arch = "wasm32")]
@@ -93,9 +93,9 @@ pub(self) fn run_program(
         let expected_trace: Vec<_> = trace
             .iter()
             .copied()
-            .map(|(pc, ap, fp)| TraceEntry { pc, ap, fp })
+            .map(|(pc, ap, fp)| RelocatedTraceEntry { pc, ap, fp })
             .collect();
-        let trace = vm.get_relocated_trace().unwrap();
+        let trace = runner.relocated_trace.as_ref().unwrap();
         assert_eq!(trace.len(), expected_trace.len());
         for (entry, expected) in trace.iter().zip(expected_trace.iter()) {
             assert_eq!(entry, expected);

vm/src/utils.rs

@@ -413,7 +413,10 @@ pub mod test_utils {
     pub(crate) use non_continuous_ids_data;
 
     #[track_caller]
-    pub(crate) fn trace_check(actual: &[TraceEntry], expected: &[(usize, usize, usize)]) {
+    pub(crate) fn trace_check(
+        actual: &[TraceEntry],
+        expected: &[(crate::utils::Relocatable, usize, usize)],
+    ) {
         assert_eq!(actual.len(), expected.len());
         for (entry, expected) in actual.iter().zip(expected.iter()) {
             assert_eq!(&(entry.pc, entry.ap, entry.fp), expected);
@@ -704,22 +707,29 @@ mod test {
     fn assert_trace() {
         let trace = vec![
             TraceEntry {
-                pc: 2,
+                pc: (0, 2).into(),
                 ap: 7,
                 fp: 1,
             },
             TraceEntry {
-                pc: 5,
+                pc: (0, 5).into(),
                 ap: 1,
                 fp: 0,
             },
             TraceEntry {
-                pc: 9,
+                pc: (0, 9).into(),
                 ap: 5,
                 fp: 7,
             },
         ];
-        trace_check(&trace, &[(2, 7, 1), (5, 1, 0), (9, 5, 7)]);
+        trace_check(
+            &trace,
+            &[
+                ((0, 2).into(), 7, 1),
+                ((0, 5).into(), 1, 0),
+                ((0, 9).into(), 5, 7),
+            ],
+        );
     }
 
     #[test]