Implement extend_additional_data for BuiltinRunner

CHANGELOG.md

@@ -2,6 +2,8 @@
 
 #### Upcoming Changes
 
+* feat: Implement `extend_additional_data` for `BuiltinRunner`[#1726](https://github.com/lambdaclass/cairo-vm/pull/1726)
+
 * refactor(BREAKING): Use `BuiltinName` enum instead of string representation [#1722](https://github.com/lambdaclass/cairo-vm/pull/1722)
   * `BuiltinName` moved from `crate::serde::deserialize_program` module to `crate::types::builtin_name`.
     * Implement `BuiltinName` methods `to_str`, `to_str_with_suffix`, `from_str` & `from_str_with_suffix`.

vm/src/vm/errors/runner_errors.rs

@@ -126,6 +126,8 @@ pub enum RunnerError {
     MissingBuiltin(BuiltinName),
     #[error("The stop pointer of the missing builtin {0} must be 0")]
     MissingBuiltinStopPtrNotZero(BuiltinName),
+    #[error("{0}: Invalid additional data")]
+    InvalidAdditionalData(BuiltinName),
 }
 
 #[cfg(test)]

vm/src/vm/runners/builtin_runner/hash.rs

@@ -1,5 +1,6 @@
 use crate::air_private_input::{PrivateInput, PrivateInputPair};
 use crate::stdlib::{cell::RefCell, prelude::*};
+use crate::types::builtin_name::BuiltinName;
 use crate::types::instance_definitions::pedersen_instance_def::CELLS_PER_HASH;
 use crate::types::relocatable::{MaybeRelocatable, Relocatable};
 use crate::vm::errors::memory_errors::MemoryError;
@@ -118,6 +119,28 @@ impl HashBuiltinRunner {
         BuiltinAdditionalData::Hash(verified_addresses)
     }
 
+    pub fn extend_additional_data(
+        &mut self,
+        additional_data: &BuiltinAdditionalData,
+    ) -> Result<(), RunnerError> {
+        let additional_data = match additional_data {
+            BuiltinAdditionalData::Hash(d) => d,
+            _ => return Err(RunnerError::InvalidAdditionalData(BuiltinName::pedersen)),
+        };
+        let mut verified_addresses = self.verified_addresses.borrow_mut();
+        for addr in additional_data {
+            if addr.segment_index != self.base as isize {
+                return Err(RunnerError::InvalidAdditionalData(BuiltinName::pedersen));
+            }
+            // Mark offset as verified
+            if addr.offset > verified_addresses.len() {
+                verified_addresses.resize(addr.offset, false);
+            }
+            verified_addresses.insert(addr.offset, true)
+        }
+        Ok(())
+    }
+
     pub fn air_private_input(&self, memory: &Memory) -> Vec<PrivateInput> {
         let mut private_inputs = vec![];
         if let Some(segment) = memory.data.get(self.base) {
@@ -442,7 +465,7 @@ mod tests {
     }
 
     #[test]
-    fn get_additional_info() {
+    fn get_additional_data() {
         let mut builtin = HashBuiltinRunner::new(Some(1), true);
         let verified_addresses = vec![Relocatable::from((0, 3)), Relocatable::from((0, 6))];
         builtin.verified_addresses =
@@ -453,6 +476,17 @@ mod tests {
         )
     }
 
+    #[test]
+    fn get_and_extend_additional_data() {
+        let mut builtin_a = HashBuiltinRunner::new(Some(1), true);
+        builtin_a.verified_addresses =
+            RefCell::new(vec![false, false, false, true, false, false, true]);
+        let additional_data = builtin_a.get_additional_data();
+        let mut builtin_b = HashBuiltinRunner::new(Some(1), true);
+        builtin_b.extend_additional_data(&additional_data).unwrap();
+        assert_eq!(builtin_a.verified_addresses, builtin_b.verified_addresses);
+    }
+
     #[test]
     #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
     fn get_air_private_input() {

vm/src/vm/runners/builtin_runner/mod.rs

@@ -508,6 +508,18 @@ impl BuiltinRunner {
         }
     }
 
+    pub fn extend_additional_data(
+        &mut self,
+        additional_data: &BuiltinAdditionalData,
+    ) -> Result<(), RunnerError> {
+        match self {
+            BuiltinRunner::Hash(builtin) => builtin.extend_additional_data(additional_data),
+            BuiltinRunner::Output(builtin) => builtin.extend_additional_data(additional_data),
+            BuiltinRunner::Signature(builtin) => builtin.extend_additional_data(additional_data),
+            _ => Ok(()),
+        }
+    }
+
     // Returns information about the builtin that should be added to the AIR private input.
     pub fn air_private_input(&self, segments: &MemorySegmentManager) -> Vec<PrivateInput> {
         match self {

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

@@ -124,6 +124,19 @@ impl OutputBuiltinRunner {
         })
     }
 
+    pub fn extend_additional_data(
+        &mut self,
+        additional_data: &BuiltinAdditionalData,
+    ) -> Result<(), RunnerError> {
+        let additional_data = match additional_data {
+            BuiltinAdditionalData::Output(d) => d,
+            _ => return Err(RunnerError::InvalidAdditionalData(BuiltinName::output)),
+        };
+        self.pages.extend(additional_data.pages.clone());
+        self.attributes.extend(additional_data.attributes.clone());
+        Ok(())
+    }
+
     pub(crate) fn set_stop_ptr_offset(&mut self, offset: usize) {
         self.stop_ptr = Some(offset)
     }
@@ -456,7 +469,7 @@ mod tests {
     }
 
     #[test]
-    fn get_additional_info_no_pages_no_attributes() {
+    fn get_additional_data_no_pages_no_attributes() {
         let builtin = OutputBuiltinRunner::new(true);
         assert_eq!(
             builtin.get_additional_data(),
@@ -599,4 +612,26 @@ mod tests {
             vec![(0, 0), (1, 0), (2, 1), (3, 1), (4, 2), (5, 2), (6, 2)]
         );
     }
+
+    #[test]
+    fn get_and_extend_additional_data() {
+        let builtin_a = OutputBuiltinRunner {
+            base: 0,
+            pages: HashMap::from([(1, PublicMemoryPage { start: 0, size: 3 })]),
+            attributes: HashMap::from([("gps_fact_topology".to_string(), vec![0, 2, 0])]),
+            stop_ptr: None,
+            included: true,
+        };
+        let additional_data = builtin_a.get_additional_data();
+        let mut builtin_b = OutputBuiltinRunner {
+            base: 0,
+            pages: Default::default(),
+            attributes: Default::default(),
+            stop_ptr: None,
+            included: true,
+        };
+        builtin_b.extend_additional_data(&additional_data).unwrap();
+        assert_eq!(builtin_a.attributes, builtin_b.attributes);
+        assert_eq!(builtin_a.pages, builtin_b.pages);
+    }
 }

vm/src/vm/runners/builtin_runner/signature.rs

@@ -2,8 +2,10 @@ use crate::air_private_input::{PrivateInput, PrivateInputSignature, SignatureInp
 use crate::math_utils::div_mod;
 use crate::stdlib::{cell::RefCell, collections::HashMap, prelude::*, rc::Rc};
 
+use crate::types::builtin_name::BuiltinName;
 use crate::types::errors::math_errors::MathError;
 use crate::types::instance_definitions::ecdsa_instance_def::CELLS_PER_SIGNATURE;
+use crate::vm::errors::runner_errors::RunnerError;
 use crate::vm::runners::cairo_pie::BuiltinAdditionalData;
 use crate::Felt252;
 use crate::{
@@ -180,6 +182,31 @@ impl SignatureBuiltinRunner {
         BuiltinAdditionalData::Signature(signatures)
     }
 
+    pub fn extend_additional_data(
+        &mut self,
+        additional_data: &BuiltinAdditionalData,
+    ) -> Result<(), RunnerError> {
+        let additional_data = match additional_data {
+            BuiltinAdditionalData::Signature(d) => d,
+            _ => return Err(RunnerError::InvalidAdditionalData(BuiltinName::ecdsa)),
+        };
+        for (addr, (r, s)) in additional_data {
+            if addr.segment_index != self.base as isize {
+                return Err(RunnerError::InvalidAdditionalData(BuiltinName::ecdsa));
+            }
+            self.signatures.borrow_mut().insert(
+                *addr,
+                Signature {
+                    r: FieldElement::from_bytes_be(&r.to_bytes_be())
+                        .map_err(|_| MathError::ByteConversionError)?,
+                    s: FieldElement::from_bytes_be(&s.to_bytes_be())
+                        .map_err(|_| MathError::ByteConversionError)?,
+                },
+            );
+        }
+        Ok(())
+    }
+
     pub fn air_private_input(&self, memory: &Memory) -> Vec<PrivateInput> {
         let mut private_inputs = vec![];
         for (addr, signature) in self.signatures.borrow().iter() {
@@ -482,7 +509,7 @@ mod tests {
     }
 
     #[test]
-    fn get_additional_info() {
+    fn get_additional_data() {
         let mut builtin = SignatureBuiltinRunner::new(Some(512), true);
         let signatures = HashMap::from([(
             Relocatable::from((4, 0)),
@@ -501,4 +528,31 @@ mod tests {
             BuiltinAdditionalData::Signature(signatures)
         )
     }
+
+    #[test]
+    fn get_and_extend_additional_data() {
+        let mut builtin_a = SignatureBuiltinRunner::new(Some(512), true);
+        let signatures = HashMap::from([(
+            Relocatable::from((0, 0)),
+            Signature {
+                r: FieldElement::from_dec_str("45678").unwrap(),
+                s: FieldElement::from_dec_str("1239").unwrap(),
+            },
+        )]);
+        builtin_a.signatures = Rc::new(RefCell::new(signatures));
+        let additional_data = builtin_a.get_additional_data();
+        let mut builtin_b = SignatureBuiltinRunner::new(Some(512), true);
+        builtin_b.extend_additional_data(&additional_data).unwrap();
+        // Signature doesn't implement PartialEq so we can't comapre the list of signatures directly
+        let signatures_a = builtin_a.signatures.borrow();
+        let signatures_b = builtin_b.signatures.borrow();
+        assert_eq!(signatures_a.len(), signatures_b.len());
+        for ((addr_a, signature_a), (addr_b, signature_b)) in
+            signatures_a.iter().zip(signatures_b.iter())
+        {
+            assert_eq!(addr_a, addr_b);
+            assert_eq!(signature_a.r, signature_b.r);
+            assert_eq!(signature_a.s, signature_b.s);
+        }
+    }
 }