starknet_os: output txs_trace from os_run (#8512)

Author: Yonatan-Starkware

crates/starknet_committer_and_os_cli/src/os_cli/commands.rs

@@ -17,7 +17,7 @@ use starknet_api::executable_transaction::{AccountTransaction, Transaction};
 use starknet_os::hint_processor::aggregator_hint_processor::AggregatorInput;
 use starknet_os::io::os_input::{OsBlockInput, OsHints, StarknetOsInput};
 use starknet_os::io::os_output::{StarknetAggregatorRunnerOutput, StarknetOsRunnerOutput};
-use starknet_os::runner::{run_aggregator, run_os_stateless};
+use starknet_os::runner::{run_aggregator, run_os_stateless, run_os_stateless_for_testing};
 use tracing::info;
 use tracing::level_filters::LevelFilter;
 use tracing_subscriber::reload::Handle;
@@ -88,6 +88,7 @@ pub(crate) fn parse_and_run_os(
     input_path: String,
     output_path: String,
     log_filter_handle: Handle<LevelFilter, Registry>,
+    include_txs_trace: bool,
 ) {
     let OsCliInput { layout, os_hints, cairo_pie_zip_path } = load_input(input_path);
     log_filter_handle
@@ -97,13 +98,21 @@ pub(crate) fn parse_and_run_os(
     validate_os_input(&os_hints.os_input);
 
     info!("Running OS...");
-    let StarknetOsRunnerOutput { cairo_pie, da_segment, metrics, unused_hints, .. } =
-        run_os_stateless(layout, os_hints)
+    let (runner_output, txs_trace) = if include_txs_trace {
+        let (output, traces) = run_os_stateless_for_testing(layout, os_hints)
             .unwrap_or_else(|err| panic!("OS run failed. Error: {err}"));
+        (output, Some(traces))
+    } else {
+        let output = run_os_stateless(layout, os_hints)
+            .unwrap_or_else(|err| panic!("OS run failed. Error: {err}"));
+        (output, None)
+    };
+
+    let StarknetOsRunnerOutput { cairo_pie, da_segment, metrics, unused_hints, .. } = runner_output;
 
     info!("Finished running OS. Serializing OS output...");
     serialize_runner_output(
-        &OsCliOutput { da_segment, metrics: metrics.into(), unused_hints },
+        &OsCliOutput { da_segment, metrics: metrics.into(), unused_hints, txs_trace },
         output_path,
         &cairo_pie,
         cairo_pie_zip_path,

crates/starknet_committer_and_os_cli/src/os_cli/run_os_cli.rs

@@ -5,6 +5,7 @@ use cairo_vm::types::relocatable::MaybeRelocatable;
 use cairo_vm::vm::runners::cairo_runner::ExecutionResources;
 use clap::{Parser, Subcommand};
 use serde::Serialize;
+use starknet_os::hint_processor::os_logger::OsTransactionTrace;
 use starknet_os::hints::enum_definition::AllHints;
 use starknet_os::metrics::OsMetrics;
 use starknet_types_core::felt::Felt;
@@ -62,6 +63,8 @@ enum Command {
     RunOsStateless {
         #[clap(flatten)]
         io_args: IoArgs,
+        #[clap(long)]
+        include_txs_trace: bool,
     },
     RunAggregator {
         #[clap(flatten)]
@@ -81,8 +84,11 @@ pub async fn run_os_cli(
         Command::PythonTest(python_test_arg) => {
             run_python_test::<OsPythonTestRunner>(python_test_arg).await;
         }
-        Command::RunOsStateless { io_args: IoArgs { input_path, output_path } } => {
-            parse_and_run_os(input_path, output_path, log_filter_handle);
+        Command::RunOsStateless {
+            io_args: IoArgs { input_path, output_path },
+            include_txs_trace,
+        } => {
+            parse_and_run_os(input_path, output_path, log_filter_handle, include_txs_trace);
         }
         Command::RunAggregator { io_args: IoArgs { input_path, output_path } } => {
             parse_and_run_aggregator(input_path, output_path, log_filter_handle);
@@ -141,6 +147,8 @@ pub(crate) struct OsCliOutput {
     pub(crate) da_segment: Option<Vec<Felt>>,
     pub(crate) metrics: OsCliMetrics,
     pub unused_hints: HashSet<AllHints>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub(crate) txs_trace: Option<Vec<OsTransactionTrace>>,
 }
 
 #[derive(Serialize)]

crates/starknet_os/src/hint_processor/os_logger.rs

@@ -9,6 +9,7 @@ use cairo_vm::types::relocatable::Relocatable;
 use cairo_vm::vm::errors::memory_errors::MemoryError;
 use cairo_vm::vm::runners::cairo_runner::ExecutionResources;
 use cairo_vm::vm::vm_core::VirtualMachine;
+use serde::Serialize;
 use starknet_api::executable_transaction::TransactionType;
 use starknet_api::transaction::TransactionHash;
 
@@ -82,6 +83,7 @@ pub trait ResourceFinalizer {
     }
 }
 
+#[cfg_attr(any(test, feature = "testing"), derive(Serialize, Debug, Clone))]
 pub struct SyscallTrace {
     selector: SyscallSelector,
     is_deprecated: bool,
@@ -142,6 +144,7 @@ impl TryFrom<&SyscallTrace> for String {
     }
 }
 
+#[cfg_attr(any(test, feature = "testing"), derive(Serialize, Debug, Clone))]
 pub struct OsTransactionTrace {
     tx_type: TransactionType,
     tx_hash: TransactionHash,
@@ -508,4 +511,9 @@ impl OsLogger {
         self.txs.push(current_tx);
         Ok(())
     }
+
+    #[cfg(any(test, feature = "testing"))]
+    pub fn get_txs(&self) -> &Vec<OsTransactionTrace> {
+        &self.txs
+    }
 }

crates/starknet_os/src/runner.rs

@@ -1,21 +1,32 @@
+use std::collections::BTreeMap;
+
 use apollo_starknet_os_program::{AGGREGATOR_PROGRAM, OS_PROGRAM};
 use blockifier::state::state_api::StateReader;
+use cairo_lang_starknet_classes::casm_contract_class::CasmContractClass;
 use cairo_vm::cairo_run::CairoRunConfig;
 use cairo_vm::hint_processor::hint_processor_definition::HintProcessor;
 use cairo_vm::types::layout_name::LayoutName;
 use cairo_vm::types::program::Program;
 use cairo_vm::vm::errors::vm_exception::VmException;
 use cairo_vm::vm::runners::cairo_pie::CairoPie;
 use cairo_vm::vm::runners::cairo_runner::CairoRunner;
-#[cfg(feature = "include_program_output")]
+use starknet_api::core::CompiledClassHash;
+use starknet_api::deprecated_contract_class::ContractClass;
 use starknet_types_core::felt::Felt;
 
 use crate::errors::StarknetOsError;
 use crate::hint_processor::aggregator_hint_processor::{AggregatorHintProcessor, AggregatorInput};
 use crate::hint_processor::common_hint_processor::CommonHintProcessor;
+use crate::hint_processor::os_logger::OsTransactionTrace;
 use crate::hint_processor::panicking_state_reader::PanickingStateReader;
 use crate::hint_processor::snos_hint_processor::SnosHintProcessor;
-use crate::io::os_input::{OsHints, StarknetOsInput};
+use crate::io::os_input::{
+    CachedStateInput,
+    OsBlockInput,
+    OsHints,
+    OsHintsConfig,
+    StarknetOsInput,
+};
 use crate::io::os_output::{StarknetAggregatorRunnerOutput, StarknetOsRunnerOutput};
 use crate::metrics::OsMetrics;
 
@@ -105,6 +116,33 @@ pub fn run_os<S: StateReader>(
     }: OsHints,
     state_readers: Vec<S>,
 ) -> Result<StarknetOsRunnerOutput, StarknetOsError> {
+    let (runner_output, snos_hint_processor) = create_hint_processor_and_run_os(
+        layout,
+        os_hints_config,
+        &os_block_inputs,
+        cached_state_inputs,
+        deprecated_compiled_classes,
+        compiled_classes,
+        state_readers,
+        public_key_x,
+        public_key_y,
+    )?;
+
+    generate_os_output(runner_output, snos_hint_processor)
+}
+
+#[allow(clippy::too_many_arguments)]
+fn create_hint_processor_and_run_os<'a, S: StateReader>(
+    layout: LayoutName,
+    os_hints_config: OsHintsConfig,
+    os_block_inputs: &'a [OsBlockInput],
+    cached_state_inputs: Vec<CachedStateInput>,
+    deprecated_compiled_classes: BTreeMap<CompiledClassHash, ContractClass>,
+    compiled_classes: BTreeMap<CompiledClassHash, CasmContractClass>,
+    state_readers: Vec<S>,
+    public_key_x: Felt,
+    public_key_y: Felt,
+) -> Result<(RunnerReturnObject, SnosHintProcessor<'a, S>), StarknetOsError> {
     // Create the hint processor.
     let mut snos_hint_processor = SnosHintProcessor::new(
         &OS_PROGRAM,
@@ -118,8 +156,16 @@ pub fn run_os<S: StateReader>(
         public_key_y,
     )?;
 
-    let mut runner_output = run_program(layout, &OS_PROGRAM, &mut snos_hint_processor)?;
+    // Run the OS program.
+    let runner_output = run_program(layout, &OS_PROGRAM, &mut snos_hint_processor)?;
+
+    Ok((runner_output, snos_hint_processor))
+}
 
+fn generate_os_output(
+    mut runner_output: RunnerReturnObject,
+    mut snos_hint_processor: SnosHintProcessor<'_, impl StateReader>,
+) -> Result<StarknetOsRunnerOutput, StarknetOsError> {
     Ok(StarknetOsRunnerOutput {
         #[cfg(feature = "include_program_output")]
         os_output: {
@@ -143,6 +189,43 @@ pub fn run_os<S: StateReader>(
     })
 }
 
+/// Runs the OS the same way as `run_os`. Returns also the transactions trace which are needed for
+/// some tests.
+#[cfg(any(test, feature = "testing"))]
+pub fn run_os_for_testing<S: StateReader>(
+    layout: LayoutName,
+    OsHints {
+        os_hints_config,
+        os_input:
+            StarknetOsInput {
+                os_block_inputs,
+                cached_state_inputs,
+                deprecated_compiled_classes,
+                compiled_classes,
+                public_key_x,
+                public_key_y,
+            },
+    }: OsHints,
+    state_readers: Vec<S>,
+) -> Result<(StarknetOsRunnerOutput, Vec<OsTransactionTrace>), StarknetOsError> {
+    let (runner_output, snos_hint_processor) = create_hint_processor_and_run_os(
+        layout,
+        os_hints_config,
+        &os_block_inputs,
+        cached_state_inputs,
+        deprecated_compiled_classes,
+        compiled_classes,
+        state_readers,
+        public_key_x,
+        public_key_y,
+    )?;
+
+    let txs_trace: Vec<OsTransactionTrace> =
+        snos_hint_processor.get_current_execution_helper().unwrap().os_logger.get_txs().clone();
+
+    Ok((generate_os_output(runner_output, snos_hint_processor)?, txs_trace))
+}
+
 /// Run the OS with a "stateless" state reader - panics if the state is accessed for data that was
 /// not pre-loaded as part of the input.
 pub fn run_os_stateless(
@@ -153,6 +236,14 @@ pub fn run_os_stateless(
     run_os(layout, os_hints, vec![PanickingStateReader; n_blocks])
 }
 
+pub fn run_os_stateless_for_testing(
+    layout: LayoutName,
+    os_hints: OsHints,
+) -> Result<(StarknetOsRunnerOutput, Vec<OsTransactionTrace>), StarknetOsError> {
+    let n_blocks = os_hints.os_input.os_block_inputs.len();
+    run_os_for_testing(layout, os_hints, vec![PanickingStateReader; n_blocks])
+}
+
 /// Run the Aggregator.
 pub fn run_aggregator(
     layout: LayoutName,