zinc_piop/

sumcheck.rs

pub mod multi_degree;
pub mod prover;
// pub mod utils;
pub mod verifier;

#[cfg(test)]
mod tests;

use self::verifier::Subclaim;
use crate::sumcheck::{
    prover::{NatEvaluatedPolyWithoutConstant, ProverMsg},
    verifier::VerifierState,
};
use crypto_primitives::{BaseFieldConfig, ProjectPrimitiveIntegersWithConfig};
use prover::ProverState;
use std::marker::PhantomData;
use thiserror::Error;
use zinc_poly::{EvaluationError, mle::DenseMultilinearExtension, utils::ArithErrors};
use zinc_transcript::traits::{ConstTranscribable, GenTranscribable, Transcribable, Transcript};
use zinc_utils::mul;

/// Sumcheck for products of multilinear polynomial, generic over the field
/// config `C`.
pub struct MLSumcheck<C>(PhantomData<C>);

/// Proof generated by the sumcheck prover.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SumcheckProof<F> {
    /// List of prover messages, one for each round.
    pub messages: Vec<ProverMsg<F>>,
    /// The claimed sum for the first round polynomial.
    pub claimed_sum: F,
}

/// The proof is transcribed as raw field elements without field metadata:
/// the field config is bound into the transcript separately, at sampling
/// time / at the top level of the surrounding proof.
impl<F: ConstTranscribable> GenTranscribable for SumcheckProof<F> {
    fn read_transcription_bytes_exact(bytes: &[u8]) -> Self {
        let (n_msgs, mut bytes) = u32::read_transcription_bytes_subset(bytes);
        let n_msgs = usize::try_from(n_msgs).expect("message count must fit into usize");

        let mut messages = Vec::with_capacity(n_msgs);
        for _ in 0..n_msgs {
            let (len, rest) = u32::read_transcription_bytes_subset(bytes);
            let len = usize::try_from(len).expect("polynomial length must fit into usize");
            bytes = rest;
            let end = mul!(len, F::NUM_BYTES);
            let tail_evaluations: Vec<F> = Vec::read_transcription_bytes_exact(&bytes[..end]);
            messages.push(ProverMsg(NatEvaluatedPolyWithoutConstant {
                tail_evaluations,
            }));
            bytes = &bytes[end..];
        }

        let claimed_sum = F::read_transcription_bytes_exact(bytes);
        Self {
            messages,
            claimed_sum,
        }
    }

    fn write_transcription_bytes_exact(&self, mut buf: &mut [u8]) {
        buf = {
            let len =
                u32::try_from(self.messages.len()).expect("messages length must fit into u32");
            len.write_transcription_bytes_exact(&mut buf[0..u32::NUM_BYTES]);
            &mut buf[u32::NUM_BYTES..]
        };
        for msg in &self.messages {
            let evals = &msg.0.tail_evaluations;
            buf = {
                let len = u32::try_from(evals.len()).expect("messages length must fit into u32");
                len.write_transcription_bytes_exact(&mut buf[0..u32::NUM_BYTES]);
                &mut buf[u32::NUM_BYTES..]
            };
            let end = mul!(evals.len(), F::NUM_BYTES);
            evals.write_transcription_bytes_exact(&mut buf[..end]);
            buf = &mut buf[end..];
        }
        self.claimed_sum.write_transcription_bytes_exact(buf);
    }
}

impl<F: ConstTranscribable> Transcribable for SumcheckProof<F> {
    #[allow(clippy::arithmetic_side_effects)]
    fn get_num_bytes(&self) -> usize {
        let n_msgs = self.messages.len();
        let total_evals: usize = self
            .messages
            .iter()
            .map(|m| m.0.tail_evaluations.len())
            .sum();
        u32::NUM_BYTES // n_msgs
            + n_msgs * u32::NUM_BYTES // n_evals for each message
            + total_evals * F::NUM_BYTES // evals
            + F::NUM_BYTES // claimed_sum
    }
}

impl<C> MLSumcheck<C>
where
    C: BaseFieldConfig + ProjectPrimitiveIntegersWithConfig,
    C::Integer: ConstTranscribable,
{
    /// Sumcheck prover main entry point.
    ///
    /// This function executes the Prover side of the Sumcheck protocol.
    /// It verifies a claim of the form:
    ///
    /// $$
    /// \sum_{x \in \{0, 1\}^{\text{nvars}}} \text{comb\\_fn}(\text{mles}(x)) =
    /// \text{claimed\\_sum}. $$
    ///
    /// It is designed to be used as a subprotocol within a larger system
    /// since it takes the FS transcript (`transcript` argument) as input
    /// and returns the **internal ProverState** alongside the final proof.
    ///
    /// The claimed sum is derived by the prover.
    ///
    /// ---
    ///
    /// # Arguments
    ///
    /// * `transcript`: A mutable reference to a Fiat-Shamir `Transcript`.
    /// * `mles`: A `Vec` of dense multilinear extension over the base field
    ///   `F`. The sumcheck polynomial is made over the combined result of these
    ///   multilinear extensions.
    /// * `nvars`: The number of variables over which the `mles` are defined.
    ///   This must be consistent across all `mles`.
    /// * `degree`: The maximum combined degree of the `mles` under the
    ///   `comb_fn`.
    /// * `comb_fn`: A closure that defines the combination function
    ///   $G(\text{mles}(x))$. It takes a slice of field elements (the
    ///   evaluations of the `mles` at a point $x$) and returns a single field
    ///   element.
    /// * `config`: The configuration for the underlying field used in the
    ///   protocol.
    ///
    /// ---
    ///
    /// # Returns
    ///
    /// A tuple containing:
    ///
    /// 1. `SumcheckProof<F>`: The final sumcheck proof.
    /// 2. `ProverState<F>`: The state of the Prover after the protocol
    ///    completes.
    ///
    /// ---
    ///
    /// # Panics
    ///
    /// * Panics if the number of variables is `0`.
    pub fn prove_as_subprotocol(
        transcript: &mut impl Transcript,
        mles: Vec<DenseMultilinearExtension<C::Element>>,
        nvars: usize,
        degree: usize,
        comb_fn: impl Fn(&[C::Element]) -> C::Element + Send + Sync,
        config: &C,
    ) -> (SumcheckProof<C::Element>, ProverState<C>) {
        if nvars == 0 {
            panic!("Attempt to prove a constant")
        }

        let mut buf = vec![0; <C::Integer as ConstTranscribable>::NUM_BYTES];
        let nvars_field = config.project(&(nvars as u64));
        let degree_field = config.project(&(degree as u64));

        transcript.absorb_field_element(config, &nvars_field, &mut buf);
        transcript.absorb_field_element(config, &degree_field, &mut buf);

        let mut prover_state = ProverState::new(mles, nvars, degree);
        let mut verifier_msg = None;
        let mut prover_msgs = Vec::with_capacity(nvars);

        for _ in 0..nvars {
            let prover_msg = prover_state.prove_round(&verifier_msg, &comb_fn, config);
            transcript.absorb_field_element_slice(config, &prover_msg.0.tail_evaluations, &mut buf);
            prover_msgs.push(prover_msg);
            let next_verifier_msg = transcript.get_field_challenge(config);
            transcript.absorb_field_element(config, &next_verifier_msg, &mut buf);

            verifier_msg = Some(next_verifier_msg);
        }
        let asserted_sum = prover_state
            .asserted_sum
            .clone()
            .expect("asserted sum should be recorded after the first prover round");
        if let Some(vmsg) = verifier_msg {
            prover_state.randomness.push(vmsg);
        }

        (
            SumcheckProof {
                messages: prover_msgs,
                claimed_sum: asserted_sum,
            },
            prover_state,
        )
    }

    /// Sumcheck verifier main entry point.
    ///
    /// This function executes the Verifier side of the Sumcheck protocol.
    /// It takes a `proof` and a `claimed_sum` and verifies the
    /// intermediate steps of the sumcheck.
    ///
    /// The sumcheck verifies the claim:
    ///
    /// $$
    /// \sum_{x \in \{0, 1\}^{\text{num\\_vars}}} G(x) = \text{claimed\\_sum}.
    /// $$
    ///
    /// It is designed to be used as a subprotocol within a larger system.
    /// If successful, it returns a **Subclaim**, a final equation that the
    /// outer protocol must satisfy for the overall proof to be valid.
    ///
    /// ---
    ///
    /// # Arguments
    ///
    /// * `transcript`: A mutable reference to a Fiat-Shamir `Transcript`.
    /// * `num_vars`: The number of variables over which the sum was originally
    ///   computed.
    /// * `degree`: The maximum combined degree of the underlying polynomial
    ///   $G(x)$. This must match the degree used by the Prover.
    /// * `proof`: A reference to the `SumcheckProof<F>` generated by the
    ///   Prover.
    /// * `config`: The configuration for the underlying field used in the
    ///   protocol.
    ///
    /// ---
    ///
    /// # Returns
    ///
    /// A `Result` which is:
    ///
    /// * `Ok(Subclaim<F>)`: If the Sumcheck protocol passes successfully, it
    ///   returns a `Subclaim`. This claim consists of:
    ///     1. The final random challenge point $r \in
    ///        \text{F}^{\text{num\\_vars}}$.
    ///     2. The expected evaluation $v$ of the combined polynomial $G(r)$ at
    ///        that point.
    ///
    /// * `Err(SumCheckError<F>)`: If any of the round checks fail during the
    ///   protocol.
    ///
    /// ---
    ///
    /// # Panics
    ///
    /// * Panics if the number of variables is `0`.
    pub fn verify_as_subprotocol(
        transcript: &mut impl Transcript,
        num_vars: usize,
        degree: usize,
        proof: &SumcheckProof<C::Element>,
        config: &C,
    ) -> Result<Subclaim<C::Element>, SumCheckError<C::Element>> {
        if num_vars == 0 {
            panic!("Attempt to verify a sumcheck claim for 0 variables")
        }

        let mut buf = vec![0; <C::Integer as ConstTranscribable>::NUM_BYTES];

        let (nvars_field, degree_field) = {
            (
                config.project(&(num_vars as u64)),
                config.project(&(degree as u64)),
            )
        };
        transcript.absorb_field_element(config, &nvars_field, &mut buf);
        transcript.absorb_field_element(config, &degree_field, &mut buf);

        if proof.messages.len() != num_vars {
            return Err(SumCheckError::InvalidProofLength {
                expected: num_vars,
                got: proof.messages.len(),
            });
        }

        let mut verifier_state = VerifierState::new(num_vars, degree, config);

        for i in 0..num_vars {
            let prover_msg = &proof.messages[i];
            transcript.absorb_field_element_slice(config, &prover_msg.0.tail_evaluations, &mut buf);
            let verifier_msg = verifier_state.verify_round(prover_msg, transcript);
            transcript.absorb_field_element(config, &verifier_msg, &mut buf);
        }

        verifier_state.check_and_generate_subclaim(proof.claimed_sum.clone())
    }
}

#[derive(Error, Debug)]
pub enum SumCheckError<F> {
    #[error("univariate polynomial evaluation error")]
    EvaluationError(ArithErrors),
    #[error("incorrect sumcheck sum at round {0}. Expected `{1}`. Received `{2}`")]
    SumCheckFailed(usize, Box<F>, Box<F>),
    #[error("max degree exceeded")]
    MaxDegreeExceeded,
    #[error("invalid proof length: expected {expected}, got {got}")]
    InvalidProofLength { expected: usize, got: usize },
    #[error("verifier failed to evaluate a round polynomial: {0}")]
    UnivariateEvaluationError(EvaluationError),
}

impl<F> From<ArithErrors> for SumCheckError<F> {
    fn from(arith_error: ArithErrors) -> Self {
        Self::EvaluationError(arith_error)
    }
}

impl<F> From<EvaluationError> for SumCheckError<F> {
    fn from(error: EvaluationError) -> Self {
        Self::UnivariateEvaluationError(error)
    }
}