Output builtin features for bootloader support

CHANGELOG.md

@@ -4,6 +4,8 @@
 
 * feat: Add cairo1-run output pretty-printing for felts, arrays/spans and dicts [#1630](https://github.com/lambdaclass/cairo-vm/pull/1630)
 
+* feat: output builtin features for bootloader support [#1580](https://github.com/lambdaclass/cairo-vm/pull/1580)
+
 #### [1.0.0-rc1] - 2024-02-23
 
 * Bump `starknet-types-core` dependency version to 0.0.9 [#1628](https://github.com/lambdaclass/cairo-vm/pull/1628)

vm/src/vm/errors/runner_errors.rs

@@ -102,6 +102,8 @@ pub enum RunnerError {
     Trace(#[from] TraceError),
     #[error("EcOp builtin: Invalid Point")]
     InvalidPoint,
+    #[error("Page ({0}) is not on the expected segment {1}")]
+    PageNotOnSegment(Relocatable, usize),
 }
 
 #[cfg(test)]

vm/src/vm/runners/builtin_runner/output.rs

@@ -2,16 +2,27 @@ use crate::stdlib::{collections::HashMap, prelude::*};
 use crate::types::relocatable::{MaybeRelocatable, Relocatable};
 use crate::vm::errors::memory_errors::MemoryError;
 use crate::vm::errors::runner_errors::RunnerError;
-use crate::vm::runners::cairo_pie::{BuiltinAdditionalData, OutputBuiltinAdditionalData};
+use crate::vm::runners::cairo_pie::{
+    Attributes, BuiltinAdditionalData, OutputBuiltinAdditionalData, Pages, PublicMemoryPage,
+};
 use crate::vm::vm_core::VirtualMachine;
 use crate::vm::vm_memory::memory::Memory;
 use crate::vm::vm_memory::memory_segments::MemorySegmentManager;
 
 use super::OUTPUT_BUILTIN_NAME;
 
+#[derive(Debug, Clone, PartialEq)]
+pub struct OutputBuiltinState {
+    pub base: usize,
+    pub pages: Pages,
+    pub attributes: Attributes,
+}
+
 #[derive(Debug, Clone)]
 pub struct OutputBuiltinRunner {
     base: usize,
+    pub(crate) pages: Pages,
+    pub(crate) attributes: Attributes,
     pub(crate) stop_ptr: Option<usize>,
     pub(crate) included: bool,
 }
@@ -20,11 +31,21 @@ impl OutputBuiltinRunner {
     pub fn new(included: bool) -> OutputBuiltinRunner {
         OutputBuiltinRunner {
             base: 0,
+            pages: HashMap::default(),
+            attributes: HashMap::default(),
             stop_ptr: None,
             included,
         }
     }
 
+    pub fn new_state(&mut self, base: usize, included: bool) {
+        self.base = base;
+        self.pages = HashMap::default();
+        self.attributes = HashMap::default();
+        self.stop_ptr = None;
+        self.included = included;
+    }
+
     pub fn initialize_segments(&mut self, segments: &mut MemorySegmentManager) {
         self.base = segments.add().segment_index as usize // segments.add() always returns a positive index
     }
@@ -110,14 +131,64 @@ impl OutputBuiltinRunner {
 
     pub fn get_additional_data(&self) -> BuiltinAdditionalData {
         BuiltinAdditionalData::Output(OutputBuiltinAdditionalData {
-            pages: HashMap::default(),
-            attributes: HashMap::default(),
+            pages: self.pages.clone(),
+            attributes: self.attributes.clone(),
         })
     }
 
     pub(crate) fn set_stop_ptr_offset(&mut self, offset: usize) {
         self.stop_ptr = Some(offset)
     }
+
+    pub fn set_state(&mut self, new_state: OutputBuiltinState) {
+        self.base = new_state.base;
+        self.pages = new_state.pages;
+        self.attributes = new_state.attributes;
+    }
+
+    pub fn get_state(&mut self) -> OutputBuiltinState {
+        OutputBuiltinState {
+            base: self.base,
+            pages: self.pages.clone(),
+            attributes: self.attributes.clone(),
+        }
+    }
+
+    pub fn add_page(
+        &mut self,
+        page_id: usize,
+        page_start: Relocatable,
+        page_size: usize,
+    ) -> Result<(), RunnerError> {
+        if page_start.segment_index as usize != self.base {
+            return Err(RunnerError::PageNotOnSegment(page_start, self.base));
+        }
+
+        self.pages.insert(
+            page_id,
+            PublicMemoryPage {
+                start: page_start.offset,
+                size: page_size,
+            },
+        );
+
+        Ok(())
+    }
+
+    pub fn get_public_memory(&self) -> Result<Vec<(usize, usize)>, RunnerError> {
+        let size = self
+            .stop_ptr
+            .ok_or(RunnerError::NoStopPointer(Box::new(OUTPUT_BUILTIN_NAME)))?;
+
+        let mut public_memory: Vec<(usize, usize)> = (0..size).map(|i| (i, 0)).collect();
+        for (page_id, page) in self.pages.iter() {
+            for index in 0..page.size {
+                public_memory[page.start + index].1 = *page_id;
+            }
+        }
+
+        Ok(public_memory)
+    }
 }
 
 impl Default for OutputBuiltinRunner {
@@ -463,4 +534,115 @@ mod tests {
         let memory = memory![((0, 0), 0), ((0, 1), 1), ((0, 2), 2), ((0, 3), 3)];
         assert!(builtin.air_private_input(&memory).is_empty());
     }
+
+    #[test]
+    fn set_state() {
+        let mut builtin = OutputBuiltinRunner::new(true);
+        assert_eq!(builtin.base, 0);
+
+        let new_state = OutputBuiltinState {
+            base: 10,
+            pages: HashMap::from([(1, PublicMemoryPage { start: 0, size: 3 })]),
+            attributes: HashMap::from([("gps_fact_topology".to_string(), vec![0, 2, 0])]),
+        };
+        builtin.set_state(new_state.clone());
+
+        assert_eq!(builtin.base, new_state.base);
+        assert_eq!(builtin.pages, new_state.pages);
+        assert_eq!(builtin.attributes, new_state.attributes);
+
+        let state = builtin.get_state();
+        assert_eq!(state, new_state);
+    }
+
+    #[test]
+    fn new_state() {
+        let mut builtin = OutputBuiltinRunner {
+            base: 10,
+            pages: HashMap::from([(1, PublicMemoryPage { start: 0, size: 3 })]),
+            attributes: HashMap::from([("gps_fact_topology".to_string(), vec![0, 2, 0])]),
+            stop_ptr: Some(10),
+            included: true,
+        };
+
+        let new_base = 11;
+        let new_included = false;
+        builtin.new_state(new_base, new_included);
+
+        assert_eq!(builtin.base, new_base);
+        assert!(builtin.pages.is_empty());
+        assert!(builtin.attributes.is_empty());
+        assert_eq!(builtin.stop_ptr, None);
+        assert_eq!(builtin.included, new_included);
+    }
+
+    #[test]
+    fn add_page() {
+        let mut builtin = OutputBuiltinRunner::new(true);
+        assert_eq!(
+            builtin.add_page(
+                1,
+                Relocatable {
+                    segment_index: builtin.base() as isize,
+                    offset: 0
+                },
+                3
+            ),
+            Ok(())
+        );
+
+        assert_eq!(
+            builtin.pages,
+            HashMap::from([(1, PublicMemoryPage { start: 0, size: 3 }),])
+        )
+    }
+
+    #[test]
+    fn add_page_wrong_segment() {
+        let mut builtin = OutputBuiltinRunner::new(true);
+        let page_start = Relocatable {
+            segment_index: 18,
+            offset: 0,
+        };
+
+        let result = builtin.add_page(1, page_start, 3);
+        assert!(
+            matches!(result, Err(RunnerError::PageNotOnSegment(relocatable, base)) if relocatable == page_start && base == builtin.base())
+        )
+    }
+
+    #[test]
+    fn get_public_memory() {
+        let mut builtin = OutputBuiltinRunner::new(true);
+
+        builtin
+            .add_page(
+                1,
+                Relocatable {
+                    segment_index: builtin.base() as isize,
+                    offset: 2,
+                },
+                2,
+            )
+            .unwrap();
+
+        builtin
+            .add_page(
+                2,
+                Relocatable {
+                    segment_index: builtin.base() as isize,
+                    offset: 4,
+                },
+                3,
+            )
+            .unwrap();
+
+        builtin.stop_ptr = Some(7);
+
+        let public_memory = builtin.get_public_memory().unwrap();
+        assert_eq!(
+            public_memory,
+            vec![(0, 0), (1, 0), (2, 1), (3, 1), (4, 2), (5, 2), (6, 2)]
+        );
+    }
 }

vm/src/vm/runners/cairo_runner.rs

@@ -55,7 +55,7 @@ use num_traits::{ToPrimitive, Zero};
 use serde::{Deserialize, Serialize};
 
 use super::{
-    builtin_runner::{KeccakBuiltinRunner, PoseidonBuiltinRunner, OUTPUT_BUILTIN_NAME},
+    builtin_runner::{KeccakBuiltinRunner, PoseidonBuiltinRunner},
     cairo_pie::{self, CairoPie, CairoPieMetadata, CairoPieVersion},
 };
 
@@ -1090,8 +1090,8 @@ impl CairoRunner {
             let (_, size) = builtin_runner
                 .get_used_cells_and_allocated_size(vm)
                 .map_err(RunnerError::FinalizeSegements)?;
-            if builtin_runner.name() == OUTPUT_BUILTIN_NAME {
-                let public_memory = (0..size).map(|i| (i, 0)).collect();
+            if let BuiltinRunner::Output(output_builtin) = builtin_runner {
+                let public_memory = output_builtin.get_public_memory()?;
                 vm.segments
                     .finalize(Some(size), builtin_runner.base(), Some(&public_memory))
             } else {