diff --git a/CHANGELOG.md b/CHANGELOG.md index 0619cbc798..7aeee15383 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,8 @@ * BREAKING: Remove `CairoRunner::add_additional_hash_builtin` & `VirtualMachine::disable_trace`[#1658](https://github.com/lambdaclass/cairo-vm/pull/1658) * feat: output builtin add_attribute method [#1691](https://github.com/lambdaclass/cairo-vm/pull/1691) + +* feat: add a method to retrieve the output builtin from the VM [#1690](https://github.com/lambdaclass/cairo-vm/pull/1690) * feat: Add zero segment [#1668](https://github.com/lambdaclass/cairo-vm/pull/1668) diff --git a/vm/src/vm/errors/vm_errors.rs b/vm/src/vm/errors/vm_errors.rs index db8f0278e4..cbccde51b5 100644 --- a/vm/src/vm/errors/vm_errors.rs +++ b/vm/src/vm/errors/vm_errors.rs @@ -81,6 +81,8 @@ pub enum VirtualMachineError { InconsistentAutoDeduction(Box<(&'static str, MaybeRelocatable, Option)>), #[error("Invalid hint encoding at pc: {0}")] InvalidHintEncoding(Box), + #[error("Expected output builtin to be present")] + NoOutputBuiltin, #[error("Expected range_check builtin to be present")] NoRangeCheckBuiltin, #[error("Expected ecdsa builtin to be present")] diff --git a/vm/src/vm/vm_core.rs b/vm/src/vm/vm_core.rs index 9a20d8887a..d087952000 100644 --- a/vm/src/vm/vm_core.rs +++ b/vm/src/vm/vm_core.rs @@ -32,7 +32,7 @@ use core::num::NonZeroUsize; use num_traits::{ToPrimitive, Zero}; use super::errors::runner_errors::RunnerError; -use super::runners::builtin_runner::OUTPUT_BUILTIN_NAME; +use super::runners::builtin_runner::{OutputBuiltinRunner, OUTPUT_BUILTIN_NAME}; const MAX_TRACEBACK_ENTRIES: u32 = 20; @@ -943,6 +943,18 @@ impl VirtualMachine { Err(VirtualMachineError::NoSignatureBuiltin) } + pub fn get_output_builtin_mut( + &mut self, + ) -> Result<&mut OutputBuiltinRunner, VirtualMachineError> { + for builtin in self.get_builtin_runners_as_mut() { + if let BuiltinRunner::Output(output_builtin) = builtin { + return Ok(output_builtin); + }; + } + + Err(VirtualMachineError::NoOutputBuiltin) + } + #[cfg(feature = "with_tracer")] pub fn relocate_segments(&self) -> Result, MemoryError> { self.segments.relocate_segments() @@ -3839,6 +3851,30 @@ mod tests { assert_eq!(builtins[1].name(), BITWISE_BUILTIN_NAME); } + #[test] + #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)] + fn test_get_output_builtin_mut() { + let mut vm = vm!(); + + assert_matches!( + vm.get_output_builtin_mut(), + Err(VirtualMachineError::NoOutputBuiltin) + ); + + let output_builtin = OutputBuiltinRunner::new(true); + vm.builtin_runners.push(output_builtin.clone().into()); + + let vm_output_builtin = vm + .get_output_builtin_mut() + .expect("Output builtin should be returned"); + + assert_eq!(vm_output_builtin.base(), output_builtin.base()); + assert_eq!(vm_output_builtin.pages, output_builtin.pages); + assert_eq!(vm_output_builtin.attributes, output_builtin.attributes); + assert_eq!(vm_output_builtin.stop_ptr, output_builtin.stop_ptr); + assert_eq!(vm_output_builtin.included, output_builtin.included); + } + #[test] #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)] fn get_range_for_continuous_memory() {