1// Copyright (c) 2023 Espresso Systems (espressosys.com)
  2// This file is part of the HyperPlonk library.
  3
  4// You should have received a copy of the MIT License
  5// along with the HyperPlonk library. If not, see <https://mit-license.org/>.
  6
  7//! Sumcheck based batch opening and verify commitment.
  8// TODO: refactoring this code to somewhere else
  9// currently IOP depends on PCS because perm check requires commitment.
 10// The sumcheck based batch opening therefore cannot stay in the PCS repo --
 11// which creates a cyclic dependency.
 12
 13use crate::{
 14    pcs::{
 15        multilinear_kzg::util::eq_eval,
 16        prelude::{Commitment, PCSError},
 17        PolynomialCommitmentScheme,
 18    },
 19    poly_iop::{prelude::SumCheck, PolyIOP},
 20    IOPProof,
 21};
 22use arithmetic::{build_eq_x_r_vec, DenseMultilinearExtension, VPAuxInfo, VirtualPolynomial};
 23use ark_ec::{msm::VariableBaseMSM, PairingEngine, ProjectiveCurve};
 24use ark_ff::PrimeField;
 25use ark_std::{end_timer, log2, start_timer, One, Zero};
 26use std::{collections::BTreeMap, iter, marker::PhantomData, ops::Deref, sync::Arc};
 27use transcript::IOPTranscript;
 28
 29#[derive(Clone, Debug, Default, PartialEq, Eq)]
 30pub struct BatchProof<E, PCS>
 31where
 32    E: PairingEngine,
 33    PCS: PolynomialCommitmentScheme<E>,
 34{
 35    /// A sum check proof proving tilde g's sum
 36    pub(crate) sum_check_proof: IOPProof<E::Fr>,
 37    /// f_i(point_i)
 38    pub f_i_eval_at_point_i: Vec<E::Fr>,
 39    /// proof for g'(a_2)
 40    pub(crate) g_prime_proof: PCS::Proof,
 41}
 42
 43/// Steps:
 44/// 1. get challenge point t from transcript
 45/// 2. build eq(t,i) for i in [0..k]
 46/// 3. build \tilde g_i(b) = eq(t, i) * f_i(b)
 47/// 4. compute \tilde eq_i(b) = eq(b, point_i)
 48/// 5. run sumcheck on \sum_i=1..k \tilde eq_i * \tilde g_i
 49/// 6. build g'(X) = \sum_i=1..k \tilde eq_i(a2) * \tilde g_i(X) where (a2) is
 50/// the sumcheck's point 7. open g'(X) at point (a2)
 51pub(crate) fn multi_open_internal<E, PCS>(
 52    prover_param: &PCS::ProverParam,
 53    polynomials: &[PCS::Polynomial],
 54    points: &[PCS::Point],
 55    evals: &[PCS::Evaluation],
 56    transcript: &mut IOPTranscript<E::Fr>,
 57) -> Result<BatchProof<E, PCS>, PCSError>
 58where
 59    E: PairingEngine,
 60    PCS: PolynomialCommitmentScheme<
 61        E,
 62        Polynomial = Arc<DenseMultilinearExtension<E::Fr>>,
 63        Point = Vec<E::Fr>,
 64        Evaluation = E::Fr,
 65    >,
 66{
 67    let open_timer = start_timer!(|| format!("multi open {} points", points.len()));
 68
 69    // TODO: sanity checks
 70    let num_var = polynomials[0].num_vars;
 71    let k = polynomials.len();
 72    let ell = log2(k) as usize;
 73
 74    // challenge point t
 75    let t = transcript.get_and_append_challenge_vectors("t".as_ref(), ell)?;
 76
 77    // eq(t, i) for i in [0..k]
 78    let eq_t_i_list = build_eq_x_r_vec(t.as_ref())?;
 79
 80    // \tilde g_i(b) = eq(t, i) * f_i(b)
 81    let timer = start_timer!(|| format!("compute tilde g for {} points", points.len()));
 82    // combine the polynomials that have same opening point first to reduce the
 83    // cost of sum check later.
 84    let point_indices = points
 85        .iter()
 86        .fold(BTreeMap::<_, _>::new(), |mut indices, point| {
 87            let idx = indices.len();
 88            indices.entry(point).or_insert(idx);
 89            indices
 90        });
 91    let deduped_points =
 92        BTreeMap::from_iter(point_indices.iter().map(|(point, idx)| (*idx, *point)))
 93            .into_values()
 94            .collect::<Vec<_>>();
 95    let merged_tilde_gs = polynomials
 96        .iter()
 97        .zip(points.iter())
 98        .zip(eq_t_i_list.iter())
 99        .fold(
100            iter::repeat_with(DenseMultilinearExtension::zero)
101                .map(Arc::new)
102                .take(point_indices.len())
103                .collect::<Vec<_>>(),
104            |mut merged_tilde_gs, ((poly, point), coeff)| {
105                *Arc::make_mut(&mut merged_tilde_gs[point_indices[point]]) +=
106                    (*coeff, poly.deref());
107                merged_tilde_gs
108            },
109        );
110    end_timer!(timer);
111
112    let timer = start_timer!(|| format!("compute tilde eq for {} points", points.len()));
113    let tilde_eqs: Vec<_> = deduped_points
114        .iter()
115        .map(|point| {
116            let eq_b_zi = build_eq_x_r_vec(point).unwrap();
117            Arc::new(DenseMultilinearExtension::from_evaluations_vec(
118                num_var, eq_b_zi,
119            ))
120        })
121        .collect();
122    end_timer!(timer);
123
124    // built the virtual polynomial for SumCheck
125    let timer = start_timer!(|| format!("sum check prove of {} variables", num_var));
126
127    let step = start_timer!(|| "add mle");
128    let mut sum_check_vp = VirtualPolynomial::new(num_var);
129    for (merged_tilde_g, tilde_eq) in merged_tilde_gs.iter().zip(tilde_eqs.into_iter()) {
130        sum_check_vp.add_mle_list([merged_tilde_g.clone(), tilde_eq], E::Fr::one())?;
131    }
132    end_timer!(step);
133
134    let proof = match <PolyIOP<E::Fr> as SumCheck<E::Fr>>::prove(&sum_check_vp, transcript) {
135        Ok(p) => p,
136        Err(_e) => {
137            // cannot wrap IOPError with PCSError due to cyclic dependency
138            return Err(PCSError::InvalidProver(
139                "Sumcheck in batch proving Failed".to_string(),
140            ));
141        },
142    };
143
144    end_timer!(timer);
145
146    // a2 := sumcheck's point
147    let a2 = &proof.point[..num_var];
148
149    // build g'(X) = \sum_i=1..k \tilde eq_i(a2) * \tilde g_i(X) where (a2) is the
150    // sumcheck's point \tilde eq_i(a2) = eq(a2, point_i)
151    let step = start_timer!(|| "evaluate at a2");
152    let mut g_prime = Arc::new(DenseMultilinearExtension::zero());
153    for (merged_tilde_g, point) in merged_tilde_gs.iter().zip(deduped_points.iter()) {
154        let eq_i_a2 = eq_eval(a2, point)?;
155        *Arc::make_mut(&mut g_prime) += (eq_i_a2, merged_tilde_g.deref());
156    }
157    end_timer!(step);
158
159    let step = start_timer!(|| "pcs open");
160    let (g_prime_proof, _g_prime_eval) = PCS::open(prover_param, &g_prime, a2.to_vec().as_ref())?;
161    // assert_eq!(g_prime_eval, tilde_g_eval);
162    end_timer!(step);
163
164    let step = start_timer!(|| "evaluate fi(pi)");
165    end_timer!(step);
166    end_timer!(open_timer);
167
168    Ok(BatchProof {
169        sum_check_proof: proof,
170        f_i_eval_at_point_i: evals.to_vec(),
171        g_prime_proof,
172    })
173}
174
175/// Steps:
176/// 1. get challenge point t from transcript
177/// 2. build g' commitment
178/// 3. ensure \sum_i eq(a2, point_i) * eq(t, <i>) * f_i_evals matches the sum
179/// via SumCheck verification 4. verify commitment
180pub(crate) fn batch_verify_internal<E, PCS>(
181    verifier_param: &PCS::VerifierParam,
182    f_i_commitments: &[Commitment<E>],
183    points: &[PCS::Point],
184    proof: &BatchProof<E, PCS>,
185    transcript: &mut IOPTranscript<E::Fr>,
186) -> Result<bool, PCSError>
187where
188    E: PairingEngine,
189    PCS: PolynomialCommitmentScheme<
190        E,
191        Polynomial = Arc<DenseMultilinearExtension<E::Fr>>,
192        Point = Vec<E::Fr>,
193        Evaluation = E::Fr,
194        Commitment = Commitment<E>,
195    >,
196{
197    let open_timer = start_timer!(|| "batch verification");
198
199    // TODO: sanity checks
200
201    let k = f_i_commitments.len();
202    let ell = log2(k) as usize;
203    let num_var = proof.sum_check_proof.point.len();
204
205    // challenge point t
206    let t = transcript.get_and_append_challenge_vectors("t".as_ref(), ell)?;
207
208    // sum check point (a2)
209    let a2 = &proof.sum_check_proof.point[..num_var];
210
211    // build g' commitment
212    let step = start_timer!(|| "build homomorphic commitment");
213    let eq_t_list = build_eq_x_r_vec(t.as_ref())?;
214
215    let mut scalars = vec![];
216    let mut bases = vec![];
217
218    for (i, point) in points.iter().enumerate() {
219        let eq_i_a2 = eq_eval(a2, point)?;
220        scalars.push((eq_i_a2 * eq_t_list[i]).into_repr());
221        bases.push(f_i_commitments[i].0);
222    }
223    let g_prime_commit = VariableBaseMSM::multi_scalar_mul(&bases, &scalars);
224    end_timer!(step);
225
226    // ensure \sum_i eq(t, <i>) * f_i_evals matches the sum via SumCheck
227    let mut sum = E::Fr::zero();
228    for (i, &e) in eq_t_list.iter().enumerate().take(k) {
229        sum += e * proof.f_i_eval_at_point_i[i];
230    }
231    let aux_info = VPAuxInfo {
232        max_degree: 2,
233        num_variables: num_var,
234        phantom: PhantomData,
235    };
236    let subclaim = match <PolyIOP<E::Fr> as SumCheck<E::Fr>>::verify(
237        sum,
238        &proof.sum_check_proof,
239        &aux_info,
240        transcript,
241    ) {
242        Ok(p) => p,
243        Err(_e) => {
244            // cannot wrap IOPError with PCSError due to cyclic dependency
245            return Err(PCSError::InvalidProver(
246                "Sumcheck in batch verification failed".to_string(),
247            ));
248        },
249    };
250    let tilde_g_eval = subclaim.expected_evaluation;
251
252    // verify commitment
253    let res = PCS::verify(
254        verifier_param,
255        &Commitment(g_prime_commit.into_affine()),
256        a2.to_vec().as_ref(),
257        &tilde_g_eval,
258        &proof.g_prime_proof,
259    )?;
260
261    end_timer!(open_timer);
262    Ok(res)
263}
264
265#[cfg(test)]
266mod tests {
267    use super::*;
268    use crate::pcs::{
269        prelude::{MultilinearKzgPCS, MultilinearUniversalParams},
270        StructuredReferenceString,
271    };
272    use arithmetic::get_batched_nv;
273    use ark_bls12_381::Bls12_381 as E;
274    use ark_ec::PairingEngine;
275    use ark_poly::{DenseMultilinearExtension, MultilinearExtension};
276    use ark_std::{
277        rand::{CryptoRng, RngCore},
278        test_rng,
279        vec::Vec,
280        UniformRand,
281    };
282
283    type Fr = <E as PairingEngine>::Fr;
284
285    fn test_multi_open_helper<R: RngCore + CryptoRng>(
286        ml_params: &MultilinearUniversalParams<E>,
287        polys: &[Arc<DenseMultilinearExtension<Fr>>],
288        rng: &mut R,
289    ) -> Result<(), PCSError> {
290        let merged_nv = get_batched_nv(polys[0].num_vars(), polys.len());
291        let (ml_ck, ml_vk) = ml_params.trim(merged_nv)?;
292
293        let mut points = Vec::new();
294        for poly in polys.iter() {
295            let point = (0..poly.num_vars())
296                .map(|_| Fr::rand(rng))
297                .collect::<Vec<Fr>>();
298            points.push(point);
299        }
300
301        let evals = polys
302            .iter()
303            .zip(points.iter())
304            .map(|(f, p)| f.evaluate(p).unwrap())
305            .collect::<Vec<_>>();
306
307        let commitments = polys
308            .iter()
309            .map(|poly| MultilinearKzgPCS::commit(&ml_ck.clone(), poly).unwrap())
310            .collect::<Vec<_>>();
311
312        let mut transcript = IOPTranscript::new("test transcript".as_ref());
313        transcript.append_field_element("init".as_ref(), &Fr::zero())?;
314
315        let batch_proof = multi_open_internal::<E, MultilinearKzgPCS<E>>(
316            &ml_ck,
317            polys,
318            &points,
319            &evals,
320            &mut transcript,
321        )?;
322
323        // good path
324        let mut transcript = IOPTranscript::new("test transcript".as_ref());
325        transcript.append_field_element("init".as_ref(), &Fr::zero())?;
326        assert!(batch_verify_internal::<E, MultilinearKzgPCS<E>>(
327            &ml_vk,
328            &commitments,
329            &points,
330            &batch_proof,
331            &mut transcript
332        )?);
333
334        Ok(())
335    }
336
337    #[test]
338    fn test_multi_open_internal() -> Result<(), PCSError> {
339        let mut rng = test_rng();
340
341        let ml_params = MultilinearUniversalParams::<E>::gen_srs_for_testing(&mut rng, 20)?;
342        for num_poly in 5..6 {
343            for nv in 15..16 {
344                let polys1: Vec<_> = (0..num_poly)
345                    .map(|_| Arc::new(DenseMultilinearExtension::rand(nv, &mut rng)))
346                    .collect();
347                test_multi_open_helper(&ml_params, &polys1, &mut rng)?;
348            }
349        }
350
351        Ok(())
352    }
353}