Merge branch 'main' of github.com:lambdaclass/cairo-rs-py into dict-manager

Author: Juan-M-V

Cargo.toml

@@ -13,7 +13,7 @@ default = ["pyo3/num-bigint", "pyo3/auto-initialize"]
 
 [dependencies]
 pyo3 = { version = "0.16.5" }
-cairo-rs = { git = "https://github.com/lambdaclass/cairo-rs.git", rev = "8c47dda53e874545895b34d675be6254878a9e7b" }
+cairo-rs = { git = "https://github.com/lambdaclass/cairo-rs.git", rev = "2ddf78e20cc25e660263a0c9c1b942780d95a0e6" }
 num-bigint = "0.4"
 lazy_static = "1.4.0"
 

src/cairo_runner.rs

@@ -541,6 +541,13 @@ impl PyCairoRunner {
             .collect::<Vec<_>>()
             .to_object(py))
     }
+
+    /// Add (or replace if already present) a custom hash builtin.
+    /// Returns a Relocatable with the new hash builtin base.
+    pub fn add_additional_hash_builtin(&self) -> PyRelocatable {
+        let mut vm = (*self.pyvm.vm).borrow_mut();
+        self.inner.add_additional_hash_builtin(&mut vm).into()
+    }
 }
 
 #[pyclass]
@@ -1221,6 +1228,10 @@ mod test {
                 segment_index: 6,
                 offset: 0,
             })],
+            vec![RelocatableValue(PyRelocatable {
+                segment_index: 7,
+                offset: 0,
+            })],
         ];
 
         Python::with_gil(|py| {
@@ -1522,4 +1533,43 @@ mod test {
             );
         });
     }
+
+    /// Test that add_additional_hash_builtin() returns successfully.
+    #[test]
+    fn add_additional_hash_builtin() {
+        Python::with_gil(|_| {
+            let program = fs::read_to_string("cairo_programs/fibonacci.json").unwrap();
+            let runner = PyCairoRunner::new(
+                program,
+                Some("main".to_string()),
+                Some("small".to_string()),
+                false,
+            )
+            .unwrap();
+
+            let expected_relocatable = PyRelocatable {
+                segment_index: 0,
+                offset: 0,
+            };
+            let relocatable = runner.add_additional_hash_builtin();
+            assert_eq!(expected_relocatable, relocatable);
+
+            assert_eq!(
+                (*runner.pyvm.vm)
+                    .borrow()
+                    .get_builtin_runners()
+                    .last()
+                    .map(|(key, _)| key.as_str()),
+                Some("hash_builtin"),
+            );
+
+            let mut vm = (*runner.pyvm.vm).borrow_mut();
+            // Check that the segment exists by writing to it.
+            vm.insert_value(
+                &Relocatable::from((0, 0)),
+                MaybeRelocatable::Int(bigint!(42)),
+            )
+            .expect("memory insert failed");
+        });
+    }
 }

src/memory.rs

@@ -1,16 +1,18 @@
 use crate::{
     relocatable::{PyMaybeRelocatable, PyRelocatable},
+    utils::to_py_error,
     vm_core::PyVM,
 };
 use cairo_rs::{
     types::relocatable::{MaybeRelocatable, Relocatable},
     vm::vm_core::VirtualMachine,
 };
+use num_bigint::BigInt;
 use pyo3::{
     exceptions::{PyTypeError, PyValueError},
     prelude::*,
 };
-use std::{cell::RefCell, rc::Rc};
+use std::{borrow::Cow, cell::RefCell, rc::Rc};
 
 const MEMORY_GET_ERROR_MSG: &str = "Failed to get value from Cairo memory";
 const MEMORY_SET_ERROR_MSG: &str = "Failed to set value to Cairo memory";
@@ -69,6 +71,18 @@ impl PyMemory {
             .collect::<Vec<PyMaybeRelocatable>>()
             .to_object(py))
     }
+
+    /// Return a continuous section of memory as a vector of integers.
+    pub fn get_range_as_ints(&self, addr: PyRelocatable, size: usize) -> PyResult<Vec<BigInt>> {
+        Ok(self
+            .vm
+            .borrow()
+            .get_integer_range(&Relocatable::from(&addr), size)
+            .map_err(to_py_error)?
+            .into_iter()
+            .map(Cow::into_owned)
+            .collect())
+    }
 }
 
 #[cfg(test)]
@@ -282,4 +296,104 @@ assert memory[ap] == fp
             assert_eq!(range.unwrap_err(), expected_error);
         });
     }
+
+    // Test that get_range_as_ints() works as intended.
+    #[test]
+    fn get_range_as_ints() {
+        let vm = PyVM::new(
+            BigInt::new(Sign::Plus, vec![1, 0, 0, 0, 0, 0, 17, 134217728]),
+            false,
+        );
+        let memory = PyMemory::new(&vm);
+
+        let addr = {
+            let mut vm = vm.vm.borrow_mut();
+            let addr = vm.add_memory_segment();
+
+            vm.load_data(
+                &MaybeRelocatable::from(&addr),
+                vec![
+                    bigint!(1).into(),
+                    bigint!(2).into(),
+                    bigint!(3).into(),
+                    bigint!(4).into(),
+                ],
+            )
+            .expect("memory insertion failed");
+
+            addr
+        };
+
+        assert_eq!(
+            memory
+                .get_range_as_ints(addr.into(), 4)
+                .expect("get_range_as_ints() failed"),
+            vec![bigint!(1), bigint!(2), bigint!(3), bigint!(4)],
+        );
+    }
+
+    // Test that get_range_as_ints() fails when not all values are integers.
+    #[test]
+    fn get_range_as_ints_mixed() {
+        let vm = PyVM::new(
+            BigInt::new(Sign::Plus, vec![1, 0, 0, 0, 0, 0, 17, 134217728]),
+            false,
+        );
+        let memory = PyMemory::new(&vm);
+
+        let addr = {
+            let mut vm = vm.vm.borrow_mut();
+            let addr = vm.add_memory_segment();
+
+            vm.load_data(
+                &MaybeRelocatable::from(&addr),
+                vec![
+                    bigint!(1).into(),
+                    bigint!(2).into(),
+                    MaybeRelocatable::RelocatableValue((1, 2).into()),
+                    bigint!(4).into(),
+                ],
+            )
+            .expect("memory insertion failed");
+
+            addr
+        };
+
+        memory
+            .get_range_as_ints(addr.into(), 4)
+            .expect_err("get_range_as_ints() succeeded (should have failed)");
+    }
+
+    // Test that get_range_as_ints() fails when the requested range is larger than the available
+    // segments.
+    #[test]
+    fn get_range_as_ints_incomplete() {
+        let vm = PyVM::new(
+            BigInt::new(Sign::Plus, vec![1, 0, 0, 0, 0, 0, 17, 134217728]),
+            false,
+        );
+        let memory = PyMemory::new(&vm);
+
+        let addr = {
+            let mut vm = vm.vm.borrow_mut();
+            let addr = vm.add_memory_segment();
+
+            vm.load_data(
+                &MaybeRelocatable::from(&addr),
+                vec![
+                    bigint!(1).into(),
+                    bigint!(2).into(),
+                    bigint!(3).into(),
+                    bigint!(4).into(),
+                ],
+            )
+            .expect("memory insertion failed");
+
+            addr
+        };
+
+        memory
+            .get_range_as_ints(addr.into(), 8)
+            .expect_err("get_range_as_ints() succeeded (should have failed)");
+    }
 }