openvm-circuit 2.0.1

use std::{
    alloc::{alloc, dealloc, handle_alloc_error, Layout},
    borrow::{Borrow, BorrowMut},
    iter::repeat_n,
    ptr::NonNull,
};

use itertools::Itertools;
use openvm_circuit_primitives_derive::AlignedBytesBorrow;
use openvm_instructions::{
    exe::{SparseMemoryImage, VmExe},
    instruction::Instruction,
    program::{Program, DEFAULT_PC_STEP},
    LocalOpcode, SystemOpcode,
};
use openvm_stark_backend::p3_field::PrimeField32;

#[cfg(feature = "tco")]
use crate::arch::Handler;
use crate::{
    arch::{
        execution_mode::{
            ExecutionCtx, ExecutionCtxTrait, MeteredCostCtx, MeteredCtx, MeteredExecutionCtxTrait,
            Segment,
        },
        ExecuteFunc, ExecutionError, Executor, ExecutorInventory, ExitCode, MeteredExecutor,
        StaticProgramError, Streams, SystemConfig, VmExecState, VmState,
    },
    system::memory::online::GuestMemory,
};

/// VM pure executor(E1/E2 executor) which doesn't consider trace generation.
/// Note: This executor doesn't hold any VM state and can be used for multiple execution.
///
/// The generic `Ctx` and constructor determine whether this supported pure execution or metered
/// execution.
// NOTE: the lifetime 'a represents the lifetime of borrowed ExecutorInventory, which must outlive
// the InterpretedInstance because `pre_compute_buf` may contain pointers to references held by
// executors.
pub struct InterpretedInstance<'a, F, Ctx> {
    system_config: &'a SystemConfig,
    // SAFETY: this is not actually dead code, but `pre_compute_insns` contains raw pointer refers
    // to this buffer.
    #[allow(dead_code)]
    pre_compute_buf: AlignedBuf,
    /// Instruction table of function pointers and pointers to the pre-computed buffer. Indexed by
    /// `pc_index = pc / DEFAULT_PC_STEP`.
    /// SAFETY: The first `pc_base / DEFAULT_PC_STEP` entries will be unreachable. We do this to
    /// avoid needing to subtract `pc_base` during runtime.
    #[cfg(not(feature = "tco"))]
    pre_compute_insns: Vec<PreComputeInstruction<F, Ctx>>,
    #[cfg(feature = "tco")]
    pre_compute_max_size: usize,
    /// Handler function pointers for tail call optimization.
    #[cfg(feature = "tco")]
    handlers: Vec<Handler<F, Ctx>>,

    pc_start: u32,

    init_memory: SparseMemoryImage,
}

#[repr(C)]
#[cfg_attr(feature = "tco", allow(dead_code))]
pub(crate) struct PreComputeInstruction<F, Ctx> {
    pub(crate) handler: ExecuteFunc<F, Ctx>,
    pub(crate) pre_compute: *const u8,
}

unsafe impl<F, Ctx> Send for PreComputeInstruction<F, Ctx> {}
unsafe impl<F, Ctx> Sync for PreComputeInstruction<F, Ctx> {}

#[derive(AlignedBytesBorrow, Clone)]
#[repr(C)]
struct TerminatePreCompute {
    exit_code: u32,
}

macro_rules! run {
    ($span:literal, $interpreter:ident, $exec_state:ident, $ctx:ident) => {{
        tracing::info_span!($span).in_scope(|| -> Result<(), ExecutionError> {
            // SAFETY:
            // - it is the responsibility of each Executor to ensure that pre_compute_insts contains
            //   valid function pointers and pre-computed data
            #[cfg(not(feature = "tco"))]
            {
                unsafe {
                    execute_trampoline(&mut $exec_state, &$interpreter.pre_compute_insns);
                }
            }
            #[cfg(feature = "tco")]
            {
                if $ctx::should_suspend(&mut $exec_state) {
                    return Ok(());
                }

                let handler = $interpreter
                    .get_handler($exec_state.pc())
                    .ok_or(ExecutionError::PcOutOfBounds($exec_state.pc()))?;
                // SAFETY:
                // - handler is generated by Executor, MeteredExecutor traits
                // - it is the responsibility of each Executor to ensure handler is safe given a
                //   valid VM state
                unsafe {
                    handler($interpreter, &mut $exec_state);
                }
            }
            Ok(())
        })?;
    }};
}

// Constructors for E1 and E2 respectively, which generate pre-computed buffers and function
// pointers
// - Generic in `Ctx`

impl<'a, F, Ctx> InterpretedInstance<'a, F, Ctx>
where
    F: PrimeField32,
    Ctx: ExecutionCtxTrait,
{
    /// Creates a new interpreter instance for pure execution.
    // (E1 execution)
    pub fn new<E>(
        inventory: &'a ExecutorInventory<E>,
        exe: &VmExe<F>,
    ) -> Result<Self, StaticProgramError>
    where
        E: Executor<F>,
    {
        let program = &exe.program;
        let pre_compute_max_size = get_pre_compute_max_size(program, inventory);
        let mut pre_compute_buf = alloc_pre_compute_buf(program, pre_compute_max_size);
        let mut split_pre_compute_buf =
            split_pre_compute_buf(program, &mut pre_compute_buf, pre_compute_max_size);
        #[cfg(not(feature = "tco"))]
        let pre_compute_insns = get_pre_compute_instructions::<F, Ctx, E>(
            program,
            inventory,
            &mut split_pre_compute_buf,
        )?;
        let pc_start = exe.pc_start;
        let init_memory = exe.init_memory.clone();
        #[cfg(feature = "tco")]
        let handlers = repeat_n(&None, get_pc_index(program.pc_base))
            .chain(program.instructions_and_debug_infos.iter())
            .zip_eq(split_pre_compute_buf.iter_mut())
            .enumerate()
            .map(
                |(pc_idx, (inst_opt, pre_compute))| -> Result<Handler<F, Ctx>, StaticProgramError> {
                    if let Some((inst, _)) = inst_opt {
                        let pc = pc_idx as u32 * DEFAULT_PC_STEP;
                        if get_system_opcode_handler::<F, Ctx>(inst, pre_compute).is_some() {
                            Ok(terminate_execute_e12_tco_handler)
                        } else {
                            // unwrap because get_pre_compute_instructions would have errored
                            // already on DisabledOperation
                            let executor = inventory.get_executor(inst.opcode).unwrap();
                            executor.handler(pc, inst, pre_compute)
                        }
                    } else {
                        Ok(unreachable_tco_handler)
                    }
                },
            )
            .collect::<Result<Vec<_>, _>>()?;

        Ok(Self {
            system_config: inventory.config(),
            pre_compute_buf,
            #[cfg(not(feature = "tco"))]
            pre_compute_insns,
            pc_start,
            init_memory,
            #[cfg(feature = "tco")]
            pre_compute_max_size,
            #[cfg(feature = "tco")]
            handlers,
        })
    }

    pub fn create_initial_vm_state(&self, inputs: impl Into<Streams<F>>) -> VmState<F> {
        VmState::initial(self.system_config, &self.init_memory, self.pc_start, inputs)
    }

    /// # Safety
    /// - This function assumes that the `pc` is within program bounds - this should be the case if
    ///   the pc is checked to be in bounds before jumping to it.
    /// - The returned slice may not be entirely initialized, but it is the job of each Executor to
    ///   initialize the parts of the buffer that the instruction handler will use.
    #[cfg(feature = "tco")]
    #[inline(always)]
    pub fn get_pre_compute(&self, pc: u32) -> *const u8 {
        let pc_idx = get_pc_index(pc);
        // SAFETY:
        // - we assume that pc is in bounds
        // - pre_compute_buf is allocated for pre_compute_max_size * program_len bytes, with each
        //   instruction getting pre_compute_max_size bytes
        // - self.pre_compute_buf.ptr is non-null
        // - initialization of the contents of the slice is the responsibility of each Executor
        debug_assert!(
            (pc_idx + 1) * self.pre_compute_max_size <= self.pre_compute_buf.layout.size()
        );
        unsafe {
            let ptr = self
                .pre_compute_buf
                .ptr
                .add(pc_idx * self.pre_compute_max_size);
            ptr
        }
    }

    #[cfg(feature = "tco")]
    #[inline(always)]
    pub fn get_handler(&self, pc: u32) -> Option<Handler<F, Ctx>> {
        let pc_idx = get_pc_index(pc);
        self.handlers.get(pc_idx).copied()
    }
}

impl<'a, F, Ctx> InterpretedInstance<'a, F, Ctx>
where
    F: PrimeField32,
    Ctx: MeteredExecutionCtxTrait,
{
    /// Creates a new interpreter instance for pure execution.
    // (E1 execution)
    pub fn new_metered<E>(
        inventory: &'a ExecutorInventory<E>,
        exe: &VmExe<F>,
        executor_idx_to_air_idx: &[usize],
    ) -> Result<Self, StaticProgramError>
    where
        E: MeteredExecutor<F>,
    {
        let program = &exe.program;
        let pre_compute_max_size = get_metered_pre_compute_max_size(program, inventory);
        let mut pre_compute_buf = alloc_pre_compute_buf(program, pre_compute_max_size);
        let mut split_pre_compute_buf =
            split_pre_compute_buf(program, &mut pre_compute_buf, pre_compute_max_size);
        #[cfg(not(feature = "tco"))]
        let pre_compute_insns = get_metered_pre_compute_instructions::<F, Ctx, E>(
            program,
            inventory,
            executor_idx_to_air_idx,
            &mut split_pre_compute_buf,
        )?;

        let pc_start = exe.pc_start;
        let init_memory = exe.init_memory.clone();
        #[cfg(feature = "tco")]
        let handlers = repeat_n(&None, get_pc_index(program.pc_base))
            .chain(program.instructions_and_debug_infos.iter())
            .zip_eq(split_pre_compute_buf.iter_mut())
            .enumerate()
            .map(
                |(pc_idx, (inst_opt, pre_compute))| -> Result<Handler<F, Ctx>, StaticProgramError> {
                    if let Some((inst, _)) = inst_opt {
                        let pc = pc_idx as u32 * DEFAULT_PC_STEP;
                        if get_system_opcode_handler::<F, Ctx>(inst, pre_compute).is_some() {
                            Ok(terminate_execute_e12_tco_handler)
                        } else {
                            // unwrap because get_pre_compute_instructions would have errored
                            // already on DisabledOperation
                            let executor_idx = inventory.instruction_lookup[&inst.opcode] as usize;
                            let executor = &inventory.executors[executor_idx];
                            let air_idx = executor_idx_to_air_idx[executor_idx];
                            executor.metered_handler(air_idx, pc, inst, pre_compute)
                        }
                    } else {
                        Ok(unreachable_tco_handler)
                    }
                },
            )
            .collect::<Result<Vec<_>, _>>()?;

        Ok(Self {
            system_config: inventory.config(),
            pre_compute_buf,
            #[cfg(not(feature = "tco"))]
            pre_compute_insns,
            pc_start,
            init_memory,
            #[cfg(feature = "tco")]
            pre_compute_max_size,
            #[cfg(feature = "tco")]
            handlers,
        })
    }
}

// Execute functions specialize to relevant Ctx types to provide more streamlines APIs

impl<'a, F> InterpretedInstance<'a, F, ExecutionCtx>
where
    F: PrimeField32,
{
    /// Pure execution, without metering, for the given `inputs`. Execution begins from the initial
    /// state specified by the `VmExe`. This function executes the program until either termination
    /// if `num_insns` is `None` or for exactly `num_insns` instructions if `num_insns` is `Some`.
    ///
    /// Returns the final VM state when execution stops.
    pub fn execute(
        &self,
        inputs: impl Into<Streams<F>>,
        num_insns: Option<u64>,
    ) -> Result<VmState<F, GuestMemory>, ExecutionError> {
        let vm_state =
            VmState::initial(self.system_config, &self.init_memory, self.pc_start, inputs);
        self.execute_from_state(vm_state, num_insns)
    }

    /// Pure execution, without metering, from the given `VmState`. This function executes the
    /// program until either termination if `num_insns` is `None` or for exactly `num_insns`
    /// instructions if `num_insns` is `Some`.
    ///
    /// Returns the final VM state when execution stops.
    pub fn execute_from_state(
        &self,
        from_state: VmState<F, GuestMemory>,
        num_insns: Option<u64>,
    ) -> Result<VmState<F, GuestMemory>, ExecutionError> {
        let ctx = ExecutionCtx::new(num_insns);
        let mut exec_state = VmExecState::new(from_state, ctx);

        #[cfg(feature = "metrics")]
        let start = std::time::Instant::now();
        #[cfg(feature = "metrics")]
        let start_instret_left = exec_state.ctx.instret_left;

        run!("execute_e1", self, exec_state, ExecutionCtx);

        #[cfg(feature = "metrics")]
        {
            let elapsed = start.elapsed();
            let insns = start_instret_left - exec_state.ctx.instret_left;
            tracing::info!("instructions_executed={insns}");
            metrics::counter!("execute_e1_insns").absolute(insns);
            metrics::gauge!("execute_e1_insn_mi/s").set(insns as f64 / elapsed.as_micros() as f64);
        }
        tracing::debug!("pc: {}", exec_state.vm_state.pc());
        tracing::debug!("interpreter exit code {:?}", exec_state.exit_code);
        tracing::debug!("num_insns {:?}", num_insns);

        if num_insns.is_some() {
            check_exit_code(exec_state.exit_code)?;
        } else {
            check_termination(exec_state.exit_code)?;
        }
        Ok(exec_state.vm_state)
    }
}

impl<'a, F> InterpretedInstance<'a, F, MeteredCtx>
where
    F: PrimeField32,
{
    /// Metered execution for the given `inputs`. Execution begins from the initial
    /// state specified by the `VmExe`. This function executes the program until termination.
    ///
    /// Returns the segmentation boundary data and the final VM state when execution stops.
    pub fn execute_metered(
        &self,
        inputs: impl Into<Streams<F>>,
        ctx: MeteredCtx,
    ) -> Result<(Vec<Segment>, VmState<F, GuestMemory>), ExecutionError> {
        let vm_state = self.create_initial_vm_state(inputs);
        self.execute_metered_from_state(vm_state, ctx)
    }

    /// Metered execution for the given `VmState`. This function executes the program until
    /// termination.
    ///
    /// Returns the segmentation boundary data and the final VM state when execution stops.
    ///
    /// The [MeteredCtx] can be constructed using either
    /// [VmExecutor::build_metered_ctx](super::VmExecutor::build_metered_ctx) or
    /// [VirtualMachine::build_metered_ctx](super::VirtualMachine::build_metered_ctx).
    pub fn execute_metered_from_state(
        &self,
        from_state: VmState<F, GuestMemory>,
        ctx: MeteredCtx,
    ) -> Result<(Vec<Segment>, VmState<F, GuestMemory>), ExecutionError> {
        let mut exec_state = VmExecState::new(from_state, ctx);

        loop {
            exec_state = self.execute_metered_until_suspend(exec_state)?;
            // The execution has terminated.
            if exec_state.exit_code.is_ok() && exec_state.exit_code.as_ref().unwrap().is_some() {
                break;
            }
            if exec_state.exit_code.is_err() {
                return Err(exec_state.exit_code.unwrap_err());
            }
        }
        check_termination(exec_state.exit_code)?;
        let VmExecState { vm_state, ctx, .. } = exec_state;
        Ok((ctx.into_segments(), vm_state))
    }
    /// Executes a metered virtual machine operation starting from a given execution state until
    /// suspension.
    ///
    /// This function resumes and continues execution of a guest virtual machine until either it:
    /// - Hits a suspension trigger (e.g. out of gas or a specific halt condition). ATTENTION: when
    ///   a suspension is triggered, the VM state is not at the boundary of the last segment.
    ///   Instead, the VM state is slightly after the segment boundary.
    /// - Completes its run based on the instructions or context provided.
    ///
    /// # Parameters
    /// - `self`: The reference to the current executor or VM context.
    /// - `exec_state`: A mutable `VmExecState<F, GuestMemory, MeteredCtx>` which represents the
    ///   execution state of the virtual machine, including its program counter (`pc`), instruction
    ///   retirement (`instret`), and execution context (`MeteredCtx`).
    ///
    /// # Returns
    /// - `Ok(VmExecState<F, GuestMemory, MeteredCtx>)`: The execution state after suspension or
    ///   normal completion.
    /// - `Err(ExecutionError)`: If there is an error during execution, such as an invalid state or
    ///   run-time error.
    pub fn execute_metered_until_suspend(
        &self,
        mut exec_state: VmExecState<F, GuestMemory, MeteredCtx>,
    ) -> Result<VmExecState<F, GuestMemory, MeteredCtx>, ExecutionError> {
        #[cfg(feature = "metrics")]
        let start = std::time::Instant::now();
        #[cfg(feature = "metrics")]
        let start_instret = exec_state.ctx.segmentation_ctx.instret;

        // Start execution
        run!("execute_metered", self, exec_state, MeteredCtx);

        #[cfg(feature = "metrics")]
        {
            let elapsed = start.elapsed();
            let insns = exec_state.ctx.segmentation_ctx.instret - start_instret;
            tracing::info!("instructions_executed={insns}");
            metrics::counter!("execute_metered_insns").absolute(insns);
            metrics::gauge!("execute_metered_insn_mi/s")
                .set(insns as f64 / elapsed.as_micros() as f64);
        }
        Ok(exec_state)
    }
}

impl<'a, F> InterpretedInstance<'a, F, MeteredCostCtx>
where
    F: PrimeField32,
{
    /// Metered cost execution for the given `inputs`. Execution begins from the initial
    /// state specified by the `VmExe`. This function executes the program until termination.
    ///
    /// Returns the trace cost and final VM state when execution stops.
    pub fn execute_metered_cost(
        &self,
        inputs: impl Into<Streams<F>>,
        ctx: MeteredCostCtx,
    ) -> Result<(MeteredCostCtx, VmState<F, GuestMemory>), ExecutionError> {
        let vm_state = self.create_initial_vm_state(inputs);
        self.execute_metered_cost_from_state(vm_state, ctx)
    }

    /// Metered cost execution for the given `VmState`. This function executes the program until
    /// termination.
    ///
    /// Returns the trace cost and final VM state when execution stops.
    pub fn execute_metered_cost_from_state(
        &self,
        from_state: VmState<F, GuestMemory>,
        ctx: MeteredCostCtx,
    ) -> Result<(MeteredCostCtx, VmState<F, GuestMemory>), ExecutionError> {
        let mut exec_state = VmExecState::new(from_state, ctx);

        #[cfg(feature = "metrics")]
        let start = std::time::Instant::now();
        #[cfg(feature = "metrics")]
        let start_instret = exec_state.ctx.instret;

        // Start execution
        run!("execute_metered_cost", self, exec_state, MeteredCostCtx);

        #[cfg(feature = "metrics")]
        {
            let elapsed = start.elapsed();
            let insns = exec_state.ctx.instret - start_instret;
            tracing::info!("instructions_executed={insns}");
            metrics::counter!("execute_metered_cost_insns").absolute(insns);
            metrics::gauge!("execute_metered_cost_insn_mi/s")
                .set(insns as f64 / elapsed.as_micros() as f64);
        }

        check_exit_code(exec_state.exit_code)?;
        let VmExecState { ctx, vm_state, .. } = exec_state;
        Ok((ctx, vm_state))
    }
}

pub(crate) fn alloc_pre_compute_buf<F>(
    program: &Program<F>,
    pre_compute_max_size: usize,
) -> AlignedBuf {
    let base_idx = get_pc_index(program.pc_base);
    let padded_program_len = base_idx + program.instructions_and_debug_infos.len();
    let buf_len = padded_program_len * pre_compute_max_size;
    AlignedBuf::uninit(buf_len, pre_compute_max_size)
}

pub(crate) fn split_pre_compute_buf<'a, F>(
    program: &Program<F>,
    pre_compute_buf: &'a mut AlignedBuf,
    pre_compute_max_size: usize,
) -> Vec<&'a mut [u8]> {
    let base_idx = get_pc_index(program.pc_base);
    let padded_program_len = base_idx + program.instructions_and_debug_infos.len();
    let buf_len = padded_program_len * pre_compute_max_size;
    // SAFETY:
    // - pre_compute_buf.ptr was allocated with exactly buf_len bytes
    // - lifetime 'a ensures the returned slices don't outlive the AlignedBuf
    let pre_compute_buf = unsafe { std::slice::from_raw_parts_mut(pre_compute_buf.ptr, buf_len) };
    pre_compute_buf
        .chunks_exact_mut(pre_compute_max_size)
        .collect()
}

/// Executes using function pointers with the trampoline (loop) approach.
///
/// # Safety
/// The `fn_ptrs` pointer to pre-computed buffers that outlive this function.
#[cfg(not(feature = "tco"))]
#[inline(always)]
unsafe fn execute_trampoline<F: PrimeField32, Ctx: ExecutionCtxTrait>(
    exec_state: &mut VmExecState<F, GuestMemory, Ctx>,
    fn_ptrs: &[PreComputeInstruction<F, Ctx>],
) {
    while exec_state
        .exit_code
        .as_ref()
        .is_ok_and(|exit_code| exit_code.is_none())
    {
        if Ctx::should_suspend(exec_state) {
            tracing::debug!("stop because of should_suspend");
            break;
        }
        let pc = exec_state.pc();
        let pc_index = get_pc_index(pc);

        if let Some(inst) = fn_ptrs.get(pc_index) {
            // SAFETY: pre_compute assumed to live long enough
            unsafe { (inst.handler)(inst.pre_compute, exec_state) };
        } else {
            exec_state.exit_code = Err(ExecutionError::PcOutOfBounds(pc));
        }
    }
}

#[inline(always)]
pub fn get_pc_index(pc: u32) -> usize {
    (pc / DEFAULT_PC_STEP) as usize
}

/// Bytes allocated according to the given Layout.
/// Careful: this struct implements Send and Sync unsafely. Don't change the underlying data after
/// initialization.git
// @dev: This is duplicate from the openvm crate, but it doesn't seem worth importing `openvm` here
// just for this.
pub(crate) struct AlignedBuf {
    pub ptr: *mut u8,
    pub layout: Layout,
}

unsafe impl Send for AlignedBuf {}
unsafe impl Sync for AlignedBuf {}

impl AlignedBuf {
    /// Allocate a new buffer whose start address is aligned to `align` bytes.
    /// *NOTE* if `len` is zero then a creates new `NonNull` that is dangling and 16-byte aligned.
    pub fn uninit(len: usize, align: usize) -> Self {
        let layout = Layout::from_size_align(len, align).unwrap();
        if layout.size() == 0 {
            return Self {
                ptr: NonNull::<u128>::dangling().as_ptr() as *mut u8,
                layout,
            };
        }
        // SAFETY: `len` is nonzero
        let ptr = unsafe { alloc(layout) };
        if ptr.is_null() {
            handle_alloc_error(layout);
        }
        AlignedBuf { ptr, layout }
    }
}

impl Drop for AlignedBuf {
    fn drop(&mut self) {
        if self.layout.size() != 0 {
            // SAFETY: self.ptr was allocated with self.layout in AlignedBuf::uninit
            unsafe {
                dealloc(self.ptr, self.layout);
            }
        }
    }
}

#[inline(always)]
unsafe fn terminate_execute_e12_impl<F: PrimeField32, CTX: ExecutionCtxTrait>(
    pre_compute: *const u8,
    exec_state: &mut VmExecState<F, GuestMemory, CTX>,
) {
    let pre_compute: &TerminatePreCompute =
        std::slice::from_raw_parts(pre_compute, size_of::<TerminatePreCompute>()).borrow();
    exec_state.exit_code = Ok(Some(pre_compute.exit_code));
    CTX::on_terminate(exec_state);
}

#[cfg(feature = "tco")]
unsafe fn terminate_execute_e12_tco_handler<F: PrimeField32, CTX: ExecutionCtxTrait>(
    interpreter: &InterpretedInstance<'_, F, CTX>,
    exec_state: &mut VmExecState<F, GuestMemory, CTX>,
) {
    let pre_compute = interpreter.get_pre_compute(exec_state.vm_state.pc());
    terminate_execute_e12_impl(pre_compute, exec_state);
}

#[cfg(feature = "tco")]
unsafe fn unreachable_tco_handler<F: PrimeField32, CTX>(
    _: &InterpretedInstance<'_, F, CTX>,
    exec_state: &mut VmExecState<F, GuestMemory, CTX>,
) {
    exec_state.exit_code = Err(ExecutionError::Unreachable(exec_state.vm_state.pc()));
}

pub(crate) fn get_pre_compute_max_size<F, E: Executor<F>>(
    program: &Program<F>,
    inventory: &ExecutorInventory<E>,
) -> usize {
    program
        .instructions_and_debug_infos
        .iter()
        .map(|inst_opt| {
            if let Some((inst, _)) = inst_opt {
                if let Some(size) = system_opcode_pre_compute_size(inst) {
                    size
                } else {
                    inventory
                        .get_executor(inst.opcode)
                        .map(|executor| executor.pre_compute_size())
                        .unwrap()
                }
            } else {
                0
            }
        })
        .max()
        .unwrap()
        .next_power_of_two()
}

pub(crate) fn get_metered_pre_compute_max_size<F, E: MeteredExecutor<F>>(
    program: &Program<F>,
    inventory: &ExecutorInventory<E>,
) -> usize {
    program
        .instructions_and_debug_infos
        .iter()
        .map(|inst_opt| {
            if let Some((inst, _)) = inst_opt {
                if let Some(size) = system_opcode_pre_compute_size(inst) {
                    size
                } else {
                    inventory
                        .get_executor(inst.opcode)
                        .map(|executor| executor.metered_pre_compute_size())
                        .unwrap()
                }
            } else {
                0
            }
        })
        .max()
        .unwrap()
        .next_power_of_two()
}

fn system_opcode_pre_compute_size<F>(inst: &Instruction<F>) -> Option<usize> {
    if inst.opcode == SystemOpcode::TERMINATE.global_opcode() {
        return Some(size_of::<TerminatePreCompute>());
    }
    None
}

#[cfg(not(feature = "tco"))]
pub(crate) fn get_pre_compute_instructions<F, Ctx, E>(
    program: &Program<F>,
    inventory: &ExecutorInventory<E>,
    pre_compute: &mut [&mut [u8]],
) -> Result<Vec<PreComputeInstruction<F, Ctx>>, StaticProgramError>
where
    F: PrimeField32,
    Ctx: ExecutionCtxTrait,
    E: Executor<F>,
{
    let unreachable_handler: ExecuteFunc<F, Ctx> = |_, exec_state| {
        exec_state.exit_code = Err(ExecutionError::Unreachable(exec_state.pc()));
    };

    repeat_n(&None, get_pc_index(program.pc_base))
        .chain(program.instructions_and_debug_infos.iter())
        .zip_eq(pre_compute.iter_mut())
        .enumerate()
        .map(|(i, (inst_opt, buf))| {
            // SAFETY: we cast to raw pointer and then borrow to remove the lifetime. This
            // is safe only in the current context because `buf` comes
            // from `pre_compute_buf` which will outlive the returned
            // `PreComputeInstruction`s.
            let buf: &mut [u8] = unsafe { &mut *(*buf as *mut [u8]) };
            let pre_inst = if let Some((inst, _)) = inst_opt {
                tracing::trace!("get_pre_compute_instruction {inst:?}");
                let pc = i as u32 * DEFAULT_PC_STEP;
                if let Some(handler) = get_system_opcode_handler(inst, buf) {
                    PreComputeInstruction {
                        handler,
                        pre_compute: buf.as_ptr(),
                    }
                } else if let Some(executor) = inventory.get_executor(inst.opcode) {
                    PreComputeInstruction {
                        handler: executor.pre_compute(pc, inst, buf)?,
                        pre_compute: buf.as_ptr(),
                    }
                } else {
                    return Err(StaticProgramError::DisabledOperation {
                        pc,
                        opcode: inst.opcode,
                    });
                }
            } else {
                // Dead instruction at this pc
                PreComputeInstruction {
                    handler: unreachable_handler,
                    pre_compute: buf.as_ptr(),
                }
            };
            Ok(pre_inst)
        })
        .collect::<Result<Vec<_>, _>>()
}

#[cfg(not(feature = "tco"))]
pub(crate) fn get_metered_pre_compute_instructions<F, Ctx, E>(
    program: &Program<F>,
    inventory: &ExecutorInventory<E>,
    executor_idx_to_air_idx: &[usize],
    pre_compute: &mut [&mut [u8]],
) -> Result<Vec<PreComputeInstruction<F, Ctx>>, StaticProgramError>
where
    F: PrimeField32,
    Ctx: MeteredExecutionCtxTrait,
    E: MeteredExecutor<F>,
{
    let unreachable_handler: ExecuteFunc<F, Ctx> = |_, exec_state| {
        exec_state.exit_code = Err(ExecutionError::Unreachable(exec_state.pc()));
    };
    repeat_n(&None, get_pc_index(program.pc_base))
        .chain(program.instructions_and_debug_infos.iter())
        .zip_eq(pre_compute.iter_mut())
        .enumerate()
        .map(|(i, (inst_opt, buf))| {
            // SAFETY: we cast to raw pointer and then borrow to remove the lifetime. This
            // is safe only in the current context because `buf` comes
            // from `pre_compute_buf` which will outlive the returned
            // `PreComputeInstruction`s.
            let buf: &mut [u8] = unsafe { &mut *(*buf as *mut [u8]) };
            let pre_inst = if let Some((inst, _)) = inst_opt {
                tracing::trace!("get_metered_pre_compute_instruction {inst:?}");
                let pc = program.pc_base + i as u32 * DEFAULT_PC_STEP;
                if let Some(handler) = get_system_opcode_handler(inst, buf) {
                    PreComputeInstruction {
                        handler,
                        pre_compute: buf.as_ptr(),
                    }
                } else if let Some(&executor_idx) = inventory.instruction_lookup.get(&inst.opcode) {
                    let executor_idx = executor_idx as usize;
                    let executor = inventory
                        .executors
                        .get(executor_idx)
                        .expect("ExecutorInventory ensures executor_idx is in bounds");
                    let air_idx = executor_idx_to_air_idx[executor_idx];
                    PreComputeInstruction {
                        handler: executor.metered_pre_compute(air_idx, pc, inst, buf)?,
                        pre_compute: buf.as_ptr(),
                    }
                } else {
                    return Err(StaticProgramError::DisabledOperation {
                        pc,
                        opcode: inst.opcode,
                    });
                }
            } else {
                PreComputeInstruction {
                    handler: unreachable_handler,
                    pre_compute: buf.as_ptr(),
                }
            };
            Ok(pre_inst)
        })
        .collect::<Result<Vec<_>, _>>()
}

fn get_system_opcode_handler<F: PrimeField32, Ctx: ExecutionCtxTrait>(
    inst: &Instruction<F>,
    buf: &mut [u8],
) -> Option<ExecuteFunc<F, Ctx>> {
    if inst.opcode == SystemOpcode::TERMINATE.global_opcode() {
        let pre_compute: &mut TerminatePreCompute = buf.borrow_mut();
        pre_compute.exit_code = inst.c.as_canonical_u32();
        return Some(terminate_execute_e12_impl);
    }
    None
}

/// Errors if exit code is either error or terminated with non-successful exit code.
fn check_exit_code(exit_code: Result<Option<u32>, ExecutionError>) -> Result<(), ExecutionError> {
    let exit_code = exit_code?;
    if let Some(exit_code) = exit_code {
        // This means execution did terminate
        if exit_code != ExitCode::Success as u32 {
            return Err(ExecutionError::FailedWithExitCode(exit_code));
        }
    }
    Ok(())
}

/// Same as [check_exit_code] but errors if program did not terminate.
pub(super) fn check_termination(
    exit_code: Result<Option<u32>, ExecutionError>,
) -> Result<(), ExecutionError> {
    let did_terminate = matches!(exit_code.as_ref(), Ok(Some(_)));
    check_exit_code(exit_code)?;
    match did_terminate {
        true => Ok(()),
        false => Err(ExecutionError::DidNotTerminate),
    }
}