zinc_piop/sumcheck/

prover.rs

//! Prover

use crypto_primitives::{SemiringConfig, SetConfig};
use itertools::Itertools;
#[cfg(feature = "parallel")]
use rayon::iter::*;
use std::slice;
use zinc_poly::mle::{DenseMultilinearExtension, MultilinearExtension};
use zinc_transcript::{delegate_transcribable, traits::ConstTranscribable};
use zinc_utils::{cfg_into_iter, cfg_iter_mut};

/// Evaluation of a polynomial on natural points without the constant term.
#[repr(transparent)]
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct NatEvaluatedPolyWithoutConstant<F> {
    /// Evaluations at 1, 2, ... (P(0) is omitted).
    pub tail_evaluations: Vec<F>,
}

impl<F> NatEvaluatedPolyWithoutConstant<F> {
    pub fn new(tail_evaluations: Vec<F>) -> Self {
        Self { tail_evaluations }
    }

    /// Maps every field element through `f`, preserving structure — used to
    /// lift elements into wire integers and to project wire integers back
    /// into elements at the (de)serialization boundary.
    pub fn try_map<T, E>(
        &self,
        f: impl FnMut(&F) -> Result<T, E> + Copy,
    ) -> Result<NatEvaluatedPolyWithoutConstant<T>, E> {
        Ok(NatEvaluatedPolyWithoutConstant {
            tail_evaluations: self.tail_evaluations.iter().map(f).try_collect()?,
        })
    }
}

impl<F> std::ops::Deref for NatEvaluatedPolyWithoutConstant<F> {
    type Target = [F];

    fn deref(&self) -> &Self::Target {
        &self.tail_evaluations
    }
}

impl<F> std::ops::DerefMut for NatEvaluatedPolyWithoutConstant<F> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.tail_evaluations
    }
}

delegate_transcribable!(NatEvaluatedPolyWithoutConstant<F> { tail_evaluations: Vec<F> }
    where F: ConstTranscribable);

#[repr(transparent)]
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ProverMsg<F>(pub NatEvaluatedPolyWithoutConstant<F>);

delegate_transcribable!(ProverMsg<F>(NatEvaluatedPolyWithoutConstant<F>)
    where F: ConstTranscribable);

impl<F> ProverMsg<F> {
    /// Maps every field element through `f`, preserving structure — used to
    /// lift elements into wire integers and to project wire integers back
    /// into elements at the (de)serialization boundary.
    pub fn try_map<T, E>(
        &self,
        f: impl FnMut(&F) -> Result<T, E> + Copy,
    ) -> Result<ProverMsg<T>, E> {
        Ok(ProverMsg(self.0.try_map(f)?))
    }
}

/// Sumcheck Prover State, generic over the field config `C`.
pub struct ProverState<C: SetConfig> {
    /// Sampled randomness given by the verifier.
    pub randomness: Vec<C::Element>,
    /// Stores the list of multilinear extensions
    /// the sumcheck polynomial is comprised of.
    pub mles: Vec<DenseMultilinearExtension<C::Element>>,
    /// Number of variables.
    pub num_vars: usize,
    /// Max degree.
    pub max_degree: usize,
    /// The current round number.
    pub round: usize,
    /// Claimed sum for the first round polynomial.
    pub asserted_sum: Option<C::Element>,
    /// When `true`, the next [`Self::prove_round`] invocation pushes the
    /// verifier challenge into `randomness` but skips the
    /// `fix_variables` fold of `mles`. Used by round-1 fast
    /// paths (see that pre-fold the MLEs as part of their setup.
    /// The flag is reset to `false` after the skipped fold.
    pub skip_next_fold: bool,
}

impl<C: SetConfig> ProverState<C> {
    /// Initialize the prover to argue for the sum of products of
    /// MLE's in {0,1}^`num_vars`.
    pub fn new(
        mles: Vec<DenseMultilinearExtension<C::Element>>,
        nvars: usize,
        degree: usize,
    ) -> Self {
        Self {
            randomness: Vec::with_capacity(nvars),
            mles,
            num_vars: nvars,
            max_degree: degree,
            round: 0,
            asserted_sum: None,
            skip_next_fold: false,
        }
    }
}

impl<C: SemiringConfig> ProverState<C> {
    /// Receive message from verifier, generate prover message, and proceed to
    /// next round.
    ///
    /// Adapted Jolt's sumcheck implementation, with some additions like round 1
    /// fast path.
    #[allow(clippy::arithmetic_side_effects)]
    pub fn prove_round(
        &mut self,
        v_msg: &Option<C::Element>,
        comb_fn: impl Fn(&[C::Element]) -> C::Element + Send + Sync,
        config: &C,
    ) -> ProverMsg<C::Element> {
        if let Some(msg) = v_msg {
            if self.round == 0 {
                panic!("first round should be prover first.");
            }
            self.randomness.push(msg.clone());

            if self.skip_next_fold {
                // A fast path already produced the post-fold MLEs
                self.skip_next_fold = false;
            } else {
                // fix the next variable at the verifier randomness for this round
                let i = self.round;
                let r = self.randomness[i - 1].clone();

                cfg_iter_mut!(self.mles).for_each(|multiplicand| {
                    multiplicand.fix_variables(config, slice::from_ref(&r));
                });
            }
        } else if self.round > 0 {
            panic!("verifier message is empty");
        }

        self.round += 1;

        if self.round > self.num_vars {
            panic!("Prover is not active");
        }

        let i = self.round;
        let nv = self.num_vars;
        let degree = self.max_degree;

        let polys = &self.mles;

        struct Scratch<R> {
            evals: Vec<R>,
            steps: Vec<R>,
            vals0: Vec<R>,
            vals1: Vec<R>,
            vals: Vec<R>,
            levals: Vec<R>,
        }
        let zero = config.zero();
        let zero_vec_deg = vec![zero.clone(); degree + 1];
        let zero_vec_poly = vec![zero.clone(); polys.len()];
        let scratch = || Scratch {
            evals: zero_vec_deg.clone(),
            steps: zero_vec_poly.clone(),
            vals0: zero_vec_poly.clone(),
            vals1: zero_vec_poly.clone(),
            vals: zero_vec_poly.clone(),
            levals: zero_vec_deg.clone(),
        };

        #[cfg(not(feature = "parallel"))]
        let zeros = scratch();
        #[cfg(feature = "parallel")]
        let zeros = scratch;

        let summer = cfg_into_iter!(0..1 << (nv - i)).fold(zeros, |mut s, b| {
            let index = b << 1;

            // TODO(Alex): Once you have benches set,
            //             could please try getting rid of vals0 and vals1 fields in the
            // structs, replacing them with
            //
            //             ```rust
            //             let vals0: Vec<_> = polys.iter().map(|poly|
            // poly[index].clone()).collect();             let vals1: Vec<_> =
            // polys.iter().map(|poly| poly[index + 1].clone()).collect();
            //             ```
            //             My bet is that it won't affect running time, but better safe than
            // sorry.

            s.vals0
                .iter_mut()
                .zip(polys.iter())
                .for_each(|(v0, poly)| *v0 = poly[index].clone());
            s.levals[0] = comb_fn(&s.vals0);

            if degree > 0 {
                s.vals1
                    .iter_mut()
                    .zip(polys.iter())
                    .for_each(|(v1, poly)| *v1 = poly[index + 1].clone());
                s.levals[1] = comb_fn(&s.vals1);

                for (i, (v1, v0)) in s.vals1.iter().zip(s.vals0.iter()).enumerate() {
                    s.steps[i] = config.sub(v1, v0);
                    s.vals[i] = v1.clone();
                }

                for eval_point in s.levals.iter_mut().take(degree + 1).skip(2) {
                    for poly_i in 0..polys.len() {
                        config.add_assign(&mut s.vals[poly_i], &s.steps[poly_i]);
                    }
                    *eval_point = comb_fn(&s.vals);
                }
            }

            // TODO(Alex): It seems that the only thing
            //             we pass around meaningfully is evals,
            //             so this loop could be reworked to map/reduce - maybe even without
            //             #[cfg(feature = "parallel")]. Would help to get benchmarks up and
            //             running first though.
            s.evals
                .iter_mut()
                .zip(s.levals.iter())
                .for_each(|(e, l)| config.add_assign(e, l));

            s
        });

        // Rayon's fold outputs an iter which still needs to be summed over
        #[cfg(feature = "parallel")]
        let evaluations = summer.map(|s| s.evals).reduce(
            || vec![zero.clone(); degree + 1],
            |mut evaluations, evals| {
                evaluations
                    .iter_mut()
                    .zip(evals)
                    .for_each(|(e, l)| config.add_assign(e, &l));
                evaluations
            },
        );

        #[cfg(not(feature = "parallel"))]
        let evaluations = summer.evals;

        // Record the claimed sum once during the first round.
        if self.round == 1 {
            let p0 = evaluations
                .first()
                .expect("evaluations should always contain the constant term");
            let sum = if degree > 0 {
                let eval = evaluations
                    .get(1)
                    .expect("degree > 0 implies evaluation at 1 is present");
                config.add(p0, eval)
            } else {
                p0.clone()
            };
            self.asserted_sum = Some(sum);
        }

        // Strip the constant term before sending, without re-allocating all elements.
        let mut tail = evaluations;
        tail.remove(0); // leaves P(0) behind; tail holds P(1..)

        ProverMsg(NatEvaluatedPolyWithoutConstant::new(tail))
    }
}