1use super::{Claim, FullProof, RoundProof, Witness};
  2use crate::gkr::{
  3    multivariate::{ComposedMultiLinears, ComposedMultiLinearsOracle},
  4    sumcheck::{reduce_claim, FinalEvaluationClaim, PartialProof, RoundClaim, RoundOutput},
  5};
  6use rayon::iter::{IntoParallelIterator, ParallelIterator};
  7use winter_crypto::{ElementHasher, RandomCoin};
  8use winter_math::{fields::f64::BaseElement, FieldElement};
  9
 10pub fn sum_check_prove<
 11    E: FieldElement<BaseField = BaseElement>,
 12    C: RandomCoin<Hasher = H, BaseField = BaseElement>,
 13    H: ElementHasher<BaseField = BaseElement>,
 14>(
 15    claim: &Claim<E>,
 16    oracle: ComposedMultiLinearsOracle<E>,
 17    witness: Witness<E>,
 18    coin: &mut C,
 19) -> FullProof<E> {
 20    // Setup first round
 21    let mut prev_claim = RoundClaim {
 22        partial_eval_point: vec![],
 23        current_claim: claim.sum_value.clone(),
 24    };
 25    let prev_proof = PartialProof { round_proofs: vec![] };
 26    let num_vars = witness.polynomial.num_variables_ml();
 27    let prev_output = RoundOutput { proof: prev_proof, witness };
 28
 29    let mut output = sumcheck_round(prev_output);
 30    let poly_evals = &output.proof.round_proofs[0].poly_evals;
 31    coin.reseed(H::hash_elements(&poly_evals));
 32
 33    for i in 1..num_vars {
 34        let round_challenge = coin.draw().unwrap();
 35        let new_claim = reduce_claim(
 36            output.proof.round_proofs.last().unwrap().clone(),
 37            prev_claim,
 38            round_challenge,
 39        );
 40        output.witness.polynomial = output.witness.polynomial.bind(round_challenge);
 41
 42        output = sumcheck_round(output);
 43        prev_claim = new_claim;
 44
 45        let poly_evals = &output.proof.round_proofs[i].poly_evals;
 46        coin.reseed(H::hash_elements(&poly_evals));
 47    }
 48
 49    let round_challenge = coin.draw().unwrap();
 50    let RoundClaim { partial_eval_point, current_claim } = reduce_claim(
 51        output.proof.round_proofs.last().unwrap().clone(),
 52        prev_claim,
 53        round_challenge,
 54    );
 55    let final_eval_claim = FinalEvaluationClaim {
 56        evaluation_point: partial_eval_point,
 57        claimed_evaluation: current_claim,
 58        polynomial: oracle,
 59    };
 60
 61    FullProof {
 62        sum_check_proof: output.proof,
 63        final_evaluation_claim: final_eval_claim,
 64    }
 65}
 66
 67fn sumcheck_round<E: FieldElement>(prev_proof: RoundOutput<E>) -> RoundOutput<E> {
 68    let RoundOutput { mut proof, witness } = prev_proof;
 69
 70    let polynomial = witness.polynomial;
 71    let num_ml = polynomial.num_ml();
 72    let num_vars = polynomial.num_variables_ml();
 73    let num_rounds = num_vars - 1;
 74
 75    let mut evals_zero = vec![E::ZERO; num_ml];
 76    let mut evals_one = vec![E::ZERO; num_ml];
 77    let mut deltas = vec![E::ZERO; num_ml];
 78    let mut evals_x = vec![E::ZERO; num_ml];
 79
 80    let total_evals = (0..1 << num_rounds).into_iter().map(|i| {
 81        for (j, ml) in polynomial.multi_linears.iter().enumerate() {
 82            evals_zero[j] = ml.evaluations[(i << 1) as usize];
 83            evals_one[j] = ml.evaluations[(i << 1) + 1];
 84        }
 85        let mut total_evals = vec![E::ZERO; polynomial.degree()];
 86        total_evals[0] = polynomial.composer.evaluate(&evals_one);
 87        evals_zero
 88            .iter()
 89            .zip(evals_one.iter().zip(deltas.iter_mut().zip(evals_x.iter_mut())))
 90            .for_each(|(a0, (a1, (delta, evx)))| {
 91                *delta = *a1 - *a0;
 92                *evx = *a1;
 93            });
 94        total_evals.iter_mut().skip(1).for_each(|e| {
 95            evals_x.iter_mut().zip(deltas.iter()).for_each(|(evx, delta)| {
 96                *evx += *delta;
 97            });
 98            *e = polynomial.composer.evaluate(&evals_x);
 99        });
100        total_evals
101    });
102    let evaluations = total_evals.fold(vec![E::ZERO; polynomial.degree()], |mut acc, evals| {
103        acc.iter_mut().zip(evals.iter()).for_each(|(a, ev)| *a += *ev);
104        acc
105    });
106    let proof_update = RoundProof { poly_evals: evaluations };
107    proof.round_proofs.push(proof_update);
108    RoundOutput { proof, witness: Witness { polynomial } }
109}