zinc_piop/

combined_poly_resolver.rs

//! Combined polynomial resolver subprotocol.

mod folder;
mod structs;

pub use structs::*;

use crate::{
    CombFn,
    combined_poly_resolver::{
        folder::ConstraintFolder,
        structs::{Proof as CprProof, ProverState as CprProverState},
    },
    ideal_check,
    projections::ProjectedScalars,
    sumcheck::{
        SumCheckError, multi_degree::MultiDegreeSumcheckGroup,
        prover::ProverState as SumcheckProverState,
    },
};
use crypto_primitives::{
    BaseFieldConfig, ProjectPrimitiveIntegersWithConfig, SetConfig, SetElement,
};
use itertools::Itertools;
#[cfg(feature = "parallel")]
use rayon::prelude::*;
use std::{marker::PhantomData, slice};
use thiserror::Error;
use zinc_poly::{
    EvaluationError,
    mle::DenseMultilinearExtension,
    univariate::dynamic::{DynamicPolynomial, HasDynamicPolynomialConfig},
    utils::{ArithErrors, build_eq_x_r, eq_eval},
};
use zinc_transcript::traits::{ConstTranscribable, Transcript};
use zinc_uair::{TraceRow, Uair, ideal::ImpossibleIdeal};
use zinc_utils::{
    UNCHECKED, add, cfg_iter,
    from_ref::FromRef,
    inner_product::{InnerProduct, NativeInnerProduct},
    powers,
};

/// Combined polynomial resolver.
///
/// Builds one sumcheck group over either the Q[X]-projected trace and the
/// $Q[X]$ family of constraints, or the $F_{q_i}[X]$ family
/// (for a specific $q_i$ prime). The `MultiDegreeSumcheck` at the protocol
/// layer batches them all together.
///
/// Food for thought:
/// The n+1 CPR groups could in principle be collapsed into a single
/// `MultiDegreeSumcheck` group sharing one evaluation-point trajectory, with
/// per-family soundness factor $q_i / q^*$. The shared-integer challenges
/// (folding $\alpha$, projecting $\psi$) are already in place; what remains is
/// fusing the per-family combination functions and writing down the
/// corresponding soundness lemma. The trade-off is loss of per-family
/// arithmetic locality — each $F_{q_i}[X]$ family currently does its
/// sumcheck work in its native (cheap) field, which a merged comb_fn would have
/// to give up or dispatch internally. Worth revisiting only if proof size /
/// verifier simplicity outweighs prover cost.
pub struct CombinedPolyResolver<C: SetConfig>(PhantomData<C>);

impl<C> CombinedPolyResolver<C>
where
    C: BaseFieldConfig + ProjectPrimitiveIntegersWithConfig + 'static,
    C::Integer: ConstTranscribable,
{
    /// Build the CPR sumcheck group for use in the multi-degree sumcheck.
    ///
    /// Pre-sumcheck half of the CPR prover. Samples the folding challenge `α`,
    /// builds the MLE vector and combination function with the constraint
    /// polynomial identity:
    ///
    /// $$
    /// \sum_{b \in H} (f_0(b, x_0[b],...,x_n[b], x_0ˆdown[b],...,x_nˆdown[b])
    ///                 + \alpha f_1(...) + ... + \alpha^k f_k(...)) = v_0 +
    ///                   \alpha * v_1 + ... + \alphaˆk * v_k,
    /// $$
    /// where $f_i(b, x_0[b],...,x_n[b], x_0ˆdown[b],...,x_nˆdown[b])
    ///         = eq(r, b) * (1 - eq(r, 1,...1))
    ///             * g_i(x_0[b],...,x_n[b], x_0ˆdown[b],...,x_nˆdown[b])$
    /// and `g_i` is a constraint polynomial given by the UAIR `U`.
    /// `v_0,...,v_k` are the claimed evaluations of the combined polynomials.
    ///
    /// # Parameters
    /// - `trace_matrix`: The trace that have been projected to F.
    /// - `bit_op_down_mles`: MLEs of the bit-op virtual columns, projected to
    ///   `F::Inner`, in `UairSignature::bit_op_specs()` order. The caller is
    ///   responsible for applying the bit-op (ROTR / SHR) entry-wise on the
    ///   *unprojected* binary_poly source column *before* projection — see
    ///   Lemma 2.3 of the Zinc+ paper. The length must equal the signature's
    ///   `bit_op_specs().len()`.
    /// - `evaluation_point`: The evaluation point for the claims.
    /// - `projected_scalars`: The UAIR scalars projected to `F`.
    /// - `family_idx`: which constraint family to fold. `0` -> $Q[X]$; `i >= 1`
    ///   -> $F_{q_{i-1}}[X]$.
    /// - `num_constraints`: The number of constraint polynomials in the UAIR
    ///   `U`.
    /// - `num_vars`: The number of variables of the trace MLEs.
    /// - `max_degree`: The degree of the UAIR `U`.
    /// - `folding_challenge`: pre-sampled batching challenge $\alpha$ used to
    ///   batch the constraint claims.
    /// - `field_cfg`: The random field config.
    #[allow(clippy::arithmetic_side_effects, clippy::too_many_arguments)]
    pub fn prepare_sumcheck_group<U>(
        trace_matrix: Vec<DenseMultilinearExtension<C::Element>>,
        bit_op_down_mles: Vec<DenseMultilinearExtension<C::Element>>,
        evaluation_point: &[C::Element],
        projected_scalars: &ProjectedScalars<U::Scalar, C::Element>,
        family_idx: usize,
        num_constraints: usize,
        num_vars: usize,
        max_degree: usize,
        folding_challenge: &C::Element,
        field_cfg: &C,
    ) -> Result<
        (MultiDegreeSumcheckGroup<C>, CprProverAncillary),
        CombinedPolyResolverError<C::Element>,
    >
    where
        U::Scalar: 'static,
        U: Uair,
    {
        debug_assert_ne!(
            num_vars, 1,
            "The protocol is not needed when the number of variables is 1 :)"
        );

        let zero = field_cfg.zero();
        let one = field_cfg.one();

        // Shifted trace: for each ShiftSpec, take the source column,
        // drop the first `shift_amount` rows, and zero-pad to the full
        // domain size so the MLE keeps the correct `num_vars`.
        // TODO consider working with pointers since down cols are virtual cols until
        // folded in sumcheck - virtual MLE trait needed in sumcheck.
        let uair_sig = U::signature();

        assert_eq!(
            bit_op_down_mles.len(),
            uair_sig.bit_op_specs().len(),
            "bit_op_down_mles count must match UairSignature::bit_op_specs().len()",
        );

        let n = 1usize << num_vars;
        let shift_mles: Vec<DenseMultilinearExtension<C::Element>> = cfg_iter!(uair_sig.shifts())
            .map(|spec| {
                let mut evals = trace_matrix[spec.source_col()][spec.shift_amount()..].to_vec();
                evals.resize(n, zero.clone());
                DenseMultilinearExtension {
                    evaluations: evals,
                    num_vars,
                }
            })
            .collect();

        // Down-row layout (see UairSignature::with_bit_op_specs):
        //
        //     [shifted_binary..., bit_op_binary...,
        //      shifted_arbitrary..., shifted_int...]
        //
        // Shifts are sorted by source_col, so binary-source shifts come first.
        // We splice the bit-op MLEs in between the binary and non-binary
        // shift groups so that the resulting `down` vector is consistent with
        // the down ColumnLayout (binary_poly_cols + arbitrary_poly_cols +
        // int_cols).
        let bit_op_down_offset = uair_sig.bit_op_down_offset();
        let num_shift_down = shift_mles.len();
        let num_bit_op_specs = bit_op_down_mles.len();
        let mut down: Vec<DenseMultilinearExtension<C::Element>> =
            Vec::with_capacity(num_shift_down + num_bit_op_specs);
        let mut shift_iter = shift_mles.into_iter();
        for _ in 0..bit_op_down_offset {
            down.push(shift_iter.next().expect("offset within shift_mles range"));
        }
        down.extend(bit_op_down_mles);
        down.extend(shift_iter);

        let eq_r = build_eq_x_r(field_cfg, evaluation_point)?;
        // To get the constraints on the last row ignored
        // we multiply each constraint polynomial
        // by the selector (1 - eq(1,...,1, x))
        let last_row_selector = DenseMultilinearExtension {
            num_vars,
            evaluations: {
                let mut evals = vec![zero; 1 << num_vars];
                evals[(1 << num_vars) - 1] = one.clone();
                evals
            },
        };

        // The batching challenge $\alpha$ is supplied by the caller. The
        // protocol layer samples one shared integer in $[0, q^*)$ once
        // and lifts it into each family's field, so the Q[X] CPR and
        // per-prime CPRs reuse the same underlying integer.
        let folding_challenge_powers: Vec<C::Element> =
            powers(field_cfg, folding_challenge, num_constraints);

        let num_cols = trace_matrix.len();
        let num_down_cols = down.len();
        let mles: Vec<DenseMultilinearExtension<C::Element>> = {
            let mut mles = Vec::with_capacity(2 + num_cols + num_down_cols);

            mles.push(last_row_selector);
            mles.push(eq_r);

            mles.extend(trace_matrix);
            mles.extend(down);

            mles
        };

        let projected_scalars = projected_scalars.clone();
        let comb_cfg = field_cfg.clone();
        let comb_fn: CombFn<C::Element> = Box::new(move |mle_values: &[C::Element]| {
            let uair_sig = U::signature();
            let up_layout = uair_sig.total_cols().as_column_layout();
            let down_layout = uair_sig.down_cols().as_column_layout();

            let selector = &mle_values[0];
            let eq_r = &mle_values[1];

            let mut folder =
                ConstraintFolder::new(family_idx, &folding_challenge_powers, &comb_cfg);

            let project = |scalar: &U::Scalar| {
                projected_scalars
                    .get(scalar)
                    .expect("all scalars should have been projected at this point")
            };

            U::constrain_general(
                &mut folder,
                &comb_cfg,
                TraceRow::from_slice_with_layout(&mle_values[2..num_cols + 2], up_layout),
                TraceRow::from_slice_with_layout(&mle_values[num_cols + 2..], down_layout),
                project,
                |x, y| Some(comb_cfg.mul(x, &project(y))),
                ImpossibleIdeal::from_ref,
                ImpossibleIdeal::from_ref,
            );

            comb_cfg.mul(
                &comb_cfg.mul(&folder.folded_constraints, &comb_cfg.sub(&one, selector)),
                eq_r,
            )
        });

        Ok((
            MultiDegreeSumcheckGroup::new(max_degree + 2, mles, comb_fn),
            CprProverAncillary {
                num_cols,
                num_down_cols: num_shift_down,
                num_vars,
            },
        ))
    }

    /// Finalize the CPR proof after the multi-degree sumcheck completes.
    ///
    /// # Parameters
    /// - `transcript`: FS-transcript (absorbs `up_evals` and `down_evals`).
    /// - `sumcheck_prover_state`: The CPR group's `ProverState` from
    ///   `MultiDegreeSumcheck::prove_as_subprotocol` (states\[0\]).
    /// - `ancillary`: Produced by [`prepare_sumcheck_group`]; carries column
    ///   counts and `num_vars` needed to split the flat eval vector.
    /// - `field_cfg`: Field configuration.
    #[allow(clippy::arithmetic_side_effects, clippy::type_complexity)]
    pub fn finalize_prover<U>(
        transcript: &mut impl Transcript,
        sumcheck_prover_state: SumcheckProverState<C>,
        ancillary: CprProverAncillary,
        field_cfg: &C,
    ) -> Result<
        (CprProof<C::Element>, CprProverState<C::Element>),
        CombinedPolyResolverError<C::Element>,
    >
    where
        U: Uair,
    {
        // Sumcheck prover stops evaluating MLEs
        // at the second-to-last challenge
        // leaving all MLEs in num_vars=1
        // state. We need to evaluate them up
        // and send to the verifier.
        debug_assert!(
            sumcheck_prover_state
                .mles
                .iter()
                .all(|mle| mle.num_vars == 1)
        );

        let last_sumcheck_challenge = sumcheck_prover_state
            .randomness
            .last()
            .expect("sumcheck could not have had 0 rounds");

        let mut mles = sumcheck_prover_state.mles;
        let evals: Vec<C::Element> = mles
            .drain(2..)
            .map(|mle| mle.evaluate(field_cfg, slice::from_ref(last_sumcheck_challenge)))
            .try_collect()?;

        let uair_sig = U::signature();
        let num_bit_op_specs = uair_sig.bit_op_specs().len();
        let bit_op_down_offset = uair_sig.bit_op_down_offset();

        debug_assert_eq!(
            evals.len(),
            ancillary.num_cols + ancillary.num_down_cols + num_bit_op_specs,
        );

        // The post-sumcheck evals are laid out as
        //   [up_evals..., shifted_binary..., bit_op_evals..., shifted_non_binary...]
        // matching the order in which `prepare_sumcheck_group` packed them
        // into the MLE vector. We split them into three on-wire vectors:
        // `up_evals`, `down_evals` (shifts only, in their UAIR-signature
        // order), and `bit_op_evals` (in `bit_op_specs()` order).
        let up_end = ancillary.num_cols;
        let bit_op_start = up_end + bit_op_down_offset;
        let bit_op_end = bit_op_start + num_bit_op_specs;

        let up_evals = evals[..up_end].to_vec();
        let bit_op_evals = evals[bit_op_start..bit_op_end].to_vec();
        let mut down_evals = Vec::with_capacity(ancillary.num_down_cols);
        down_evals.extend_from_slice(&evals[up_end..bit_op_start]);
        down_evals.extend_from_slice(&evals[bit_op_end..]);

        let mut transcription_buf: Vec<u8> = vec![0; <C::Integer as ConstTranscribable>::NUM_BYTES];
        transcript.absorb_field_element_slice(field_cfg, &up_evals, &mut transcription_buf);
        transcript.absorb_field_element_slice(field_cfg, &down_evals, &mut transcription_buf);
        transcript.absorb_field_element_slice(field_cfg, &bit_op_evals, &mut transcription_buf);
        Ok((
            CprProof {
                up_evals,
                down_evals,
                bit_op_evals,
            },
            CprProverState {
                evaluation_point: sumcheck_prover_state.randomness,
            },
        ))
    }

    /// Pre-sumcheck half of the CPR verifier.
    ///
    /// Must run before [`MultiDegreeSumcheck::verify_as_subprotocol`].
    /// Pure check: takes the pre-sampled $\psi$ projecting element and the
    /// pre-sampled batching challenge $\alpha$; does not touch the transcript.
    ///
    /// # Parameters
    /// - `proof`: The CPR proof (`up_evals`, `down_evals`).
    /// - `claimed_sum`: The claimed sum from
    ///   `combined_sumcheck.claimed_sums()[0]`.
    /// - `ic_check_subclaim`: Subclaim from the ideal check; provides the
    ///   evaluation point and claimed values used to verify the sumcheck sum.
    /// - `num_constraints`: Number of constraint polynomials in `U`.
    /// - `num_vars`: Number of variables of the trace MLEs.
    /// - `projecting_element`: The random challenge used to project `F[X] → F`.
    /// - `folding_challenge`: pre-sampled batching challenge $\alpha$ used to
    ///   batch the constraint claims.
    /// - `field_cfg`: Field configuration.
    #[allow(clippy::too_many_arguments)]
    pub fn prepare_verifier<U>(
        proof: &CprProof<C::Element>,
        claimed_sum: C::Element,
        ic_check_subclaim: &ideal_check::VerifierSubclaim<C::Element>,
        num_constraints: usize,
        num_vars: usize,
        projecting_element: &C::Element,
        folding_challenge: &C::Element,
        field_cfg: &C,
    ) -> Result<CprVerifierAncillary<C::Element>, CombinedPolyResolverError<C::Element>>
    where
        U: Uair,
    {
        let uair_sig = U::signature();
        proof.validate_evaluation_sizes(
            uair_sig.total_cols().cols(),
            uair_sig.shifts().len(),
            uair_sig.bit_op_specs().len(),
        )?;

        let zero = field_cfg.zero();
        let poly_cfg = field_cfg.dyn_poly_cfg();

        // Precompute powers of the projecting element for batch evaluation.
        let projection_powers: Vec<C::Element> = {
            let max_coeffs_len = ic_check_subclaim
                .values
                .iter()
                .map(|poly| poly_cfg.degree(poly).map_or(0, |d| add!(d, 1)))
                .max()
                .unwrap_or(0)
                .max(1);
            powers(field_cfg, projecting_element, max_coeffs_len)
        };

        // The batching challenge $\alpha$ is supplied by the caller; see
        // [`prepare_sumcheck_group`] for the shared-challenge rationale.
        let folding_challenge_powers: Vec<C::Element> =
            powers(field_cfg, folding_challenge, num_constraints);

        // TODO(Alex): investigate if parallelising this is beneficial.
        // Compute v_0 + \alpha * v_1 + ... + \alpha ^ k * v_k.
        let expected_sum = ic_check_subclaim
            .values
            .iter()
            .zip(&folding_challenge_powers)
            .map(|(claimed_value, random_coeff)| {
                let deg = poly_cfg.degree(claimed_value).map_or(0, |d| add!(d, 1));
                let value = NativeInnerProduct::inner_product::<UNCHECKED>(
                    field_cfg,
                    &claimed_value.coeffs[..deg],
                    &projection_powers[..deg],
                    zero.clone(),
                )
                .expect("inner product cannot fail here");
                field_cfg.mul(&value, random_coeff)
            })
            .fold(zero.clone(), |acc, term| field_cfg.add(&acc, &term));

        if claimed_sum != expected_sum {
            return Err(CombinedPolyResolverError::WrongSumcheckSum {
                got: claimed_sum,
                expected: expected_sum,
            });
        }

        Ok(CprVerifierAncillary {
            folding_challenge_powers,
            ic_evaluation_point: ic_check_subclaim.evaluation_point.clone(),
            num_vars,
        })
    }

    /// Post-sumcheck half of the CPR verifier.
    ///
    /// Runs after [`MultiDegreeSumcheck::verify_as_subprotocol`] produces the
    /// shared evaluation point.
    ///
    /// # Parameters
    /// - `transcript`: FS-transcript (absorbs `up_evals` and `down_evals`).
    /// - `proof`: The CPR proof (consumed to produce the subclaim).
    /// - `shared_point`: The shared evaluation point `r*` from the multi-degree
    ///   sumcheck.
    /// - `expected_evaluation`: `md_subclaims.expected_evaluations()[0]` — the
    ///   expected value of the CPR combination function at `r*`.
    /// - `ancillary`: Produced by [`prepare_verifier`]; carries folding
    ///   challenge powers, ideal-check evaluation point, and `num_vars`.
    /// - `projected_scalars`: UAIR scalars projected to `F`.
    /// - `family_idx`: which constraint family. `0` -> $Q[X]$; `i >= 1` ->
    ///   $F_{q_{i-1}}[X]$.
    /// - `field_cfg`: Field configuration.
    #[allow(clippy::too_many_arguments)]
    pub fn finalize_verifier<U>(
        transcript: &mut impl Transcript,
        proof: CprProof<C::Element>,
        shared_point: Vec<C::Element>,
        expected_evaluation: C::Element,
        ancillary: CprVerifierAncillary<C::Element>,
        projected_scalars: &ProjectedScalars<U::Scalar, C::Element>,
        family_idx: usize,
        field_cfg: &C,
    ) -> Result<VerifierSubclaim<C::Element>, CombinedPolyResolverError<C::Element>>
    where
        U: Uair,
    {
        let uair_sig = U::signature();
        let down_layout = uair_sig.down_cols().as_column_layout();
        let one = field_cfg.one();

        let eq_r_value = eq_eval(field_cfg, &shared_point, &ancillary.ic_evaluation_point)?;
        let selector_value = eq_eval(
            field_cfg,
            &shared_point,
            &vec![one.clone(); ancillary.num_vars],
        )?;

        let mut folder =
            ConstraintFolder::new(family_idx, &ancillary.folding_challenge_powers, field_cfg);

        let project = |scalar: &U::Scalar| {
            projected_scalars
                .get(scalar)
                .expect("all scalars should have been projected at this point")
        };

        // Build the full down-row eval vec by splicing bit-op evals into the
        // binary_poly slot, matching the down ColumnLayout enforced by
        // `UairSignature::with_bit_op_specs`:
        //   [shifted_binary..., bit_op_evals..., shifted_arbitrary..., shifted_int...]
        let bit_op_down_offset = uair_sig.bit_op_down_offset();
        let mut full_down_evals =
            Vec::with_capacity(add!(proof.down_evals.len(), proof.bit_op_evals.len()));
        full_down_evals.extend_from_slice(&proof.down_evals[..bit_op_down_offset]);
        full_down_evals.extend_from_slice(&proof.bit_op_evals);
        full_down_evals.extend_from_slice(&proof.down_evals[bit_op_down_offset..]);

        U::constrain_general(
            &mut folder,
            field_cfg,
            TraceRow::from_slice_with_layout(
                &proof.up_evals,
                uair_sig.total_cols().as_column_layout(),
            ),
            TraceRow::from_slice_with_layout(&full_down_evals, down_layout),
            project,
            |x, y| Some(field_cfg.mul(x, &project(y))),
            ImpossibleIdeal::from_ref,
            ImpossibleIdeal::from_ref,
        );

        // eq * (1 - selector) * folded
        let expected_claim_value = field_cfg.mul(
            &field_cfg.mul(&eq_r_value, &field_cfg.sub(&one, &selector_value)),
            &folder.folded_constraints,
        );

        if expected_claim_value != expected_evaluation {
            return Err(CombinedPolyResolverError::ClaimValueDoesNotMatch {
                got: expected_evaluation,
                expected: expected_claim_value,
            });
        }

        let mut transcription_buf: Vec<u8> = vec![0; <C::Integer as ConstTranscribable>::NUM_BYTES];
        transcript.absorb_field_element_slice(field_cfg, &proof.up_evals, &mut transcription_buf);
        transcript.absorb_field_element_slice(field_cfg, &proof.down_evals, &mut transcription_buf);
        transcript.absorb_field_element_slice(
            field_cfg,
            &proof.bit_op_evals,
            &mut transcription_buf,
        );

        Ok(VerifierSubclaim {
            up_evals: proof.up_evals,
            down_evals: proof.down_evals,
            bit_op_evals: proof.bit_op_evals,
            evaluation_point: shared_point,
        })
    }
}

#[derive(Debug, Error)]
pub enum CombinedPolyResolverError<F: SetElement> {
    #[error("failed to build eq_r: {0}")]
    EqrError(ArithErrors),
    #[error("error evaluating MLE: {0}")]
    MleEvaluationError(EvaluationError),
    #[error("error projecting polynomial {0:?} by point {1:?}: {2}")]
    ProjectionError(DynamicPolynomial<F>, F, EvaluationError),
    #[error("wrong trace columns evaluations number: got {got}, expected {expected}")]
    WrongUpEvalsNumber { got: usize, expected: usize },
    #[error("wrong shifted trace columns evaluations number: got {got}, expected {expected}")]
    WrongDownEvalsNumber { got: usize, expected: usize },
    #[error("wrong bit-op virtual columns evaluations number: got {got}, expected {expected}")]
    WrongBitOpEvalsNumber { got: usize, expected: usize },
    #[error("sumcheck verification failed: {0}")]
    SumcheckError(SumCheckError<F>),
    #[error("wrong sumcheck claimed sum: received {got:?}, expected {expected:?}")]
    WrongSumcheckSum { got: F, expected: F },
    #[error("resulting claim value does not match: received {got:?}, expected {expected:?}")]
    ClaimValueDoesNotMatch { got: F, expected: F },
}

impl<F: SetElement> From<EvaluationError> for CombinedPolyResolverError<F> {
    fn from(eval_error: EvaluationError) -> Self {
        Self::MleEvaluationError(eval_error)
    }
}

impl<F: SetElement> From<ArithErrors> for CombinedPolyResolverError<F> {
    fn from(arith_error: ArithErrors) -> Self {
        Self::EqrError(arith_error)
    }
}

impl<F: SetElement> From<SumCheckError<F>> for CombinedPolyResolverError<F> {
    fn from(sumcheck_error: SumCheckError<F>) -> Self {
        Self::SumcheckError(sumcheck_error)
    }
}

#[cfg(test)]
#[allow(
    clippy::arithmetic_side_effects,
    clippy::cast_possible_truncation,
    clippy::cast_precision_loss,
    clippy::cast_sign_loss,
    clippy::clone_on_copy,
    clippy::redundant_clone
)]
mod tests {
    use super::*;
    use crate::{
        ideal_check::IdealCheckProtocol,
        projections::{ProjectedTrace, evaluate_trace_to_column_mles, project_scalars_to_field},
        sumcheck::multi_degree::MultiDegreeSumcheck,
        test_utils::{LIMBS, run_ideal_check_prover_combined, test_config},
    };
    use crypto_primitives::{
        crypto_bigint_int::Int,
        crypto_bigint_monty::{MontyField, MontyFieldElement},
        crypto_bigint_uint::Uint,
    };
    use rand::rng;
    use zinc_poly::univariate::{dense::DensePolynomial, dynamic::DynamicPolynomialConfig};
    use zinc_test_uair::{
        GenerateRandomTrace, TestUairNoMultiplication, TestUairSimpleMultiplication,
    };
    use zinc_transcript::Blake3Transcript;
    use zinc_uair::{
        constraint_counter::count_constraints,
        degree_counter::count_max_degree,
        ideal::{DegreeOneIdeal, Ideal, IdealCheck},
        ideal_collector::IdealOrZero,
    };

    // TODO(Ilia): These tests are absolute joke.
    //             Once we have time we need to create a comprehensive test suite
    //             akin to the one we have for the PCS or the sumcheck.

    fn test_successful_verification_generic<
        U,
        IdealOverF,
        IdealOverFFromRef,
        const DEGREE_PLUS_ONE: usize,
    >(
        num_vars: usize,
        prime_idx: Option<usize>,
        ideal_over_f_from_ref: IdealOverFFromRef,
    ) where
        U: Uair<Scalar = DensePolynomial<Int<5>, DEGREE_PLUS_ONE>>
            + GenerateRandomTrace<DEGREE_PLUS_ONE, PolyCoeff = Int<5>, Int = Int<5>>,
        IdealOverF: Ideal + for<'a> IdealCheck<DynamicPolynomialConfig<'a, MontyField<LIMBS>>>,
        IdealOverFFromRef: Fn(&IdealOrZero<U::Ideal>) -> IdealOverF,
    {
        let mut rng = rng();
        let field_cfg = test_config();

        let mut prover_transcript = Blake3Transcript::new();
        let mut verifier_transcript = prover_transcript.clone();

        let trace = U::generate_random_trace(num_vars, &mut rng);

        let (ic_proof, evaluation_point, projected_scalars, projected_trace) =
            run_ideal_check_prover_combined::<U, DEGREE_PLUS_ONE>(
                num_vars,
                &trace,
                prime_idx,
                &mut prover_transcript,
            );

        let num_constraints = count_constraints::<U>();

        // Mirror the prover's evaluation-point sampling:
        // the test helper (`run_ideal_check_prover_combined`) squeezed `num_vars` field
        // challenges from the prover transcript before invoking `prove_combined`, so do
        // the same here for the verifier.
        let ic_evaluation_point: Vec<MontyFieldElement<LIMBS>> =
            verifier_transcript.get_field_challenges(num_vars, &field_cfg);

        let ic_check_subclaim = IdealCheckProtocol::<U>::verify_as_subprotocol(
            &mut verifier_transcript,
            ic_proof,
            /* family_idx = */ 0,
            num_constraints.q,
            &ic_evaluation_point,
            ideal_over_f_from_ref,
            |_| panic!("F_q[X] not supported here!"),
            &field_cfg,
        )
        .expect("Verification failed");

        let max_degree = count_max_degree::<U>();

        let projecting_element: MontyFieldElement<4> =
            prover_transcript.get_field_challenge(&field_cfg);

        let projected_scalars =
            project_scalars_to_field(&field_cfg, projected_scalars, &projecting_element).unwrap();

        // Prover: prepare -> MultiDegreeSumcheck → finalize.
        // $\alpha$ is sampled by the caller per the shared-challenge API.
        let folding_challenge: MontyFieldElement<LIMBS> =
            prover_transcript.get_field_challenge(&field_cfg);
        let (cpr_group, cpr_ancillary) = CombinedPolyResolver::prepare_sumcheck_group::<U>(
            evaluate_trace_to_column_mles(
                &field_cfg,
                &ProjectedTrace::RowMajor(projected_trace),
                &projecting_element,
            ),
            Vec::new(),
            &evaluation_point,
            &projected_scalars,
            /* family_idx = */ 0,
            num_constraints.q,
            num_vars,
            max_degree,
            &folding_challenge,
            &field_cfg,
        )
        .expect("CPR prepare failed");

        let mut sumcheck_outputs = MultiDegreeSumcheck::prove_as_subprotocol(
            &mut prover_transcript,
            vec![(vec![cpr_group], &field_cfg)],
            num_vars,
            &field_cfg,
        );
        let (md_proof, states) = sumcheck_outputs.pop().expect("single family");

        let (proof, _) = CombinedPolyResolver::finalize_prover::<U>(
            &mut prover_transcript,
            states.into_iter().next().unwrap(),
            cpr_ancillary,
            &field_cfg,
        )
        .expect("CPR finalize failed");

        let projecting_element: MontyFieldElement<LIMBS> =
            verifier_transcript.get_field_challenge(&field_cfg);
        let folding_challenge: MontyFieldElement<LIMBS> =
            verifier_transcript.get_field_challenge(&field_cfg);

        // Verifier: prepare → MultiDegreeSumcheck → finalize
        let cpr_verifier_ancillary = CombinedPolyResolver::prepare_verifier::<U>(
            &proof,
            md_proof.claimed_sums()[0].clone(),
            &ic_check_subclaim,
            num_constraints.q,
            num_vars,
            &projecting_element,
            &folding_challenge,
            &field_cfg,
        )
        .expect("CPR prepare_verifier failed");

        let md_subclaims = MultiDegreeSumcheck::verify_as_subprotocol(
            &mut verifier_transcript,
            num_vars,
            &[(&md_proof, &field_cfg)],
            &field_cfg,
        )
        .expect("MultiDegreeSumcheck verify failed")
        .pop()
        .expect("single family");

        assert!(
            CombinedPolyResolver::finalize_verifier::<U>(
                &mut verifier_transcript,
                proof,
                md_subclaims.point().to_vec(),
                md_subclaims.expected_evaluations()[0].clone(),
                cpr_verifier_ancillary,
                &projected_scalars,
                /* family_idx = */ 0,
                &field_cfg,
            )
            .is_ok()
        );
    }

    #[test]
    fn test_successful_verification() {
        let field_cfg = test_config();

        let num_vars = 2;

        test_successful_verification_generic::<
            TestUairNoMultiplication<Int<5>, Uint<LIMBS>>,
            _,
            _,
            32,
        >(num_vars, None, |ideal_over_ring| {
            ideal_over_ring.map(|i| DegreeOneIdeal::project(&field_cfg, i))
        });
        test_successful_verification_generic::<
            TestUairSimpleMultiplication<Int<5>, Uint<LIMBS>>,
            _,
            _,
            32,
        >(num_vars, None, |_ideal_over_ring| {
            IdealOrZero::<DegreeOneIdeal<_>>::zero()
        });
    }
}