cairo-lang-sierra 2.19.2

#[expect(clippy::disallowed_types)]
use std::collections::HashMap;

use cairo_lang_utils::ordered_hash_map::OrderedHashMap;
use itertools::izip;
use thiserror::Error;

use self::value::CoreValue;
use crate::edit_state::{EditState, EditStateError};
use crate::extensions::core::{CoreConcreteLibfunc, CoreLibfunc, CoreType};
use crate::ids::{FunctionId, VarId};
use crate::program::{Program, Statement, StatementIdx};
use crate::program_registry::{ProgramRegistry, ProgramRegistryError};

pub mod core;
#[cfg(test)]
mod test;
pub mod value;

/// Error occurring while simulating a libfunc.
#[derive(Error, Debug, Eq, PartialEq)]
pub enum LibfuncSimulationError {
    #[error("Expected different number of arguments")]
    WrongNumberOfArgs,
    #[error("Expected a different type of an argument")]
    WrongArgType,
    #[error("Could not resolve requested symbol value")]
    UnresolvedStatementGasInfo,
    #[error("Error occurred during user function call")]
    FunctionSimulationError(FunctionId, Box<SimulationError>),
}

/// Error occurring while simulating a program function.
#[derive(Error, Debug, Eq, PartialEq)]
pub enum SimulationError {
    #[error("error from the program registry")]
    ProgramRegistryError(#[from] Box<ProgramRegistryError>),
    #[error("error from editing a variable state")]
    EditStateError(EditStateError, StatementIdx),
    #[error("error from simulating a libfunc")]
    LibfuncSimulationError(LibfuncSimulationError, StatementIdx),
    #[error("jumped out of bounds during simulation")]
    StatementOutOfBounds(StatementIdx),
    #[error("unexpected number of arguments to function")]
    FunctionArgumentCountMismatch { function_id: FunctionId, expected: usize, actual: usize },
    #[error("identifiers left at function return")]
    FunctionDidNotConsumeAllArgs(FunctionId, StatementIdx),
}

/// Runs a function from the program with the given inputs.
#[expect(clippy::disallowed_types)]
pub fn run(
    program: &Program,
    statement_gas_info: &HashMap<StatementIdx, i64>,
    function_id: &FunctionId,
    inputs: Vec<CoreValue>,
) -> Result<Vec<CoreValue>, SimulationError> {
    let context = SimulationContext {
        program,
        statement_gas_info,
        registry: &ProgramRegistry::new(program)?,
    };
    context.simulate_function(function_id, inputs)
}

/// Helper class for running the simulation.
#[expect(clippy::disallowed_types)]
struct SimulationContext<'a> {
    pub program: &'a Program,
    pub statement_gas_info: &'a HashMap<StatementIdx, i64>,
    pub registry: &'a ProgramRegistry<CoreType, CoreLibfunc>,
}
impl SimulationContext<'_> {
    /// Simulates the run of a function, even recursively.
    fn simulate_function(
        &self,
        function_id: &FunctionId,
        inputs: Vec<CoreValue>,
    ) -> Result<Vec<CoreValue>, SimulationError> {
        let func = self.registry.get_function(function_id)?;
        let mut current_statement_id = func.entry_point;
        if func.params.len() != inputs.len() {
            return Err(SimulationError::FunctionArgumentCountMismatch {
                function_id: func.id.clone(),
                expected: func.params.len(),
                actual: inputs.len(),
            });
        }
        let mut state = OrderedHashMap::<VarId, CoreValue>::from_iter(
            izip!(func.params.iter(), inputs).map(|(param, input)| (param.id.clone(), input)),
        );
        loop {
            let statement = self
                .program
                .get_statement(current_statement_id)
                .ok_or(SimulationError::StatementOutOfBounds(current_statement_id))?;
            match statement {
                Statement::Return(ids) => {
                    let outputs = state.take_vars(ids.iter()).map_err(|error| {
                        SimulationError::EditStateError(error, current_statement_id)
                    })?;
                    return if state.is_empty() {
                        Ok(outputs)
                    } else {
                        Err(SimulationError::FunctionDidNotConsumeAllArgs(
                            func.id.clone(),
                            current_statement_id,
                        ))
                    };
                }
                Statement::Invocation(invocation) => {
                    let inputs = state.take_vars(invocation.args.iter()).map_err(|error| {
                        SimulationError::EditStateError(error, current_statement_id)
                    })?;
                    let libfunc = self.registry.get_libfunc(&invocation.libfunc_id)?;
                    let (outputs, chosen_branch) =
                        self.simulate_libfunc(current_statement_id, libfunc, inputs)?;
                    let branch_info = &invocation.branches[chosen_branch];
                    state.put_vars(izip!(branch_info.results.iter(), outputs)).map_err(
                        |error| SimulationError::EditStateError(error, current_statement_id),
                    )?;
                    current_statement_id = current_statement_id.next(branch_info.target);
                }
            }
        }
    }
    /// Simulates the run of libfuncs. Returns the memory representations of the outputs given the
    /// inputs.
    fn simulate_libfunc(
        &self,
        statement_id: StatementIdx,
        libfunc: &CoreConcreteLibfunc,
        inputs: Vec<CoreValue>,
    ) -> Result<(Vec<CoreValue>, usize), SimulationError> {
        core::simulate(
            libfunc,
            inputs,
            || self.statement_gas_info.get(&statement_id).copied(),
            |function_id, inputs| {
                self.simulate_function(function_id, inputs).map_err(|error| {
                    LibfuncSimulationError::FunctionSimulationError(
                        function_id.clone(),
                        Box::new(error),
                    )
                })
            },
        )
        .map_err(|error| SimulationError::LibfuncSimulationError(error, statement_id))
    }
}