starknet_in_rust 0.4.0

use crate::execution::CallResult;
use crate::syscalls::syscall_handler::HintProcessorPostRun;
use crate::transaction::error::TransactionError;
use cairo_lang_starknet::casm_contract_class::CasmContractClass;
use cairo_vm::felt::Felt252;
use cairo_vm::hint_processor::hint_processor_definition::HintProcessor;
use cairo_vm::serde::deserialize_program::BuiltinName;
use cairo_vm::types::errors::math_errors::MathError;
use cairo_vm::{
    types::relocatable::{MaybeRelocatable, Relocatable},
    vm::{
        runners::{
            builtin_runner::BuiltinRunner,
            cairo_runner::{CairoArg, CairoRunner, ExecutionResources},
        },
        vm_core::VirtualMachine,
    },
};
use num_traits::{ToPrimitive, Zero};
use std::{borrow::Cow, collections::HashMap};

/// Returns a vector that holds the names of the builtins that a contract class uses
/// ## Parameters
/// - contract_class: A casm Contract Class generated by cairo 1 compiler of the contract to be executed.
/// - entrypoint_offset: offset of the function that will be executed.
pub fn get_casm_contract_builtins(
    contract_class: &CasmContractClass,
    entrypoint_offset: usize,
) -> Vec<BuiltinName> {
    contract_class
        .entry_points_by_type
        .external
        .iter()
        .find(|e| e.offset == entrypoint_offset)
        .unwrap()
        .builtins
        .iter()
        .map(|n| format!("{n}_builtin"))
        .map(|s| match &*s {
            cairo_vm::vm::runners::builtin_runner::OUTPUT_BUILTIN_NAME => BuiltinName::output,
            cairo_vm::vm::runners::builtin_runner::RANGE_CHECK_BUILTIN_NAME => {
                BuiltinName::range_check
            }
            cairo_vm::vm::runners::builtin_runner::HASH_BUILTIN_NAME => BuiltinName::pedersen,
            cairo_vm::vm::runners::builtin_runner::SIGNATURE_BUILTIN_NAME => BuiltinName::ecdsa,
            cairo_vm::vm::runners::builtin_runner::KECCAK_BUILTIN_NAME => BuiltinName::keccak,
            cairo_vm::vm::runners::builtin_runner::BITWISE_BUILTIN_NAME => BuiltinName::bitwise,
            cairo_vm::vm::runners::builtin_runner::EC_OP_BUILTIN_NAME => BuiltinName::ec_op,
            cairo_vm::vm::runners::builtin_runner::POSEIDON_BUILTIN_NAME => BuiltinName::poseidon,
            cairo_vm::vm::runners::builtin_runner::SEGMENT_ARENA_BUILTIN_NAME => {
                BuiltinName::segment_arena
            }
            _ => panic!("Invalid builtin {s}"),
        })
        .collect()
}

/// Creates a wrapper over CairoRunner, the CairoVM and a Hint Processor
pub(crate) struct StarknetRunner<H>
where
    H: HintProcessor + HintProcessorPostRun,
{
    pub(crate) cairo_runner: CairoRunner,
    pub(crate) vm: VirtualMachine,
    pub(crate) hint_processor: H,
}

impl<H> StarknetRunner<H>
where
    H: HintProcessor + HintProcessorPostRun,
{
    pub const fn new(cairo_runner: CairoRunner, vm: VirtualMachine, hint_processor: H) -> Self {
        StarknetRunner {
            cairo_runner,
            vm,
            hint_processor,
        }
    }

    /// Executes the entry point in the cairo vm.
    /// ## Parameters
    /// - entrypoint: the offset of the function that will be executed.
    /// - args: arguments of the function.
    /// - program_segment_size: the size of the segment that stores the program
    pub fn run_from_entrypoint(
        &mut self,
        entrypoint: usize,
        args: &[&CairoArg],
        program_segment_size: Option<usize>,
    ) -> Result<(), TransactionError> {
        let verify_secure = true;
        let args: Vec<&CairoArg> = args.iter().map(ToOwned::to_owned).collect();

        self.cairo_runner.run_from_entrypoint(
            entrypoint,
            &args,
            verify_secure,
            program_segment_size,
            &mut self.vm,
            &mut self.hint_processor,
        )?;
        Ok(())
    }

    /// Creates the data structures required to execute the call on the cairo vm according to the entry_point_offset provided
    /// ## Parameters:
    /// - contract_class: A casm Contract Class generated by cairo 1 compiler of the contract to be executed.
    /// - entrypoint_offset: offset of the function that will be executed.
    /// - args: paramenters of the entry point that will be executed.
    #[allow(dead_code)]
    pub fn run_from_cairo1_entrypoint(
        &mut self,
        contract_class: &CasmContractClass,
        entrypoint_offset: usize,
        args: &[MaybeRelocatable],
    ) -> Result<(), TransactionError> {
        let program_builtins = get_casm_contract_builtins(contract_class, entrypoint_offset);

        self.cairo_runner
            .initialize_function_runner_cairo_1(&mut self.vm, &program_builtins)?;

        // Load builtin costs
        let builtin_costs: Vec<MaybeRelocatable> =
            vec![0.into(), 0.into(), 0.into(), 0.into(), 0.into()];
        let builtin_costs_ptr = self.vm.add_memory_segment();
        self.vm.load_data(builtin_costs_ptr, &builtin_costs)?;

        // Load extra data
        let core_program_end_ptr = (self
            .cairo_runner
            .program_base
            .ok_or(TransactionError::NotAFelt)?
            + self.cairo_runner.get_program().data_len())?;
        let program_extra_data: Vec<MaybeRelocatable> =
            vec![0x208B7FFF7FFF7FFE.into(), builtin_costs_ptr.into()];
        self.vm
            .load_data(core_program_end_ptr, &program_extra_data)?;

        // Load calldata
        let calldata_start = self.vm.add_memory_segment();
        let calldata_end = self.vm.load_data(calldata_start, &args.to_vec())?;

        // Create entrypoint_args
        let mut entrypoint_args: Vec<CairoArg> =
            args.iter().map(|m| CairoArg::from(m.clone())).collect();
        entrypoint_args.extend([
            MaybeRelocatable::from(calldata_start).into(),
            MaybeRelocatable::from(calldata_end).into(),
        ]);
        let entrypoint_args: Vec<&CairoArg> = entrypoint_args.iter().collect();

        // Once we have all the entrypoint args in place we can run it
        self.cairo_runner.run_from_entrypoint(
            entrypoint_offset,
            &entrypoint_args,
            true,
            Some(self.cairo_runner.get_program().data_len() + program_extra_data.len()),
            &mut self.vm,
            &mut self.hint_processor,
        )?;

        Ok(())
    }

    /// Returns and ExecutionResources struct that contains the resources used by the contract being execute.
    pub fn get_execution_resources(&self) -> Result<ExecutionResources, TransactionError> {
        Ok(self.cairo_runner.get_execution_resources(&self.vm)?)
    }

    /// Return a vector that holds the data and pointers used to build the CallResult
    pub fn get_return_values(&self) -> Result<Vec<Felt252>, TransactionError> {
        let ret_data = self.vm.get_return_values(2)?;

        let n_rets = ret_data[0]
            .get_int_ref()
            .ok_or(TransactionError::NotAFelt)?;

        let ret_ptr = ret_data[1]
            .get_relocatable()
            .ok_or(TransactionError::NotARelocatableValue)?;

        let ret_data = self.vm.get_integer_range(
            ret_ptr,
            n_rets
                .to_usize()
                .ok_or_else(|| MathError::Felt252ToUsizeConversion(Box::new(n_rets.clone())))?,
        )?;
        Ok(ret_data.into_iter().map(Cow::into_owned).collect())
    }

    /// returns a CallResult that holds the gas consumed, if the execution was succesfull and the retdata of the call.
    /// ## Parameters
    /// - initial_gas: The amount of gas the caller has available.
    pub fn get_call_result(&self, initial_gas: u128) -> Result<CallResult, TransactionError> {
        let return_values = self.vm.get_return_values(5)?;
        let remaining_gas = return_values[0]
            .get_int_ref()
            .and_then(ToPrimitive::to_u128)
            .ok_or(TransactionError::NotAFelt)?;
        let is_success = return_values[2]
            .get_int_ref()
            .ok_or(TransactionError::NotAFelt)?
            .is_zero();
        let retdata_start = return_values[3]
            .get_relocatable()
            .ok_or(TransactionError::NotARelocatableValue)?;
        let retdata_end = return_values[4]
            .get_relocatable()
            .ok_or(TransactionError::NotARelocatableValue)?;
        let size = (retdata_end - retdata_start)?;
        let retdata: Vec<MaybeRelocatable> = self
            .vm
            .get_continuous_range(retdata_start, size)?
            .iter()
            .map(Clone::clone)
            .collect();
        Ok(CallResult {
            gas_consumed: initial_gas.saturating_sub(remaining_gas),
            is_success,
            retdata,
        })
    }

    /// Returns a vector of pointers to the initial stack of the builtins invoked by the contract besides the gas and the syscall_segment pointer.
    /// ## Parameters
    /// - CairoRunner: An instance of a cairo runner that will execute the contract.
    /// - vm: An instance of the cairo virutal machine that will execute the contract.
    /// - gas: The amount of gas that the caller has available.
    pub fn prepare_os_context_cairo1(
        cairo_runner: &CairoRunner,
        vm: &mut VirtualMachine,
        gas: Felt252,
    ) -> Vec<MaybeRelocatable> {
        let mut os_context = vec![];
        // first, add for each builtin, its initial stack to the os_context
        let builtin_runners = vm
            .get_builtin_runners()
            .iter()
            .map(|runner| (runner.name(), runner.clone()))
            .collect::<HashMap<&str, BuiltinRunner>>();

        cairo_runner
            .get_program_builtins()
            .iter()
            .for_each(|builtin| {
                if builtin_runners.contains_key(builtin.name()) {
                    let b_runner = builtin_runners.get(builtin.name()).unwrap();
                    let stack = b_runner.initial_stack();
                    os_context.extend(stack);
                }
            });

        // add the gas
        os_context.push(gas.into());

        // finally add the syscall segment
        let syscall_segment = vm.add_memory_segment();
        os_context.push(syscall_segment.into());
        os_context
    }

    /// Returns a vector of pointers to the initial stack of the builtins invoked by the contract
    /// ## Parameters
    /// - CairoRunner: An instance of a cairo runner that will execute the contract.
    /// - vm: An instance of the cairo virutal machine that will execute the contract.
    pub fn prepare_os_context_cairo0(
        cairo_runner: &CairoRunner,
        vm: &mut VirtualMachine,
    ) -> Vec<MaybeRelocatable> {
        let syscall_segment = vm.add_memory_segment();
        let mut os_context = [syscall_segment.into()].to_vec();
        let builtin_runners = vm
            .get_builtin_runners()
            .iter()
            .map(|runner| (runner.name(), runner))
            .collect::<HashMap<&str, &BuiltinRunner>>();

        cairo_runner
            .get_program_builtins()
            .iter()
            .for_each(|builtin| {
                if builtin_runners.contains_key(builtin.name()) {
                    let b_runner = builtin_runners.get(builtin.name()).unwrap();
                    let stack = b_runner.initial_stack();
                    os_context.extend(stack);
                }
            });
        os_context
    }

    /// Returns the base and stop ptr of the OS-designated segment that starts at ptr_offset.
    /// ## Parameters
    /// - ptr_offset: A pointer that points where the base pointer is stored in memory.
    /// - os_context: The Os context used to fetch the stop pointer segment from memory.
    pub(crate) fn get_os_segment_ptr_range(
        &self,
        ptr_offset: usize,
        os_context: Vec<MaybeRelocatable>,
    ) -> Result<(MaybeRelocatable, MaybeRelocatable), TransactionError> {
        if ptr_offset != 0 {
            return Err(TransactionError::IllegalOsPtrOffset);
        }

        let os_context_end = (self.vm.get_ap() - 2)?;
        let final_os_context_ptr = (os_context_end - os_context.len())?;
        let os_context_ptr = os_context
            .get(ptr_offset)
            .ok_or(TransactionError::InvalidPtrFetch)?
            .to_owned();

        let addr = (final_os_context_ptr + ptr_offset)?;
        let ptr_fetch_from_memory = self
            .vm
            .get_maybe(&addr)
            .ok_or(TransactionError::InvalidPtrFetch)?;

        Ok((os_context_ptr, ptr_fetch_from_memory))
    }

    pub(crate) fn validate_segment_pointers(
        &self,
        segment_base_ptr: &MaybeRelocatable,
        segment_stop_ptr: &MaybeRelocatable,
    ) -> Result<(), TransactionError> {
        let seg_base_ptr = match segment_base_ptr {
            MaybeRelocatable::RelocatableValue(val) => {
                if val.offset != 0 {
                    return Err(TransactionError::InvalidSegBasePtrOffset(val.offset));
                }
                val
            }
            _ => return Err(TransactionError::NotARelocatableValue),
        };

        let expected_stop_ptr = seg_base_ptr
            + self
                .vm
                .get_segment_used_size(seg_base_ptr.segment_index as usize)
                .ok_or(TransactionError::InvalidSegmentSize)?;

        let seg_stop_ptr: Relocatable = match segment_stop_ptr {
            MaybeRelocatable::RelocatableValue(val) => *val,
            _ => return Err(TransactionError::NotARelocatableValue),
        };

        if expected_stop_ptr != seg_stop_ptr {
            return Err(TransactionError::InvalidStopPointer(
                expected_stop_ptr,
                seg_stop_ptr,
            ));
        }
        Ok(())
    }

    /// Validates and processes an OS context that was returned by a transaction.
    /// Returns the syscall processor object containing the accumulated syscall information.
    pub(crate) fn validate_and_process_os_context(
        &mut self,
        initial_os_context: Vec<MaybeRelocatable>,
    ) -> Result<(), TransactionError> {
        // The returned values are os_context, retdata_size, retdata_ptr.
        let os_context_end = (self.vm.get_ap() - 5)?;

        let stack_ptr = self
            .cairo_runner
            .get_builtins_final_stack(&mut self.vm, os_context_end)?;

        let final_os_context_ptr = (stack_ptr - 2)?;

        if final_os_context_ptr + initial_os_context.len() != Ok(os_context_end) {
            return Err(TransactionError::OsContextPtrNotEqual);
        }

        // Validate system calls
        let syscall_base_ptr = initial_os_context
            .last()
            .ok_or(TransactionError::EmptyOsContext)?;
        // Stack ends with: syscall_ptr, vairant_selector, retdata_start, retdata_end.
        let syscall_stop_ptr = self
            .vm
            .get_maybe(&(self.vm.get_ap() - 4)?)
            .ok_or(TransactionError::InvalidPtrFetch)?;

        self.validate_segment_pointers(syscall_base_ptr, &syscall_stop_ptr)?;

        self.hint_processor
            .post_run(&mut self.vm, syscall_stop_ptr.try_into()?)?;

        Ok(())
    }

    /// Validates and processes an OS context that was returned by a transaction.
    /// Returns the syscall processor object containing the accumulated syscall information.
    pub(crate) fn validate_and_process_os_context_for_version0_class(
        &mut self,
        initial_os_context: Vec<MaybeRelocatable>,
    ) -> Result<(), TransactionError> {
        // The returned values are os_context, retdata_size, retdata_ptr.
        let os_context_end = (self.vm.get_ap() - 2)?;

        let stack_ptr = self
            .cairo_runner
            .get_builtins_final_stack(&mut self.vm, os_context_end)?;

        let final_os_context_ptr = (stack_ptr - 1)?;

        if final_os_context_ptr + initial_os_context.len() != Ok(os_context_end) {
            return Err(TransactionError::OsContextPtrNotEqual);
        }

        // Validate system calls
        let (syscall_base_ptr, syscall_stop_ptr) =
            self.get_os_segment_ptr_range(0, initial_os_context)?;

        self.validate_segment_pointers(&syscall_base_ptr, &syscall_stop_ptr)?;

        self.hint_processor
            .post_run(&mut self.vm, syscall_stop_ptr.try_into()?)?;

        Ok(())
    }
}

#[cfg(test)]
mod test {
    use super::StarknetRunner;
    use crate::{
        state::cached_state::CachedState,
        state::in_memory_state_reader::InMemoryStateReader,
        syscalls::{
            deprecated_business_logic_syscall_handler::DeprecatedBLSyscallHandler,
            deprecated_syscall_handler::DeprecatedSyscallHintProcessor,
            syscall_handler::SyscallHintProcessor,
        },
        transaction::error::TransactionError,
    };
    use cairo_vm::{
        types::relocatable::{MaybeRelocatable, Relocatable},
        vm::{
            runners::cairo_runner::{CairoRunner, RunResources},
            vm_core::VirtualMachine,
        },
    };

    #[test]
    fn prepare_os_context_test() {
        let program = cairo_vm::types::program::Program::default();
        let cairo_runner = CairoRunner::new(&program, "starknet", false).unwrap();
        let mut vm = VirtualMachine::new(true);

        let os_context = StarknetRunner::<SyscallHintProcessor<CachedState::<InMemoryStateReader>>>::prepare_os_context_cairo0(&cairo_runner, &mut vm);

        // is expected to return a pointer to the first segment as there is nothing more in the vm
        let expected = Vec::from([MaybeRelocatable::from((0, 0))]);

        assert_eq!(os_context, expected);
    }

    #[test]
    fn run_from_entrypoint_should_fail_with_no_exec_base() {
        let program = cairo_vm::types::program::Program::default();
        let cairo_runner = CairoRunner::new(&program, "starknet", false).unwrap();
        let vm = VirtualMachine::new(true);

        let mut state = CachedState::<InMemoryStateReader>::default();
        let hint_processor = DeprecatedSyscallHintProcessor::new(
            DeprecatedBLSyscallHandler::default_with(&mut state),
            RunResources::default(),
        );

        let mut runner = StarknetRunner::new(cairo_runner, vm, hint_processor);
        assert!(runner.run_from_entrypoint(1, &[], None).is_err())
    }

    #[test]
    fn get_os_segment_ptr_range_should_fail_when_ptr_offset_is_not_zero() {
        let program = cairo_vm::types::program::Program::default();
        let cairo_runner = CairoRunner::new(&program, "starknet", false).unwrap();
        let vm = VirtualMachine::new(true);

        let mut state = CachedState::<InMemoryStateReader>::default();
        let hint_processor = DeprecatedSyscallHintProcessor::new(
            DeprecatedBLSyscallHandler::default_with(&mut state),
            RunResources::default(),
        );

        let runner = StarknetRunner::new(cairo_runner, vm, hint_processor);
        assert_matches!(
            runner.get_os_segment_ptr_range(1, vec![]).unwrap_err(),
            TransactionError::IllegalOsPtrOffset
        );
    }

    #[test]
    fn validate_segment_pointers_should_fail_when_offset_is_not_zero() {
        let program = cairo_vm::types::program::Program::default();
        let cairo_runner = CairoRunner::new(&program, "starknet", false).unwrap();
        let vm = VirtualMachine::new(true);

        let mut state = CachedState::<InMemoryStateReader>::default();
        let hint_processor = DeprecatedSyscallHintProcessor::new(
            DeprecatedBLSyscallHandler::default_with(&mut state),
            RunResources::default(),
        );

        let runner = StarknetRunner::new(cairo_runner, vm, hint_processor);
        let relocatable = MaybeRelocatable::RelocatableValue((0, 1).into());
        assert_matches!(
            runner
                .validate_segment_pointers(&relocatable, &relocatable)
                .unwrap_err(),
            TransactionError::InvalidSegBasePtrOffset(1)
        );
    }

    #[test]
    fn validate_segment_pointers_should_fail_when_base_is_not_a_value() {
        let program = cairo_vm::types::program::Program::default();
        let cairo_runner = CairoRunner::new(&program, "starknet", false).unwrap();
        let vm = VirtualMachine::new(true);

        let mut state = CachedState::<InMemoryStateReader>::default();
        let hint_processor = DeprecatedSyscallHintProcessor::new(
            DeprecatedBLSyscallHandler::default_with(&mut state),
            RunResources::default(),
        );

        let runner = StarknetRunner::new(cairo_runner, vm, hint_processor);
        let relocatable = MaybeRelocatable::Int((1).into());
        assert_matches!(
            runner
                .validate_segment_pointers(&relocatable, &relocatable)
                .unwrap_err(),
            TransactionError::NotARelocatableValue
        );
    }

    #[test]
    fn validate_segment_pointers_should_fail_with_invalid_segment_size() {
        let program = cairo_vm::types::program::Program::default();
        let cairo_runner = CairoRunner::new(&program, "starknet", false).unwrap();
        let vm = VirtualMachine::new(true);

        let mut state = CachedState::<InMemoryStateReader>::default();
        let hint_processor = DeprecatedSyscallHintProcessor::new(
            DeprecatedBLSyscallHandler::default_with(&mut state),
            RunResources::default(),
        );

        let runner = StarknetRunner::new(cairo_runner, vm, hint_processor);
        let base = MaybeRelocatable::RelocatableValue((0, 0).into());
        assert_matches!(
            runner.validate_segment_pointers(&base, &base).unwrap_err(),
            TransactionError::InvalidSegmentSize
        );
    }

    #[test]
    fn validate_segment_pointers_should_fail_when_stop_is_not_a_value() {
        let program = cairo_vm::types::program::Program::default();
        let cairo_runner = CairoRunner::new(&program, "starknet", false).unwrap();
        let mut vm = VirtualMachine::new(true);
        vm.add_memory_segment();
        vm.compute_segments_effective_sizes();

        let mut state = CachedState::<InMemoryStateReader>::default();
        let hint_processor = DeprecatedSyscallHintProcessor::new(
            DeprecatedBLSyscallHandler::default_with(&mut state),
            RunResources::default(),
        );

        let runner = StarknetRunner::new(cairo_runner, vm, hint_processor);
        let base = MaybeRelocatable::RelocatableValue((0, 0).into());
        let stop = MaybeRelocatable::Int((1).into());
        assert_matches!(
            runner.validate_segment_pointers(&base, &stop).unwrap_err(),
            TransactionError::NotARelocatableValue
        );
    }

    #[test]
    fn validate_segment_pointers_should_fail_with_invalid_stop_pointer() {
        let program = cairo_vm::types::program::Program::default();
        let cairo_runner = CairoRunner::new(&program, "starknet", false).unwrap();
        let mut vm = VirtualMachine::new(true);
        vm.add_memory_segment();
        vm.compute_segments_effective_sizes();

        let mut state = CachedState::<InMemoryStateReader>::default();
        let hint_processor = DeprecatedSyscallHintProcessor::new(
            DeprecatedBLSyscallHandler::default_with(&mut state),
            RunResources::default(),
        );

        let runner = StarknetRunner::new(cairo_runner, vm, hint_processor);
        let base = MaybeRelocatable::RelocatableValue((0, 0).into());
        let stop = MaybeRelocatable::RelocatableValue((0, 1).into());
        assert_matches!(
            runner.validate_segment_pointers(&base, &stop).unwrap_err(),
            TransactionError::InvalidStopPointer(
                Relocatable {
                    segment_index: 0,
                    offset: 0,
                },
                Relocatable {
                    segment_index: 0,
                    offset: 1
                },
            )
        );
    }
}