spartan2 0.2.0

//! This module implements the sum-check protocol used in the Spartan SNARK.
//!
//! The sum-check protocol allows a prover to convince a verifier that a claimed sum
//! over a multivariate polynomial equals a specific value, without the verifier
//! needing to compute the sum directly.

use crate::{
  errors::SpartanError,
  polys::{
    multilinear::MultilinearPolynomial,
    univariate::{CompressedUniPoly, UniPoly},
  },
  start_span,
  traits::{Engine, transcript::TranscriptEngineTrait},
};
use ff::Field;
use rayon::prelude::*;
use serde::{Deserialize, Serialize};
use std::time::Instant;
use tracing::{info, info_span};

/// 4 k elements is a good cut-off on a 16-core machine.
const PAR_THRESHOLD: usize = 4 << 10; // 4096

/// Adaptive parallel-for helper.
/// Falls back to a plain `for` loop when the slice is small **or**
/// we are already inside a Rayon worker, avoiding nested pools.
/// `par_for` – run `map(i)` for `i = 0..len` and fold the
/// results with `reduce`, starting from `identity()`.
///
/// When `len` is small or we’re already on a Rayon thread it executes
/// serially; otherwise it uses `into_par_iter`.
pub fn par_for<R, Map, Red, Id>(len: usize, map: Map, reduce: Red, identity: Id) -> R
where
  R: Send, // result must cross Rayon threads
  Map: Fn(usize) -> R + Sync + Send,
  Red: Fn(R, R) -> R + Sync + Send,
  Id: Fn() -> R + Sync + Send,
{
  // Fast-path for empty ranges
  if len == 0 {
    return identity();
  }

  // Are we *already* running inside a Rayon worker thread?
  let in_rayon_ctx = rayon::current_thread_index().is_some();

  if len < PAR_THRESHOLD || in_rayon_ctx {
    // ---------- serial fallback ----------
    let mut acc = identity();
    for i in 0..len {
      let v = map(i);
      acc = reduce(acc, v);
    }
    acc
  } else {
    // ---------- true parallel path ----------
    (0..len).into_par_iter().map(map).reduce(identity, reduce)
  }
}

/// A proof generated by the sum-check protocol.
///
/// This struct contains the compressed univariate polynomials that constitute
/// the prover's messages in each round of the sum-check protocol.
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(bound = "")]
pub struct SumcheckProof<E: Engine> {
  compressed_polys: Vec<CompressedUniPoly<E::Scalar>>,
}

impl<E: Engine> SumcheckProof<E> {
  /// Verifies a sum-check proof.
  ///
  /// # Arguments
  /// * `claim` - The claimed sum that the prover asserts
  /// * `num_rounds` - The number of rounds in the sum-check protocol
  /// * `degree_bound` - The maximum degree of univariate polynomials in each round
  /// * `transcript` - The transcript for the interactive proof
  ///
  /// # Returns
  /// A tuple containing the final evaluation point and the sequence of verifier challenges,
  /// or an error if verification fails.
  pub fn verify(
    &self,
    claim: E::Scalar,
    num_rounds: usize,
    degree_bound: usize,
    transcript: &mut E::TE,
  ) -> Result<(E::Scalar, Vec<E::Scalar>), SpartanError> {
    let (_verify_span, verify_t) = start_span!("sumcheck_verify");
    let mut e = claim;
    let mut r: Vec<E::Scalar> = Vec::new();

    // verify that there is a univariate polynomial for each round
    if self.compressed_polys.len() != num_rounds {
      return Err(SpartanError::InvalidSumcheckProof);
    }

    for i in 0..self.compressed_polys.len() {
      let (_round_span, round_t) = start_span!("sumcheck_verify_round", round = i);
      let poly = self.compressed_polys[i].decompress(&e);

      // verify degree bound
      if poly.degree() != degree_bound {
        return Err(SpartanError::InvalidSumcheckProof);
      }

      // we do not need to check if poly(0) + poly(1) = e, as
      // decompress() call above already ensures that holds
      debug_assert_eq!(poly.eval_at_zero() + poly.eval_at_one(), e);

      // append the prover's message to the transcript
      transcript.absorb(b"p", &poly);

      //derive the verifier's challenge for the next round
      let r_i = transcript.squeeze(b"c")?;

      r.push(r_i);

      // evaluate the claimed degree-ell polynomial at r_i
      e = poly.evaluate(&r_i);

      if round_t.elapsed().as_millis() > 0 {
        info!(elapsed_ms = %round_t.elapsed().as_millis(), "sumcheck_verify_round");
      }
    }

    info!(elapsed_ms = %verify_t.elapsed().as_millis(), "sumcheck_verify");
    Ok((e, r))
  }

  #[inline]
  fn compute_eval_points_quad<F>(
    poly_A: &MultilinearPolynomial<E::Scalar>,
    poly_B: &MultilinearPolynomial<E::Scalar>,
    comb_func: &F,
  ) -> (E::Scalar, E::Scalar)
  where
    F: Fn(&E::Scalar, &E::Scalar) -> E::Scalar + Sync,
  {
    let len = poly_A.Z.len() / 2;

    // Using `par_for` keeps the map-reduce logic identical
    // but avoids Rayon overhead on tiny slices.
    par_for(
      len,
      // map-closure (returned pair is the per-index contribution)
      |i| {
        let a_low = poly_A[i];
        let a_high = poly_A[len + i];
        let b_low = poly_B[i];
        let b_high = poly_B[len + i];

        // eval 0:   A(low)
        let eval0 = comb_func(&a_low, &b_low);

        // eval 2:  −A(low) + 2·A(high)   (same for B)
        let a_bound = a_high + a_high - a_low;
        let b_bound = b_high + b_high - b_low;
        let eval2 = comb_func(&a_bound, &b_bound);

        (eval0, eval2)
      },
      // reduce-closure (pairwise accumulation)
      |mut acc, val| {
        acc.0 += val.0;
        acc.1 += val.1;
        acc
      },
      // identity value
      || (E::Scalar::ZERO, E::Scalar::ZERO),
    )
  }

  /// Generates a sum-check proof for a quadratic combination of two multilinear polynomials.
  ///
  /// # Arguments
  /// * `claim` - The claimed sum over the hypercube
  /// * `num_rounds` - The number of variables/rounds in the sum-check
  /// * `poly_A` - First multilinear polynomial (mutable, will be bound during protocol)
  /// * `poly_B` - Second multilinear polynomial (mutable, will be bound during protocol)  
  /// * `comb_func` - Function that combines evaluations of the two polynomials
  /// * `transcript` - The transcript for generating randomness
  ///
  /// # Returns
  /// A tuple containing the sum-check proof, the sequence of verifier challenges,
  /// and the final evaluations of the polynomials.
  pub fn prove_quad<F>(
    claim: &E::Scalar,
    num_rounds: usize,
    poly_A: &mut MultilinearPolynomial<E::Scalar>,
    poly_B: &mut MultilinearPolynomial<E::Scalar>,
    comb_func: F,
    transcript: &mut E::TE,
  ) -> Result<(Self, Vec<E::Scalar>, Vec<E::Scalar>), SpartanError>
  where
    F: Fn(&E::Scalar, &E::Scalar) -> E::Scalar + Sync,
  {
    let mut r: Vec<E::Scalar> = Vec::new();
    let mut polys: Vec<CompressedUniPoly<E::Scalar>> = Vec::new();
    let mut claim_per_round = *claim;
    for round in 0..num_rounds {
      let (_round_span, round_t) = start_span!("sumcheck_quad_round", round = round);

      let poly = {
        let (_eval_span, eval_t) = start_span!("compute_eval_points_quad");
        let (eval_point_0, eval_point_2) =
          Self::compute_eval_points_quad(poly_A, poly_B, &comb_func);
        if eval_t.elapsed().as_millis() > 0 {
          info!(elapsed_ms = %eval_t.elapsed().as_millis(), "compute_eval_points_quad");
        }

        let evals = vec![eval_point_0, claim_per_round - eval_point_0, eval_point_2];
        UniPoly::from_evals(&evals)?
      };

      // append the prover's message to the transcript
      transcript.absorb(b"p", &poly);

      //derive the verifier's challenge for the next round
      let r_i = transcript.squeeze(b"c")?;
      r.push(r_i);
      polys.push(poly.compress());

      // Set up next round
      claim_per_round = poly.evaluate(&r_i);

      // bind all tables to the verifier's challenge
      let (_bind_span, bind_t) = start_span!("bind_poly_vars_quad");
      rayon::join(
        || poly_A.bind_poly_var_top(&r_i),
        || poly_B.bind_poly_var_top(&r_i),
      );
      info!(elapsed_ms = %bind_t.elapsed().as_millis(), "bind_poly_vars_quad");
      info!(elapsed_ms = %round_t.elapsed().as_millis(), round = round, "sumcheck_quad_round");
    }

    Ok((
      SumcheckProof {
        compressed_polys: polys,
      },
      r,
      vec![poly_A[0], poly_B[0]],
    ))
  }

  #[inline]
  /// Computes evaluation points for a cubic polynomial with additive term.
  ///
  /// This function computes three evaluation points (at 0, 2, and 3) for a univariate
  /// polynomial that represents the sum over a hypercube edge in the sum-check protocol
  /// for a cubic combination of four multilinear polynomials.
  ///
  /// # Arguments
  /// * `poly_A` - First multilinear polynomial
  /// * `poly_B` - Second multilinear polynomial  
  /// * `poly_C` - Third multilinear polynomial
  /// * `poly_D` - Fourth multilinear polynomial
  /// * `comb_func` - Function that combines evaluations of the four polynomials
  ///
  /// # Returns
  /// A tuple containing the evaluations at points 0, 2, and 3.
  pub fn compute_eval_points_cubic_with_additive_term<F>(
    poly_A: &MultilinearPolynomial<E::Scalar>,
    poly_B: &MultilinearPolynomial<E::Scalar>,
    poly_C: &MultilinearPolynomial<E::Scalar>,
    poly_D: &MultilinearPolynomial<E::Scalar>,
    comb_func: &F,
  ) -> (E::Scalar, E::Scalar, E::Scalar)
  where
    F: Fn(&E::Scalar, &E::Scalar, &E::Scalar, &E::Scalar) -> E::Scalar + Sync,
  {
    let len = poly_A.Z.len() / 2;
    par_for(
      len,
      |i| {
        let a_low = poly_A[i];
        let a_high = poly_A[i + len];
        let b_low = poly_B[i];
        let b_high = poly_B[i + len];
        let c_low = poly_C[i];
        let c_high = poly_C[i + len];
        let d_low = poly_D[i];
        let d_high = poly_D[i + len];

        // eval 0: bound_func is A(low)
        let eval_point_0 = comb_func(&a_low, &b_low, &c_low, &d_low);

        // eval 2: bound_func is -A(low) + 2*A(high)
        let poly_A_bound_point = a_high + a_high - a_low;
        let poly_B_bound_point = b_high + b_high - b_low;
        let poly_C_bound_point = c_high + c_high - c_low;
        let poly_D_bound_point = d_high + d_high - d_low;
        let eval_point_2 = comb_func(
          &poly_A_bound_point,
          &poly_B_bound_point,
          &poly_C_bound_point,
          &poly_D_bound_point,
        );

        // eval 3: bound_func is -2A(low) + 3A(high); computed incrementally with bound_func applied to eval(2)
        let poly_A_bound_point = poly_A_bound_point + a_high - a_low;
        let poly_B_bound_point = poly_B_bound_point + b_high - b_low;
        let poly_C_bound_point = poly_C_bound_point + c_high - c_low;
        let poly_D_bound_point = poly_D_bound_point + d_high - d_low;
        let eval_point_3 = comb_func(
          &poly_A_bound_point,
          &poly_B_bound_point,
          &poly_C_bound_point,
          &poly_D_bound_point,
        );
        (eval_point_0, eval_point_2, eval_point_3)
      },
      |mut acc, val| {
        acc.0 += val.0;
        acc.1 += val.1;
        acc.2 += val.2;
        acc
      },
      || (E::Scalar::ZERO, E::Scalar::ZERO, E::Scalar::ZERO),
    )
  }

  /// Generates a sum-check proof for a cubic combination with additive term of four multilinear polynomials.
  ///
  /// # Arguments
  /// * `claim` - The claimed sum over the hypercube
  /// * `num_rounds` - The number of variables/rounds in the sum-check
  /// * `poly_A` - First multilinear polynomial (mutable, will be bound during protocol)
  /// * `poly_B` - Second multilinear polynomial (mutable, will be bound during protocol)
  /// * `poly_C` - Third multilinear polynomial (mutable, will be bound during protocol)
  /// * `poly_D` - Fourth multilinear polynomial (mutable, will be bound during protocol)
  /// * `comb_func` - Function that combines evaluations of the four polynomials
  /// * `transcript` - The transcript for generating randomness
  ///
  /// # Returns
  /// A tuple containing the sum-check proof, the sequence of verifier challenges,
  /// and the final evaluations of the polynomials.
  #[allow(clippy::too_many_arguments)]
  pub fn prove_cubic_with_additive_term<F>(
    claim: &E::Scalar,
    num_rounds: usize,
    poly_A: &mut MultilinearPolynomial<E::Scalar>,
    poly_B: &mut MultilinearPolynomial<E::Scalar>,
    poly_C: &mut MultilinearPolynomial<E::Scalar>,
    poly_D: &mut MultilinearPolynomial<E::Scalar>,
    comb_func: F,
    transcript: &mut E::TE,
  ) -> Result<(Self, Vec<E::Scalar>, Vec<E::Scalar>), SpartanError>
  where
    F: Fn(&E::Scalar, &E::Scalar, &E::Scalar, &E::Scalar) -> E::Scalar + Sync,
  {
    let mut r: Vec<E::Scalar> = Vec::new();
    let mut polys: Vec<CompressedUniPoly<E::Scalar>> = Vec::new();
    let mut claim_per_round = *claim;

    for round in 0..num_rounds {
      let (_round_span, round_t) = start_span!("sumcheck_round", round = round);

      let poly = {
        // Make an iterator returning the contributions to the evaluations
        let (_eval_span, eval_t) = start_span!("compute_eval_points");
        let (eval_point_0, eval_point_2, eval_point_3) =
          Self::compute_eval_points_cubic_with_additive_term(
            poly_A, poly_B, poly_C, poly_D, &comb_func,
          );
        if eval_t.elapsed().as_millis() > 0 {
          info!(elapsed_ms = %eval_t.elapsed().as_millis(), "compute_eval_points");
        }
        let evals = vec![
          eval_point_0,
          claim_per_round - eval_point_0,
          eval_point_2,
          eval_point_3,
        ];
        UniPoly::from_evals(&evals)?
      };

      // append the prover's message to the transcript
      transcript.absorb(b"p", &poly);

      //derive the verifier's challenge for the next round
      let r_i = transcript.squeeze(b"c")?;
      r.push(r_i);
      polys.push(poly.compress());

      // Set up next round
      claim_per_round = poly.evaluate(&r_i);

      // bound all tables to the verifier's challenge
      let (_bind_span, bind_t) = start_span!("bind_poly_vars");
      rayon::join(
        || {
          rayon::join(
            || poly_A.bind_poly_var_top(&r_i),
            || poly_B.bind_poly_var_top(&r_i),
          )
        },
        || {
          rayon::join(
            || poly_C.bind_poly_var_top(&r_i),
            || poly_D.bind_poly_var_top(&r_i),
          )
        },
      );
      info!(elapsed_ms = %bind_t.elapsed().as_millis(), "bind_poly_vars");
      info!(elapsed_ms = %round_t.elapsed().as_millis(), round = round, "sumcheck_round");
    }

    Ok((
      SumcheckProof {
        compressed_polys: polys,
      },
      r,
      vec![poly_A[0], poly_B[0], poly_C[0], poly_D[0]],
    ))
  }
}