spartan2 0.5.0

//! This module implements the sum-check protocol used in Spartan.
//!
//! 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::{
  CommitmentKey,
  bellpepper::{
    r1cs::{MultiRoundSpartanWitness, MultiRoundState},
    solver::SatisfyingAssignment,
  },
  errors::SpartanError,
  polys::{
    multilinear::MultilinearPolynomial,
    univariate::{CompressedUniPoly, UniPoly},
  },
  r1cs::SplitMultiRoundR1CSShape,
  start_span,
  traits::{Engine, transcript::TranscriptEngineTrait},
  zk::{NeutronNovaVerifierCircuit, SpartanVerifierCircuit},
};
use ff::Field;
use rayon::prelude::*;
use serde::{Deserialize, Serialize};
use tracing::info;

/// 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`.
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]],
    ))
  }

  /// Proves a batch of quadratic combinations of two multilinear polynomials.
  pub fn prove_quad_batched<F>(
    challenge: &E::Scalar,
    claims: &[E::Scalar; 2],
    num_rounds: usize,
    poly_A_0: &mut MultilinearPolynomial<E::Scalar>,
    poly_A_1: &mut MultilinearPolynomial<E::Scalar>,
    poly_B_0: &mut MultilinearPolynomial<E::Scalar>,
    poly_B_1: &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();

    // compute the joint claim for both sum-check instances
    let mut claim_per_round = claims[0] + *challenge * claims[1];

    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_0, eval_point_2_0) =
          Self::compute_eval_points_quad(poly_A_0, poly_B_0, &comb_func);
        let (eval_point_0_1, eval_point_2_1) =
          Self::compute_eval_points_quad(poly_A_1, poly_B_1, &comb_func);

        let eval_point_0 = eval_point_0_0 + *challenge * eval_point_0_1;
        let eval_point_2 = eval_point_2_0 + *challenge * eval_point_2_1;

        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(
        || {
          rayon::join(
            || poly_A_0.bind_poly_var_top(&r_i),
            || poly_B_0.bind_poly_var_top(&r_i),
          )
        },
        || {
          rayon::join(
            || poly_A_1.bind_poly_var_top(&r_i),
            || poly_B_1.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[0], poly_A_1[0], poly_B_0[0], poly_B_1[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.
  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.
  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]],
    ))
  }

  /// Executes the **outer** cubic-with-additive-term sum-check in
  /// Zero-knowledge outer sum-check for the cubic-with-additive-term case.
  pub fn prove_cubic_with_additive_term_zk(
    num_rounds: usize,
    poly_tau: &mut MultilinearPolynomial<E::Scalar>,
    poly_Az: &mut MultilinearPolynomial<E::Scalar>,
    poly_Bz: &mut MultilinearPolynomial<E::Scalar>,
    poly_Cz: &mut MultilinearPolynomial<E::Scalar>,
    verifier_circuit: &mut SpartanVerifierCircuit<E>,
    state: &mut MultiRoundState<E>,
    vc_shape: &SplitMultiRoundR1CSShape<E>,
    vc_ck: &CommitmentKey<E>,
    transcript: &mut E::TE,
  ) -> Result<Vec<E::Scalar>, SpartanError> {
    let mut r_x: Vec<E::Scalar> = Vec::with_capacity(num_rounds);
    let mut claim_outer_round = E::Scalar::ZERO;

    for i in 0..num_rounds {
      // -------- interpolate coefficients --------

      let comb = |a: &E::Scalar, b: &E::Scalar, c: &E::Scalar, d: &E::Scalar| -> E::Scalar {
        *a * (*b * *c - *d)
      };
      let (eval0, eval2, eval3) = Self::compute_eval_points_cubic_with_additive_term(
        poly_tau, poly_Az, poly_Bz, poly_Cz, &comb,
      );
      let evals = vec![eval0, claim_outer_round - eval0, eval2, eval3];
      let poly = UniPoly::from_evals(&evals)?;
      verifier_circuit.outer_polys[i] = [
        poly.coeffs[0],
        poly.coeffs[1],
        poly.coeffs[2],
        poly.coeffs[3],
      ];

      // -------- transcript / witness handling --------
      let chals = SatisfyingAssignment::<E>::process_round(
        state,
        vc_shape,
        vc_ck,
        verifier_circuit,
        i,
        transcript,
      )?;
      r_x.push(chals[0]);

      // -------- advance claim and bind polys --------
      claim_outer_round = poly.evaluate(&chals[0]);

      rayon::join(
        || poly_tau.bind_poly_var_top(&chals[0]),
        || {
          rayon::join(
            || poly_Az.bind_poly_var_top(&chals[0]),
            || poly_Bz.bind_poly_var_top(&chals[0]),
          );
          poly_Cz.bind_poly_var_top(&chals[0]);
        },
      );
    }

    Ok(r_x)
  }

  /// Executes a **quadratic** sum-check in zero-knowledge mode and returns the
  /// Zero-knowledge quadratic sum-check used for the inner round.
  pub fn prove_quad_zk(
    claim: &E::Scalar,
    num_rounds: usize,
    poly_ABC: &mut MultilinearPolynomial<E::Scalar>,
    poly_z: &mut MultilinearPolynomial<E::Scalar>,
    verifier_circuit: &mut SpartanVerifierCircuit<E>,
    state: &mut MultiRoundState<E>,
    vc_shape: &SplitMultiRoundR1CSShape<E>,
    vc_ck: &CommitmentKey<E>,
    transcript: &mut E::TE,
    start_round: usize,
  ) -> Result<(Vec<E::Scalar>, Vec<E::Scalar>), SpartanError> {
    let mut r_y: Vec<E::Scalar> = Vec::with_capacity(num_rounds);
    let mut claim_current_round = *claim;

    for j in 0..num_rounds {
      // -------- interpolate coeffs --------
      let comb = |a: &E::Scalar, b: &E::Scalar| -> E::Scalar { *a * *b };
      let (eval0, eval2) = Self::compute_eval_points_quad(poly_ABC, poly_z, &comb);
      let evals = vec![eval0, claim_current_round - eval0, eval2];
      let poly = UniPoly::from_evals(&evals)?;

      verifier_circuit.inner_polys[j] = [poly.coeffs[0], poly.coeffs[1], poly.coeffs[2]];

      // -------- transcript / witness handling --------
      let chals = SatisfyingAssignment::<E>::process_round(
        state,
        vc_shape,
        vc_ck,
        verifier_circuit,
        start_round + j,
        transcript,
      )?;
      r_y.push(chals[0]);

      // -------- bind polys --------
      rayon::join(
        || poly_ABC.bind_poly_var_top(&chals[0]),
        || poly_z.bind_poly_var_top(&chals[0]),
      );

      // -------- advance claim for next round --------
      claim_current_round = poly.evaluate(&chals[0]);
    }

    Ok((r_y, vec![poly_ABC[0], poly_z[0]]))
  }

  /// Generates a batched sum-check proof for a cubic combination with additive term of four multilinear polynomials.
  pub fn prove_cubic_with_additive_term_batched<F>(
    challenge: &E::Scalar,
    claims: &[E::Scalar; 2],
    num_rounds: usize,
    poly_A: &mut MultilinearPolynomial<E::Scalar>,
    poly_B_0: &mut MultilinearPolynomial<E::Scalar>,
    poly_B_1: &mut MultilinearPolynomial<E::Scalar>,
    poly_C_0: &mut MultilinearPolynomial<E::Scalar>,
    poly_C_1: &mut MultilinearPolynomial<E::Scalar>,
    poly_D_0: &mut MultilinearPolynomial<E::Scalar>,
    poly_D_1: &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();

    // compute the joint claim for both sum-check instances
    let mut claim_per_round = claims[0] + *challenge * claims[1];

    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_0, eval_point_2_0, eval_point_3_0) =
          Self::compute_eval_points_cubic_with_additive_term(
            poly_A, poly_B_0, poly_C_0, poly_D_0, &comb_func,
          );
        let (eval_point_0_1, eval_point_2_1, eval_point_3_1) =
          Self::compute_eval_points_cubic_with_additive_term(
            poly_A, poly_B_1, poly_C_1, poly_D_1, &comb_func,
          );

        let eval_point_0 = eval_point_0_0 + *challenge * eval_point_0_1;
        let eval_point_2 = eval_point_2_0 + *challenge * eval_point_2_1;
        let eval_point_3 = eval_point_3_0 + *challenge * eval_point_3_1;

        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_0.bind_poly_var_top(&r_i),
          );
        },
        || {
          rayon::join(
            || {
              rayon::join(
                || poly_B_1.bind_poly_var_top(&r_i),
                || poly_C_0.bind_poly_var_top(&r_i),
              );
            },
            || {
              rayon::join(
                || poly_C_1.bind_poly_var_top(&r_i),
                || {
                  rayon::join(
                    || poly_D_0.bind_poly_var_top(&r_i),
                    || poly_D_1.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[0],
        poly_B_1[0],
        poly_C_0[0],
        poly_C_1[0],
        poly_D_0[0],
        poly_D_1[0],
      ],
    ))
  }

  /// Executes a **quadratic** batched sum-check in zero-knowledge mode and returns the
  /// sequence of verifier challenges used for the inner round.
  pub fn prove_quad_batched_zk(
    claims: &[E::Scalar; 2],
    num_rounds: usize,
    poly_A_0: &mut MultilinearPolynomial<E::Scalar>,
    poly_A_1: &mut MultilinearPolynomial<E::Scalar>,
    poly_B_0: &mut MultilinearPolynomial<E::Scalar>,
    poly_B_1: &mut MultilinearPolynomial<E::Scalar>,
    verifier_circuit: &mut NeutronNovaVerifierCircuit<E>,
    state: &mut MultiRoundState<E>,
    vc_shape: &SplitMultiRoundR1CSShape<E>,
    vc_ck: &CommitmentKey<E>,
    transcript: &mut E::TE,
    start_round: usize,
  ) -> Result<(Vec<E::Scalar>, Vec<E::Scalar>), SpartanError> {
    let mut r_y: Vec<E::Scalar> = Vec::with_capacity(num_rounds);
    // Maintain separate claims for step and core branches
    let mut claim_step_round = claims[0];
    let mut claim_core_round = claims[1];

    for j in 0..num_rounds {
      // -------- interpolate coeffs --------
      let comb = |a: &E::Scalar, b: &E::Scalar| -> E::Scalar { *a * *b };

      let ((eval0_s, eval2_s), (eval0_c, eval2_c)) = rayon::join(
        || Self::compute_eval_points_quad(poly_A_0, poly_B_0, &comb),
        || Self::compute_eval_points_quad(poly_A_1, poly_B_1, &comb),
      );

      // step branch
      let evals_s = vec![eval0_s, claim_step_round - eval0_s, eval2_s];
      let poly_s = UniPoly::from_evals(&evals_s)?;
      let coeffs_step = [poly_s.coeffs[0], poly_s.coeffs[1], poly_s.coeffs[2]];

      // core branch
      let evals_c = vec![eval0_c, claim_core_round - eval0_c, eval2_c];
      let poly_c = UniPoly::from_evals(&evals_c)?;
      let coeffs_core = [poly_c.coeffs[0], poly_c.coeffs[1], poly_c.coeffs[2]];

      verifier_circuit.inner_polys_step[j] = coeffs_step;
      verifier_circuit.inner_polys_core[j] = coeffs_core;

      // -------- transcript / witness handling --------
      let chals = SatisfyingAssignment::<E>::process_round(
        state,
        vc_shape,
        vc_ck,
        verifier_circuit,
        start_round + j,
        transcript,
      )?;
      let r_j = chals[0];
      r_y.push(r_j);

      // -------- bind polys --------
      rayon::join(
        || {
          rayon::join(
            || poly_A_0.bind_poly_var_top(&r_j),
            || poly_B_0.bind_poly_var_top(&r_j),
          );
        },
        || {
          rayon::join(
            || poly_A_1.bind_poly_var_top(&r_j),
            || poly_B_1.bind_poly_var_top(&r_j),
          );
        },
      );

      // -------- advance claim for next round --------
      claim_step_round = poly_s.evaluate(&r_j);
      claim_core_round = poly_c.evaluate(&r_j);
    }

    Ok((
      r_y,
      vec![poly_A_0[0], poly_A_1[0], poly_B_0[0], poly_B_1[0]],
    ))
  }

  /// Executes a **cubic-with-additive-term** batched outer sum-check in zero-knowledge mode
  /// and returns the sequence of verifier challenges.
  pub fn prove_cubic_with_additive_term_batched_zk(
    num_rounds: usize,
    poly_A: &mut MultilinearPolynomial<E::Scalar>,
    poly_B_step: &mut MultilinearPolynomial<E::Scalar>,
    poly_B_core: &mut MultilinearPolynomial<E::Scalar>,
    poly_C_step: &mut MultilinearPolynomial<E::Scalar>,
    poly_C_core: &mut MultilinearPolynomial<E::Scalar>,
    poly_D_step: &mut MultilinearPolynomial<E::Scalar>,
    poly_D_core: &mut MultilinearPolynomial<E::Scalar>,
    verifier_circuit: &mut NeutronNovaVerifierCircuit<E>,
    state: &mut MultiRoundState<E>,
    vc_shape: &SplitMultiRoundR1CSShape<E>,
    vc_ck: &CommitmentKey<E>,
    transcript: &mut E::TE,
    start_round: usize,
  ) -> Result<Vec<E::Scalar>, SpartanError> {
    let mut r_x: Vec<E::Scalar> = Vec::with_capacity(num_rounds);

    let mut claim_step = verifier_circuit.t_out_step;
    let mut claim_core = E::Scalar::ZERO;

    for i in 0..num_rounds {
      let comb = |a: &E::Scalar, b: &E::Scalar, c: &E::Scalar, d: &E::Scalar| -> E::Scalar {
        *a * (*b * *c - *d)
      };

      // step branch
      let ((eval0_s, eval2_s, eval3_s), (eval0_c, eval2_c, eval3_c)) = rayon::join(
        || {
          Self::compute_eval_points_cubic_with_additive_term(
            poly_A,
            poly_B_step,
            poly_C_step,
            poly_D_step,
            &comb,
          )
        },
        || {
          Self::compute_eval_points_cubic_with_additive_term(
            poly_A,
            poly_B_core,
            poly_C_core,
            poly_D_core,
            &comb,
          )
        },
      );

      let evals_s = vec![eval0_s, claim_step - eval0_s, eval2_s, eval3_s];
      let poly_s = UniPoly::from_evals(&evals_s)?;
      let coeffs_step = [
        poly_s.coeffs[0],
        poly_s.coeffs[1],
        poly_s.coeffs[2],
        poly_s.coeffs[3],
      ];

      let evals_c = vec![eval0_c, claim_core - eval0_c, eval2_c, eval3_c];
      let poly_c = UniPoly::from_evals(&evals_c)?;
      let coeffs_core = [
        poly_c.coeffs[0],
        poly_c.coeffs[1],
        poly_c.coeffs[2],
        poly_c.coeffs[3],
      ];

      verifier_circuit.outer_polys_step[i] = coeffs_step;
      verifier_circuit.outer_polys_core[i] = coeffs_core;

      // -------- transcript / witness handling --------
      let chals = SatisfyingAssignment::<E>::process_round(
        state,
        vc_shape,
        vc_ck,
        verifier_circuit,
        start_round + i,
        transcript,
      )?;
      let r_i = chals[0];
      r_x.push(r_i);

      // -------- advance claim and bind polys --------
      claim_step = poly_s.evaluate(&r_i);
      claim_core = poly_c.evaluate(&r_i);

      // bind polynomials to the verifier's challenge
      rayon::join(
        || poly_A.bind_poly_var_top(&r_i),
        || {
          rayon::join(
            || {
              rayon::join(
                || poly_B_step.bind_poly_var_top(&r_i),
                || poly_B_core.bind_poly_var_top(&r_i),
              );
            },
            || {
              rayon::join(
                || {
                  rayon::join(
                    || poly_C_step.bind_poly_var_top(&r_i),
                    || poly_C_core.bind_poly_var_top(&r_i),
                  );
                },
                || {
                  rayon::join(
                    || poly_D_step.bind_poly_var_top(&r_i),
                    || poly_D_core.bind_poly_var_top(&r_i),
                  );
                },
              );
            },
          );
        },
      );
    }

    Ok(r_x)
  }
}