starknet_os: split out FullOsOutputs parsing logic (#10359)

Author: dorimedini-starkware

crates/starknet_os/src/hints/hint_implementation/aggregator/implementation.rs

@@ -3,17 +3,15 @@ use cairo_vm::hint_processor::builtin_hint_processor::hint_utils::{
     insert_value_from_var_name,
     insert_value_into_ap,
 };
-use cairo_vm::hint_processor::hint_processor_utils::felt_to_usize;
 use starknet_types_core::felt::Felt;
 
 use crate::hint_processor::aggregator_hint_processor::{AggregatorHintProcessor, DataAvailability};
 use crate::hint_processor::common_hint_processor::CommonHintProcessor;
 use crate::hints::error::{OsHintError, OsHintResult};
-use crate::hints::hint_implementation::aggregator::utils::FullOsOutputs;
+use crate::hints::hint_implementation::aggregator::utils::FullOsOutputsData;
 use crate::hints::hint_implementation::output::load_public_keys_into_memory;
 use crate::hints::types::HintArgs;
 use crate::hints::vars::Ids;
-use crate::io::os_output::{wrap_missing, FullOsOutput, OsOutput};
 use crate::io::os_output_types::TryFromOutputIter;
 use crate::vm_utils::LoadIntoVmMemory;
 
@@ -51,32 +49,13 @@ pub(crate) fn get_os_output_for_inner_blocks(
         .take()
         .ok_or(OsHintError::AggregatorBootloaderInputAlreadyConsumed)?
         .into_iter();
-    let n_outputs = wrap_missing(bootloader_iter.next(), "n_output")?;
-    let n_outputs_usize = felt_to_usize(&n_outputs)?;
-    assert!(n_outputs_usize > 0, "No tasks found in the bootloader output.");
-
-    let mut program_hash: Option<Felt> = None;
-    let mut outputs = FullOsOutputs(Vec::<FullOsOutput>::with_capacity(n_outputs_usize));
-    for _ in 0..n_outputs_usize {
-        wrap_missing(bootloader_iter.next(), "output_size")?;
-        let current_output_program_hash = wrap_missing(bootloader_iter.next(), "program_hash")?;
-        assert_eq!(
-            program_hash.get_or_insert(current_output_program_hash),
-            &current_output_program_hash
-        );
-        // Regardless of encryption mode: OS outputs that reach the aggregator are never encrypted.
-        // No need to pass any decryption data here.
-        outputs.0.push(OsOutput::try_from_output_iter(&mut bootloader_iter, None)?.try_into()?);
-    }
 
-    insert_value_from_var_name(
-        Ids::OsProgramHash.into(),
-        program_hash.expect("n_outputs > 0 but program hash wasn't initialized."),
-        vm,
-        ids_data,
-        ap_tracking,
-    )?;
+    // Regardless of encryption mode: OS outputs that reach the aggregator are never
+    // encrypted. No need to pass any decryption data here.
+    let FullOsOutputsData { outputs, n_outputs, program_hash } =
+        FullOsOutputsData::try_from_output_iter(&mut bootloader_iter, None)?;
 
+    insert_value_from_var_name(Ids::OsProgramHash.into(), program_hash, vm, ids_data, ap_tracking)?;
     insert_value_from_var_name(Ids::NTasks.into(), n_outputs, vm, ids_data, ap_tracking)?;
 
     let os_output_ptr = get_ptr_from_var_name(Ids::OsOutputs.into(), vm, ids_data, ap_tracking)?;

crates/starknet_os/src/hints/hint_implementation/aggregator/utils.rs

@@ -1,6 +1,7 @@
 use std::collections::hash_map::Entry;
 use std::collections::HashMap;
 
+use cairo_vm::hint_processor::hint_processor_utils::felt_to_usize;
 use cairo_vm::types::errors::math_errors::MathError;
 use cairo_vm::types::relocatable::{MaybeRelocatable, Relocatable};
 use cairo_vm::vm::vm_core::VirtualMachine;
@@ -10,15 +11,19 @@ use starknet_types_core::felt::Felt;
 
 use crate::hint_processor::state_update_pointers::{StateEntryPtr, StoragePtr};
 use crate::io::os_output::{
+    wrap_missing,
     FullOsOutput,
     MessageToL2,
+    OsOutput,
+    OsOutputError,
     MESSAGE_TO_L1_CONST_FIELD_SIZE,
     MESSAGE_TO_L2_CONST_FIELD_SIZE,
 };
 use crate::io::os_output_types::{
     FullCompiledClassHashUpdate,
     FullContractChanges,
     FullContractStorageUpdate,
+    TryFromOutputIter,
 };
 use crate::vm_utils::{LoadIntoVmMemory, VmUtilsResult};
 
@@ -300,3 +305,37 @@ impl LoadIntoVmMemory for FullOsOutputs {
         Ok(os_output_ptr)
     }
 }
+
+pub struct FullOsOutputsData {
+    pub outputs: FullOsOutputs,
+    pub n_outputs: usize,
+    pub program_hash: Felt,
+}
+
+impl TryFromOutputIter for FullOsOutputsData {
+    fn try_from_output_iter<It: Iterator<Item = Felt>>(
+        iter: &mut It,
+        private_keys: Option<&Vec<Felt>>,
+    ) -> Result<Self, OsOutputError> {
+        let n_outputs = wrap_missing(iter.next(), "n_output")?;
+        let n_outputs_usize = felt_to_usize(&n_outputs)?;
+        assert!(n_outputs_usize > 0, "No tasks found in the bootloader output.");
+
+        let mut program_hash: Option<Felt> = None;
+        let mut outputs = FullOsOutputs(Vec::<FullOsOutput>::with_capacity(n_outputs_usize));
+        for _ in 0..n_outputs_usize {
+            wrap_missing(iter.next(), "output_size")?;
+            let current_output_program_hash = wrap_missing(iter.next(), "program_hash")?;
+            assert_eq!(
+                program_hash.get_or_insert(current_output_program_hash),
+                &current_output_program_hash
+            );
+            outputs.0.push(OsOutput::try_from_output_iter(iter, private_keys)?.try_into()?);
+        }
+        Ok(Self {
+            outputs,
+            n_outputs: n_outputs_usize,
+            program_hash: program_hash.expect("n_outputs > 0 but program hash wasn't initialized."),
+        })
+    }
+}

crates/starknet_os/src/io/os_output.rs

@@ -1,6 +1,7 @@
 use std::sync::LazyLock;
 
 use cairo_vm::types::builtin_name::BuiltinName;
+use cairo_vm::types::errors::math_errors::MathError;
 use cairo_vm::types::relocatable::{MaybeRelocatable, Relocatable};
 use cairo_vm::vm::errors::memory_errors::MemoryError;
 use cairo_vm::vm::errors::vm_errors::VirtualMachineError;
@@ -38,6 +39,8 @@ pub static STARKNET_OS_CONFIG_HASH_VERSION: LazyLock<Felt> = LazyLock::new(|| {
 
 #[derive(Debug, thiserror::Error)]
 pub enum OsOutputError {
+    #[error(transparent)]
+    Math(#[from] MathError),
     #[error("Missing expected field: {0}.")]
     MissingFieldInOutput(String),
     #[error("Invalid output in field: {value_name}. Val: {val}. Error: {message}")]