miden-processor 0.22.1

use alloc::vec::Vec;
#[cfg(any(test, feature = "testing"))]
use core::ops::Range;

#[cfg(feature = "std")]
use miden_air::trace::PADDED_TRACE_WIDTH;
use miden_air::{
    AirWitness, AuxBuilder, ProcessorAir, PublicInputs, debug,
    trace::{
        DECODER_TRACE_OFFSET, MainTrace, STACK_TRACE_OFFSET, TRACE_WIDTH,
        decoder::{NUM_USER_OP_HELPERS, USER_OP_HELPERS_OFFSET},
    },
};

use crate::{
    AdviceProvider, Felt, MIN_STACK_DEPTH, ProgramInfo, StackInputs, StackOutputs, Word, ZERO,
    fast::ExecutionOutput,
    field::{ExtensionField, QuadFelt},
    precompile::{PrecompileRequest, PrecompileTranscript},
    utils::{ColMatrix, Matrix, RowMajorMatrix},
};

pub(crate) mod utils;
use miden_air::trace::Challenges;
use utils::{AuxColumnBuilder, TraceFragment};

pub mod chiplets;
pub mod execution_tracer;

mod decoder;
mod parallel;
mod range;
mod stack;
mod trace_state;

#[cfg(test)]
mod tests;

// RE-EXPORTS
// ================================================================================================

pub use miden_air::trace::RowIndex;
pub use parallel::{CORE_TRACE_WIDTH, build_trace, build_trace_with_max_len};
pub use utils::{ChipletsLengths, TraceLenSummary};

// VM EXECUTION TRACE
// ================================================================================================

/// Execution trace which is generated when a program is executed on the VM.
///
/// The trace consists of the following components:
/// - Main traces of System, Decoder, Operand Stack, Range Checker, and Chiplets.
/// - Auxiliary trace builders.
/// - Information about the program (program hash and the kernel).
/// - Information about execution outputs (stack state, advice provider, and precompile transcript).
/// - Summary of trace lengths of the main trace components.
#[derive(Debug)]
pub struct ExecutionTrace {
    main_trace: MainTrace,
    aux_trace_builders: AuxTraceBuilders,
    program_info: ProgramInfo,
    stack_outputs: StackOutputs,
    advice: AdviceProvider,
    final_pc_transcript: PrecompileTranscript,
    trace_len_summary: TraceLenSummary,
}

impl ExecutionTrace {
    // CONSTRUCTOR
    // --------------------------------------------------------------------------------------------

    pub fn new_from_parts(
        program_info: ProgramInfo,
        execution_output: ExecutionOutput,
        main_trace: MainTrace,
        aux_trace_builders: AuxTraceBuilders,
        trace_len_summary: TraceLenSummary,
    ) -> Self {
        Self {
            main_trace,
            aux_trace_builders,
            program_info,
            stack_outputs: execution_output.stack,
            advice: execution_output.advice,
            final_pc_transcript: execution_output.final_pc_transcript,
            trace_len_summary,
        }
    }

    // PUBLIC ACCESSORS
    // --------------------------------------------------------------------------------------------

    /// Returns the program info of this execution trace.
    pub fn program_info(&self) -> &ProgramInfo {
        &self.program_info
    }

    /// Returns hash of the program execution of which resulted in this execution trace.
    pub fn program_hash(&self) -> &Word {
        self.program_info.program_hash()
    }

    /// Returns outputs of the program execution which resulted in this execution trace.
    pub fn stack_outputs(&self) -> &StackOutputs {
        &self.stack_outputs
    }

    /// Returns the public inputs for this execution trace.
    pub fn public_inputs(&self) -> PublicInputs {
        PublicInputs::new(
            self.program_info.clone(),
            self.init_stack_state(),
            self.stack_outputs,
            self.final_pc_transcript.state(),
        )
    }

    /// Returns the public values for this execution trace.
    pub fn to_public_values(&self) -> Vec<Felt> {
        self.public_inputs().to_elements()
    }

    /// Returns a clone of the auxiliary trace builders.
    pub fn aux_trace_builders(&self) -> AuxTraceBuilders {
        self.aux_trace_builders.clone()
    }

    /// Returns a reference to the main trace.
    pub fn main_trace(&self) -> &MainTrace {
        &self.main_trace
    }

    /// Returns a mutable reference to the main trace.
    pub fn main_trace_mut(&mut self) -> &mut MainTrace {
        &mut self.main_trace
    }

    /// Returns the precompile requests generated during program execution.
    pub fn precompile_requests(&self) -> &[PrecompileRequest] {
        self.advice.precompile_requests()
    }

    /// Moves all accumulated precompile requests out of the trace, leaving it empty.
    ///
    /// Intended for proof packaging, where requests are serialized into the proof and no longer
    /// needed in the trace after consumption.
    pub fn take_precompile_requests(&mut self) -> Vec<PrecompileRequest> {
        self.advice.take_precompile_requests()
    }

    /// Returns the final precompile transcript after executing all precompile requests.
    pub fn final_precompile_transcript(&self) -> PrecompileTranscript {
        self.final_pc_transcript
    }

    /// Returns the initial state of the top 16 stack registers.
    pub fn init_stack_state(&self) -> StackInputs {
        let mut result = [ZERO; MIN_STACK_DEPTH];
        for (i, result) in result.iter_mut().enumerate() {
            *result = self.main_trace.get_column(i + STACK_TRACE_OFFSET)[0];
        }
        result.into()
    }

    /// Returns the final state of the top 16 stack registers.
    pub fn last_stack_state(&self) -> StackOutputs {
        let last_step = self.last_step();
        let mut result = [ZERO; MIN_STACK_DEPTH];
        for (i, result) in result.iter_mut().enumerate() {
            *result = self.main_trace.get_column(i + STACK_TRACE_OFFSET)[last_step];
        }
        result.into()
    }

    /// Returns helper registers state at the specified `clk` of the VM
    pub fn get_user_op_helpers_at(&self, clk: u32) -> [Felt; NUM_USER_OP_HELPERS] {
        let mut result = [ZERO; NUM_USER_OP_HELPERS];
        for (i, result) in result.iter_mut().enumerate() {
            *result = self.main_trace.get_column(DECODER_TRACE_OFFSET + USER_OP_HELPERS_OFFSET + i)
                [clk as usize];
        }
        result
    }

    /// Returns the trace length.
    pub fn get_trace_len(&self) -> usize {
        self.main_trace.num_rows()
    }

    /// Returns the length of the trace (number of rows in the main trace).
    pub fn length(&self) -> usize {
        self.get_trace_len()
    }

    /// Returns a summary of the lengths of main, range and chiplet traces.
    pub fn trace_len_summary(&self) -> &TraceLenSummary {
        &self.trace_len_summary
    }

    /// Returns the final advice provider state.
    pub fn advice_provider(&self) -> &AdviceProvider {
        &self.advice
    }

    /// Destructures this execution trace into the process’s final stack and advice states.
    pub fn into_outputs(self) -> (StackOutputs, AdviceProvider) {
        (self.stack_outputs, self.advice)
    }

    // DEBUG CONSTRAINT CHECKING
    // --------------------------------------------------------------------------------------------

    /// Validates this execution trace against all AIR constraints without generating a STARK
    /// proof.
    ///
    /// This is the recommended way to test trace correctness. It is much faster than full STARK
    /// proving and provides better error diagnostics (panics on the first constraint violation
    /// with the instance and row index).
    ///
    /// # Panics
    ///
    /// Panics if any AIR constraint evaluates to nonzero.
    pub fn check_constraints(&self) {
        let public_inputs = self.public_inputs();
        let trace_matrix = self.to_row_major_matrix();

        let (public_values, kernel_felts) = public_inputs.to_air_inputs();
        let var_len_public_inputs: &[&[Felt]] = &[&kernel_felts];

        let aux_builder = self.aux_trace_builders();

        // Derive deterministic challenges by hashing public values with Poseidon2.
        // The 4-element digest maps directly to 2 QuadFelt challenges.
        let digest = crate::crypto::hash::Poseidon2::hash_elements(&public_values);
        let challenges =
            [QuadFelt::new([digest[0], digest[1]]), QuadFelt::new([digest[2], digest[3]])];

        let witness = AirWitness::new(&trace_matrix, &public_values, var_len_public_inputs);
        debug::check_constraints(&ProcessorAir, witness, &aux_builder, &challenges);
    }

    /// Converts the main trace from column-major to row-major format.
    ///
    /// Only includes the first [`TRACE_WIDTH`] columns (excluding padding columns added for
    /// Poseidon2 rate alignment), which is the width expected by the AIR.
    // TODO: the padding columns can be removed once we use the lifted-stark's virtual trace
    // alignment, which pads to the required rate width without materializing extra columns.
    pub fn to_row_major_matrix(&self) -> RowMajorMatrix<Felt> {
        let trace_len = self.get_trace_len();
        let mut col_major_data = Vec::with_capacity(TRACE_WIDTH * trace_len);
        for col_idx in 0..TRACE_WIDTH {
            col_major_data.extend_from_slice(self.main_trace.get_column(col_idx));
        }
        let col_major_matrix = RowMajorMatrix::new(col_major_data, trace_len);
        col_major_matrix.transpose()
    }

    // HELPER METHODS
    // --------------------------------------------------------------------------------------------

    /// Returns the index of the last row in the trace.
    fn last_step(&self) -> usize {
        self.length() - 1
    }

    // TEST HELPERS
    // --------------------------------------------------------------------------------------------
    #[cfg(feature = "std")]
    pub fn print(&self) {
        use miden_air::trace::TRACE_WIDTH;

        let mut row = [ZERO; PADDED_TRACE_WIDTH];
        for i in 0..self.length() {
            self.main_trace.read_row_into(i, &mut row);
            std::println!(
                "{:?}",
                row.iter().take(TRACE_WIDTH).map(|v| v.as_canonical_u64()).collect::<Vec<_>>()
            );
        }
    }

    #[cfg(any(test, feature = "testing"))]
    pub fn get_column_range(&self, range: Range<usize>) -> Vec<Vec<Felt>> {
        self.main_trace.get_column_range(range)
    }

    pub fn build_aux_trace<E>(&self, rand_elements: &[E]) -> Option<ColMatrix<E>>
    where
        E: ExtensionField<Felt>,
    {
        let aux_columns =
            self.aux_trace_builders.build_aux_columns(&self.main_trace, rand_elements);

        Some(ColMatrix::new(aux_columns))
    }
}

// AUX TRACE BUILDERS
// ================================================================================================

#[derive(Debug, Clone)]
pub struct AuxTraceBuilders {
    pub(crate) decoder: decoder::AuxTraceBuilder,
    pub(crate) stack: stack::AuxTraceBuilder,
    pub(crate) range: range::AuxTraceBuilder,
    pub(crate) chiplets: chiplets::AuxTraceBuilder,
}

impl AuxTraceBuilders {
    /// Builds auxiliary columns for all trace segments given the main trace and challenges.
    ///
    /// This is the internal column-major version used by the processor.
    pub fn build_aux_columns<E>(&self, main_trace: &MainTrace, challenges: &[E]) -> Vec<Vec<E>>
    where
        E: ExtensionField<Felt>,
    {
        // Expand raw challenges (alpha, beta) into coefficient array once, then pass
        // the expanded challenges to all sub-builders.
        let challenges = Challenges::<E>::new(challenges[0], challenges[1]);

        let decoder_cols = self.decoder.build_aux_columns(main_trace, &challenges);
        let stack_cols = self.stack.build_aux_columns(main_trace, &challenges);
        let range_cols = self.range.build_aux_columns(main_trace, &challenges);
        let chiplets_cols = self.chiplets.build_aux_columns(main_trace, &challenges);

        decoder_cols
            .into_iter()
            .chain(stack_cols)
            .chain(range_cols)
            .chain(chiplets_cols)
            .collect()
    }
}

// PLONKY3 AUX TRACE BUILDER
// ================================================================================================
//
// Implements the upstream `AuxBuilder` trait from `p3_miden_lifted_air` directly on
// `AuxTraceBuilders`. Plonky3 uses row-major matrices while our existing aux trace building logic
// uses column-major format. This impl adapts between the two by converting the main trace from
// row-major to column-major, delegating to the existing logic, and converting the result back.

impl<EF: ExtensionField<Felt>> AuxBuilder<Felt, EF> for AuxTraceBuilders {
    fn build_aux_trace(
        &self,
        main: &RowMajorMatrix<Felt>,
        challenges: &[EF],
    ) -> (RowMajorMatrix<EF>, Vec<EF>) {
        let _span = tracing::info_span!("build_aux_trace").entered();

        // Transpose the row-major main trace into column-major `MainTrace` needed by the
        // auxiliary trace builders. The last program row is the point where the clock
        // (column 0) stops incrementing.
        let main_trace_col_major = {
            let num_rows = main.height();
            // Detect last program row from row-major layout using column 0 (clock).
            let last_program_row = (1..num_rows)
                .find(|&i| {
                    main.get(i, 0).expect("valid indices")
                        != main.get(i - 1, 0).expect("valid indices") + Felt::ONE
                })
                .map_or(num_rows - 1, |i| i - 1);
            let transposed = main.transpose();
            let columns: Vec<Vec<Felt>> = transposed.row_slices().map(|row| row.to_vec()).collect();
            MainTrace::new(ColMatrix::new(columns), last_program_row.into())
        };

        // Build auxiliary columns in column-major format.
        let aux_columns = self.build_aux_columns(&main_trace_col_major, challenges);
        assert!(!aux_columns.is_empty(), "aux columns should not be empty");

        // Flatten column-major aux columns into a contiguous buffer, then transpose
        // to the row-major layout expected by the lifted prover.
        let trace_len = main.height();
        let num_ef_cols = aux_columns.len();
        let mut col_major_data = Vec::with_capacity(trace_len * num_ef_cols);
        for col in aux_columns {
            col_major_data.extend_from_slice(&col);
        }
        let aux_trace = RowMajorMatrix::new(col_major_data, trace_len).transpose();

        // Extract the last row from the row-major aux trace for Fiat-Shamir.
        let last_row = aux_trace
            .row_slice(trace_len - 1)
            .expect("aux trace has at least one row")
            .to_vec();

        (aux_trace, last_row)
    }
}