p3-sumcheck 0.8.0

//! Honest-verifier zero-knowledge sumcheck prover, generic over the binding mode.

use alloc::vec::Vec;
use core::marker::PhantomData;

use p3_challenger::fs::TranscriptField;
use p3_challenger::{CanObserve, FieldChallenger, GrindingChallenger};
use p3_commit::Mmcs;
use p3_field::{ExtensionField, Field, TwoAdicField};
use p3_matrix::Matrix;
use p3_multilinear_util::point::Point;
use p3_zk_codes::{ZkEncoding, ZkEncodingWithRandomness};
use rand::CryptoRng;

use super::common::{mask_endpoints, sample_masks};
use super::layout::ZkLayout;
use super::round::{MaskedRounds, PlainPiece};
use crate::extrapolate_01inf;
use crate::layout::{PrefixProver, SuffixProver};
use crate::strategy::SumcheckProver;
use crate::table::{OpeningEvals, OpeningRequest};
use crate::zk::data::{ZkSumcheckData, ZkSumcheckHandoff};
use crate::zk::transcript::{ZkProverTranscript, ZkSumcheckShape};

/// Honest-verifier zero-knowledge sumcheck prover.
///
/// Plain accumulator math is unchanged.
/// The overlay adds mask sampling, mask commits, and the round formula below.
///
/// # Per-round polynomial
///
/// At round `j` with past challenges `gamma_1, ..., gamma_{j-1}`:
///
/// ```text
///     h_j(X) = 2^{k-j}   * s_j(X)                        (live mask)
///            + 2^{k-j}   * sum_{l<j} s_l(gamma_l)        (past masks)
///            + 2^{k-j-1} * sum_{l>j} ( s_l(0) + s_l(1) ) (future ends)
///            + eps       * plain_j(X)                    (plain piece)
/// ```
///
/// Combined degree is `max(ell_zk - 1, 2)`.
/// The mask piece is degree `ell_zk - 1`, the plain piece is degree `2`.
///
/// # `mu_tilde` closed form
///
/// ```text
///     mu_tilde = sum_{b in {0,1}^k} ( s_1(b_1) + ... + s_k(b_k) )
///              = 2^{k-1} * sum_l ( s_l(0) + s_l(1) )
/// ```
///
/// # Residual handoff
///
/// After `k` rounds the residual claim is scaled by `eps`.
/// The factor is folded into the residual polynomial during compression, so downstream consumers see the scaling automatically.
///
/// # Binding mode
///
/// The generic layout type selects the binding direction.
/// Two ready-made aliases cover the two supported modes; see [`ZkPrefixProver`] and [`ZkSuffixProver`].
pub struct ZkProver<F, EF, Enc, M, L>
where
    F: Field,
    EF: ExtensionField<F>,
    Enc: ZkEncoding<EF>,
    M: Mmcs<EF>,
{
    /// Plain stacked-layout prover.
    inner: L,

    /// Zero-knowledge code used to encode the mask polynomials.
    encoding: Enc,

    /// Merkle commitment scheme used to commit each encoded mask.
    mmcs: M,

    /// Marker tying `F`/`EF` to the storage type without inflating the runtime layout.
    _marker: PhantomData<(F, EF)>,
}

/// HVZK prover for the prefix-binding sumcheck.
pub type ZkPrefixProver<F, EF, Enc, M> = ZkProver<F, EF, Enc, M, PrefixProver<F, EF>>;

/// HVZK prover for the suffix-binding sumcheck.
pub type ZkSuffixProver<F, EF, Enc, M> = ZkProver<F, EF, Enc, M, SuffixProver<F, EF>>;

impl<F, EF, Enc, M, L> ZkProver<F, EF, Enc, M, L>
where
    F: TwoAdicField,
    EF: ExtensionField<F>,
    Enc: ZkEncodingWithRandomness<EF>,
    M: Mmcs<EF>,
    L: ZkLayout<F, EF>,
{
    /// Wraps a plain layout with the HVZK ingredients.
    pub const fn new(inner: L, encoding: Enc, mmcs: M) -> Self {
        Self {
            inner,
            encoding,
            mmcs,
            _marker: PhantomData,
        }
    }

    /// Returns the folding factor of the wrapped inner prover.
    pub fn folding(&self) -> usize {
        self.inner.folding()
    }

    /// Returns the variable count of the stacked polynomial.
    pub fn num_variables(&self) -> usize {
        self.inner.num_variables()
    }

    /// Records opening claims at the current points and at their repeat-last successor points on the inner prover.
    ///
    /// # Returns
    ///
    /// The current-point evaluations first, then the successor-point evaluations.
    pub fn eval<Ch>(
        &mut self,
        table_idx: usize,
        batch: &OpeningRequest,
        challenger: &mut Ch,
    ) -> OpeningEvals<EF>
    where
        F: TranscriptField,
        Ch: FieldChallenger<F> + GrindingChallenger<Witness = F>,
    {
        // Delegate; the HVZK overlay carries no extra state at claim time.
        self.inner.eval(table_idx, batch, challenger)
    }

    /// Records a virtual opening claim on the inner prover.
    pub fn add_virtual_eval<Ch>(&mut self, challenger: &mut Ch) -> EF
    where
        F: TranscriptField,
        Ch: FieldChallenger<F> + GrindingChallenger<Witness = F>,
    {
        // Same delegation pattern as concrete openings.
        self.inner.add_virtual_eval(challenger)
    }

    /// Runs the HVZK sumcheck and returns the residual claim plus the mask oracles.
    ///
    /// # Phases
    ///
    /// 1. Plain-piece preamble (alpha, accumulators, plain_sum).
    /// 2. Sample, encode and commit the masks (Construction 6.3 step 1).
    /// 3. Compute `mu_tilde` from the mask endpoints (step 2).
    /// 4. Bind the oracle and `mu_tilde`, then draw the combining challenge `eps` (step 3).
    /// 5. Per-round sumcheck (step 4).
    /// 6. Residual handoff scaled by `eps` (mode-specific compression).
    ///
    /// # Returns
    ///
    /// - Residual sumcheck prover over the mode-specific product polynomial, scaled by `eps`.
    /// - Vector of per-round challenges `gamma_1, ..., gamma_k`.
    /// - Combining challenge `eps`, made explicit for code-switch composition.
    /// - Plain mask messages and one mask oracle per round, in round order.
    ///
    /// # Randomness
    ///
    /// Every mask comes from the supplied generator.
    ///
    /// The reveals hide the witness only as long as that stream is unpredictable.
    ///
    /// The trait bound rejects generators that are known not to be cryptographically secure.
    ///
    /// Callers still have to seed from real entropy.
    ///
    /// That bound is load-bearing, so the example below pins it.
    ///
    /// A generator built on a fast non-cryptographic algorithm must fail to compile.
    ///
    /// ```compile_fail
    /// use p3_challenger::fs::TranscriptField;
    /// use p3_challenger::{CanObserve, FieldChallenger, GrindingChallenger};
    /// use p3_commit::Mmcs;
    /// use p3_field::{ExtensionField, TwoAdicField};
    /// use p3_matrix::Matrix;
    /// use p3_sumcheck::zk::{ZkLayout, ZkProver};
    /// use p3_zk_codes::ZkEncodingWithRandomness;
    /// use rand::rngs::SmallRng;
    ///
    /// fn masks_from_xoshiro<F, EF, Enc, M, L, Ch>()
    /// where
    ///     F: TranscriptField + TwoAdicField,
    ///     EF: ExtensionField<F> + TwoAdicField,
    ///     Enc: ZkEncodingWithRandomness<EF>,
    ///     Enc::Codeword: Matrix<EF>,
    ///     M: Mmcs<EF>,
    ///     L: ZkLayout<F, EF>,
    ///     Ch: FieldChallenger<F> + GrindingChallenger<Witness = F> + CanObserve<M::Commitment>,
    /// {
    ///     // Xoshiro256++ is fast and well distributed, but not a CSPRNG.
    ///     // Its state is recoverable from a short output run, so it cannot mask a witness.
    ///     let _ = ZkProver::<F, EF, Enc, M, L>::into_sumcheck::<SmallRng, Ch>;
    /// }
    /// ```
    ///
    /// # Panics
    ///
    /// - The configuration cannot describe a masked batch.
    /// - Folding factor exceeds the polynomial's arity.
    #[tracing::instrument(skip_all)]
    pub fn into_sumcheck<R, Ch>(
        self,
        zk_data: &mut ZkSumcheckData<F, EF>,
        pow_bits: usize,
        challenger: &mut Ch,
        rng: &mut R,
    ) -> ZkSumcheckHandoff<F, EF, M>
    where
        F: TranscriptField,
        EF: TwoAdicField,
        Enc::Codeword: Matrix<EF>,
        R: CryptoRng,
        Ch: FieldChallenger<F> + GrindingChallenger<Witness = F> + CanObserve<M::Commitment>,
    {
        // Protocol shape resolved from the inner prover + mask encoding.
        let k = self.inner.folding();
        let ell_zk = self.encoding.message_len();
        let n_vars = self.inner.num_variables();

        // The batch derives its claim from the claims the layout recorded.
        let shape = ZkSumcheckShape::new_batching(k, ell_zk, pow_bits);

        // Lemma 6.4 hypotheses, plus the one bound the transcript never sees.
        shape
            .validate::<F>()
            .expect("a prover's own configuration must describe a masked batch");
        assert!(
            k <= n_vars,
            "folding_factor must be <= poly.num_variables()",
        );

        // Every step from here on is played through the description.
        let mut transcript = ZkProverTranscript::<Ch, F, EF>::new(challenger, shape);

        // Phase 1: plain-piece preamble (setup, precedes Construction 6.3).

        // `alpha` is the per-claim batching base: powers a^0, a^1, ... weight the claim accumulators below.
        let alpha: EF = transcript.batching_challenge();

        let accumulators = self.inner.claims().batched_accumulators(alpha);

        // Plain sumcheck claim `mu`, batched by the alphas.
        let mut plain_sum = self.inner.batched_sum(alpha);

        // Phase 2: sample, encode and commit the masks (Construction 6.3 step 1).
        //
        // The encoder draws zero-knowledge padding randomness from the same rng.
        let (masks, mask_randomness, mask_oracle) =
            sample_masks::<EF, _, _, _>(k, &self.encoding, &self.mmcs, rng);

        // Phase 3: mu_tilde via the closed form (Construction 6.3 step 2).
        //
        // The helper also returns the running future-mask endpoint budget for
        // the per-round loop.
        let (mu_tilde, sum_endpoints_init) = mask_endpoints::<EF>(&masks, k);
        zk_data.mu_tilde = mu_tilde;

        // Phase 4: bind the oracle and mu_tilde, then draw `eps` (Construction 6.3 step 3).
        //
        // The construction is instantiated over `EF`: the masks, `eps`, and the
        // round polynomials all live in `EF`, so Lemma 6.4 applies with `F := EF`
        // and the per-round polynomial is uniform over the full extension field.
        let eps: EF = transcript.masks(mask_oracle.0.clone(), mu_tilde);

        // Phase 5: per-round sumcheck (Construction 6.3 step 4).

        // Per-round challenges.
        let mut rs: Vec<EF> = Vec::with_capacity(k);

        // Mask side of every round: endpoints still ahead, past evaluations, wire, and proof.
        let mut rounds = MaskedRounds::new(&masks, ell_zk, sum_endpoints_init, eps);

        for _ in 0..k {
            // Plain `(c_0, c_inf)` of every recorded claim at the challenges so far.
            let (plain_c0, plain_c_inf) = accumulators.round_coefficients(&rs);

            // Assemble h_j, bind its wire, and record it; see the round module for the formula
            // and the in-place affine-consistency cross-check.
            //
            // The linear coefficient is not passed: it is dropped from the wire
            // and reconstructed by the verifier from the affine identity.
            let gamma_j = rounds.play(
                PlainPiece {
                    c0: plain_c0,
                    c_inf: plain_c_inf,
                },
                &mut transcript,
                zk_data,
            );

            // Advance the plain sumcheck claim via quadratic extrapolation through (0, 1, inf).
            plain_sum = extrapolate_01inf(plain_c0, plain_sum - plain_c0, plain_c_inf, gamma_j);
            rs.push(gamma_j);
        }

        // Every described step has been played, so the sponge goes back to the caller.
        transcript.finish();

        // Phase 6: residual handoff scaled by eps.
        //
        // The mode-specific compression lives behind the layout trait.
        // The shared invariant is `prod_poly.dot_product() == eps * plain_residual_sum`.
        let rs = Point::new(rs);
        let prod_poly = self.inner.zk_residual_handoff(&rs, alpha, eps);

        let residual_sum = eps * plain_sum;
        debug_assert_eq!(
            prod_poly.dot_product(),
            residual_sum,
            "residual product polynomial dot product must equal eps * plain_residual_sum",
        );

        ZkSumcheckHandoff {
            residual_prover: SumcheckProver::new(prod_poly, residual_sum),
            randomness: rs,
            eps,
            mask_messages: masks,
            mask_randomness,
            mask_oracle,
        }
    }
}

#[cfg(test)]
mod tests {
    use alloc::vec;
    use alloc::vec::Vec;

    use proptest::prelude::*;
    use rand::rngs::StdRng;
    use rand::{RngExt, SeedableRng};

    use crate::layout::{PrefixProver, SuffixProver};
    use crate::strategy::VariableOrder;
    use crate::table::OpeningBatch;
    use crate::tests::transcript_fingerprint;
    use crate::zk::test_helpers::{
        EF, F, MyChallenger, MyMmcs, build_prover_verifier, make_setup, run_roundtrip,
    };
    use crate::zk::{ZkLayout, ZkSumcheckData};

    #[test]
    fn prover_verifier_roundtrip_prefix() {
        // Invariant: honest prover-verifier roundtrip on the prefix path with both claim kinds present accepts.
        //
        // Fixture state:
        //
        //     n_vars       = 8       (witness has 2^8 = 256 evaluations)
        //     folding      = 3       (three sumcheck rounds)
        //     ell_zk       = 4       (mask polynomial degree 3)
        //     num_concrete = 1       (drives the inner accumulator branch)
        //     num_virtual  = 1       (drives the virtual accumulator branch)
        //
        // Both claim kinds force the alpha-power split to do real work in both branches.
        // This catches a wrong skip count in the per-round accumulator assembly.
        run_roundtrip(VariableOrder::Prefix, 8, 3, 4, 1, 1, 0)
            .expect("honest roundtrip should accept");
    }

    #[test]
    fn prover_verifier_roundtrip_prefix_with_no_residual_variables() {
        // Folding all three variables exercises the scalar handoff on SIMD builds.
        run_roundtrip(VariableOrder::Prefix, 3, 3, 4, 1, 1, 0)
            .expect("fully folded prefix roundtrip should accept");
    }

    #[test]
    fn prover_verifier_roundtrip_suffix() {
        // Invariant: honest prover-verifier roundtrip on the suffix path with both claim kinds present accepts.
        //
        // Fixture mirrors the prefix-mode pin so a regression in one mode surfaces with the same fingerprint as the other.
        // The residual handoff routes the combining challenge through suffix-specific compression that the prefix overlay does not exercise.
        run_roundtrip(VariableOrder::Suffix, 8, 3, 4, 1, 1, 0)
            .expect("honest roundtrip should accept");
    }

    #[test]
    fn long_mask_horner_path_prefix() {
        // Invariant: per-round Horner handles long masks identically to short ones on the prefix path.
        //
        // Fixture state: ell_zk = 32, matching a soundness-realistic mask setup.
        // The proptests below cap ell_zk at 5 to keep runtime small.
        // This case pins the long-mask arithmetic path the proptests never reach.
        run_roundtrip(VariableOrder::Prefix, 8, 3, 32, 1, 1, 0)
            .expect("honest roundtrip should accept");
    }

    #[test]
    fn long_mask_horner_path_suffix() {
        // Invariant: per-round Horner handles long masks identically to short ones on the suffix path.
        //
        // Mirrors the prefix long-mask pin.
        // A regression that mishandles the combining factor on a long mask passes the prefix pin and trips here.
        run_roundtrip(VariableOrder::Suffix, 8, 3, 32, 1, 1, 0)
            .expect("honest roundtrip should accept");
    }

    #[test]
    fn prover_verifier_roundtrip_mixed_current_next_batch() {
        // Invariant: the zero-knowledge prefix wrapper accepts a batch that mixes
        // a current opening point with its repeat-last successor point.
        // Acceptance means the prover and verifier draw identical randomness.

        // Fixture state: 8 variables, fold 3 at a time, mask width 4, no grinding.
        let n_vars = 8;
        let folding_factor = 3;
        let ell_zk = 4;
        let seed = 7u64;
        let pow_bits = 0;

        // Shared permutation, Merkle scheme, and encoding for both parties.
        let (perm, mmcs, encoding) = make_setup(seed, ell_zk);
        // Draw a random multilinear over the boolean cube of 2^8 = 256 points.
        let mut data_rng = StdRng::seed_from_u64(seed.wrapping_add(1));
        let evals: Vec<F> = (0..(1usize << n_vars)).map(|_| data_rng.random()).collect();
        let (mut prover, mut verifier, _) =
            build_prover_verifier::<PrefixProver<F, EF>>(evals, folding_factor, encoding, mmcs);

        // Two challengers seeded identically so the transcripts stay in lockstep.
        let mut prover_challenger = MyChallenger::new(perm.clone());
        let mut verifier_challenger = MyChallenger::new(perm);

        // Open column 0 at one current point and at its successor point.
        let batch = OpeningBatch::new(vec![0], vec![0]);
        // Prover records the claims and returns one evaluation per requested point.
        let evals = prover.eval(0, &batch, &mut prover_challenger);
        assert_eq!(evals.len(), batch.len());
        // Verifier mirrors the same draws against the returned evaluations.
        verifier
            .add_claim(0, &batch, &evals, &mut verifier_challenger)
            .unwrap();

        // Add the masking claim that hides the witness in the final sumcheck.
        let virtual_eval = prover.add_virtual_eval(&mut prover_challenger);
        verifier.add_virtual_eval(virtual_eval, &mut verifier_challenger);

        // Prover commits to the mask and hands off the sumcheck transcript.
        let mut zk_data = ZkSumcheckData::<F, EF>::default();
        let mut prover_rng = StdRng::seed_from_u64(seed.wrapping_add(2));
        let prover_handoff = prover.into_sumcheck(
            &mut zk_data,
            pow_bits,
            &mut prover_challenger,
            &mut prover_rng,
        );
        // The mask commitment is the only extra public input the verifier needs.
        let mask_commitment = prover_handoff.mask_oracle.0.clone();

        // Verifier replays the sumcheck and must accept the mixed batch.
        let verifier_handoff = verifier
            .into_sumcheck::<MyMmcs, _>(
                &zk_data,
                &mask_commitment,
                ell_zk,
                folding_factor,
                pow_bits,
                &mut verifier_challenger,
            )
            .expect("verifier should accept mixed current/Next batch");

        // Both parties must have folded along the same challenge sequence.
        assert_eq!(
            prover_handoff
                .randomness
                .iter()
                .copied()
                .collect::<Vec<_>>(),
            verifier_handoff
                .randomness
                .iter()
                .copied()
                .collect::<Vec<_>>(),
        );
    }

    proptest! {
        #![proptest_config(ProptestConfig::with_cases(16))]

        #[test]
        fn prop_completeness_prefix(
            n_vars in 3usize..=8,
            ell_zk in 3usize..=5,
            num_concrete in 0usize..=2,
            num_virtual in 0usize..=2,
            seed in 0u64..1024,
        ) {
            // Invariant: every honest prover output verifies across the parameter cube on the prefix path.
            //
            // Sweeping concrete and virtual independently exercises both accumulator branches and the alpha-power split between them.

            // The protocol needs at least one claim for a meaningful mu.
            prop_assume!(num_concrete + num_virtual >= 1);

            // Cover every legal prefix folding depth.
            let folding_factor = 1 + (seed as usize % n_vars);

            prop_assert!(
                run_roundtrip(VariableOrder::Prefix, n_vars, folding_factor, ell_zk, num_concrete, num_virtual, seed)
                    .is_ok()
            );
        }

        #[test]
        fn prop_completeness_suffix(
            n_vars in 3usize..=8,
            ell_zk in 3usize..=5,
            num_concrete in 0usize..=2,
            num_virtual in 0usize..=2,
            seed in 0u64..1024,
        ) {
            // Invariant: every honest prover output verifies across the parameter cube on the suffix path.

            // The protocol needs at least one claim for a meaningful mu.
            prop_assume!(num_concrete + num_virtual >= 1);

            // Suffix mode has no packed-lane reservation, so folding goes up to n_vars - 1.
            // One variable stays for the residual sumcheck.
            let folding_factor = 1 + (seed as usize % (n_vars - 1));

            prop_assert!(
                run_roundtrip(VariableOrder::Suffix, n_vars, folding_factor, ell_zk, num_concrete, num_virtual, seed)
                    .is_ok()
            );
        }
    }

    #[test]
    fn hiding_batch_transcripts_are_pinned() {
        // Invariant: every message the HVZK batch sends, and every challenge it
        // draws, stays the same value over this fixed run.
        //
        // Fixture state: a batch that mixes a current-and-successor opening, a
        // current-only opening, and two virtual claims, folded to a constant and
        // one further challenge drawn from the residual handoff.
        //
        // Grinding stays off: under `--features parallel`, a PoW search may return
        // any valid witness, so a pinned value would be flaky with grinding on.
        //
        // The constants below depend on the seeded `StdRng` streams the setup, the
        // witness, and the masks draw from.
        fn run<L: ZkLayout<F, EF>>() -> [u32; 4] {
            let (perm, mmcs, encoding) = make_setup(11, 4);
            let mut data_rng = StdRng::seed_from_u64(12);
            let evals: Vec<F> = (0..256).map(|_| data_rng.random()).collect();
            let (mut prover, _verifier, _n_vars) =
                build_prover_verifier::<L>(evals, 3, encoding, mmcs);
            let mut ch = MyChallenger::new(perm);

            prover.eval(0, &OpeningBatch::new(vec![0], vec![0]), &mut ch);
            prover.eval(0, &OpeningBatch::new(vec![0], Vec::new()), &mut ch);
            let _ = prover.add_virtual_eval(&mut ch);
            let _ = prover.add_virtual_eval(&mut ch);

            let mut zk_data = ZkSumcheckData::<F, EF>::default();
            let mut prover_rng = StdRng::seed_from_u64(13);
            let mut handoff = prover.into_sumcheck(&mut zk_data, 0, &mut ch, &mut prover_rng);
            transcript_fingerprint(&mut handoff.residual_prover, &mut ch)
        }
        assert_eq!(
            run::<PrefixProver<F, EF>>(),
            [45553658, 1785166884, 1370940203, 990586681]
        );
        assert_eq!(
            run::<SuffixProver<F, EF>>(),
            [1049900532, 349798432, 102124702, 1192330144]
        );
    }
}