miden-processor 0.32.1

//! End-to-end collection-phase smoke test for the prover-side LogUp pipeline.
//!
//! Runs a tiny MASM basic block through `build_trace_from_ops`, materialises the resulting
//! main trace as a [`RowMajorMatrix<Felt>`], and pipes it through [`build_lookup_fractions`]
//! + [`accumulate`]. The test validates:
//!
//! 1. **Shape-const drift**: every bus emitter's declared `MAX_INTERACTIONS_PER_ROW` is large
//!    enough to accommodate real trace data (the `debug_assert!` inside
//!    `ProverLookupBuilder::column` panics on overflow).
//! 2. **Zero-denominator bugs**: every encoded `LookupMessage` evaluates to a non-zero
//!    extension-field element, so per-fraction `try_inverse` inside the accumulator does not panic.
//! 3. **Pipeline plumbing**: row slicing with wraparound, per-row periodic composition, `RowWindow`
//!    construction over a real matrix, and the dense `LookupFractions` buffer all line up.
//! 4. **Constraint agreement**: the trace and its prover-built auxiliary columns satisfy the AIR.

use alloc::{boxed::Box, vec::Vec};
use core::{
    borrow::{Borrow, BorrowMut},
    mem::size_of,
};

use miden_air::{
    BaseAir, ChipletCols, ControllerCols, MidenAir, MidenMultiAir, ProverStatement, StarkConfig,
    Statement, config, debug,
    logup::{BusId, HasherPermLinkMsg, MIDEN_MAX_MESSAGE_WIDTH},
    lookup::{Challenges, LookupMessage, accumulate, build_lookup_fractions},
    trace::CHIPLET_CONTROLLER_OFFSET,
};
use miden_core::{
    field::QuadFelt,
    utils::{Matrix, RowMajorMatrix},
};

use super::{Felt, VmTrace, build_trace_from_ops, rand_array};
use crate::operation::Operation;

const CONTROLLER_OFFSET: usize = CHIPLET_CONTROLLER_OFFSET;
const CONTROLLER_WIDTH: usize = size_of::<ControllerCols<u8>>();
static PANIC_HOOK_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());

/// Pad/Add/Mul/Drop inside a span — same flavour of ops the decoder/stack tests use, with
/// enough variety to exercise decoder, stack, and range-check bus emitters.
fn tiny_span() -> Vec<Operation> {
    vec![
        Operation::Pad,
        Operation::Pad,
        Operation::Add,
        Operation::Pad,
        Operation::Mul,
        Operation::Drop,
    ]
}

fn perm_link_fractions(
    chip_matrix: &RowMajorMatrix<Felt>,
    poseidon2_matrix: &RowMajorMatrix<Felt>,
    challenges: &Challenges<QuadFelt>,
) -> Vec<(Felt, QuadFelt)> {
    let chip_periodic = MidenAir::Chiplets.periodic_columns();
    let poseidon2_periodic = MidenAir::Poseidon2Permutation.periodic_columns();
    let chip_fractions =
        build_lookup_fractions(&MidenAir::Chiplets, chip_matrix, &chip_periodic, challenges);
    let poseidon2_fractions = build_lookup_fractions(
        &MidenAir::Poseidon2Permutation,
        poseidon2_matrix,
        &poseidon2_periodic,
        challenges,
    );

    chip_fractions
        .fractions()
        .iter()
        .chain(poseidon2_fractions.fractions())
        .copied()
        .collect()
}

fn net_multiplicity(fractions: &[(Felt, QuadFelt)], denom: QuadFelt) -> Felt {
    let mut net = Felt::ZERO;
    for &(multiplicity, encoded) in fractions {
        if encoded == denom {
            net += multiplicity;
        }
    }
    net
}

fn chiplet_row(matrix: &RowMajorMatrix<Felt>, row: usize) -> &ChipletCols<Felt> {
    let width = matrix.width();
    matrix.values[row * width..(row + 1) * width].borrow()
}

fn controller_row(matrix: &RowMajorMatrix<Felt>, row: usize) -> &ControllerCols<Felt> {
    chiplet_row(matrix, row).controller()
}

fn controller_row_mut(matrix: &mut RowMajorMatrix<Felt>, row: usize) -> &mut ControllerCols<Felt> {
    let width = matrix.width();
    let start = row * width + CONTROLLER_OFFSET;
    matrix.values[start..start + CONTROLLER_WIDTH].borrow_mut()
}

pub(super) fn assert_trace_constraints_reject(
    trace: &VmTrace,
    core_matrix: RowMajorMatrix<Felt>,
    chip_matrix: RowMajorMatrix<Felt>,
    poseidon2_matrix: RowMajorMatrix<Felt>,
) {
    let (public_values, aux_inputs) = trace.public_inputs().to_air_inputs();
    let statement =
        Statement::<Felt, QuadFelt, _>::new(MidenMultiAir::new(), public_values, aux_inputs)
            .expect("valid statement inputs");
    let prover_statement =
        ProverStatement::new(statement, vec![core_matrix, chip_matrix, poseidon2_matrix])
            .expect("valid trace shapes");

    let config = config::poseidon2_config(config::pcs_params(), config::RELATION_DIGEST);
    let _guard = PANIC_HOOK_LOCK.lock().expect("panic hook lock poisoned");
    let panic_hook = std::panic::take_hook();
    std::panic::set_hook(Box::new(|_| {}));
    let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
        debug::check_constraints(&prover_statement, config.challenger());
    }));
    std::panic::set_hook(panic_hook);
    assert!(result.is_err(), "mutated trace should violate AIR constraints");
}

#[test]
fn build_lookup_fractions_runs_on_execution_trace() {
    let trace = build_trace_from_ops(tiny_span(), &[]);

    let (core_matrix, chip_matrix, poseidon2_matrix) = trace.main_trace().to_air_matrices();
    let chip_periodic = MidenAir::Chiplets.periodic_columns();
    let poseidon2_periodic = MidenAir::Poseidon2Permutation.periodic_columns();

    // QuadFelt challenges for LogUp, built from 4 random Felts (QuadFelt itself doesn't
    // implement Randomizable, so we draw base-field elements and pair them).
    let raw = rand_array::<Felt, 4>();
    let alpha = QuadFelt::new([raw[0], raw[1]]);
    let beta = QuadFelt::new([raw[2], raw[3]]);
    let challenges =
        Challenges::<QuadFelt>::new(alpha, beta, MIDEN_MAX_MESSAGE_WIDTH, BusId::COUNT);

    // --- Core ---
    let core_fractions = build_lookup_fractions(&MidenAir::Core, &core_matrix, &[], &challenges);
    assert!(
        !core_fractions.fractions().is_empty(),
        "no Core fractions collected — trace is degenerate or emitters are broken",
    );
    let _ = accumulate(&core_fractions);

    // --- Chiplets ---
    let chip_fractions =
        build_lookup_fractions(&MidenAir::Chiplets, &chip_matrix, &chip_periodic, &challenges);
    assert!(
        !chip_fractions.fractions().is_empty(),
        "no Chiplets fractions collected — trace is degenerate or emitters are broken",
    );
    let _ = accumulate(&chip_fractions);

    // --- Poseidon2 permutation ---
    let poseidon2_fractions = build_lookup_fractions(
        &MidenAir::Poseidon2Permutation,
        &poseidon2_matrix,
        &poseidon2_periodic,
        &challenges,
    );
    assert!(
        !poseidon2_fractions.fractions().is_empty(),
        "no Poseidon2 fractions collected; trace is degenerate or emitters are broken",
    );
    let _ = accumulate(&poseidon2_fractions);

    trace.check_constraints();
}

#[test]
fn u32div_trace_satisfies_air() {
    let operations = vec![Operation::U32div, Operation::Drop, Operation::Drop, Operation::U32div];
    let trace = build_trace_from_ops(operations, &[0x0008_000b, 0x003b_0051, 3, 0x0003_0004]);
    trace.check_constraints();
}

#[test]
fn perm_link_rejects_swapped_controller_outputs() {
    let trace =
        build_trace_from_ops(vec![Operation::HPerm, Operation::HPerm], &[8, 7, 6, 5, 4, 3, 2, 1]);
    let (core_matrix, chip_matrix, poseidon2_matrix) = trace.main_trace().to_air_matrices();

    let output_rows: Vec<_> = (0..chip_matrix.height())
        .filter(|&row| {
            let chiplet = chiplet_row(&chip_matrix, row);
            let ctrl = chiplet.controller();
            chiplet.chiplet_selectors()[0] == Felt::ZERO
                && ctrl.s0 == Felt::ZERO
                && ctrl.s1 == Felt::ZERO
                && ctrl.s2 == Felt::ONE
        })
        .take(2)
        .collect();
    assert_eq!(output_rows.len(), 2, "expected two controller output rows");

    let ctrl_a = controller_row(&chip_matrix, output_rows[0]);
    let ctrl_b = controller_row(&chip_matrix, output_rows[1]);
    let state_a = ctrl_a.state;
    let state_b = ctrl_b.state;
    assert_ne!(state_a, state_b, "test needs two distinct permutation outputs");

    let perm_id_a = ctrl_a.perm_id;
    let perm_id_b = ctrl_b.perm_id;
    assert_ne!(perm_id_a, perm_id_b, "test needs two distinct permutation ids");

    let challenges = Challenges::<QuadFelt>::new(
        QuadFelt::new([Felt::new_unchecked(7), Felt::ZERO]),
        QuadFelt::new([Felt::new_unchecked(11), Felt::ZERO]),
        MIDEN_MAX_MESSAGE_WIDTH,
        BusId::COUNT,
    );

    let honest_fractions = perm_link_fractions(&chip_matrix, &poseidon2_matrix, &challenges);
    for msg in [
        HasherPermLinkMsg::Output { perm_id: perm_id_a, state: state_a },
        HasherPermLinkMsg::Output { perm_id: perm_id_b, state: state_b },
    ] {
        assert_eq!(
            net_multiplicity(&honest_fractions, msg.encode(&challenges)),
            Felt::ZERO,
            "honest controller output link is balanced"
        );
    }

    let mut state_swapped_chip_matrix = chip_matrix.clone();
    controller_row_mut(&mut state_swapped_chip_matrix, output_rows[0]).state = state_b;
    controller_row_mut(&mut state_swapped_chip_matrix, output_rows[1]).state = state_a;

    let swapped_fractions =
        perm_link_fractions(&state_swapped_chip_matrix, &poseidon2_matrix, &challenges);
    for msg in [
        HasherPermLinkMsg::Output { perm_id: perm_id_a, state: state_b },
        HasherPermLinkMsg::Output { perm_id: perm_id_b, state: state_a },
    ] {
        assert_eq!(
            net_multiplicity(&swapped_fractions, msg.encode(&challenges)),
            Felt::ONE,
            "swapped controller output leaves an unmatched perm-link addition"
        );
    }

    let mut tuple_swapped_chip_matrix = chip_matrix;
    {
        let row = controller_row_mut(&mut tuple_swapped_chip_matrix, output_rows[0]);
        row.state = state_b;
        row.perm_id = perm_id_b;
    }
    {
        let row = controller_row_mut(&mut tuple_swapped_chip_matrix, output_rows[1]);
        row.state = state_a;
        row.perm_id = perm_id_a;
    }

    let balanced_fractions =
        perm_link_fractions(&tuple_swapped_chip_matrix, &poseidon2_matrix, &challenges);
    for msg in [
        HasherPermLinkMsg::Output { perm_id: perm_id_a, state: state_a },
        HasherPermLinkMsg::Output { perm_id: perm_id_b, state: state_b },
    ] {
        assert_eq!(
            net_multiplicity(&balanced_fractions, msg.encode(&challenges)),
            Felt::ZERO,
            "swapping output tuples keeps the perm-link bus balanced"
        );
    }

    assert_trace_constraints_reject(
        &trace,
        core_matrix,
        tuple_swapped_chip_matrix,
        poseidon2_matrix,
    );
}