Skip to content

Commit 7a834dd

Browse files
committed
mpcs e2e batch in one proof
1 parent c6c7461 commit 7a834dd

File tree

9 files changed

+422
-237
lines changed

9 files changed

+422
-237
lines changed

ceno_zkvm/src/structs.rs

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -71,9 +71,6 @@ pub enum RAMType {
7171
Memory,
7272
}
7373

74-
/// A point is a vector of num_var length
75-
pub type Point<F> = Vec<F>;
76-
7774
/// A point and the evaluation of this point.
7875
#[derive(Clone, Debug, PartialEq)]
7976
pub struct PointAndEval<F> {

mpcs/src/basefold.rs

Lines changed: 95 additions & 115 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
use crate::{
2-
Error, Evaluation, PolynomialCommitmentScheme,
2+
Error, Evaluation, Point, PolynomialCommitmentScheme,
33
util::{
44
arithmetic::inner_product,
55
ext_to_usize,
@@ -19,7 +19,7 @@ use sumcheck::macros::{entered_span, exit_span};
1919
use transcript::Transcript;
2020
use witness::RowMajorMatrix;
2121

22-
use itertools::Itertools;
22+
use itertools::{Itertools, izip};
2323
use serde::{Serialize, de::DeserializeOwned};
2424

2525
use multilinear_extensions::{mle::FieldType, virtual_poly::build_eq_x_r_vec};
@@ -35,14 +35,15 @@ pub use structure::{
3535
BasefoldProverParams, BasefoldRSParams, BasefoldVerifierParams,
3636
};
3737
mod commit_phase;
38-
use commit_phase::simple_batch_commit_phase;
38+
use commit_phase::{batch_commit_phase, simple_batch_commit_phase};
3939
mod encoding;
4040
use multilinear_extensions::virtual_poly::ArcMultilinearExtension;
4141

4242
mod query_phase;
4343
// This sumcheck module is different from the mpcs::sumcheck module, in that
4444
// it deals only with the special case of the form \sum eq(r_i)f_i().
4545

46+
// make it pure error
4647
pub enum PolyEvalsCodeword<E: ExtensionField> {
4748
Normal(Box<DenseMatrix<E::BaseField>>),
4849
TooSmall(Box<DenseMatrix<E::BaseField>>), // The polynomial is too small to apply FRI
@@ -213,84 +214,50 @@ where
213214

214215
fn batch_commit(
215216
pp: &Self::ProverParam,
216-
rmm: witness::RowMajorMatrix<<E as ff_ext::ExtensionField>::BaseField>,
217+
rmm: Vec<witness::RowMajorMatrix<<E as ff_ext::ExtensionField>::BaseField>>,
217218
) -> Result<Self::CommitmentWithWitness, Error> {
218-
let span = entered_span!("to_mles", profiling_3 = true);
219-
let polys = rmm.to_mles();
220-
exit_span!(span);
221219
// assumptions
222220
// 1. there must be at least one polynomial
223221
// 2. all polynomials must exist in the same field type
224222
// (TODO: eliminate this assumption by supporting commiting
225223
// and opening mixed-type polys)
226224
// 3. all polynomials must have the same number of variables
227225

228-
if polys.is_empty() {
226+
if rmm.is_empty() {
229227
return Err(Error::InvalidPcsParam(
230228
"cannot batch commit to zero polynomials".to_string(),
231229
));
232230
}
233231

234-
let num_vars = polys[0].num_vars();
235-
let num_polys = polys.len();
236-
237-
let is_base = match polys[0].evaluations() {
238-
FieldType::Ext(_) => false,
239-
FieldType::Base(_) => true,
240-
_ => unreachable!(),
241-
};
232+
let span = entered_span!("to_mles", profiling_3 = true);
233+
let polys: Vec<Vec<ArcMultilinearExtension<E>>> = rmm
234+
.iter()
235+
.map(|rmm| rmm.to_mles().into_iter().map(|p| p.into()).collect_vec())
236+
.collect_vec();
237+
exit_span!(span);
242238

243-
if !polys.iter().map(|poly| poly.num_vars()).all_equal() {
244-
return Err(Error::InvalidPcsParam(
245-
"cannot batch commit to polynomials with different number of variables".to_string(),
246-
));
247-
}
248239
let timer = start_timer!(|| "Basefold::batch commit");
249-
let encode_timer = start_timer!(|| "Basefold::batch commit::encoding and interpolations");
250240
let span = entered_span!("encode_codeword_and_mle", profiling_3 = true);
251-
252-
let evals_codewords = Spec::EncodingScheme::encode(&pp.encoding_params, rmm);
253-
241+
let evals_codewords = rmm
242+
.into_iter()
243+
.map(|rmm| Spec::EncodingScheme::encode(&pp.encoding_params, rmm))
244+
.collect::<Result<Vec<DenseMatrix<E::BaseField>>, _>>()?;
254245
exit_span!(span);
255-
end_timer!(encode_timer);
256246

257247
let span = entered_span!("build mt", profiling_3 = true);
258-
// build merkle tree
259-
let ret = match evals_codewords {
260-
PolyEvalsCodeword::Normal(codewords) => {
261-
let mmcs = poseidon2_merkle_tree::<E>();
262-
let (comm, codeword) = mmcs.commit_matrix(*codewords);
263-
let polys: Vec<ArcMultilinearExtension<E>> =
264-
polys.into_iter().map(|poly| poly.into()).collect_vec();
265-
Self::CommitmentWithWitness {
266-
pi_d_digest: comm,
267-
codeword,
268-
polynomials_bh_evals: polys,
269-
num_vars,
270-
is_base,
271-
num_polys,
272-
}
273-
}
274-
PolyEvalsCodeword::TooSmall(rmm_padded) => {
275-
let mmcs = poseidon2_merkle_tree::<E>();
276-
let (comm, codeword) = mmcs.commit_matrix(*rmm_padded);
277-
let polys: Vec<ArcMultilinearExtension<E>> =
278-
polys.into_iter().map(|poly| poly.into()).collect_vec();
279-
Self::CommitmentWithWitness {
280-
pi_d_digest: comm,
281-
codeword,
282-
polynomials_bh_evals: polys,
283-
num_vars,
284-
is_base,
285-
num_polys,
286-
}
287-
}
288-
PolyEvalsCodeword::TooBig(_) => return Err(Error::PolynomialTooLarge(num_vars)),
289-
};
248+
let mmcs = poseidon2_merkle_tree::<E>();
249+
let (comm, codeword) = mmcs.commit(evals_codewords);
290250
exit_span!(span);
291-
end_timer!(timer);
292-
293-
Ok(ret)
251+
let meta_info = polys
252+
.iter()
253+
.map(|polys| (polys[0].num_vars(), polys.len()))
254+
.collect_vec();
255+
Ok(Self::CommitmentWithWitness {
256+
pi_d_digest: comm,
257+
codeword,
258+
polynomials_bh_evals: polys,
259+
meta_info,
260+
})
294261
}
295262

296263
fn write_commitment(
@@ -320,75 +287,72 @@ where
320287
}
321288

322289
/// Open a batch of polynomial commitments at several points.
323-
/// The current version only supports one polynomial per commitment.
324-
/// Because otherwise it is complex to match the polynomials and
325-
/// the commitments, and because currently this high flexibility is
326-
/// not very useful in ceno.
327290
fn batch_open(
328-
_pp: &Self::ProverParam,
329-
_polys: &[ArcMultilinearExtension<E>],
330-
_comms: &[Self::CommitmentWithWitness],
331-
_points: &[Vec<E>],
332-
_evals: &[Evaluation<E>],
333-
_transcript: &mut impl Transcript<E>,
334-
) -> Result<Self::Proof, Error> {
335-
unimplemented!()
336-
}
337-
338-
/// This is a simple version of batch open:
339-
/// 1. Open at one point
340-
/// 2. All the polynomials share the same commitment and have the same
341-
/// number of variables.
342-
/// 3. The point is already a random point generated by a sum-check.
343-
fn simple_batch_open(
344291
pp: &Self::ProverParam,
345-
polys: &[ArcMultilinearExtension<E>],
346-
comm: &Self::CommitmentWithWitness,
347-
point: &[E],
348-
evals: &[E],
292+
comms: &Self::CommitmentWithWitness,
293+
points: &[Point<E>],
294+
evals: &[Vec<E>],
349295
transcript: &mut impl Transcript<E>,
350296
) -> Result<Self::Proof, Error> {
351-
let timer = start_timer!(|| "Basefold::batch_open");
352-
let num_vars = polys[0].num_vars();
353-
354-
if comm.is_trivial::<Spec>() {
355-
let mmcs = poseidon2_merkle_tree::<E>();
356-
return Ok(Self::Proof::trivial(
357-
mmcs.get_matrices(&comm.codeword)[0].clone(),
358-
));
359-
}
360-
361-
polys
362-
.iter()
363-
.for_each(|poly| assert_eq!(poly.num_vars(), num_vars));
364-
assert!(num_vars >= Spec::get_basecode_msg_size_log());
365-
assert_eq!(comm.num_polys, polys.len());
366-
assert_eq!(comm.num_polys, evals.len());
297+
let span = entered_span!("Basefold::batch_open");
298+
299+
// sanity check
300+
// number of point match with commitment length, assuming each commitment are opening under same point
301+
assert_eq!(points.len(), comms.polynomials_bh_evals.len());
302+
303+
assert!(izip!(&comms.polynomials_bh_evals, &comms.meta_info).all(
304+
|(polynomials_bh_evals, meta_info)| {
305+
let (num_var, num_polys) = meta_info;
306+
// check num_vars & num_poly match
307+
polynomials_bh_evals
308+
.iter()
309+
.all(|p| p.num_vars() == *num_var)
310+
&& polynomials_bh_evals.len() == *num_polys
311+
},
312+
));
313+
314+
let (min_num_vars, max_num_vars) = (
315+
*comms
316+
.meta_info
317+
.iter()
318+
.map(|(num_var, _)| num_var)
319+
.min()
320+
.unwrap(),
321+
*comms
322+
.meta_info
323+
.iter()
324+
.map(|(num_var, _)| num_var)
325+
.max()
326+
.unwrap(),
327+
);
328+
assert!(min_num_vars >= Spec::get_basecode_msg_size_log());
367329

368330
if cfg!(feature = "sanity-check") {
369-
evals
370-
.iter()
371-
.zip(polys)
372-
.for_each(|(eval, poly)| assert_eq!(&poly.evaluate(point), eval))
331+
assert!(izip!(comms.polynomials_bh_evals, points, evals).all(
332+
|(polys, point, evals)| {
333+
izip!(polys, evals).all(|(poly, eval)| poly.evaluate(point) == *eval)
334+
}
335+
));
373336
}
374-
// evals.len() is the batch size, i.e., how many polynomials are being opened together
375-
let batch_coeffs = &transcript
376-
.sample_and_append_challenge_pows(evals.len(), b"batch coeffs")[0..evals.len()];
377-
let _target_sum = inner_product(evals, batch_coeffs);
337+
let total_num_polys = evals.iter().map(|evals| evals.len()).sum();
338+
// evals.len() is the number of polynomials
339+
let batch_coeffs =
340+
&transcript.sample_and_append_challenge_pows(total_num_polys, b"batch coeffs");
341+
// let _target_sum = inner_product(evals, batch_coeffs);
378342

379343
// Now the verifier has obtained the new target sum, and is able to compute the random
380344
// linear coefficients.
381345
// The remaining tasks for the prover is to prove that
382346
// sum_i coeffs[i] poly_evals[i] is equal to
383347
// the new target sum, where coeffs is computed as follows
384-
let (trees, commit_phase_proof) = simple_batch_commit_phase::<E, Spec>(
348+
let (trees, commit_phase_proof) = batch_commit_phase::<E, Spec>(
385349
&pp.encoding_params,
386-
point,
350+
points,
387351
batch_coeffs,
388-
comm,
352+
comms,
389353
transcript,
390-
num_vars,
391-
num_vars - Spec::get_basecode_msg_size_log(),
354+
(min_num_vars, max_num_vars),
355+
max_num_vars - Spec::get_basecode_msg_size_log(),
392356
);
393357

394358
let query_timer = start_timer!(|| "Basefold::open::query_phase");
@@ -398,7 +362,7 @@ where
398362
simple_batch_prover_query_phase(transcript, comm, &trees, Spec::get_number_queries());
399363
end_timer!(query_timer);
400364

401-
end_timer!(timer);
365+
end_timer!(span);
402366
Ok(Self::Proof {
403367
sumcheck_messages: commit_phase_proof.sumcheck_messages,
404368
roots: commit_phase_proof.roots,
@@ -409,6 +373,22 @@ where
409373
})
410374
}
411375

376+
/// This is a simple version of batch open:
377+
/// 1. Open at one point
378+
/// 2. All the polynomials share the same commitment and have the same
379+
/// number of variables.
380+
/// 3. The point is already a random point generated by a sum-check.
381+
fn simple_batch_open(
382+
pp: &Self::ProverParam,
383+
polys: &[ArcMultilinearExtension<E>],
384+
comm: &Self::CommitmentWithWitness,
385+
point: &[E],
386+
evals: &[E],
387+
transcript: &mut impl Transcript<E>,
388+
) -> Result<Self::Proof, Error> {
389+
unimplemented!()
390+
}
391+
412392
fn verify(
413393
_vp: &Self::VerifierParam,
414394
_comm: &Self::Commitment,

0 commit comments

Comments
 (0)