p3-sumcheck 0.8.0

//! Per-round prover state and the default sumcheck driver.

use alloc::vec::Vec;

use p3_challenger::fs::TranscriptField;
use p3_challenger::{FieldChallenger, GrindingChallenger};
use p3_field::ExtensionField;
use p3_multilinear_util::point::Point;

use super::proof::GenericDegreeProof;
use super::transcript::{GenericDegreeShape, ProverTranscript};

/// Per-round callback used to drive the prover loop.
///
/// # Contract
///
/// - Carries the prover's mutable polynomial state across rounds.
/// - Exposes two operations called in fixed order by the driver.
pub trait RoundProver<EF> {
    /// Bind the most recently active variable to the verifier's challenge.
    ///
    /// - Called exactly once between consecutive round-polynomial requests.
    /// - In-place updates are recommended.
    fn fold(&mut self, r: EF);

    /// Return the current round polynomial as a list of evaluations.
    ///
    /// This reads the current state without mutating it; binding happens only in [`Self::fold`].
    ///
    /// # Returns
    ///
    /// Vector of length `degree` containing `h(0), h(2), h(3), ..., h(degree)`.
    fn round_poly(&self) -> Vec<EF>;

    /// Run the prover side of a generic-degree sumcheck driven by this state.
    ///
    /// # Arguments
    ///
    /// - `challenger`: Fiat-Shamir transcript shared with the verifier.
    /// - `num_rounds`: number of variables to bind.
    /// - `degree`: per-variable degree of the polynomial being summed.
    /// - `pow_bits`: grinding difficulty per round, or `0` to skip.
    /// - `claimed_sum`: claimed value of the sum at the start of round zero.
    ///
    /// # Returns
    ///
    /// - The transcript record, which carries the claimed sum.
    /// - The vector of challenges sampled by the verifier.
    ///
    /// # Transcript
    ///
    /// The claimed sum is bound here, not by the caller.
    /// The verifier binds it at the same point.
    /// Both sides therefore stay in step with nothing for the caller to mirror.
    ///
    /// # Panics
    ///
    /// - When the per-variable degree is zero.
    /// - When a round polynomial carries a different number of evaluations.
    fn prove<F, Challenger>(
        &mut self,
        challenger: &mut Challenger,
        num_rounds: usize,
        degree: usize,
        pow_bits: usize,
        claimed_sum: EF,
    ) -> (GenericDegreeProof<F, EF>, Point<EF>)
    where
        F: TranscriptField,
        EF: ExtensionField<F>,
        Challenger: FieldChallenger<F> + GrindingChallenger<Witness = F>,
    {
        // A degree-zero polynomial would carry no information.
        assert!(degree > 0, "generic-degree sumcheck: degree must be > 0");

        let mut proof = GenericDegreeProof {
            claimed_sum,
            round_polys: Vec::with_capacity(num_rounds),
            pow_witnesses: Vec::with_capacity(if pow_bits > 0 { num_rounds } else { 0 }),
        };
        let mut challenges = Vec::with_capacity(num_rounds);

        // Seeding binds the shape and the claimed sum before the first round.
        let shape = GenericDegreeShape::new(num_rounds, degree, pow_bits);
        let mut transcript =
            ProverTranscript::<Challenger, F, EF>::new(challenger, shape, claimed_sum);

        for _ in 0..num_rounds {
            let evals = self.round_poly();

            // Bind the polynomial, grind, and draw this round's challenge.
            //
            // Binding compares the evaluation count against the described width.
            let (challenge, witness) = transcript.round(&evals);

            // Store what the round produced alongside what it bound.
            proof.round_polys.push(evals);
            proof.pow_witnesses.extend(witness);

            self.fold(challenge);
            challenges.push(challenge);
        }

        // Require that every described step was played.
        transcript.finish();

        (proof, Point::new(challenges))
    }
}

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

    use p3_baby_bear::{BabyBear, Poseidon2BabyBear};
    use p3_challenger::DuplexChallenger;
    use p3_field::PrimeCharacteristicRing;
    use p3_field::extension::BinomialExtensionField;
    use p3_multilinear_util::poly::Poly;
    use rand::rngs::SmallRng;
    use rand::{RngExt, SeedableRng};

    use super::*;
    use crate::generic_degree::GenericDegreeError;

    type F = BabyBear;
    type EF = BinomialExtensionField<F, 4>;
    type Perm = Poseidon2BabyBear<16>;
    type Ch = DuplexChallenger<F, Perm, 16, 8>;

    fn fresh_challenger() -> Ch {
        // Fixed seed so prover and verifier transcripts match exactly.
        let mut rng = SmallRng::seed_from_u64(0xDEADBEEF);
        let perm = Perm::new_from_rng_128(&mut rng);
        Ch::new(perm)
    }

    /// Test prover for a product of three multilinears.
    ///
    /// The summed polynomial is the pointwise product `a(x) * b(x) * c(x)` over the `log_m`-variable hypercube.
    struct TripleProductProver {
        /// Three multilinear factors, all over the same variable space.
        factors: [Poly<EF>; 3],
    }

    impl RoundProver<EF> for TripleProductProver {
        fn fold(&mut self, r: EF) {
            // Bind the most recently active variable in every factor.
            for f in &mut self.factors {
                f.fix_prefix_var_mut(r);
            }
        }

        fn round_poly(&self) -> Vec<EF> {
            // Return h(0), h(2), h(3) for the current round polynomial.
            let h0 = round_poly_product(&self.factors, EF::ZERO);
            let h2 = round_poly_product(&self.factors, EF::from_u64(2));
            let h3 = round_poly_product(&self.factors, EF::from_u64(3));
            vec![h0, h2, h3]
        }
    }

    fn round_poly_product(factors: &[Poly<EF>; 3], node: EF) -> EF {
        // Step 1: for each factor, interpolate between the two halves at X=node.
        //
        //     factor(node, x') = factor(0, x') + (factor(1, x') - factor(0, x')) * node
        let interpolated: Vec<Vec<EF>> = factors
            .iter()
            .map(|p| {
                let half = p.num_evals() / 2;
                (0..half)
                    .map(|i| {
                        let a0 = p.as_slice()[i];
                        let a1 = p.as_slice()[i + half];
                        a0 + (a1 - a0) * node
                    })
                    .collect()
            })
            .collect();

        // Step 2: sum the pointwise product across the unbound cube.
        let n_remaining = interpolated[0].len();
        let mut acc = EF::ZERO;
        for j in 0..n_remaining {
            let mut prod = EF::ONE;
            for v in &interpolated {
                prod *= v[j];
            }
            acc += prod;
        }
        acc
    }

    /// Build the three random factors and their ground-truth product sum.
    fn random_instance(rng: &mut SmallRng, m: usize) -> ([Vec<F>; 3], EF) {
        let a: Vec<F> = (0..m).map(|_| rng.random()).collect();
        let b: Vec<F> = (0..m).map(|_| rng.random()).collect();
        let c: Vec<F> = (0..m).map(|_| rng.random()).collect();
        let claimed_sum: EF = (0..m)
            .map(|i| EF::from(a[i]) * EF::from(b[i]) * EF::from(c[i]))
            .sum();
        ([a, b, c], claimed_sum)
    }

    fn prover_from(columns: &[Vec<F>; 3]) -> TripleProductProver {
        TripleProductProver {
            factors: [
                Poly::new(columns[0].iter().copied().map(EF::from).collect()),
                Poly::new(columns[1].iter().copied().map(EF::from).collect()),
                Poly::new(columns[2].iter().copied().map(EF::from).collect()),
            ],
        }
    }

    #[test]
    fn end_to_end_prove_verify_degree_3() {
        // Fixture state:
        //
        //     log_m = 6, so the hypercube has 64 points
        //     three random base-field columns a, b, c
        //     claimed sum = sum_x a(x) * b(x) * c(x)
        //
        // Invariant:
        //
        //     prover and verifier sample the same challenges
        //     the final claimed sum equals a(r) * b(r) * c(r) at the challenge point
        let mut rng = SmallRng::seed_from_u64(123);
        let log_m = 6usize;
        let (columns, claimed_sum) = random_instance(&mut rng, 1 << log_m);

        // Prover side.
        let mut prover_state = prover_from(&columns);
        let mut p_ch = fresh_challenger();
        let (proof, prover_challenges) =
            prover_state.prove::<F, _>(&mut p_ch, log_m, 3, 0, claimed_sum);

        // The proof carries the claimed sum the verifier will consume.
        assert_eq!(proof.claimed_sum, claimed_sum);

        // After all rounds each factor has been bound down to a single value.
        let final_a = prover_state.factors[0].as_slice()[0];
        let final_b = prover_state.factors[1].as_slice()[0];
        let final_c = prover_state.factors[2].as_slice()[0];

        // Verifier side replays the transcript and recovers the final claim.
        let mut v_ch = fresh_challenger();
        let (verifier_challenges, final_sum) = proof.verify(&mut v_ch, log_m, 3, 0).unwrap();

        // Both sides must observe the same Fiat-Shamir transcript.
        assert_eq!(prover_challenges, verifier_challenges);
        // The final sumcheck equation is the product of multilinear evaluations.
        assert_eq!(final_sum, final_a * final_b * final_c);
    }

    #[test]
    fn end_to_end_prove_verify_with_grinding() {
        // Same product sumcheck, now with per-round proof-of-work enabled.
        // This exercises the grind / check-witness path that the degree-3 test skips.
        //
        // Fixture state: log_m = 4, pow_bits = 4 to keep grinding fast.
        let mut rng = SmallRng::seed_from_u64(321);
        let log_m = 4usize;
        let pow_bits = 4usize;
        let (columns, claimed_sum) = random_instance(&mut rng, 1 << log_m);

        let mut prover_state = prover_from(&columns);
        let mut p_ch = fresh_challenger();
        let (proof, prover_challenges) =
            prover_state.prove::<F, _>(&mut p_ch, log_m, 3, pow_bits, claimed_sum);

        // Grinding emits exactly one witness per round.
        assert_eq!(proof.pow_witnesses.len(), log_m);

        let final_a = prover_state.factors[0].as_slice()[0];
        let final_b = prover_state.factors[1].as_slice()[0];
        let final_c = prover_state.factors[2].as_slice()[0];

        // The verifier re-checks each round's witness while replaying the transcript.
        let mut v_ch = fresh_challenger();
        let (verifier_challenges, final_sum) = proof.verify(&mut v_ch, log_m, 3, pow_bits).unwrap();

        assert_eq!(prover_challenges, verifier_challenges);
        assert_eq!(final_sum, final_a * final_b * final_c);
    }

    #[test]
    fn verify_rejects_tampered_pow_witness() {
        // Invariant: a corrupted PoW witness must fail the per-round grinding check.
        //
        // Fixture state: a valid grinded proof over log_m = 4 rounds.
        //
        // Mutation: corrupt the last round's witness.
        //   - rounds 0..2 keep valid witnesses
        //   - the transcript therefore stays in sync up to the final round
        //   - the final round then fails its grinding check
        let mut rng = SmallRng::seed_from_u64(99);
        let log_m = 4usize;
        let pow_bits = 4usize;
        let (columns, claimed_sum) = random_instance(&mut rng, 1 << log_m);

        let mut prover_state = prover_from(&columns);
        let mut p_ch = fresh_challenger();
        let (mut proof, _) = prover_state.prove::<F, _>(&mut p_ch, log_m, 3, pow_bits, claimed_sum);

        let last_round = log_m - 1;
        let valid_witness = proof.pow_witnesses[last_round];

        // Why an offset search instead of a single `+1` tamper:
        //   - a corrupted witness still clears the grind by chance with probability 2^-pow_bits
        //   - parallel grinding may return any valid witness, not a fixed smallest one
        //
        // Walking successive offsets pins the rejection regardless of the witness grinding picked.
        let mut offset = F::ONE;
        loop {
            // Replace the honest witness with a nearby, almost-certainly-invalid value.
            proof.pow_witnesses[last_round] = valid_witness + offset;

            let mut v_ch = fresh_challenger();
            match proof.verify(&mut v_ch, log_m, 3, pow_bits) {
                // The failure is pinned to the corrupted round.
                Err(GenericDegreeError::InvalidPowWitness { round }) => {
                    assert_eq!(round, last_round);
                    break;
                }
                // Rare coincidental grind hit: advance to the next offset.
                Ok(_) => offset += F::ONE,
                Err(other) => panic!("expected an invalid pow witness error, got {other:?}"),
            }
        }
    }
}

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

    use p3_binary_field::{BinaryChallenger, BinaryField128, TowerLevel};
    use p3_challenger::{CanSample, HashChallenger};
    use p3_field::Field;
    use p3_keccak::Keccak256Hash;
    use p3_multilinear_util::poly::Poly;

    use super::RoundProver;
    use crate::generic_degree::transcript::GenericDegreeShape;

    type F = BinaryField128;
    type Ch = BinaryChallenger<F, HashChallenger<u8, Keccak256Hash, 32>>;

    fn fresh_challenger() -> Ch {
        Ch::from_hasher(Vec::new(), Keccak256Hash)
    }

    struct ProductProver([Poly<F>; 3]);

    impl RoundProver<F> for ProductProver {
        fn fold(&mut self, r: F) {
            for factor in &mut self.0 {
                factor.fix_prefix_var_mut(r);
            }
        }

        fn round_poly(&self) -> Vec<F> {
            [0, 2, 3]
                .into_iter()
                .map(|i| {
                    let node = F::interpolation_node(i);
                    let half = self.0[0].num_evals() / 2;
                    (0..half)
                        .map(|j| {
                            self.0
                                .iter()
                                .map(|factor| {
                                    let values = factor.as_slice();
                                    values[j] + (values[j + half] - values[j]) * node
                                })
                                .product::<F>()
                        })
                        .sum()
                })
                .collect()
        }
    }

    #[test]
    fn binary_sumcheck_round_trips_with_and_without_grinding() {
        let factors: [Poly<F>; 3] = core::array::from_fn(|column| {
            Poly::new(
                (0..16)
                    .map(|row| {
                        F::from_repr(((row + 1) as u128) << (column * 40) | (column + 1) as u128)
                    })
                    .collect::<Vec<_>>(),
            )
        });
        let claimed_sum: F = (0..16)
            .map(|row| factors.iter().map(|f| f.as_slice()[row]).product::<F>())
            .sum();
        for pow_bits in [0, 4] {
            let mut prover = ProductProver(factors.clone());
            let mut p_ch = fresh_challenger();
            let (proof, p_point) = prover.prove::<F, _>(&mut p_ch, 4, 3, pow_bits, claimed_sum);
            assert_eq!(proof.pow_witnesses.len(), if pow_bits == 0 { 0 } else { 4 });
            let mut v_ch = fresh_challenger();
            let (v_point, final_sum) = proof.verify(&mut v_ch, 4, 3, pow_bits).unwrap();
            assert_eq!(p_point, v_point);
            assert_eq!(
                final_sum,
                factors.iter().map(|f| f.eval_base(&v_point)).product::<F>()
            );
            assert_eq!(
                CanSample::<F>::sample(&mut p_ch),
                CanSample::<F>::sample(&mut v_ch)
            );
        }
    }

    #[test]
    fn binary_sumcheck_seed_binds_rounds_degree_and_grinding() {
        let sample = |rounds, degree, pow_bits| {
            let mut challenger = fresh_challenger();
            GenericDegreeShape::new(rounds, degree, pow_bits)
                .domain_separator::<F, F>()
                .seed(&mut challenger);
            CanSample::<F>::sample(&mut challenger)
        };
        let baseline = sample(4, 3, 0);
        assert_ne!(baseline, sample(5, 3, 0));
        assert_ne!(baseline, sample(4, 4, 0));
        assert_ne!(baseline, sample(4, 3, 4));
        assert_ne!(sample(4, 3, 4), sample(4, 3, 5));
    }
}