p3-miden-prover 0.4.2

//! See `prover.rs` for an overview of the protocol and a more detailed soundness analysis.
use alloc::vec;
use alloc::vec::Vec;

use core::marker::PhantomData;

use itertools::Itertools;
use p3_challenger::{CanObserve, FieldChallenger};
use p3_commit::{Pcs, PolynomialSpace};
use p3_field::{BasedVectorSpace, Field, PrimeCharacteristicRing, TwoAdicField};
use p3_matrix::dense::RowMajorMatrixView;
use p3_matrix::stack::VerticalPair;
use p3_miden_air::{BusType, MidenAir};
use p3_util::zip_eq::zip_eq;
use tracing::{debug_span, instrument};

use crate::periodic_tables::evaluate_periodic_at_point;
use crate::symbolic_builder::get_log_quotient_degree;
use crate::util::verifier_row_to_ext;
use crate::{
    AirWithBoundaryConstraints, Domain, PcsError, Proof, StarkGenericConfig, Val,
    VerifierConstraintFolder,
};

/// Recomposes the quotient polynomial from its chunks evaluated at a point.
///
/// Given quotient chunks and their domains, this computes the Lagrange
/// interpolation coefficients (zps) and reconstructs quotient(zeta).
pub fn recompose_quotient_from_chunks<SC>(
    quotient_chunks_domains: &[Domain<SC>],
    quotient_chunks: &[Vec<SC::Challenge>],
    zeta: SC::Challenge,
) -> SC::Challenge
where
    SC: StarkGenericConfig,
{
    let zps = quotient_chunks_domains
        .iter()
        .enumerate()
        .map(|(i, domain)| {
            quotient_chunks_domains
                .iter()
                .enumerate()
                .filter(|(j, _)| *j != i)
                .map(|(_, other_domain)| {
                    other_domain.vanishing_poly_at_point(zeta)
                        * other_domain
                            .vanishing_poly_at_point(domain.first_point())
                            .inverse()
                })
                .product::<SC::Challenge>()
        })
        .collect_vec();

    quotient_chunks
        .iter()
        .enumerate()
        .map(|(ch_i, ch)| {
            // We checked in valid_shape the length of "ch" is equal to
            // <SC::Challenge as BasedVectorSpace<Val<SC>>>::DIMENSION. Hence
            // the unwrap() will never panic.
            zps[ch_i]
                * ch.iter()
                    .enumerate()
                    .map(|(e_i, &c)| SC::Challenge::ith_basis_element(e_i).unwrap() * c)
                    .sum::<SC::Challenge>()
        })
        .sum::<SC::Challenge>()
}

/// Verifies that the folded constraints match the quotient polynomial at zeta.
///
/// This evaluates the AIR constraints at the out-of-domain point and checks
/// that constraints(zeta) / Z_H(zeta) = quotient(zeta).
#[allow(clippy::too_many_arguments)]
pub fn verify_constraints<SC, A, PcsErr>(
    air: &A,
    trace_local: &[SC::Challenge],
    trace_next: &[SC::Challenge],
    preprocessed_local: Option<&[SC::Challenge]>,
    preprocessed_next: Option<&[SC::Challenge]>,
    aux_local: Option<&[SC::Challenge]>,
    aux_next: Option<&[SC::Challenge]>,
    randomness: &[SC::Challenge],
    aux_bus_boundary_values: &[SC::Challenge],
    public_values: &[Val<SC>],
    trace_domain: Domain<SC>,
    zeta: SC::Challenge,
    alpha: SC::Challenge,
    quotient: SC::Challenge,
) -> Result<(), VerificationError<PcsErr>>
where
    SC: StarkGenericConfig + Sync,
    A: MidenAir<Val<SC>, SC::Challenge>,
    Val<SC>: TwoAdicField,
{
    let sels = trace_domain.selectors_at_point(zeta);

    // =====================================
    // Periodic entries
    // =====================================
    let periodic_values: Vec<SC::Challenge> = evaluate_periodic_at_point::<Val<SC>, SC::Challenge>(
        air.periodic_table(),
        trace_domain,
        zeta,
    );

    // =====================================
    // Main trace
    // =====================================
    let main = VerticalPair::new(
        RowMajorMatrixView::new_row(trace_local),
        RowMajorMatrixView::new_row(trace_next),
    );

    // =====================================
    // Preprocessed trace
    // =====================================
    let preprocessed = match (preprocessed_local, preprocessed_next) {
        (Some(local), Some(next)) => Some(VerticalPair::new(
            RowMajorMatrixView::new_row(local),
            RowMajorMatrixView::new_row(next),
        )),
        _ => None,
    };

    // =====================================
    // Aux trace
    // =====================================
    // Aux trace is committed as flattened base limbs. Recompose into EF values.
    let aux_local_ext;
    let aux_next_ext;
    let aux = match (aux_local, aux_next) {
        (Some(local), Some(next)) => {
            aux_local_ext = verifier_row_to_ext::<Val<SC>, SC::Challenge>(local)
                .ok_or(VerificationError::InvalidProofShape)?;
            aux_next_ext = verifier_row_to_ext::<Val<SC>, SC::Challenge>(next)
                .ok_or(VerificationError::InvalidProofShape)?;

            VerticalPair::new(
                RowMajorMatrixView::new_row(&aux_local_ext),
                RowMajorMatrixView::new_row(&aux_next_ext),
            )
        }
        _ => {
            // Create an empty ViewPair with zero width
            let empty: &[SC::Challenge] = &[];
            VerticalPair::new(
                RowMajorMatrixView::new_row(empty),
                RowMajorMatrixView::new_row(empty),
            )
        }
    };
    let mut folder: VerifierConstraintFolder<'_, SC> = VerifierConstraintFolder {
        main,
        aux,
        randomness,
        aux_bus_boundary_values,
        preprocessed,
        public_values,
        periodic_values: &periodic_values,
        is_first_row: sels.is_first_row,
        is_last_row: sels.is_last_row,
        is_transition: sels.is_transition,
        alpha,
        accumulator: SC::Challenge::ZERO,
    };
    air.eval(&mut folder);
    let folded_constraints = folder.accumulator;

    // Check that constraints(zeta) / Z_H(zeta) = quotient(zeta)
    if folded_constraints * sels.inv_vanishing != quotient {
        return Err(VerificationError::OodEvaluationMismatch { index: None });
    }

    Ok(())
}

/// Validates and commits the preprocessed trace if present.
/// Returns the preprocessed width and its commitment hash (available iff width > 0).
#[allow(clippy::type_complexity)]
fn process_preprocessed_trace<SC, A>(
    air: &A,
    opened_values: &crate::proof::OpenedValues<SC::Challenge>,
    pcs: &SC::Pcs,
    trace_domain: <SC::Pcs as Pcs<SC::Challenge, SC::Challenger>>::Domain,
    is_zk: usize,
) -> Result<
    (
        usize,
        Option<<SC::Pcs as Pcs<SC::Challenge, SC::Challenger>>::Commitment>,
    ),
    VerificationError<PcsError<SC>>,
>
where
    SC: StarkGenericConfig,
    A: MidenAir<Val<SC>, SC::Challenge>,
    Val<SC>: TwoAdicField,
{
    // If verifier asked for preprocessed trace, then proof should have it
    let preprocessed = air.preprocessed_trace();
    let preprocessed_width = preprocessed.as_ref().map(|m| m.width).unwrap_or(0);
    let preprocessed_local_len = opened_values
        .preprocessed_local
        .as_ref()
        .map_or(0, |v| v.len());
    let preprocessed_next_len = opened_values
        .preprocessed_next
        .as_ref()
        .map_or(0, |v| v.len());
    if preprocessed_width != preprocessed_local_len || preprocessed_width != preprocessed_next_len {
        // Verifier expects preprocessed trace while proof does not have it, or vice versa
        return Err(VerificationError::InvalidProofShape);
    }

    if preprocessed_width > 0 {
        assert_eq!(is_zk, 0); // TODO: preprocessed columns not supported in zk mode
        let height = preprocessed.as_ref().unwrap().values.len() / preprocessed_width;
        assert_eq!(
            height,
            trace_domain.size(),
            "Verifier's preprocessed trace height must be equal to trace domain size"
        );
        let (preprocessed_commit, _) = debug_span!("process preprocessed trace")
            .in_scope(|| pcs.commit([(trace_domain, preprocessed.unwrap())]));
        Ok((preprocessed_width, Some(preprocessed_commit)))
    } else {
        Ok((preprocessed_width, None))
    }
}

#[instrument(skip_all)]
pub fn verify<SC, A>(
    config: &SC,
    air: &A,
    proof: &Proof<SC>,
    public_values: &[Val<SC>],
    var_length_public_inputs: &[&[&[Val<SC>]]],
) -> Result<(), VerificationError<PcsError<SC>>>
where
    SC: StarkGenericConfig + Sync,
    A: MidenAir<Val<SC>, SC::Challenge>,
    Val<SC>: TwoAdicField,
{
    let air = &AirWithBoundaryConstraints {
        inner: air,
        phantom: PhantomData::<SC>,
    };

    let Proof {
        commitments,
        opened_values,
        opening_proof,
        aux_finals,
        degree_bits,
    } = proof;

    let pcs = config.pcs();
    let degree = 1 << degree_bits;
    let aux_width = air.aux_width();
    let num_randomness = air.num_randomness();

    let trace_domain = pcs.natural_domain_for_degree(degree);
    // TODO: allow moving preprocessed commitment to preprocess time, if known in advance
    let (preprocessed_width, preprocessed_commit) =
        process_preprocessed_trace::<SC, _>(air, opened_values, pcs, trace_domain, config.is_zk())?;

    let log_quotient_degree = get_log_quotient_degree::<Val<SC>, SC::Challenge, _>(
        air,
        preprocessed_width,
        public_values.len(),
        config.is_zk(),
        aux_width,
        num_randomness,
    );
    let quotient_degree = 1 << (log_quotient_degree + config.is_zk());
    let mut challenger = config.initialise_challenger();
    let init_trace_domain = pcs.natural_domain_for_degree(degree >> (config.is_zk()));

    let quotient_domain =
        trace_domain.create_disjoint_domain(1 << (degree_bits + log_quotient_degree));
    let quotient_chunks_domains = quotient_domain.split_domains(quotient_degree);

    let randomized_quotient_chunks_domains = quotient_chunks_domains
        .iter()
        .map(|domain| pcs.natural_domain_for_degree(domain.size() << (config.is_zk())))
        .collect_vec();
    // Check that the random commitments are/are not present depending on the ZK setting.
    // - If ZK is enabled, the prover should have random commitments.
    // - If ZK is not enabled, the prover should not have random commitments.
    if (opened_values.random.is_some() != SC::Pcs::ZK)
        || (commitments.random.is_some() != SC::Pcs::ZK)
    {
        return Err(VerificationError::RandomizationError);
    }

    // Observe the instance.
    challenger.observe(Val::<SC>::from_usize(proof.degree_bits));
    challenger.observe(Val::<SC>::from_usize(proof.degree_bits - config.is_zk()));
    challenger.observe(Val::<SC>::from_usize(preprocessed_width));
    // TODO: Might be best practice to include other instance data here in the transcript, like some
    // encoding of the AIR. This protects against transcript collisions between distinct instances.
    // Practically speaking though, the only related known attack is from failing to include public
    // values. It's not clear if failing to include other instance data could enable a transcript
    // collision, since most such changes would completely change the set of satisfying witnesses.

    challenger.observe(commitments.trace.clone());
    if preprocessed_width > 0 {
        challenger.observe(preprocessed_commit.as_ref().unwrap().clone());
    }
    challenger.observe_slice(public_values);

    // begin processing aux trace (optional)
    let num_randomness = air.num_randomness();

    let air_width = air.width();
    let bus_types = air.bus_types();
    let valid_shape = opened_values.trace_local.len() == air_width
        && opened_values.trace_next.len() == air_width
        && opened_values.quotient_chunks.len() == quotient_degree
        && opened_values
            .quotient_chunks
            .iter()
            .all(|qc| qc.len() == SC::Challenge::DIMENSION)
        // We've already checked that opened_values.random is present if and only if ZK is enabled.
        && opened_values.random.as_ref().is_none_or(|r_comm| r_comm.len() == SC::Challenge::DIMENSION)
        // Check aux trace shape
        && if num_randomness > 0 {
            let aux_width_base = aux_width * SC::Challenge::DIMENSION;
            // Note: bus_types length is not matched against aux_width, to allow for more generic aux traces.
            match (&opened_values.aux_trace_local, &opened_values.aux_trace_next) {
                (Some(l), Some(n)) => l.len() == aux_width_base
                    && n.len() == aux_width_base
                    && aux_finals.len() == aux_width,
                _ => false,
            }
        } else {
            opened_values.aux_trace_local.is_none() && opened_values.aux_trace_next.is_none() && aux_finals.is_empty()
        };
    if !valid_shape {
        return Err(VerificationError::InvalidProofShape);
    }
    let randomness = if num_randomness != 0 {
        let randomness: Vec<SC::Challenge> = (0..num_randomness)
            .map(|_| challenger.sample_algebra_element())
            .collect();

        if let Some(aux_commit) = &commitments.aux {
            challenger.observe(aux_commit.clone());
        } else {
            return Err(VerificationError::InvalidProofShape);
        }
        for aux_final in aux_finals {
            challenger.observe_algebra_element(*aux_final);
        }

        randomness
    } else {
        // No aux trace expected
        if commitments.aux.is_some() {
            return Err(VerificationError::InvalidProofShape);
        }
        vec![]
    };

    // Get the first Fiat Shamir challenge which will be used to combine all constraint polynomials
    // into a single polynomial.
    //
    // Soundness Error: n/|EF| where n is the number of constraints.
    let alpha = challenger.sample_algebra_element();

    challenger.observe(commitments.quotient_chunks.clone());

    // We've already checked that commitments.random is present if and only if ZK is enabled.
    // Observe the random commitment if it is present.
    if let Some(r_commit) = commitments.random.clone() {
        challenger.observe(r_commit);
    }

    // Get an out-of-domain point to open our values at.
    //
    // Soundness Error: dN/|EF| where `N` is the trace length and our constraint polynomial has degree `d`.
    let zeta = challenger.sample_algebra_element();
    let zeta_next = init_trace_domain
        .next_point(zeta)
        .ok_or(VerificationError::NextPointUnavailable)?;

    // We've already checked that commitments.random and opened_values.random are present if and only if ZK is enabled.
    let mut coms_to_verify = if let Some(random_commit) = &commitments.random {
        let random_values = opened_values
            .random
            .as_ref()
            .ok_or(VerificationError::RandomizationError)?;
        vec![(
            random_commit.clone(),
            vec![(trace_domain, vec![(zeta, random_values.clone())])],
        )]
    } else {
        vec![]
    };
    coms_to_verify.extend(vec![
        (
            commitments.trace.clone(),
            vec![(
                trace_domain,
                vec![
                    (zeta, opened_values.trace_local.clone()),
                    (zeta_next, opened_values.trace_next.clone()),
                ],
            )],
        ),
        (
            commitments.quotient_chunks.clone(),
            // Check the commitment on the randomized domains.
            zip_eq(
                randomized_quotient_chunks_domains.iter(),
                &opened_values.quotient_chunks,
                VerificationError::InvalidProofShape,
            )?
            .map(|(domain, values)| (*domain, vec![(zeta, values.clone())]))
            .collect_vec(),
        ),
    ]);

    // Add aux trace verification if present
    if let Some(aux_commit) = &commitments.aux {
        let aux_local = opened_values
            .aux_trace_local
            .as_ref()
            .ok_or(VerificationError::InvalidProofShape)?;
        let aux_next = opened_values
            .aux_trace_next
            .as_ref()
            .ok_or(VerificationError::InvalidProofShape)?;
        coms_to_verify.push((
            aux_commit.clone(),
            vec![(
                trace_domain,
                vec![(zeta, aux_local.clone()), (zeta_next, aux_next.clone())],
            )],
        ));
    }

    // Add preprocessed commitment verification if present
    if preprocessed_width > 0 {
        let preprocessed_local = opened_values
            .preprocessed_local
            .as_ref()
            .ok_or(VerificationError::InvalidProofShape)?;
        let preprocessed_next = opened_values
            .preprocessed_next
            .as_ref()
            .ok_or(VerificationError::InvalidProofShape)?;

        coms_to_verify.push((
            preprocessed_commit.unwrap(),
            vec![(
                trace_domain,
                vec![
                    (zeta, preprocessed_local.clone()),
                    (zeta_next, preprocessed_next.clone()),
                ],
            )],
        ));
    }

    pcs.verify(coms_to_verify, opening_proof, &mut challenger)
        .map_err(VerificationError::InvalidOpeningArgument)?;

    let quotient = recompose_quotient_from_chunks::<SC>(
        &quotient_chunks_domains,
        &opened_values.quotient_chunks,
        zeta,
    );

    // Verify the aux trace final values match the expected values if the aux trace contains buses (one bus per aux column)
    // Note: if no buses are defined (bus_types.is_empty), the boundary values of the aux_trace are not checked against the provided variable-length public inputs.
    for (idx, (bus_type, aux_final)) in bus_types.iter().zip(aux_finals).enumerate() {
        let public_inputs_for_bus = *var_length_public_inputs
            .get(idx)
            .ok_or(VerificationError::InvalidProofShape)?;
        let expected_final = match bus_type {
            BusType::Multiset => bus_multiset_boundary_varlen::<_, SC>(
                &randomness,
                public_inputs_for_bus.iter().copied(),
            ),
            BusType::Logup => bus_logup_boundary_varlen::<_, SC>(
                &randomness,
                public_inputs_for_bus.iter().copied(),
            ),
        };

        if *aux_final != expected_final {
            return Err(VerificationError::InvalidBusBoundaryValues);
        }
    }

    verify_constraints::<SC, _, PcsError<SC>>(
        air,
        &opened_values.trace_local,
        &opened_values.trace_next,
        opened_values.preprocessed_local.as_deref(),
        opened_values.preprocessed_next.as_deref(),
        opened_values.aux_trace_local.as_deref(),
        opened_values.aux_trace_next.as_deref(),
        &randomness,
        aux_finals,
        public_values,
        init_trace_domain,
        zeta,
        alpha,
        quotient,
    )?;

    Ok(())
}

/// Computes the final value for a multiset bus given variable-length public inputs.
pub fn bus_multiset_boundary_varlen<
    'a,
    I: IntoIterator<Item = &'a [Val<SC>]>,
    SC: StarkGenericConfig,
>(
    randomness: &[SC::Challenge],
    public_inputs: I,
) -> SC::Challenge {
    let mut bus_p_last = SC::Challenge::ONE;
    let rand = randomness;
    for row in public_inputs {
        let mut p_last = rand[0];
        for (c, p_i) in row.iter().enumerate() {
            p_last += SC::Challenge::from(*p_i) * rand[c + 1];
        }
        bus_p_last *= p_last;
    }
    bus_p_last
}

/// Computes the final value for a logup bus boundary constraint given variable-length public inputs.
pub fn bus_logup_boundary_varlen<
    'a,
    I: IntoIterator<Item = &'a [Val<SC>]>,
    SC: StarkGenericConfig,
>(
    randomness: &[SC::Challenge],
    public_inputs: I,
) -> SC::Challenge {
    let mut bus_q_last = SC::Challenge::ZERO;
    let rand = randomness;
    for row in public_inputs {
        let mut q_last = rand[0];
        for (c, p_i) in row.iter().enumerate() {
            let p_i = *p_i;
            q_last += SC::Challenge::from(p_i) * rand[c + 1];
        }
        bus_q_last += q_last.inverse();
    }
    bus_q_last
}

#[derive(Debug)]
pub enum VerificationError<PcsErr> {
    InvalidProofShape,
    /// An error occurred while verifying the claimed openings.
    InvalidOpeningArgument(PcsErr),
    /// Out-of-domain evaluation mismatch, i.e. `constraints(zeta)` did not match
    /// `quotient(zeta) Z_H(zeta)`.
    OodEvaluationMismatch {
        index: Option<usize>,
    },
    /// The FRI batch randomization does not correspond to the ZK setting.
    RandomizationError,
    /// The domain does not support computing the next point algebraically.
    NextPointUnavailable,
    /// The expected bus boundary final values do not match the opened values.
    InvalidBusBoundaryValues,
}