openvm-circuit 2.0.1

use std::{collections::BTreeMap, mem};

use backtrace::Backtrace;
use cycle_tracker::CycleTracker;
#[cfg(feature = "perf-metrics")]
use itertools::Itertools;
use metrics::counter;
use openvm_instructions::{
    exe::{FnBound, FnBounds},
    program::ProgramDebugInfo,
};
use openvm_stark_backend::prover::{DeviceMultiStarkProvingKey, ProverBackend};

use crate::{
    arch::{
        execution_mode::PreflightCtx, interpreter_preflight::PcEntry, Arena, PreflightExecutor,
        VmExecState,
    },
    system::memory::online::TracingMemory,
};

pub mod cycle_tracker;

#[derive(Clone, Debug, Default)]
pub struct VmMetrics {
    // Static info
    pub air_names: Vec<String>,
    pub debug_infos: ProgramDebugInfo,
    #[cfg(feature = "perf-metrics")]
    pub(crate) num_sys_airs: usize,
    #[cfg(feature = "perf-metrics")]
    pub(crate) main_widths: Vec<usize>,
    #[cfg(feature = "perf-metrics")]
    pub(crate) total_widths: Vec<usize>,

    // Dynamic stats
    /// Maps (dsl_ir, opcode) to number of times opcode was executed
    pub counts: BTreeMap<(Option<String>, String), usize>,
    /// Maps (dsl_ir, opcode, air_name) to number of trace cells generated by opcode
    pub trace_cells: BTreeMap<(Option<String>, String, String), usize>,
    /// Metric collection tools. Only collected when "perf-metrics" feature is enabled.
    pub cycle_tracker: CycleTracker,

    #[cfg(feature = "perf-metrics")]
    pub(crate) current_trace_cells: Vec<usize>,

    /// Backtrace for guest debug panic display
    pub prev_backtrace: Option<Backtrace>,
    #[allow(dead_code)]
    pub(crate) fn_bounds: FnBounds,
    /// Cycle span by function if function start/end addresses are available
    #[allow(dead_code)]
    pub(crate) current_fn: FnBound,
}

/// We assume this will be called after execute_instruction, so less error-handling is needed.
#[allow(unused_variables)]
#[inline(always)]
pub fn update_instruction_metrics<F, RA, Executor>(
    state: &mut VmExecState<F, TracingMemory, PreflightCtx<RA>>,
    executor: &Executor,
    prev_pc: u32, // the pc of the instruction executed, state.pc is next pc
    pc_entry: &PcEntry<F>,
) where
    F: Clone + Send + Sync,
    RA: Arena,
    Executor: PreflightExecutor<F, RA>,
{
    #[cfg(all(feature = "metrics", any(debug_assertions, feature = "perf-metrics")))]
    {
        let pc = state.pc();
        state.metrics.update_backtrace(pc);
    }

    #[cfg(feature = "perf-metrics")]
    {
        use std::iter::zip;

        let pc = state.pc();
        let opcode = pc_entry.insn.opcode;
        let opcode_name = executor.get_opcode_name(opcode.as_usize());

        let debug_info = state.metrics.debug_infos.get(prev_pc);
        let dsl_instr = debug_info.as_ref().map(|info| info.dsl_instruction.clone());

        let now_trace_heights: Vec<usize> = state
            .ctx
            .arenas
            .iter()
            .map(|arena| arena.current_trace_height())
            .collect();
        let now_trace_cells = zip(&state.metrics.main_widths, &now_trace_heights)
            .map(|(main_width, h)| main_width * h)
            .collect_vec();
        state
            .metrics
            .update_trace_cells(now_trace_cells, opcode_name, dsl_instr);

        state.metrics.update_current_fn(pc);
    }
}

// We clear the current trace cell counts so there aren't negative diffs at the start of the next
// segment.
#[cfg(feature = "perf-metrics")]
pub fn end_segment_metrics<F, RA>(state: &mut VmExecState<F, TracingMemory, PreflightCtx<RA>>)
where
    F: Clone + Send + Sync,
    RA: Arena,
{
    state.metrics.current_trace_cells.fill(0);
}

#[cfg(feature = "metrics")]
pub fn emit_opcode_counts(metrics: &VmMetrics, counts: BTreeMap<(usize, String), u64>) {
    for ((air_idx, opcode), count) in counts {
        let Some(air_name) = metrics.air_names.get(air_idx) else {
            continue;
        };
        let labels = [
            ("air_name", air_name.clone()),
            ("air_id", air_idx.to_string()),
            ("opcode", opcode),
        ];
        counter!("opcode_count", &labels).absolute(count);
    }
}

impl VmMetrics {
    #[cfg(feature = "metrics")]
    pub fn set_pk_air_names<PB: ProverBackend>(&mut self, pk: &DeviceMultiStarkProvingKey<PB>) {
        self.air_names = pk.per_air.iter().map(|pk| pk.air_name.clone()).collect();
    }

    #[cfg(feature = "perf-metrics")]
    pub fn set_pk_trace_info<PB: ProverBackend>(&mut self, pk: &DeviceMultiStarkProvingKey<PB>) {
        let (main_widths, total_widths): (Vec<_>, Vec<_>) = pk
            .per_air
            .iter()
            .map(|pk| {
                let width = &pk.vk.params.width;
                (width.main_width(), width.total_width())
            })
            .unzip();
        self.main_widths = main_widths;
        self.total_widths = total_widths;
        self.current_trace_cells = vec![0; self.air_names.len()];
    }

    #[cfg(feature = "perf-metrics")]
    pub fn update_trace_cells(
        &mut self,
        now_trace_cells: Vec<usize>,
        opcode_name: String,
        dsl_instr: Option<String>,
    ) {
        let key = (dsl_instr, opcode_name.clone());
        self.cycle_tracker.increment_opcode(&key);
        *self.counts.entry(key.clone()).or_insert(0) += 1;

        for (air_name, now_value, prev_value) in
            itertools::izip!(&self.air_names, &now_trace_cells, &self.current_trace_cells)
        {
            if prev_value != now_value {
                let cells_used = now_value - prev_value;
                let key = (key.0.clone(), key.1.clone(), air_name.to_owned());
                self.cycle_tracker.increment_cells_used(&key, cells_used);
                *self.trace_cells.entry(key).or_insert(0) += cells_used;
            }
        }
        self.current_trace_cells = now_trace_cells;
    }

    /// Take the cycle tracker and fn bounds information for use in
    /// next segment. Leave the rest of the metrics for recording purposes.
    pub fn partial_take(&mut self) -> Self {
        Self {
            cycle_tracker: mem::take(&mut self.cycle_tracker),
            fn_bounds: mem::take(&mut self.fn_bounds),
            current_fn: mem::take(&mut self.current_fn),
            ..Default::default()
        }
    }

    /// Clear statistics that are local to a segment
    // Important: chip and cycle count metrics should start over for SegmentationStrategy,
    // but we need to carry over the cycle tracker so spans can cross segments
    pub fn clear(&mut self) {
        *self = self.partial_take();
    }

    #[cfg(all(feature = "metrics", any(debug_assertions, feature = "perf-metrics")))]
    pub fn update_backtrace(&mut self, pc: u32) {
        if let Some(info) = self.debug_infos.get(pc) {
            if let Some(trace) = &info.trace {
                self.prev_backtrace = Some(trace.clone());
            }
        }
    }

    #[cfg(feature = "perf-metrics")]
    pub(super) fn update_current_fn(&mut self, pc: u32) {
        if self.fn_bounds.is_empty() {
            return;
        }
        if pc < self.current_fn.start || pc > self.current_fn.end {
            self.current_fn = self
                .fn_bounds
                .range(..=pc)
                .next_back()
                .map(|(_, func)| (*func).clone())
                .unwrap();
            if pc == self.current_fn.start {
                self.cycle_tracker.start(self.current_fn.name.clone());
            } else {
                while let Some(name) = self.cycle_tracker.top() {
                    if name == &self.current_fn.name {
                        break;
                    }
                    self.cycle_tracker.force_end();
                }
            }
        };
    }

    pub fn emit(&self) {
        for ((dsl_ir, opcode), value) in self.counts.iter() {
            let labels = [
                ("dsl_ir", dsl_ir.clone().unwrap_or_else(String::new)),
                ("opcode", opcode.clone()),
            ];
            counter!("frequency", &labels).absolute(*value as u64);
        }

        for ((dsl_ir, opcode, air_name), value) in self.trace_cells.iter() {
            let labels = [
                ("dsl_ir", dsl_ir.clone().unwrap_or_else(String::new)),
                ("opcode", opcode.clone()),
                ("air_name", air_name.clone()),
            ];
            counter!("cells_used", &labels).absolute(*value as u64);
        }
    }
}