rialo-s-program-runtime 0.4.0-alpha.1

// Copyright (c) Subzero Labs, Inc.
// SPDX-License-Identifier: Apache-2.0
// This file is either (a) original to Subzero Labs, Inc. or (b) derived from the Anza codebase and modified by Subzero Labs, Inc.

use std::{
    alloc::Layout,
    cell::RefCell,
    collections::HashSet,
    fmt::{self, Debug},
    rc::Rc,
};

use rialo_hash::Hash;
use rialo_s_account::{create_account_shared_data_for_test, AccountSharedData, StoredAccount};
use rialo_s_clock::Slot;
use rialo_s_compute_budget::compute_budget::ComputeBudget;
use rialo_s_epoch_schedule::EpochSchedule;
use rialo_s_feature_set::{
    lift_cpi_caller_restriction, move_precompile_verification_to_svm,
    remove_accounts_executable_flag_checks, FeatureSet,
};
use rialo_s_instruction::{error::InstructionError, AccountMeta};
use rialo_s_log_collector::{ic_msg, LogCollector};
use rialo_s_measure::measure::Measure;
use rialo_s_precompiles::Precompile;
use rialo_s_pubkey::Pubkey;
use rialo_s_sdk_ids::{bpf_loader_deprecated, native_loader, sysvar};
use rialo_s_stable_layout::stable_instruction::StableInstruction;
use rialo_s_timings::{ExecuteDetailsTimings, ExecuteTimings};
use rialo_s_transaction_context::{
    IndexOfAccount, InstructionAccount, TransactionAccount, TransactionContext,
};
use rialo_s_type_overrides::sync::{atomic::Ordering, Arc};
pub use rialo_stake_cache_interface::{StakeCacheData, StakesHandle, ValidatorAccount};
use solana_sbpf::{
    ebpf::MM_HEAP_START,
    error::{EbpfError, ProgramResult},
    memory_region::MemoryMapping,
    program::{BuiltinFunction, SBPFVersion},
    vm::{Config, ContextObject, EbpfVm},
};

use crate::{
    loaded_programs::{ProgramCacheEntryType, ProgramCacheForTx, ProgramRuntimeEnvironments},
    stable_log,
    sysvar_cache::SysvarCache,
};

pub type BuiltinFunctionWithContext = BuiltinFunction<InvokeContext<'static, 'static>>;

/// Adapter so we can unify the interfaces of built-in programs and syscalls
#[macro_export]
macro_rules! declare_process_instruction {
    ($process_instruction:ident, $cu_to_consume:expr, |$invoke_context:ident| $inner:tt) => {
        $crate::solana_sbpf::declare_builtin_function!(
            $process_instruction,
            fn rust(
                invoke_context: &mut $crate::invoke_context::InvokeContext<'_, '_>,
                _arg0: u64,
                _arg1: u64,
                _arg2: u64,
                _arg3: u64,
                _arg4: u64,
                _memory_mapping: &mut $crate::solana_sbpf::memory_region::MemoryMapping<'_>,
            ) -> std::result::Result<u64, Box<dyn std::error::Error>> {
                fn process_instruction_inner(
                    $invoke_context: &mut $crate::invoke_context::InvokeContext<'_, '_>,
                ) -> std::result::Result<(), $crate::__private::InstructionError>
                    $inner

                let consumption_result = if $cu_to_consume > 0
                {
                    invoke_context.consume_checked($cu_to_consume)
                } else {
                    Ok(())
                };
                consumption_result
                    .and_then(|_| {
                        process_instruction_inner(invoke_context)
                            .map(|_| 0)
                            .map_err(|err| Box::new(err) as Box<dyn std::error::Error>)
                    })
                    .into()
            }
        );
    };
}

/// Account insert error.
pub enum AccountInsertError {
    /// Too many accounts have been loaded.
    TooManyAccounts,

    /// Account has already been loaded.
    AccountAlreadyLoaded,
}

/// Runtime account loader.
///
/// Used by the RISC-V VM to load accounts into VM during runtime of the program.
#[allow(clippy::result_unit_err)]
pub trait RuntimeAccountLoader {
    /// Attempt to load the account info of `pubkey`.
    ///
    /// Returns `Ok(Some(AccountSharedData))` on success.
    /// Returns `Ok(None)` if the account was not found.
    /// Returns `Err(())` if there was an error with account loading that should abort the program.
    fn load_account(&self, pubkey: &Pubkey) -> Result<Option<StoredAccount>, ()>;
}

impl ContextObject for InvokeContext<'_, '_> {
    fn trace(&mut self, state: [u64; 12]) {
        self.syscall_context
            .last_mut()
            .unwrap()
            .as_mut()
            .unwrap()
            .trace_log
            .push(state);
    }

    fn consume(&mut self, amount: u64) {
        // 1 to 1 instruction to compute unit mapping
        // ignore overflow, Ebpf will bail if exceeded
        let mut compute_meter = self.compute_meter.borrow_mut();
        *compute_meter = compute_meter.saturating_sub(amount);
    }

    fn get_remaining(&self) -> u64 {
        *self.compute_meter.borrow()
    }
}

#[derive(Clone, PartialEq, Eq, Debug)]
pub struct AllocErr;
impl fmt::Display for AllocErr {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str("Error: Memory allocation failed")
    }
}

pub struct BpfAllocator {
    len: u64,
    pos: u64,
}

impl BpfAllocator {
    pub fn new(len: u64) -> Self {
        Self { len, pos: 0 }
    }

    pub fn alloc(&mut self, layout: Layout) -> Result<u64, AllocErr> {
        let bytes_to_align = (self.pos as *const u8).align_offset(layout.align()) as u64;
        if self
            .pos
            .saturating_add(bytes_to_align)
            .saturating_add(layout.size() as u64)
            <= self.len
        {
            self.pos = self.pos.saturating_add(bytes_to_align);
            let addr = MM_HEAP_START.saturating_add(self.pos);
            self.pos = self.pos.saturating_add(layout.size() as u64);
            Ok(addr)
        } else {
            Err(AllocErr)
        }
    }
}

pub struct EnvironmentConfig<'a> {
    pub blockhash: Hash,
    pub blockhash_kelvins_per_signature: u64,
    get_epoch_vote_account_stake_callback: &'a dyn Fn(&'a Pubkey) -> u64,
    pub feature_set: Arc<FeatureSet>,
    sysvar_cache: &'a SysvarCache,
    pub random_seed: u64,
    /// Handle to freeze the stake cache during instruction execution.
    /// Builtin programs can call `freeze_stakes()` to perform an O(1) swap
    /// of the pending stake cache data to the frozen history. The handle also
    /// contains the pending stake cache data for read-only access, including
    /// the next epoch number.
    ///
    /// The handle also contains the `epoch_rewards_exists_fn` callback for
    /// checking if an EpochRewards account exists for a given epoch.
    stakes_handle: StakesHandle,
}
impl<'a> EnvironmentConfig<'a> {
    #[allow(clippy::too_many_arguments)]
    pub fn new(
        blockhash: Hash,
        blockhash_kelvins_per_signature: u64,
        get_epoch_vote_account_stake_callback: &'a dyn Fn(&'a Pubkey) -> u64,
        feature_set: Arc<FeatureSet>,
        sysvar_cache: &'a SysvarCache,
        random_seed: u64,
        stakes_handle: StakesHandle,
    ) -> Self {
        Self {
            blockhash,
            blockhash_kelvins_per_signature,
            get_epoch_vote_account_stake_callback,
            feature_set,
            sysvar_cache,
            random_seed,
            stakes_handle,
        }
    }
}

pub struct SyscallContext {
    pub allocator: BpfAllocator,
    pub accounts_metadata: Vec<SerializedAccountMetadata>,
    pub trace_log: Vec<[u64; 12]>,
}

#[derive(Debug, Clone)]
pub struct SerializedAccountMetadata {
    pub original_data_len: usize,
    pub vm_data_addr: u64,
    pub vm_key_addr: u64,
    pub vm_kelvins_addr: u64,
    pub vm_owner_addr: u64,
}

/// Main pipeline from runtime to program execution.
pub struct InvokeContext<'a, 'b> {
    /// Information about the currently executing transaction.
    pub transaction_context: &'a mut TransactionContext,
    /// The local program cache for the transaction batch.
    pub program_cache_for_tx_batch: &'a mut ProgramCacheForTx<'b>,
    /// Runtime configurations used to provision the invocation environment.
    pub environment_config: EnvironmentConfig<'a>,
    /// The compute budget for the current invocation.
    compute_budget: ComputeBudget,
    /// Instruction compute meter, for tracking compute units consumed against
    /// the designated compute budget during program execution.
    compute_meter: RefCell<u64>,
    log_collector: Option<Rc<RefCell<LogCollector>>>,
    /// Latest measurement not yet accumulated in [ExecuteDetailsTimings::execute_us]
    pub execute_time: Option<Measure>,
    pub timings: ExecuteDetailsTimings,
    pub syscall_context: Vec<Option<SyscallContext>>,
    traces: Vec<Vec<[u64; 12]>>,
    pub account_loader: Option<Box<dyn RuntimeAccountLoader + 'a>>,
    /// Map of accounts already loaded in the VM.
    ///
    /// This is shared across CPI calls of the program and is used to prevent programs from reloading accounts
    /// and/or loading accounts other programs have already loaded.
    loaded_accounts: HashSet<Pubkey>,
    /// How many writable accounts can be loaded.
    num_account_loads: usize,
}

impl<'a, 'b> InvokeContext<'a, 'b> {
    #[allow(clippy::too_many_arguments)]
    pub fn new(
        transaction_context: &'a mut TransactionContext,
        program_cache_for_tx_batch: &'a mut ProgramCacheForTx<'b>,
        environment_config: EnvironmentConfig<'a>,
        log_collector: Option<Rc<RefCell<LogCollector>>>,
        compute_budget: ComputeBudget,
    ) -> Self {
        let loaded_accounts = transaction_context.account_keys().copied().collect();

        Self {
            transaction_context,
            program_cache_for_tx_batch,
            environment_config,
            log_collector,
            compute_budget,
            compute_meter: RefCell::new(compute_budget.compute_unit_limit),
            execute_time: None,
            timings: ExecuteDetailsTimings::default(),
            syscall_context: Vec::new(),
            traces: Vec::new(),
            account_loader: None,
            loaded_accounts,
            num_account_loads: 0usize,
        }
    }

    pub fn new_with_account_loader(
        transaction_context: &'a mut TransactionContext,
        program_cache_for_tx_batch: &'a mut ProgramCacheForTx<'b>,
        environment_config: EnvironmentConfig<'a>,
        log_collector: Option<Rc<RefCell<LogCollector>>>,
        compute_budget: ComputeBudget,
        account_loader: Box<dyn RuntimeAccountLoader + 'a>,
        num_account_locks: usize,
        num_writable_accounts: usize,
    ) -> Self {
        let loaded_accounts = transaction_context.account_keys().copied().collect();

        Self {
            transaction_context,
            program_cache_for_tx_batch,
            environment_config,
            log_collector,
            compute_budget,
            compute_meter: RefCell::new(compute_budget.compute_unit_limit),
            execute_time: None,
            timings: ExecuteDetailsTimings::default(),
            syscall_context: Vec::new(),
            traces: Vec::new(),
            account_loader: Some(account_loader),
            loaded_accounts,
            num_account_loads: num_account_locks.saturating_sub(num_writable_accounts),
        }
    }

    /// Attempt to loaded account `pubkey` into `InvokeContext`.
    ///
    /// Returns `Ok(())` on success
    /// Returns `Err(AccountInsertError::TooManyAccounts)` if too many accounts have been loaded.
    /// Returns `Err(AccountInsertError::AccountAlreadyLoaded)` if account has already been loaded.
    pub fn add_loaded_account(&mut self, pubkey: Pubkey) -> Result<(), AccountInsertError> {
        match self.num_account_loads.checked_sub(1) {
            Some(value) => self.num_account_loads = value,
            None => return Err(AccountInsertError::TooManyAccounts),
        }

        if self.loaded_accounts.insert(pubkey) {
            return Ok(());
        }

        Err(AccountInsertError::AccountAlreadyLoaded)
    }

    pub fn get_environments_for_slot(
        &self,
        effective_slot: Slot,
    ) -> Result<&ProgramRuntimeEnvironments, InstructionError> {
        let epoch_schedule = self.environment_config.sysvar_cache.get_epoch_schedule()?;
        let epoch = epoch_schedule.get_epoch(effective_slot);
        Ok(self
            .program_cache_for_tx_batch
            .get_environments_for_epoch(epoch))
    }

    /// Push a stack frame onto the invocation stack
    pub fn push(&mut self) -> Result<(), InstructionError> {
        let instruction_context = self
            .transaction_context
            .get_instruction_context_at_index_in_trace(
                self.transaction_context.get_instruction_trace_length(),
            )?;
        let program_id = instruction_context
            .get_last_program_key(self.transaction_context)
            .map_err(|_| InstructionError::UnsupportedProgramId)?;
        if self
            .transaction_context
            .get_instruction_context_stack_height()
            != 0
        {
            let contains = (0..self
                .transaction_context
                .get_instruction_context_stack_height())
                .any(|level| {
                    self.transaction_context
                        .get_instruction_context_at_nesting_level(level)
                        .and_then(|instruction_context| {
                            instruction_context
                                .try_borrow_last_program_account(self.transaction_context)
                        })
                        .map(|program_account| program_account.get_key() == program_id)
                        .unwrap_or(false)
                });
            let is_last = self
                .transaction_context
                .get_current_instruction_context()
                .and_then(|instruction_context| {
                    instruction_context.try_borrow_last_program_account(self.transaction_context)
                })
                .map(|program_account| program_account.get_key() == program_id)
                .unwrap_or(false);
            if contains && !is_last {
                // Reentrancy not allowed unless caller is calling itself
                return Err(InstructionError::ReentrancyNotAllowed);
            }
        }

        self.syscall_context.push(None);
        self.transaction_context.push()
    }

    /// Pop a stack frame from the invocation stack
    fn pop(&mut self) -> Result<(), InstructionError> {
        if let Some(Some(syscall_context)) = self.syscall_context.pop() {
            self.traces.push(syscall_context.trace_log);
        }
        self.transaction_context.pop()
    }

    /// Current height of the invocation stack, top level instructions are height
    /// `rialo_s_instruction::TRANSACTION_LEVEL_STACK_HEIGHT`
    pub fn get_stack_height(&self) -> usize {
        self.transaction_context
            .get_instruction_context_stack_height()
    }

    /// Entrypoint for a cross-program invocation from a builtin program
    pub fn native_invoke(
        &mut self,
        instruction: StableInstruction,
        signers: &[Pubkey],
    ) -> Result<(), InstructionError> {
        let (instruction_accounts, program_indices) =
            self.prepare_instruction(&instruction, signers)?;
        let mut compute_units_consumed = 0;
        self.process_instruction(
            &instruction.data,
            &instruction_accounts,
            &program_indices,
            &mut compute_units_consumed,
            &mut ExecuteTimings::default(),
        )?;
        Ok(())
    }

    /// Helper to prepare for process_instruction()
    #[allow(clippy::type_complexity)]
    pub fn prepare_instruction(
        &mut self,
        instruction: &StableInstruction,
        signers: &[Pubkey],
    ) -> Result<(Vec<InstructionAccount>, Vec<IndexOfAccount>), InstructionError> {
        self.prepare_instruction_inner(instruction.program_id, &instruction.accounts, signers)
    }

    /// Helper to prepare for process_instruction()
    pub fn prepare_cpi_instruction(
        &mut self,
        program_id: Pubkey,
        account_metas: &[AccountMeta],
        signers: &[Pubkey],
    ) -> Result<(Vec<InstructionAccount>, Vec<IndexOfAccount>), InstructionError> {
        self.prepare_instruction_inner(program_id, account_metas, signers)
    }

    pub fn prepare_instruction_inner(
        &mut self,
        callee_program_id: Pubkey,
        account_metas: &[AccountMeta],
        signers: &[Pubkey],
    ) -> Result<(Vec<InstructionAccount>, Vec<IndexOfAccount>), InstructionError> {
        // Finds the index of each account in the instruction by its pubkey.
        // Then normalizes / unifies the privileges of duplicate accounts.
        // Note: This is an O(n^2) algorithm,
        // but performed on a very small slice and requires no heap allocations.
        let instruction_context = self.transaction_context.get_current_instruction_context()?;
        let mut deduplicated_instruction_accounts: Vec<InstructionAccount> = Vec::new();
        let mut duplicate_indicies = Vec::with_capacity(account_metas.len());
        for (instruction_account_index, account_meta) in account_metas.iter().enumerate() {
            let index_in_transaction = self
                .transaction_context
                .find_index_of_account(&account_meta.pubkey)
                .ok_or_else(|| {
                    ic_msg!(
                        self,
                        "Instruction references an unknown account {}",
                        account_meta.pubkey,
                    );
                    InstructionError::MissingAccount
                })?;
            if let Some(duplicate_index) =
                deduplicated_instruction_accounts
                    .iter()
                    .position(|instruction_account| {
                        instruction_account.index_in_transaction == index_in_transaction
                    })
            {
                duplicate_indicies.push(duplicate_index);
                let instruction_account = deduplicated_instruction_accounts
                    .get_mut(duplicate_index)
                    .ok_or(InstructionError::NotEnoughAccountKeys)?;
                instruction_account.is_signer |= account_meta.is_signer;
                instruction_account.is_writable |= account_meta.is_writable;
            } else {
                let index_in_caller = instruction_context
                    .find_index_of_instruction_account(
                        self.transaction_context,
                        &account_meta.pubkey,
                    )
                    .ok_or_else(|| {
                        ic_msg!(
                            self,
                            "Instruction references an unknown account {}",
                            account_meta.pubkey,
                        );
                        InstructionError::MissingAccount
                    })?;
                duplicate_indicies.push(deduplicated_instruction_accounts.len());
                deduplicated_instruction_accounts.push(InstructionAccount {
                    index_in_transaction,
                    index_in_caller,
                    index_in_callee: instruction_account_index as IndexOfAccount,
                    is_signer: account_meta.is_signer,
                    is_writable: account_meta.is_writable,
                });
            }
        }
        for instruction_account in deduplicated_instruction_accounts.iter() {
            let borrowed_account = instruction_context.try_borrow_instruction_account(
                self.transaction_context,
                instruction_account.index_in_caller,
            )?;

            // Readonly in caller cannot become writable in callee
            if instruction_account.is_writable && !borrowed_account.is_writable() {
                ic_msg!(
                    self,
                    "{}'s writable privilege escalated",
                    borrowed_account.get_key(),
                );
                return Err(InstructionError::PrivilegeEscalation);
            }

            // To be signed in the callee,
            // it must be either signed in the caller or by the program
            if instruction_account.is_signer
                && !(borrowed_account.is_signer() || signers.contains(borrowed_account.get_key()))
            {
                ic_msg!(
                    self,
                    "{}'s signer privilege escalated",
                    borrowed_account.get_key()
                );
                return Err(InstructionError::PrivilegeEscalation);
            }
        }
        let instruction_accounts = duplicate_indicies
            .into_iter()
            .map(|duplicate_index| {
                deduplicated_instruction_accounts
                    .get(duplicate_index)
                    .cloned()
                    .ok_or(InstructionError::NotEnoughAccountKeys)
            })
            .collect::<Result<Vec<InstructionAccount>, InstructionError>>()?;

        // Find and validate executables / program accounts
        let program_account_index = if self
            .get_feature_set()
            .is_active(&lift_cpi_caller_restriction::id())
        {
            match (
                self.transaction_context
                    .find_index_of_program_account(&callee_program_id),
                &self.account_loader,
            ) {
                (Some(index), _) => index,
                (None, Some(account_loader)) => {
                    match account_loader.load_account(&callee_program_id) {
                        Ok(Some(account)) => {
                            self.transaction_context
                                .add_account((callee_program_id, account.clone()));

                            // program must not exist in `loaded_accounts` since it was not found in transaction context
                            assert!(self.loaded_accounts.insert(callee_program_id));

                            // program must exist since it was just added into transaction context
                            self.transaction_context
                                .find_index_of_program_account(&callee_program_id)
                                .expect("program to exist")
                        }
                        Ok(None) => return Err(InstructionError::MissingAccount),
                        Err(()) => return Err(InstructionError::GenericError),
                    }
                }
                (None, None) => {
                    ic_msg!(self, "Unknown program {}", callee_program_id);
                    return Err(InstructionError::MissingAccount);
                }
            }
        } else {
            let program_account_index = instruction_context
                .find_index_of_instruction_account(self.transaction_context, &callee_program_id)
                .ok_or_else(|| {
                    ic_msg!(self, "Unknown program {}", callee_program_id);
                    InstructionError::MissingAccount
                })?;
            let borrowed_program_account = instruction_context
                .try_borrow_instruction_account(self.transaction_context, program_account_index)?;
            #[allow(deprecated)]
            if !self
                .get_feature_set()
                .is_active(&remove_accounts_executable_flag_checks::id())
                && !borrowed_program_account.is_executable()
            {
                ic_msg!(self, "Account {} is not executable", callee_program_id);
                return Err(InstructionError::AccountNotExecutable);
            }
            borrowed_program_account.get_index_in_transaction()
        };

        Ok((instruction_accounts, vec![program_account_index]))
    }

    /// Processes an instruction and returns how many compute units were used
    pub fn process_instruction(
        &mut self,
        instruction_data: &[u8],
        instruction_accounts: &[InstructionAccount],
        program_indices: &[IndexOfAccount],
        compute_units_consumed: &mut u64,
        timings: &mut ExecuteTimings,
    ) -> Result<(), InstructionError> {
        *compute_units_consumed = 0;
        self.transaction_context
            .get_next_instruction_context()?
            .configure(program_indices, instruction_accounts, instruction_data);
        self.push()?;
        self.process_executable_chain(compute_units_consumed, timings)
            // MUST pop if and only if `push` succeeded, independent of `result`.
            // Thus, the `.and()` instead of an `.and_then()`.
            .and(self.pop())
    }

    /// Processes a precompile instruction
    pub fn process_precompile<'ix_data>(
        &mut self,
        precompile: &Precompile,
        instruction_data: &[u8],
        instruction_accounts: &[InstructionAccount],
        program_indices: &[IndexOfAccount],
        message_instruction_datas_iter: impl Iterator<Item = &'ix_data [u8]>,
    ) -> Result<(), InstructionError> {
        self.transaction_context
            .get_next_instruction_context()?
            .configure(program_indices, instruction_accounts, instruction_data);
        self.push()?;

        let feature_set = self.get_feature_set();
        let move_precompile_verification_to_svm =
            feature_set.is_active(&move_precompile_verification_to_svm::id());
        if move_precompile_verification_to_svm {
            let instruction_datas: Vec<_> = message_instruction_datas_iter.collect();
            precompile
                .verify(instruction_data, &instruction_datas, feature_set)
                .map_err(InstructionError::from)
                .and(self.pop())
        } else {
            self.pop()
        }
    }

    /// Calls the instruction's program entrypoint method
    fn process_executable_chain(
        &mut self,
        compute_units_consumed: &mut u64,
        timings: &mut ExecuteTimings,
    ) -> Result<(), InstructionError> {
        let instruction_context = self.transaction_context.get_current_instruction_context()?;
        let process_executable_chain_time = Measure::start("process_executable_chain_time");

        let builtin_id = {
            debug_assert!(instruction_context.get_number_of_program_accounts() <= 1);
            let borrowed_root_account = instruction_context
                .try_borrow_program_account(self.transaction_context, 0)
                .map_err(|_| InstructionError::UnsupportedProgramId)?;
            let owner_id = borrowed_root_account.get_owner();
            if native_loader::check_id(owner_id) {
                *borrowed_root_account.get_key()
            } else {
                *owner_id
            }
        };

        // The Murmur3 hash value (used by RBPF) of the string "entrypoint"
        const ENTRYPOINT_KEY: u32 = 0x71E3CF81;
        let entry = self
            .program_cache_for_tx_batch
            .find(&builtin_id)
            .ok_or(InstructionError::UnsupportedProgramId)?;
        let function = match &entry.program {
            ProgramCacheEntryType::Builtin(program) => program
                .get_function_registry()
                .lookup_by_key(ENTRYPOINT_KEY)
                .map(|(_name, function)| function),
            _ => None,
        }
        .ok_or(InstructionError::UnsupportedProgramId)?;
        entry.ix_usage_counter.fetch_add(1, Ordering::Relaxed);

        let program_id = *instruction_context.get_last_program_key(self.transaction_context)?;
        self.transaction_context
            .set_return_data(program_id, Vec::new())?;
        let logger = self.get_log_collector();
        stable_log::program_invoke(&logger, &program_id, self.get_stack_height());
        let pre_remaining_units = self.get_remaining();
        // In program-runtime v2 we will create this VM instance only once per transaction.
        // `program_runtime_environment_v2.get_config()` will be used instead of `mock_config`.
        // For now, only built-ins are invoked from here, so the VM and its Config are irrelevant.
        let mock_config = Config::default();
        let empty_memory_mapping =
            MemoryMapping::new(Vec::new(), &mock_config, SBPFVersion::V0).unwrap();
        let mut vm = EbpfVm::new(
            self.program_cache_for_tx_batch
                .environments
                .program_runtime_v2
                .clone(),
            SBPFVersion::V0,
            // Removes lifetime tracking
            unsafe {
                std::mem::transmute::<&mut InvokeContext<'_, '_>, &mut InvokeContext<'_, '_>>(self)
            },
            empty_memory_mapping,
            0,
        );
        vm.invoke_function(function);
        let result = match vm.program_result {
            ProgramResult::Ok(_) => {
                stable_log::program_success(&logger, &program_id);
                Ok(())
            }
            ProgramResult::Err(ref err) => {
                if let EbpfError::SyscallError(syscall_error) = err {
                    if let Some(instruction_err) = syscall_error.downcast_ref::<InstructionError>()
                    {
                        stable_log::program_failure(&logger, &program_id, instruction_err);
                        Err(instruction_err.clone())
                    } else {
                        stable_log::program_failure(&logger, &program_id, syscall_error);
                        Err(InstructionError::ProgramFailedToComplete)
                    }
                } else {
                    stable_log::program_failure(&logger, &program_id, err);
                    Err(InstructionError::ProgramFailedToComplete)
                }
            }
        };
        let post_remaining_units = self.get_remaining();
        *compute_units_consumed = pre_remaining_units.saturating_sub(post_remaining_units);

        if builtin_id == program_id && result.is_ok() && *compute_units_consumed == 0 {
            return Err(InstructionError::BuiltinProgramsMustConsumeComputeUnits);
        }

        timings
            .execute_accessories
            .process_instructions
            .process_executable_chain_us += process_executable_chain_time.end_as_us();
        result
    }

    /// Get this invocation's LogCollector
    pub fn get_log_collector(&self) -> Option<Rc<RefCell<LogCollector>>> {
        self.log_collector.clone()
    }

    /// Consume compute units
    pub fn consume_checked(&self, amount: u64) -> Result<(), Box<dyn std::error::Error>> {
        let mut compute_meter = self.compute_meter.borrow_mut();
        let exceeded = *compute_meter < amount;
        *compute_meter = compute_meter.saturating_sub(amount);
        if exceeded {
            return Err(Box::new(InstructionError::ComputationalBudgetExceeded));
        }
        Ok(())
    }

    /// Set compute units
    ///
    /// Only use for tests and benchmarks
    pub fn mock_set_remaining(&self, remaining: u64) {
        *self.compute_meter.borrow_mut() = remaining;
    }

    /// Get this invocation's compute budget
    pub fn get_compute_budget(&self) -> &ComputeBudget {
        &self.compute_budget
    }

    /// Get the current feature set.
    pub fn get_feature_set(&self) -> &FeatureSet {
        &self.environment_config.feature_set
    }

    /// Set feature set.
    ///
    /// Only use for tests and benchmarks.
    pub fn mock_set_feature_set(&mut self, feature_set: Arc<FeatureSet>) {
        self.environment_config.feature_set = feature_set;
    }

    /// Get cached sysvars
    pub fn get_sysvar_cache(&self) -> &SysvarCache {
        self.environment_config.sysvar_cache
    }

    /// Get cached stake for the epoch vote account.
    pub fn get_epoch_vote_account_stake(&self, pubkey: &'a Pubkey) -> u64 {
        (self
            .environment_config
            .get_epoch_vote_account_stake_callback)(pubkey)
    }

    /// Get the CURRENT epoch (the epoch that is actively running).
    ///
    /// This is the epoch to use when checking stake activation/deactivation status,
    /// unbonding periods, lockup periods, etc.
    ///
    /// # Implementation Note
    /// The stake cache stores the NEXT epoch (the epoch being prepared for when
    /// stake changes become effective). We subtract 1 to get the current epoch.
    /// `saturating_sub` handles the edge case of epoch 0.
    pub fn get_current_epoch(&self) -> u64 {
        self.environment_config
            .stakes_handle
            .pending_epoch()
            .saturating_sub(1)
    }

    /// Get the NEXT epoch (the epoch being prepared in the stake cache).
    ///
    /// This is the epoch that accumulated stake changes will become effective in
    /// AFTER the next FreezeStakes call. Most instruction logic should use
    /// `get_current_epoch()` instead; this method is primarily for internal use
    /// and the FreezeStakes instruction.
    pub fn get_next_epoch(&self) -> u64 {
        self.environment_config.stakes_handle.pending_epoch()
    }

    /// Get the timestamp of the last FreezeStakes call (last epoch boundary).
    ///
    /// This is the timestamp from `history.back()` which represents when the most
    /// recent epoch boundary occurred. Use this for **state transition checks**:
    /// - A stake is "activating" while `activation_requested >= last_freeze_timestamp`
    /// - A stake becomes "activated" when `activation_requested < last_freeze_timestamp`
    /// - A stake is "deactivating" while `deactivation_requested >= last_freeze_timestamp`
    /// - A stake becomes "deactivated" when `deactivation_requested < last_freeze_timestamp`
    ///
    /// Returns 0 if no history exists (should not happen in normal operation since
    /// `StakingState::new_with_initial_history()` guarantees at least one entry).
    ///
    /// **Note:** For duration enforcement (lockup, unbonding), use the current timestamp
    /// from Clock sysvar instead, to provide real-time guarantees.
    pub fn get_last_freeze_timestamp(&self) -> u64 {
        self.environment_config
            .stakes_handle
            .last_frozen_timestamp()
            .unwrap_or(0)
    }

    /// Set the last freeze timestamp for testing.
    ///
    /// This is the primary mock function for timestamp-based tests. It sets up the
    /// stake cache so that `get_last_freeze_timestamp()` returns the specified timestamp.
    ///
    /// The function:
    /// 1. Sets the pending epoch to 1 (to avoid underflow when computing current epoch)
    /// 2. Pushes a `StakeCacheData` entry to history with epoch=0 and the specified timestamp
    ///
    /// Only use for tests and benchmarks.
    pub fn mock_set_last_freeze_timestamp(&mut self, timestamp: u64) {
        // Set pending epoch to 1 to avoid underflow when computing current epoch (pending - 1)
        self.environment_config.stakes_handle.set_pending_epoch(1);

        // Push a history entry with the specified timestamp
        let history_entry = StakeCacheData {
            epoch: 0,
            timestamp,
            ..Default::default()
        };
        self.environment_config
            .stakes_handle
            .push_frozen(history_entry);
    }

    /// Insert a stake account into the pending cache for testing.
    ///
    /// This allows unit tests to set up stake accounts that reference validators,
    /// which is needed to test the validator withdrawal blocking logic.
    ///
    /// Only use for tests and benchmarks.
    pub fn mock_insert_stake_account(
        &mut self,
        pubkey: rialo_s_pubkey::Pubkey,
        account: rialo_stake_cache_interface::StakeAccount,
    ) {
        self.environment_config
            .stakes_handle
            .insert_stake_account(pubkey, account);
    }

    /// Freeze the pending stake cache data.
    ///
    /// This performs an O(1) swap of the pending stake cache data using `std::mem::take()`
    /// and pushes it to the back of the frozen queue. This is typically called by the
    /// ValidatorRegistry program's FreezeStakes instruction to capture the validator set
    /// at a specific point.
    ///
    /// This method also sets the `epoch_stakes_frozen` flag to signal the Bank to
    /// propagate `new_commission_rate` → `commission_rate`, etc.
    ///
    /// To access the frozen validator data after calling this method, use
    /// `get_all_validator_accounts_from_last_frozen()`.
    pub fn freeze_stakes(&self) {
        self.environment_config.stakes_handle.freeze_stakes();
    }

    /// Signal that FreezeStakes was requested without performing the swap.
    ///
    /// The actual pending → frozen swap is deferred to `finalize_impl()` to ensure
    /// pending modifications are persisted first. This only sets the `epoch_stakes_frozen`
    /// flag, which is read by `apply_pending_validator_changes_if_needed()` and consumed
    /// by `finalize_impl()`.
    pub fn signal_freeze_stakes(&self) {
        self.environment_config
            .stakes_handle
            .set_epoch_stakes_frozen();
    }

    /// Get all validator accounts from the last frozen epoch (current epoch state).
    ///
    /// Returns a vector of `(pubkey, ValidatorAccount)` pairs for all validators,
    /// sorted by pubkey. This performs a layered lookup: frozen (newest to oldest) → baseline,
    /// skipping pending (next epoch changes).
    ///
    /// This method allows builtin programs (like TokenomicsGovernance) to access
    /// the frozen validator data directly without needing CPI return data.
    pub fn get_all_validator_accounts_from_last_frozen(&self) -> Vec<(Pubkey, ValidatorAccount)> {
        self.environment_config
            .stakes_handle
            .get_all_validator_accounts_from_last_frozen()
    }

    /// Get the length of the frozen stake history.
    ///
    /// Used by DistributeRewards to check if there are enough frozen snapshots to merge.
    /// We require at least 2 entries to allow popping (to preserve the current epoch snapshot).
    pub fn frozen_stake_history_len(&self) -> usize {
        self.environment_config.stakes_handle.frozen_len()
    }

    /// Request epoch rewards initialization.
    ///
    /// Called by DistributeRewards instruction to signal the Bank to create an EpochRewards
    /// account for tracking partitioned reward distribution. The Bank will detect this signal
    /// after transaction execution and create the account outside of instruction processing.
    ///
    /// # Arguments
    /// * `epoch` - The epoch number for which rewards are being distributed
    /// * `total_rewards` - The total amount of rewards to be distributed this epoch
    pub fn request_epoch_rewards_init(&self, epoch: u64, total_rewards: u64) {
        self.environment_config
            .stakes_handle
            .request_epoch_rewards_init(epoch, total_rewards);
    }

    /// Get the epoch of the oldest frozen snapshot (front of the queue).
    ///
    /// This is the epoch that will be popped by `pop_front_and_merge_to_baseline()`.
    /// Used by DistributeRewards to know which epoch rewards are being distributed.
    ///
    /// Returns `None` if no frozen snapshots exist.
    pub fn front_frozen_epoch(&self) -> Option<u64> {
        self.environment_config.stakes_handle.front_frozen_epoch()
    }

    /// Check if any stake account references the given validator whose unbonding
    /// period is NOT yet complete.
    ///
    /// This performs an O(n) search over all stake accounts using a layered lookup
    /// (pending → frozen → baseline) to determine if the validator is referenced
    /// by an active or still-unbonding stake.
    ///
    /// Used by ValidatorRegistry's Withdraw handler to prevent withdrawal below
    /// rent-exempt minimum while stake accounts still reference the validator.
    ///
    /// A stake is considered to "reference" the validator if:
    /// - It is delegated to this validator, AND
    /// - Either active (no deactivation) OR still unbonding (unbonding not complete)
    ///
    /// Stake accounts whose unbonding is complete are NOT counted, since they can
    /// be fully withdrawn or reactivated to another validator.
    ///
    /// # Arguments
    /// * `validator` - The validator pubkey to check
    /// * `validator_info` - The validator's info (used to compute unbonding end via `end_of_unbonding`)
    /// * `last_freeze_timestamp` - Timestamp from last FreezeStakes (epoch boundary)
    /// * `current_timestamp` - Current block timestamp from Clock sysvar
    pub fn is_validator_referenced(
        &self,
        validator: &Pubkey,
        validator_info: &rialo_stake_cache_interface::ValidatorInfo,
        last_freeze_timestamp: u64,
        current_timestamp: u64,
    ) -> bool {
        self.environment_config
            .stakes_handle
            .is_validator_referenced(
                validator,
                validator_info,
                last_freeze_timestamp,
                current_timestamp,
            )
    }

    /// Check if any stake account delegated to the given validator is still within
    /// its lockup period.
    ///
    /// Delegates to `StakesHandle::has_locked_stakers()` which performs an O(n) scan
    /// over all stake accounts. A staker is "locked" if their
    /// `activation_requested + lockup_period > current_timestamp`.
    ///
    /// # Arguments
    /// * `validator` - The validator pubkey to check
    /// * `lockup_period` - The validator's lockup period in milliseconds
    /// * `current_timestamp` - Current block timestamp from Clock sysvar (in milliseconds)
    ///
    /// # Returns
    /// `true` if at least one staker is still within their lockup period.
    pub fn has_locked_stakers(
        &self,
        validator: &Pubkey,
        lockup_period: u64,
        current_timestamp: u64,
    ) -> bool {
        self.environment_config.stakes_handle.has_locked_stakers(
            validator,
            lockup_period,
            current_timestamp,
        )
    }

    /// Check if a validator is referenced by any stake accounts (excluding the self-bond).
    ///
    /// This variant excludes the self-bond PDA from the check to prevent circular logic
    /// where the self-bond cannot be deactivated because its existence always makes
    /// is_validator_referenced() return true.
    pub fn is_validator_referenced_excluding_self_bond(
        &self,
        validator_pubkey: &Pubkey,
        validator_info: &rialo_stake_cache_interface::ValidatorInfo,
        last_freeze_timestamp: u64,
        current_timestamp: u64,
    ) -> bool {
        self.environment_config
            .stakes_handle
            .is_validator_referenced_excluding_self_bond(
                validator_pubkey,
                validator_info,
                last_freeze_timestamp,
                current_timestamp,
            )
    }

    /// Look up a stake account by pubkey from the pending layer (next epoch state).
    ///
    /// Searches: pending → frozen (newest to oldest) → baseline.
    /// Returns `Some(StakeAccount)` if found, `None` if not found or tombstoned.
    pub fn get_stake_account_from_pending(
        &self,
        pubkey: &Pubkey,
    ) -> Option<rialo_stake_cache_interface::StakeAccount> {
        self.environment_config
            .stakes_handle
            .get_stake_account_from_pending(pubkey)
    }

    /// Check if an epoch rewards initialization request is pending for this block.
    ///
    /// Used by DistributeRewards to prevent multiple reward distributions in the same block.
    /// If this returns true, the DistributeRewards instruction should fail with InvalidArgument.
    pub fn is_epoch_rewards_init_pending(&self) -> bool {
        self.environment_config
            .stakes_handle
            .is_epoch_rewards_init_pending()
    }

    /// Get completed frozen epochs (excludes the last/current epoch).
    ///
    /// Returns epochs in order from oldest to newest. These are the epochs that have
    /// completed and are eligible for reward distribution. We exclude the last frozen
    /// epoch because it represents the current epoch that is still accumulating rewards.
    pub fn completed_frozen_epochs(&self) -> Vec<u64> {
        self.environment_config
            .stakes_handle
            .completed_frozen_epochs()
    }

    /// Check if an EpochRewards account exists for the given epoch.
    ///
    /// Delegates to the StakesHandle which contains the callback provided by Bank
    /// to query the StateStore. This allows DistributeRewards to find the first
    /// completed frozen epoch that doesn't yet have an EpochRewards account.
    pub fn epoch_rewards_exists(&self, epoch: u64) -> bool {
        self.environment_config
            .stakes_handle
            .epoch_rewards_exists(epoch)
    }

    // Should alignment be enforced during user pointer translation
    pub fn get_check_aligned(&self) -> bool {
        self.transaction_context
            .get_current_instruction_context()
            .and_then(|instruction_context| {
                let program_account =
                    instruction_context.try_borrow_last_program_account(self.transaction_context);
                debug_assert!(program_account.is_ok());
                program_account
            })
            .map(|program_account| *program_account.get_owner() != bpf_loader_deprecated::id())
            .unwrap_or(true)
    }

    // Set this instruction syscall context
    pub fn set_syscall_context(
        &mut self,
        syscall_context: SyscallContext,
    ) -> Result<(), InstructionError> {
        *self
            .syscall_context
            .last_mut()
            .ok_or(InstructionError::CallDepth)? = Some(syscall_context);
        Ok(())
    }

    // Get this instruction's SyscallContext
    pub fn get_syscall_context(&self) -> Result<&SyscallContext, InstructionError> {
        self.syscall_context
            .last()
            .and_then(std::option::Option::as_ref)
            .ok_or(InstructionError::CallDepth)
    }

    // Get this instruction's SyscallContext
    pub fn get_syscall_context_mut(&mut self) -> Result<&mut SyscallContext, InstructionError> {
        self.syscall_context
            .last_mut()
            .and_then(|syscall_context| syscall_context.as_mut())
            .ok_or(InstructionError::CallDepth)
    }

    /// Return a references to traces
    pub fn get_traces(&self) -> &Vec<Vec<[u64; 12]>> {
        &self.traces
    }
}

#[macro_export]
macro_rules! with_mock_invoke_context {
    (
        $invoke_context:ident,
        $transaction_context:ident,
        $entry:expr,
        $transaction_accounts:expr $(,)?
    ) => {
        use rialo_s_compute_budget::compute_budget::ComputeBudget;
        use rialo_s_feature_set::FeatureSet;
        use rialo_s_log_collector::LogCollector;
        use rialo_s_type_overrides::sync::Arc;
        use $crate::{
            __private::{Hash, ReadableAccount, Rent, TransactionContext},
            invoke_context::{EnvironmentConfig, InvokeContext},
            loaded_programs::{ProgramCacheEntry, ProgramCacheForTx, ProgramCacheForTxBatch},
            sysvar_cache::SysvarCache,
        };
        let compute_budget = ComputeBudget::default();
        let mut $transaction_context = TransactionContext::new(
            $transaction_accounts,
            Rent::default(),
            compute_budget.max_instruction_stack_depth,
            compute_budget.max_instruction_trace_length,
        );
        let mut sysvar_cache = SysvarCache::default();
        sysvar_cache.fill_missing_entries(|pubkey, callback| {
            for index in 0..$transaction_context.get_number_of_accounts() {
                if $transaction_context
                    .get_key_of_account_at_index(index)
                    .unwrap()
                    == pubkey
                {
                    callback(
                        $transaction_context
                            .get_account_at_index(index)
                            .unwrap()
                            .borrow()
                            .data(),
                    );
                }
            }
        });
        // Create a stakes handle with default epoch (0) - tests that need specific epochs
        // should set up the stake cache data appropriately.
        // The default StakesHandle includes an epoch_rewards_exists_fn that returns false.
        let mock_stakes_handle = $crate::invoke_context::StakesHandle::default();
        let environment_config = EnvironmentConfig::new(
            Hash::default(),
            0,
            &|_| 0,
            Arc::new(FeatureSet::all_enabled()),
            &sysvar_cache,
            0,
            mock_stakes_handle,
        );
        let mut program_cache_for_tx_batch = ProgramCacheForTxBatch::default();
        if let Some((loader_id, builtin)) = $entry {
            program_cache_for_tx_batch.replenish(
                loader_id,
                Arc::new(ProgramCacheEntry::new_builtin(0, 0, builtin)),
            );
        }
        let mut program_cache_for_tx = ProgramCacheForTx::from_cache(&program_cache_for_tx_batch);
        let mut $invoke_context = InvokeContext::new(
            &mut $transaction_context,
            &mut program_cache_for_tx,
            environment_config,
            Some(LogCollector::new_ref()),
            compute_budget,
        );
    };
    (
        $invoke_context:ident,
        $transaction_context:ident,
        $transaction_accounts:expr $(,)?
    ) => {
        with_mock_invoke_context!(
            $invoke_context,
            $transaction_context,
            None,
            $transaction_accounts
        )
    };
}

pub fn mock_process_instruction<
    T: Into<StoredAccount>,
    F: FnMut(&mut InvokeContext<'_, '_>),
    G: FnMut(&mut InvokeContext<'_, '_>),
>(
    loader_id: &Pubkey,
    mut program_indices: Vec<IndexOfAccount>,
    instruction_data: &[u8],
    transaction_accounts: Vec<(Pubkey, T)>,
    instruction_account_metas: Vec<AccountMeta>,
    expected_result: Result<(), InstructionError>,
    builtin_function: BuiltinFunctionWithContext,
    mut pre_adjustments: F,
    mut post_adjustments: G,
) -> Vec<StoredAccount> {
    // Convert to StoredAccount upfront since we need to mutate the vector
    let mut transaction_accounts: Vec<TransactionAccount> = transaction_accounts
        .into_iter()
        .map(|(key, account)| (key, account.into()))
        .collect();
    let mut instruction_accounts: Vec<InstructionAccount> =
        Vec::with_capacity(instruction_account_metas.len());
    for (instruction_account_index, account_meta) in instruction_account_metas.iter().enumerate() {
        let index_in_transaction = transaction_accounts
            .iter()
            .position(|(key, _account)| *key == account_meta.pubkey)
            .unwrap_or(transaction_accounts.len())
            as IndexOfAccount;
        let index_in_callee = instruction_accounts
            .get(0..instruction_account_index)
            .unwrap()
            .iter()
            .position(|instruction_account| {
                instruction_account.index_in_transaction == index_in_transaction
            })
            .unwrap_or(instruction_account_index) as IndexOfAccount;
        instruction_accounts.push(InstructionAccount {
            index_in_transaction,
            index_in_caller: index_in_transaction,
            index_in_callee,
            is_signer: account_meta.is_signer,
            is_writable: account_meta.is_writable,
        });
    }
    if program_indices.is_empty() {
        program_indices.insert(0, transaction_accounts.len() as IndexOfAccount);
        let processor_account =
            StoredAccount::Data(AccountSharedData::new(0, 0, &native_loader::id()));
        transaction_accounts.push((*loader_id, processor_account));
    }
    let pop_epoch_schedule_account = if !transaction_accounts
        .iter()
        .any(|(key, _)| *key == sysvar::epoch_schedule::id())
    {
        transaction_accounts.push((
            sysvar::epoch_schedule::id(),
            create_account_shared_data_for_test(&EpochSchedule::default()),
        ));
        true
    } else {
        false
    };
    with_mock_invoke_context!(
        invoke_context,
        transaction_context,
        Some((*loader_id, builtin_function)),
        transaction_accounts
    );
    pre_adjustments(&mut invoke_context);
    let result = invoke_context.process_instruction(
        instruction_data,
        &instruction_accounts,
        &program_indices,
        &mut 0,
        &mut ExecuteTimings::default(),
    );
    assert_eq!(result, expected_result);
    post_adjustments(&mut invoke_context);
    drop(invoke_context);
    let mut transaction_accounts = transaction_context.deconstruct_without_keys().unwrap();
    if pop_epoch_schedule_account {
        transaction_accounts.pop();
    }
    transaction_accounts.pop();
    transaction_accounts
}

#[cfg(test)]
mod tests {
    use rialo_s_compute_budget::compute_budget_limits;
    use rialo_s_instruction::Instruction;
    use rialo_s_rent::Rent;
    use serde::{Deserialize, Serialize};

    use super::*;

    #[derive(Debug, Serialize, Deserialize)]
    enum MockInstruction {
        NoopSuccess,
        NoopFail,
        ModifyOwned,
        ModifyNotOwned,
        ModifyReadonly,
        UnbalancedPush,
        UnbalancedPop,
        ConsumeComputeUnits {
            compute_units_to_consume: u64,
            desired_result: Result<(), InstructionError>,
        },
        Resize {
            new_len: u64,
        },
    }

    const MOCK_BUILTIN_COMPUTE_UNIT_COST: u64 = 1;

    declare_process_instruction!(
        MockBuiltin,
        MOCK_BUILTIN_COMPUTE_UNIT_COST,
        |invoke_context| {
            let transaction_context = &invoke_context.transaction_context;
            let instruction_context = transaction_context.get_current_instruction_context()?;
            let instruction_data = instruction_context.get_instruction_data();
            let program_id = instruction_context.get_last_program_key(transaction_context)?;
            let instruction_accounts = (0..4)
                .map(|instruction_account_index| InstructionAccount {
                    index_in_transaction: instruction_account_index,
                    index_in_caller: instruction_account_index,
                    index_in_callee: instruction_account_index,
                    is_signer: false,
                    is_writable: false,
                })
                .collect::<Vec<_>>();
            assert_eq!(
                program_id,
                instruction_context
                    .try_borrow_instruction_account(transaction_context, 0)?
                    .get_owner()
            );
            assert_ne!(
                instruction_context
                    .try_borrow_instruction_account(transaction_context, 1)?
                    .get_owner(),
                instruction_context
                    .try_borrow_instruction_account(transaction_context, 0)?
                    .get_key()
            );

            if let Ok(instruction) = bincode::deserialize(instruction_data) {
                match instruction {
                    MockInstruction::NoopSuccess => (),
                    MockInstruction::NoopFail => return Err(InstructionError::GenericError),
                    MockInstruction::ModifyOwned => instruction_context
                        .try_borrow_instruction_account(transaction_context, 0)?
                        .set_data_from_slice(&[1])?,
                    MockInstruction::ModifyNotOwned => instruction_context
                        .try_borrow_instruction_account(transaction_context, 1)?
                        .set_data_from_slice(&[1])?,
                    MockInstruction::ModifyReadonly => instruction_context
                        .try_borrow_instruction_account(transaction_context, 2)?
                        .set_data_from_slice(&[1])?,
                    MockInstruction::UnbalancedPush => {
                        instruction_context
                            .try_borrow_instruction_account(transaction_context, 0)?
                            .checked_add_kelvins(1)?;
                        let program_id = *transaction_context.get_key_of_account_at_index(3)?;
                        let metas = vec![
                            AccountMeta::new_readonly(
                                *transaction_context.get_key_of_account_at_index(0)?,
                                false,
                            ),
                            AccountMeta::new_readonly(
                                *transaction_context.get_key_of_account_at_index(1)?,
                                false,
                            ),
                        ];
                        let inner_instruction = Instruction::new_with_bincode(
                            program_id,
                            &MockInstruction::NoopSuccess,
                            metas,
                        );
                        invoke_context
                            .transaction_context
                            .get_next_instruction_context()
                            .unwrap()
                            .configure(&[3], &instruction_accounts, &[]);
                        let result = invoke_context.push();
                        assert_eq!(result, Err(InstructionError::UnbalancedInstruction));
                        result?;
                        invoke_context
                            .native_invoke(inner_instruction.into(), &[])
                            .and(invoke_context.pop())?;
                    }
                    MockInstruction::UnbalancedPop => instruction_context
                        .try_borrow_instruction_account(transaction_context, 0)?
                        .checked_add_kelvins(1)?,
                    MockInstruction::ConsumeComputeUnits {
                        compute_units_to_consume,
                        desired_result,
                    } => {
                        invoke_context
                            .consume_checked(compute_units_to_consume)
                            .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
                        return desired_result;
                    }
                    MockInstruction::Resize { new_len } => instruction_context
                        .try_borrow_instruction_account(transaction_context, 0)?
                        .set_data(vec![0; new_len as usize])?,
                }
            } else {
                return Err(InstructionError::InvalidInstructionData);
            }
            Ok(())
        }
    );

    #[test]
    fn test_instruction_stack_height() {
        let one_more_than_max_depth = ComputeBudget::default()
            .max_instruction_stack_depth
            .saturating_add(1);
        let mut invoke_stack = vec![];
        let mut transaction_accounts = vec![];
        let mut instruction_accounts = vec![];
        for index in 0..one_more_than_max_depth {
            invoke_stack.push(rialo_s_pubkey::new_rand());
            transaction_accounts.push((
                rialo_s_pubkey::new_rand(),
                AccountSharedData::new(index as u64, 1, invoke_stack.get(index).unwrap()),
            ));
            instruction_accounts.push(InstructionAccount {
                index_in_transaction: index as IndexOfAccount,
                index_in_caller: index as IndexOfAccount,
                index_in_callee: instruction_accounts.len() as IndexOfAccount,
                is_signer: false,
                is_writable: true,
            });
        }
        for (index, program_id) in invoke_stack.iter().enumerate() {
            transaction_accounts.push((
                *program_id,
                AccountSharedData::new(1, 1, &rialo_s_pubkey::Pubkey::default()),
            ));
            instruction_accounts.push(InstructionAccount {
                index_in_transaction: index as IndexOfAccount,
                index_in_caller: index as IndexOfAccount,
                index_in_callee: index as IndexOfAccount,
                is_signer: false,
                is_writable: false,
            });
        }
        let transaction_accounts: Vec<(Pubkey, StoredAccount)> = transaction_accounts
            .into_iter()
            .map(|(k, v)| (k, v.into()))
            .collect();
        with_mock_invoke_context!(invoke_context, transaction_context, transaction_accounts);

        // Check call depth increases and has a limit
        let mut depth_reached = 0;
        for _ in 0..invoke_stack.len() {
            invoke_context
                .transaction_context
                .get_next_instruction_context()
                .unwrap()
                .configure(
                    &[one_more_than_max_depth.saturating_add(depth_reached) as IndexOfAccount],
                    &instruction_accounts,
                    &[],
                );
            if Err(InstructionError::CallDepth) == invoke_context.push() {
                break;
            }
            depth_reached = depth_reached.saturating_add(1);
        }
        assert_ne!(depth_reached, 0);
        assert!(depth_reached < one_more_than_max_depth);
    }

    #[test]
    fn test_max_instruction_trace_length() {
        const MAX_INSTRUCTIONS: usize = 8;
        let mut transaction_context = TransactionContext::new::<StoredAccount>(
            Vec::new(),
            Rent::default(),
            1,
            MAX_INSTRUCTIONS,
        );
        for _ in 0..MAX_INSTRUCTIONS {
            transaction_context.push().unwrap();
            transaction_context.pop().unwrap();
        }
        assert_eq!(
            transaction_context.push(),
            Err(InstructionError::MaxInstructionTraceLengthExceeded)
        );
    }

    #[test]
    fn test_process_instruction() {
        let callee_program_id = rialo_s_pubkey::new_rand();
        let owned_account = AccountSharedData::new(42, 1, &callee_program_id);
        let not_owned_account = AccountSharedData::new(84, 1, &rialo_s_pubkey::new_rand());
        let readonly_account = AccountSharedData::new(168, 1, &rialo_s_pubkey::new_rand());
        let loader_account = AccountSharedData::new(0, 1, &native_loader::id());
        let program_account = AccountSharedData::new(1, 1, &native_loader::id());
        let transaction_accounts = vec![
            (rialo_s_pubkey::new_rand(), owned_account),
            (rialo_s_pubkey::new_rand(), not_owned_account),
            (rialo_s_pubkey::new_rand(), readonly_account),
            (callee_program_id, program_account),
            (rialo_s_pubkey::new_rand(), loader_account),
        ];
        let metas = vec![
            AccountMeta::new(transaction_accounts.first().unwrap().0, false),
            AccountMeta::new(transaction_accounts.get(1).unwrap().0, false),
            AccountMeta::new_readonly(transaction_accounts.get(2).unwrap().0, false),
        ];
        let instruction_accounts = (0..4)
            .map(|instruction_account_index| InstructionAccount {
                index_in_transaction: instruction_account_index,
                index_in_caller: instruction_account_index,
                index_in_callee: instruction_account_index,
                is_signer: false,
                is_writable: instruction_account_index < 2,
            })
            .collect::<Vec<_>>();
        let transaction_accounts: Vec<(Pubkey, StoredAccount)> = transaction_accounts
            .into_iter()
            .map(|(k, v)| (k, v.into()))
            .collect();
        with_mock_invoke_context!(
            invoke_context,
            transaction_context,
            Some((callee_program_id, MockBuiltin::vm)),
            transaction_accounts
        );

        // Account modification tests
        let cases = vec![
            (MockInstruction::NoopSuccess, Ok(())),
            (
                MockInstruction::NoopFail,
                Err(InstructionError::GenericError),
            ),
            (MockInstruction::ModifyOwned, Ok(())),
            (
                MockInstruction::ModifyNotOwned,
                Err(InstructionError::ExternalAccountDataModified),
            ),
            (
                MockInstruction::ModifyReadonly,
                Err(InstructionError::ReadonlyDataModified),
            ),
            (
                MockInstruction::UnbalancedPush,
                Err(InstructionError::UnbalancedInstruction),
            ),
            (
                MockInstruction::UnbalancedPop,
                Err(InstructionError::UnbalancedInstruction),
            ),
        ];
        for case in cases {
            invoke_context
                .transaction_context
                .get_next_instruction_context()
                .unwrap()
                .configure(&[4], &instruction_accounts, &[]);
            invoke_context.push().unwrap();
            let inner_instruction =
                Instruction::new_with_bincode(callee_program_id, &case.0, metas.clone());
            let result = invoke_context
                .native_invoke(inner_instruction.into(), &[])
                .and(invoke_context.pop());
            assert_eq!(result, case.1);
        }

        // Compute unit consumption tests
        let compute_units_to_consume = 10;
        let expected_results = vec![Ok(()), Err(InstructionError::GenericError)];
        for expected_result in expected_results {
            invoke_context
                .transaction_context
                .get_next_instruction_context()
                .unwrap()
                .configure(&[4], &instruction_accounts, &[]);
            invoke_context.push().unwrap();
            let inner_instruction = Instruction::new_with_bincode(
                callee_program_id,
                &MockInstruction::ConsumeComputeUnits {
                    compute_units_to_consume,
                    desired_result: expected_result.clone(),
                },
                metas.clone(),
            );
            let inner_instruction = StableInstruction::from(inner_instruction);
            let (inner_instruction_accounts, program_indices) = invoke_context
                .prepare_instruction(&inner_instruction, &[])
                .unwrap();

            let mut compute_units_consumed = 0;
            let result = invoke_context.process_instruction(
                &inner_instruction.data,
                &inner_instruction_accounts,
                &program_indices,
                &mut compute_units_consumed,
                &mut ExecuteTimings::default(),
            );

            // Because the instruction had compute cost > 0, then regardless of the execution result,
            // the number of compute units consumed should be a non-default which is something greater
            // than zero.
            assert!(compute_units_consumed > 0);
            assert_eq!(
                compute_units_consumed,
                compute_units_to_consume.saturating_add(MOCK_BUILTIN_COMPUTE_UNIT_COST),
            );
            assert_eq!(result, expected_result);

            invoke_context.pop().unwrap();
        }
    }

    #[test]
    fn test_invoke_context_compute_budget() {
        let transaction_accounts = vec![(rialo_s_pubkey::new_rand(), AccountSharedData::default())];
        let transaction_accounts: Vec<(Pubkey, StoredAccount)> = transaction_accounts
            .into_iter()
            .map(|(k, v)| (k, v.into()))
            .collect();

        with_mock_invoke_context!(invoke_context, transaction_context, transaction_accounts);
        invoke_context.compute_budget = ComputeBudget::new(
            compute_budget_limits::DEFAULT_INSTRUCTION_COMPUTE_UNIT_LIMIT as u64,
        );

        invoke_context
            .transaction_context
            .get_next_instruction_context()
            .unwrap()
            .configure(&[0], &[], &[]);
        invoke_context.push().unwrap();
        assert_eq!(
            *invoke_context.get_compute_budget(),
            ComputeBudget::new(
                compute_budget_limits::DEFAULT_INSTRUCTION_COMPUTE_UNIT_LIMIT as u64
            )
        );
        invoke_context.pop().unwrap();
    }

    #[test]
    fn test_process_instruction_accounts_resize_delta() {
        let program_key = Pubkey::new_unique();
        let user_account_data_len = 123u64;
        let user_account =
            AccountSharedData::new(100, user_account_data_len as usize, &program_key);
        let dummy_account = AccountSharedData::new(10, 0, &program_key);
        let program_account = AccountSharedData::new(500, 500, &native_loader::id());
        let transaction_accounts = vec![
            (Pubkey::new_unique(), user_account),
            (Pubkey::new_unique(), dummy_account),
            (program_key, program_account),
        ];
        let instruction_accounts = [
            InstructionAccount {
                index_in_transaction: 0,
                index_in_caller: 0,
                index_in_callee: 0,
                is_signer: false,
                is_writable: true,
            },
            InstructionAccount {
                index_in_transaction: 1,
                index_in_caller: 1,
                index_in_callee: 1,
                is_signer: false,
                is_writable: false,
            },
        ];
        let transaction_accounts: Vec<(Pubkey, StoredAccount)> = transaction_accounts
            .into_iter()
            .map(|(k, v)| (k, v.into()))
            .collect();
        with_mock_invoke_context!(
            invoke_context,
            transaction_context,
            Some((program_key, MockBuiltin::vm)),
            transaction_accounts
        );

        // Test: Resize the account to *the same size*, so not consuming any additional size; this must succeed
        {
            let resize_delta: i64 = 0;
            let new_len = (user_account_data_len as i64).saturating_add(resize_delta) as u64;
            let instruction_data =
                bincode::serialize(&MockInstruction::Resize { new_len }).unwrap();

            let result = invoke_context.process_instruction(
                &instruction_data,
                &instruction_accounts,
                &[2],
                &mut 0,
                &mut ExecuteTimings::default(),
            );

            assert!(result.is_ok());
            assert_eq!(
                invoke_context
                    .transaction_context
                    .accounts_resize_delta()
                    .unwrap(),
                resize_delta
            );
        }

        // Test: Resize the account larger; this must succeed
        {
            let resize_delta: i64 = 1;
            let new_len = (user_account_data_len as i64).saturating_add(resize_delta) as u64;
            let instruction_data =
                bincode::serialize(&MockInstruction::Resize { new_len }).unwrap();

            let result = invoke_context.process_instruction(
                &instruction_data,
                &instruction_accounts,
                &[2],
                &mut 0,
                &mut ExecuteTimings::default(),
            );

            assert!(result.is_ok());
            assert_eq!(
                invoke_context
                    .transaction_context
                    .accounts_resize_delta()
                    .unwrap(),
                resize_delta
            );
        }

        // Test: Resize the account smaller; this must succeed
        {
            let resize_delta: i64 = -1;
            let new_len = (user_account_data_len as i64).saturating_add(resize_delta) as u64;
            let instruction_data =
                bincode::serialize(&MockInstruction::Resize { new_len }).unwrap();

            let result = invoke_context.process_instruction(
                &instruction_data,
                &instruction_accounts,
                &[2],
                &mut 0,
                &mut ExecuteTimings::default(),
            );

            assert!(result.is_ok());
            assert_eq!(
                invoke_context
                    .transaction_context
                    .accounts_resize_delta()
                    .unwrap(),
                resize_delta
            );
        }
    }
}