Merge Cairo Pie extra segments into one segment (#1960)

* implement merge_extra_segments

* revert changes in CairoPieMemory

* format

* changelog

* add reference

* serialize memory according to the segment_offsets values

* format

* fature std for test

* change for

* clippy

* clippy

* add merge_extra_segments as argument

* docs

* typo in changelog

* revert making self mut in write_zip_file method

* forgot to update metadata in write_zip_file

* make merge_extra_segments method private

* clippy

* reviews

* format

* reviews

* format

* reviews

* clippy

Author: FrancoGiachetta

CHANGELOG.md

@@ -2,6 +2,9 @@
 
 #### Upcoming Changes
 
+* fix: Fix Cairo Pie limiting the number of segments to 2^16 [#1960](https://github.com/lambdaclass/cairo-vm/pull/1960)
+  * Implement `merge_extra_segments`
+
 * feat: implement an opcode that computes QM31 arithmetics (add, sub, mul, div) in the VM [#1938](https://github.com/lambdaclass/cairo-vm/pull/1938)
 
 * feat: add functions that compute packed reduced qm31 arithmetics to `math_utils` [#1944](https://github.com/lambdaclass/cairo-vm/pull/1944)

cairo-vm-cli/src/main.rs

@@ -69,6 +69,8 @@ struct Args {
         conflicts_with_all = ["proof_mode", "air_private_input", "air_public_input"]
     )]
     cairo_pie_output: Option<String>,
+    #[structopt(long = "merge_extra_segments")]
+    merge_extra_segments: bool,
     #[structopt(long = "allow_missing_builtins")]
     allow_missing_builtins: Option<bool>,
     #[structopt(long = "tracer")]
@@ -273,7 +275,7 @@ fn run(args: impl Iterator<Item = String>) -> Result<(), Error> {
         cairo_runner
             .get_cairo_pie()
             .map_err(CairoRunError::Runner)?
-            .write_zip_file(file_path)?
+            .write_zip_file(file_path, args.merge_extra_segments)?
     }
 
     Ok(())

cairo1-run/src/main.rs

@@ -48,6 +48,8 @@ struct Args {
         conflicts_with_all = ["proof_mode", "air_private_input", "air_public_input"]
     )]
     cairo_pie_output: Option<PathBuf>,
+    #[clap(long = "merge_extra_segments", value_parser)]
+    merge_extra_segments: bool,
     // Arguments should be spaced, with array elements placed between brackets
     // For example " --args '1 2 [1 2 3]'" will yield 3 arguments, with the last one being an array of 3 elements
     #[clap(long = "args", default_value = "", value_parser=process_args, conflicts_with = "args_file")]
@@ -234,7 +236,9 @@ fn run(args: impl Iterator<Item = String>) -> Result<Option<String>, Error> {
     }
 
     if let Some(ref file_path) = args.cairo_pie_output {
-        runner.get_cairo_pie()?.write_zip_file(file_path)?
+        runner
+            .get_cairo_pie()?
+            .write_zip_file(file_path, args.merge_extra_segments)?
     }
 
     if let Some(trace_path) = args.trace_file {

vm/src/vm/runners/cairo_pie.rs

@@ -301,17 +301,35 @@ impl CairoPie {
     }
 
     #[cfg(feature = "std")]
-    pub fn write_zip_file(&self, file_path: &Path) -> Result<(), std::io::Error> {
+    pub fn write_zip_file(
+        &self,
+        file_path: &Path,
+        merge_extra_segments: bool,
+    ) -> Result<(), std::io::Error> {
+        let mut metadata = self.metadata.clone();
+
+        let segment_offsets = if merge_extra_segments {
+            if let Some((segment, segment_offsets)) = self.merge_extra_segments() {
+                metadata.extra_segments = vec![segment];
+                Some(segment_offsets)
+            } else {
+                None
+            }
+        } else {
+            None
+        };
+
         let file = File::create(file_path)?;
         let mut zip_writer = ZipWriter::new(file);
         let options =
             zip::write::FileOptions::default().compression_method(zip::CompressionMethod::Deflated);
+
         zip_writer.start_file("version.json", options)?;
         serde_json::to_writer(&mut zip_writer, &self.version)?;
         zip_writer.start_file("metadata.json", options)?;
-        serde_json::to_writer(&mut zip_writer, &self.metadata)?;
+        serde_json::to_writer(&mut zip_writer, &metadata)?;
         zip_writer.start_file("memory.bin", options)?;
-        zip_writer.write_all(&self.memory.to_bytes())?;
+        zip_writer.write_all(&self.memory.to_bytes(segment_offsets))?;
         zip_writer.start_file("additional_data.json", options)?;
         serde_json::to_writer(&mut zip_writer, &self.additional_data)?;
         zip_writer.start_file("execution_resources.json", options)?;
@@ -372,6 +390,48 @@ impl CairoPie {
 
         Self::from_zip_archive(zip)
     }
+
+    // Heavily inspired in:
+    // https://github.com/starkware-libs/cairo-lang/blob/8276ac35830148a397e1143389f23253c8b80e93/src/starkware/cairo/lang/vm/cairo_pie.py#L286-L306
+    /// Merges `extra_segments` to a single segment.
+    ///
+    /// Returns a tuple with the new `extra_segments` (containing just the merged segment)
+    /// and a HashMap with the old segment indices mapped to their new offset in the new segment
+    #[cfg(feature = "std")]
+    fn merge_extra_segments(&self) -> Option<(SegmentInfo, HashMap<usize, Relocatable>)> {
+        if self.metadata.extra_segments.is_empty() {
+            return None;
+        }
+
+        let new_index = self.metadata.extra_segments[0].index;
+        let mut accumulated_size = 0;
+        let offsets: HashMap<usize, Relocatable> = self
+            .metadata
+            .extra_segments
+            .iter()
+            .map(|seg| {
+                let value = (
+                    seg.index as usize,
+                    Relocatable {
+                        segment_index: new_index,
+                        offset: accumulated_size,
+                    },
+                );
+
+                accumulated_size += seg.size;
+
+                value
+            })
+            .collect();
+
+        Some((
+            SegmentInfo {
+                index: new_index,
+                size: accumulated_size,
+            },
+            offsets,
+        ))
+    }
 }
 
 pub(super) mod serde_impl {
@@ -595,26 +655,50 @@ pub(super) mod serde_impl {
     }
 
     impl CairoPieMemory {
-        pub fn to_bytes(&self) -> Vec<u8> {
+        /// Relocates a `Relocatable` value, which represented by its
+        /// index and offset, according to a given segment offsets
+        fn relocate_value(
+            index: usize,
+            offset: usize,
+            segment_offsets: &Option<HashMap<usize, Relocatable>>,
+        ) -> (usize, usize) {
+            segment_offsets
+                .as_ref()
+                .and_then(|offsets| offsets.get(&index))
+                .map(|relocatable| {
+                    (
+                        relocatable.segment_index as usize,
+                        relocatable.offset + offset,
+                    )
+                })
+                .unwrap_or((index, offset))
+        }
+
+        pub fn to_bytes(&self, seg_offsets: Option<HashMap<usize, Relocatable>>) -> Vec<u8> {
             // Missing segment and memory holes can be ignored
             // as they can be inferred by the address on the prover side
             let values = &self.0;
             let mem_cap = values.len() * ADDR_BYTE_LEN + values.len() * FIELD_BYTE_LEN;
             let mut res = Vec::with_capacity(mem_cap);
 
             for ((segment, offset), value) in values.iter() {
-                let mem_addr = ADDR_BASE + *segment as u64 * OFFSET_BASE + *offset as u64;
+                let (segment, offset) = Self::relocate_value(*segment, *offset, &seg_offsets);
+                let mem_addr = ADDR_BASE + segment as u64 * OFFSET_BASE + offset as u64;
                 res.extend_from_slice(mem_addr.to_le_bytes().as_ref());
                 match value {
                     // Serializes RelocatableValue(little endian):
                     // 1bit |   SEGMENT_BITS |   OFFSET_BITS
                     // 1    |     segment    |   offset
                     MaybeRelocatable::RelocatableValue(rel_val) => {
+                        let (segment, offset) = Self::relocate_value(
+                            rel_val.segment_index as usize,
+                            rel_val.offset,
+                            &seg_offsets,
+                        );
                         let reloc_base = BigUint::from_str_radix(RELOCATE_BASE, 16).unwrap();
                         let reloc_value = reloc_base
-                            + BigUint::from(rel_val.segment_index as usize)
-                                * BigUint::from(OFFSET_BASE)
-                            + BigUint::from(rel_val.offset);
+                            + BigUint::from(segment) * BigUint::from(OFFSET_BASE)
+                            + BigUint::from(offset);
                         res.extend_from_slice(reloc_value.to_bytes_le().as_ref());
                     }
                     // Serializes Int(little endian):
@@ -780,7 +864,14 @@ pub(super) mod serde_impl {
 #[cfg(test)]
 mod test {
     #[cfg(feature = "std")]
-    use rstest::rstest;
+    use {
+        crate::{
+            cairo_run::CairoRunConfig,
+            hint_processor::builtin_hint_processor::builtin_hint_processor_definition::BuiltinHintProcessor,
+            types::layout_name::LayoutName,
+        },
+        rstest::rstest,
+    };
 
     use super::*;
 
@@ -858,11 +949,6 @@ mod test {
     #[case(include_bytes!("../../../../cairo_programs/bitwise_output.json"), "bitwise")]
     #[case(include_bytes!("../../../../cairo_programs/value_beyond_segment.json"), "relocate_beyond")]
     fn read_write_pie_zip(#[case] program_content: &[u8], #[case] identifier: &str) {
-        use crate::{
-            cairo_run::CairoRunConfig,
-            hint_processor::builtin_hint_processor::builtin_hint_processor_definition::BuiltinHintProcessor,
-            types::layout_name::LayoutName,
-        };
         // Run a program to obtain the CairoPie
         let cairo_pie = {
             let cairo_run_config = CairoRunConfig {
@@ -880,12 +966,157 @@ mod test {
         // Serialize the CairoPie into a zip file
         let filename = format!("temp_file_{}", identifier); // Identifier used to avoid name clashes
         let file_path = Path::new(&filename);
-        cairo_pie.write_zip_file(file_path).unwrap();
+        cairo_pie.write_zip_file(file_path, false).unwrap();
         // Deserialize the zip file
         let deserialized_pie = CairoPie::read_zip_file(file_path).unwrap();
         // Check that both pies are equal
         assert_eq!(cairo_pie, deserialized_pie);
         // Remove zip file created by the test
         std::fs::remove_file(file_path).unwrap();
     }
+
+    #[test]
+    #[cfg(feature = "std")]
+    fn cairo_pie_with_extra_segments() {
+        let program_content = include_bytes!("../../../../cairo_programs/fibonacci.json");
+        let mut cairo_pie = {
+            let cairo_run_config = CairoRunConfig {
+                layout: LayoutName::starknet_with_keccak,
+                ..Default::default()
+            };
+            let runner = crate::cairo_run::cairo_run(
+                program_content,
+                &cairo_run_config,
+                &mut BuiltinHintProcessor::new_empty(),
+            )
+            .unwrap();
+            runner.get_cairo_pie().unwrap()
+        };
+
+        cairo_pie.metadata.extra_segments = vec![
+            SegmentInfo { index: 8, size: 10 },
+            SegmentInfo { index: 9, size: 20 },
+        ];
+        let memory = CairoPieMemory(vec![
+            (
+                (3, 4),
+                MaybeRelocatable::RelocatableValue(Relocatable {
+                    segment_index: 6,
+                    offset: 7,
+                }),
+            ),
+            (
+                (8, 0),
+                MaybeRelocatable::RelocatableValue(Relocatable {
+                    segment_index: 8,
+                    offset: 4,
+                }),
+            ),
+            (
+                (9, 3),
+                MaybeRelocatable::RelocatableValue(Relocatable {
+                    segment_index: 9,
+                    offset: 7,
+                }),
+            ),
+        ]);
+
+        cairo_pie.memory = memory;
+
+        let file_path = Path::new("merge_extra_segments_test");
+
+        cairo_pie.write_zip_file(file_path, true).unwrap();
+
+        let result_cairo_pie = CairoPie::read_zip_file(file_path).unwrap();
+
+        std::fs::remove_file(file_path).unwrap();
+
+        assert_eq!(
+            result_cairo_pie.metadata.extra_segments,
+            vec![SegmentInfo { index: 8, size: 30 }]
+        );
+        assert_eq!(
+            result_cairo_pie.memory,
+            CairoPieMemory(vec![
+                (
+                    (3, 4),
+                    MaybeRelocatable::RelocatableValue(Relocatable {
+                        segment_index: 6,
+                        offset: 7
+                    })
+                ),
+                (
+                    (8, 0),
+                    MaybeRelocatable::RelocatableValue(Relocatable {
+                        segment_index: 8,
+                        offset: 4
+                    })
+                ),
+                (
+                    (8, 13),
+                    MaybeRelocatable::RelocatableValue(Relocatable {
+                        segment_index: 8,
+                        offset: 17
+                    })
+                ),
+            ])
+        )
+    }
+
+    #[test]
+    #[cfg(feature = "std")]
+    fn cairo_pie_without_extra_segments() {
+        let program_content = include_bytes!("../../../../cairo_programs/fibonacci.json");
+        let mut cairo_pie = {
+            let cairo_run_config = CairoRunConfig {
+                layout: LayoutName::starknet_with_keccak,
+                ..Default::default()
+            };
+            let runner = crate::cairo_run::cairo_run(
+                program_content,
+                &cairo_run_config,
+                &mut BuiltinHintProcessor::new_empty(),
+            )
+            .unwrap();
+            runner.get_cairo_pie().unwrap()
+        };
+
+        cairo_pie.metadata.extra_segments = vec![];
+        let memory = CairoPieMemory(vec![
+            (
+                (3, 4),
+                MaybeRelocatable::RelocatableValue(Relocatable {
+                    segment_index: 6,
+                    offset: 7,
+                }),
+            ),
+            (
+                (8, 0),
+                MaybeRelocatable::RelocatableValue(Relocatable {
+                    segment_index: 8,
+                    offset: 4,
+                }),
+            ),
+            (
+                (9, 3),
+                MaybeRelocatable::RelocatableValue(Relocatable {
+                    segment_index: 9,
+                    offset: 7,
+                }),
+            ),
+        ]);
+
+        cairo_pie.memory = memory.clone();
+
+        let file_path = Path::new("merge_without_extra_segments_test");
+
+        cairo_pie.write_zip_file(file_path, true).unwrap();
+
+        let result_cairo_pie = CairoPie::read_zip_file(file_path).unwrap();
+
+        std::fs::remove_file(file_path).unwrap();
+
+        assert_eq!(result_cairo_pie.metadata.extra_segments, vec![]);
+        assert_eq!(result_cairo_pie.memory, memory)
+    }
 }