1use alloc::collections::{BTreeMap, BTreeSet};
4use alloc::vec;
5use alloc::vec::Vec;
6
7use p3_commit::{BatchOpening, BatchOpeningRef, Mmcs};
8use p3_field::{Field, PackedValue};
9use p3_matrix::{Dimensions, Matrix};
10use serde::{Deserialize, Serialize};
11use thiserror::Error;
12
13#[derive(Clone, Debug)]
21pub struct GroupedCodewordMmcs<Inner> {
22 inner: Inner,
23 group_size: usize,
24 log_inv_rate: Option<usize>,
25}
26
27impl<Inner> GroupedCodewordMmcs<Inner> {
28 pub fn new(inner: Inner, group_size: usize) -> Self {
33 assert!(
34 group_size.is_power_of_two(),
35 "group size must be a power of two"
36 );
37 Self {
38 inner,
39 group_size,
40 log_inv_rate: None,
41 }
42 }
43
44 pub const fn for_folding(inner: Inner, config: &crate::BinaryPcsConfig) -> Self {
47 Self::with_group_size(inner, config, 1 << config.log_folding_factor())
48 }
49
50 pub const fn with_group_size(
60 inner: Inner,
61 config: &crate::BinaryPcsConfig,
62 group_size: usize,
63 ) -> Self {
64 assert!(
65 group_size.is_power_of_two(),
66 "group size must be a power of two"
67 );
68 Self {
69 inner,
70 group_size,
71 log_inv_rate: Some(config.log_inv_rate()),
72 }
73 }
74
75 pub fn group_size_at(&self, height: usize) -> Option<usize> {
84 let message_len = height >> self.log_inv_rate.unwrap_or(0);
85 (message_len != 0).then(|| self.group_size.min(message_len))
86 }
87}
88
89#[derive(Clone, Debug)]
91pub struct GroupedCodeword<M> {
92 matrix: M,
93 group_size: usize,
94}
95
96impl<F: Field, M: Matrix<F>> Matrix<F> for GroupedCodeword<M> {
97 fn width(&self) -> usize {
98 self.group_size
99 }
100
101 fn height(&self) -> usize {
102 self.matrix.height() / self.group_size
103 }
104
105 unsafe fn row_subseq_unchecked(
106 &self,
107 r: usize,
108 start: usize,
109 end: usize,
110 ) -> impl IntoIterator<Item = F, IntoIter = impl Iterator<Item = F> + Send + Sync> {
111 (start..end).map(move |lane| {
112 unsafe { self.matrix.get_unchecked(r * self.group_size + lane, 0) }
115 })
116 }
117
118 #[inline]
119 fn vertically_packed_row<P>(&self, r: usize) -> impl Iterator<Item = P>
120 where
121 F: Copy,
122 P: PackedValue<Value = F>,
123 {
124 let matrix = &self.matrix;
125 let group_size = self.group_size;
126 let height = self.height();
127 let row = r % height;
128 let no_wrap = P::WIDTH != 1 && r + P::WIDTH <= height;
129
130 (0..group_size).map(move |lane| {
134 if P::WIDTH == 1 {
135 unsafe { P::broadcast(matrix.get_unchecked(row * group_size + lane, 0)) }
137 } else if no_wrap {
138 P::from_fn(|i| unsafe { matrix.get_unchecked((r + i) * group_size + lane, 0) })
140 } else {
141 P::from_fn(|i| {
142 let row = (r + i) % height;
143 unsafe { matrix.get_unchecked(row * group_size + lane, 0) }
145 })
146 }
147 })
148 }
149}
150
151#[derive(Clone, Serialize, Deserialize)]
153#[serde(bound(
154 serialize = "F: Serialize, P: Serialize",
155 deserialize = "F: Deserialize<'de>, P: Deserialize<'de>"
156))]
157pub struct GroupedCodewordProof<F, P> {
158 inner: P,
159 missing_symbols: Vec<F>,
161}
162
163#[derive(Debug, Error)]
165pub enum GroupedCodewordError<E> {
166 #[error("expected one width-1 codeword of power-of-two length")]
168 Dimensions,
169 #[error("malformed grouped codeword opening")]
171 OpeningShape,
172 #[error("conflicting duplicate symbol openings")]
174 ConflictingDuplicate,
175 #[error("grouped codeword authentication failed: {0:?}")]
177 Inner(#[source] E),
178}
179
180fn group_indices(indices: &[usize], group_size: usize) -> Vec<usize> {
181 let mut groups: Vec<_> = indices.iter().map(|i| i / group_size).collect();
182 groups.sort_unstable();
183 groups.dedup();
184 groups
185}
186
187impl<F: Field, Inner: Mmcs<F>> Mmcs<F> for GroupedCodewordMmcs<Inner> {
188 type ProverData<M> = Inner::ProverData<GroupedCodeword<M>>;
189 type Commitment = Inner::Commitment;
190 type Proof = GroupedCodewordProof<F, Inner::MultiProof>;
191 type MultiProof = GroupedCodewordProof<F, Inner::MultiProof>;
192 type Error = GroupedCodewordError<Inner::Error>;
193
194 fn commit<M: Matrix<F>>(&self, mut inputs: Vec<M>) -> (Self::Commitment, Self::ProverData<M>) {
195 assert_eq!(inputs.len(), 1, "expected exactly one codeword");
196 let matrix = inputs.pop().unwrap();
197 assert_eq!(matrix.width(), 1, "expected a width-1 codeword");
198 assert!(
199 matrix.height().is_power_of_two(),
200 "expected power-of-two length"
201 );
202 let group_size = self
203 .group_size_at(matrix.height())
204 .expect("codeword shorter than rate expansion");
205 self.inner
206 .commit_matrix(GroupedCodeword { matrix, group_size })
207 }
208
209 fn get_matrices<'a, M: Matrix<F>>(&self, data: &'a Self::ProverData<M>) -> Vec<&'a M> {
210 self.inner
211 .get_matrices(data)
212 .into_iter()
213 .map(|view| &view.matrix)
214 .collect()
215 }
216
217 fn open_batch<M: Matrix<F>>(
218 &self,
219 index: usize,
220 data: &Self::ProverData<M>,
221 ) -> BatchOpening<F, Self> {
222 let (mut rows, proof) = self.open_multi_batch(&[index], data);
223 BatchOpening::new(rows.pop().unwrap(), proof)
224 }
225
226 fn verify_batch(
227 &self,
228 commitment: &Self::Commitment,
229 dimensions: &[Dimensions],
230 index: usize,
231 opening: BatchOpeningRef<'_, F, Self>,
232 ) -> Result<(), Self::Error> {
233 let rows = vec![opening.opened_values.iter().map(Vec::as_slice).collect()];
234 self.verify_multi_batch(
235 commitment,
236 dimensions,
237 &[index],
238 &rows,
239 opening.opening_proof,
240 )
241 }
242
243 fn open_multi_batch<M: Matrix<F>>(
244 &self,
245 indices: &[usize],
246 data: &Self::ProverData<M>,
247 ) -> (Vec<Vec<Vec<F>>>, Self::MultiProof) {
248 let views = self.inner.get_matrices(data);
249 let view = views[0];
250 assert!(
251 indices.iter().all(|&i| i < view.matrix.height()),
252 "symbol index out of range"
253 );
254 let group_size = view.group_size;
255 let groups = group_indices(indices, group_size);
256 let (group_rows, inner) = self.inner.open_multi_batch(&groups, data);
257 let requested: BTreeSet<_> = indices.iter().copied().collect();
258 let mut missing_symbols = Vec::new();
259 for (&group, rows) in groups.iter().zip(group_rows) {
260 for (lane, &value) in rows[0].iter().enumerate() {
261 if !requested.contains(&(group * group_size + lane)) {
262 missing_symbols.push(value);
263 }
264 }
265 }
266 let rows = indices
267 .iter()
268 .map(|&i| vec![vec![view.matrix.get(i, 0).unwrap()]])
269 .collect();
270 (
271 rows,
272 GroupedCodewordProof {
273 inner,
274 missing_symbols,
275 },
276 )
277 }
278
279 fn verify_multi_batch<R: AsRef<[F]> + PartialEq>(
280 &self,
281 commitment: &Self::Commitment,
282 dimensions: &[Dimensions],
283 indices: &[usize],
284 opened_values: &[Vec<R>],
285 proof: &Self::MultiProof,
286 ) -> Result<(), Self::Error> {
287 use GroupedCodewordError::{
288 ConflictingDuplicate, Dimensions as BadDimensions, OpeningShape,
289 };
290
291 let [dim] = dimensions else {
292 return Err(BadDimensions);
293 };
294 if dim.width != 1 || !dim.height.is_power_of_two() {
295 return Err(BadDimensions);
296 }
297 let group_size = self.group_size_at(dim.height).ok_or(BadDimensions)?;
298 if indices.len() != opened_values.len() {
299 return Err(OpeningShape);
300 }
301 let mut requested = BTreeMap::new();
302 for (&index, rows) in indices.iter().zip(opened_values) {
303 if index >= dim.height || rows.len() != 1 || rows[0].as_ref().len() != 1 {
304 return Err(OpeningShape);
305 }
306 let value = rows[0].as_ref()[0];
307 if requested
308 .insert(index, value)
309 .is_some_and(|old| old != value)
310 {
311 return Err(ConflictingDuplicate);
312 }
313 }
314 let groups = group_indices(indices, group_size);
315 let missing_count = groups.len() * group_size - requested.len();
318 if proof.missing_symbols.len() != missing_count {
319 return Err(OpeningShape);
320 }
321 let mut missing = proof.missing_symbols.iter();
322 let rows: Vec<_> = groups
323 .iter()
324 .map(|&group| {
325 let row: Vec<_> = (0..group_size)
326 .map(|lane| {
327 *requested
328 .get(&(group * group_size + lane))
329 .unwrap_or_else(|| {
330 missing.next().expect("missing symbol count checked above")
331 })
332 })
333 .collect();
334 vec![row]
335 })
336 .collect();
337 let grouped_dimensions = [Dimensions {
338 width: group_size,
339 height: dim.height / group_size,
340 }];
341 self.inner
342 .verify_multi_batch(
343 commitment,
344 &grouped_dimensions,
345 &groups,
346 &rows,
347 &proof.inner,
348 )
349 .map_err(GroupedCodewordError::Inner)
350 }
351}
352
353#[cfg(test)]
354mod tests {
355 use alloc::vec;
356 use alloc::vec::Vec;
357
358 use p3_binary_field::{BinaryField128 as F, TowerLevel};
359 use p3_commit::{BatchOpeningRef, Mmcs};
360 use p3_field::{FieldArray, PrimeCharacteristicRing};
361 use p3_matrix::dense::RowMajorMatrix;
362 use p3_matrix::{Dimensions, Matrix};
363
364 use super::{GroupedCodeword, GroupedCodewordMmcs};
365 use crate::test_util::mmcs;
366 use crate::{BinaryPcsConfig, BinaryPcsParams};
367
368 type Packed = FieldArray<F, 3>;
370
371 struct DefaultPackedGroupedCodeword<M> {
374 matrix: M,
375 group_size: usize,
376 }
377
378 impl<M: Matrix<F>> Matrix<F> for DefaultPackedGroupedCodeword<M> {
379 fn width(&self) -> usize {
380 self.group_size
381 }
382
383 fn height(&self) -> usize {
384 self.matrix.height() / self.group_size
385 }
386
387 unsafe fn row_subseq_unchecked(
388 &self,
389 r: usize,
390 start: usize,
391 end: usize,
392 ) -> impl IntoIterator<Item = F, IntoIter = impl Iterator<Item = F> + Send + Sync> {
393 (start..end).map(move |lane| {
394 unsafe { self.matrix.get_unchecked(r * self.group_size + lane, 0) }
396 })
397 }
398 }
399
400 #[test]
401 fn groups_match_an_explicitly_reshaped_commitment() {
402 let values: Vec<_> = (0..32).map(F::from_repr).collect();
403 for group_size in [1, 2, 4, 8, 16, 64] {
404 let grouped = GroupedCodewordMmcs::new(mmcs(), group_size);
405 let (root, data) = grouped.commit_matrix(RowMajorMatrix::new(values.clone(), 1));
406 let (expected, _) =
407 mmcs().commit_matrix(RowMajorMatrix::new(values.clone(), group_size.min(32)));
408 assert_eq!(root, expected);
409 assert_eq!(grouped.get_matrices(&data)[0].values, values);
410 assert_eq!(grouped.get_matrices(&data)[0].width, 1);
411
412 let indices = [31, 0, 1, 17, 0];
413 let (rows, proof) = grouped.open_multi_batch(&indices, &data);
414 for (&index, row) in indices.iter().zip(&rows) {
415 assert_eq!(row, &vec![vec![values[index]]]);
416 }
417 let dims = [Dimensions {
418 width: 1,
419 height: 32,
420 }];
421 grouped
422 .verify_multi_batch(&root, &dims, &indices, &rows, &proof)
423 .unwrap();
424
425 let opening = grouped.open_batch(17, &data);
426 grouped
427 .verify_batch(&root, &dims, 17, (&opening).into())
428 .unwrap();
429 }
430 }
431
432 #[test]
433 fn pairs_need_no_extra_symbols_and_larger_groups_only_send_missing_symbols() {
434 let values: Vec<_> = (0..16).map(F::from_repr).collect();
435 let indices = [0, 1, 8, 9, 0, 1];
436 for (size, missing) in [(2, vec![]), (4, vec![2, 3, 10, 11])] {
437 let grouped = GroupedCodewordMmcs::new(mmcs(), size);
438 let (_, data) = grouped.commit_matrix(RowMajorMatrix::new(values.clone(), 1));
439 let (_, proof) = grouped.open_multi_batch(&indices, &data);
440 assert_eq!(
441 proof.missing_symbols,
442 missing.into_iter().map(F::from_repr).collect::<Vec<_>>()
443 );
444 }
445 }
446
447 #[test]
448 fn folding_schedule_sizes_leaves_for_the_next_actual_batch() {
449 let config = BinaryPcsConfig::try_new::<F, F>(
450 6,
451 BinaryPcsParams {
452 log_inv_rate: 2,
453 security_level: 100,
454 pow_bits: 0,
455 },
456 )
457 .unwrap()
458 .try_with_folding(4)
459 .unwrap();
460 let grouped = GroupedCodewordMmcs::for_folding(mmcs(), &config);
461 for (len, width) in [(256, 16), (16, 4)] {
462 let values: Vec<_> = (0..len).map(|i| F::from_repr(i as u128)).collect();
463 let (root, data) = grouped.commit_matrix(RowMajorMatrix::new(values.clone(), 1));
464 let (expected, _) = mmcs().commit_matrix(RowMajorMatrix::new(values, width));
465 assert_eq!(root, expected);
466 let opening = grouped.open_batch(len - 1, &data);
467 grouped
468 .verify_batch(
469 &root,
470 &[Dimensions {
471 height: len,
472 width: 1,
473 }],
474 len - 1,
475 (&opening).into(),
476 )
477 .unwrap();
478 }
479 }
480
481 #[test]
482 fn a_leaf_wider_than_the_fold_coset_carries_the_symbols_no_query_asked_for() {
483 let config = BinaryPcsConfig::try_new::<F, F>(
484 6,
485 BinaryPcsParams {
486 log_inv_rate: 2,
487 security_level: 100,
488 pow_bits: 0,
489 },
490 )
491 .unwrap()
492 .try_with_folding(2)
493 .unwrap();
494 let values: Vec<_> = (0..256).map(|i| F::from_repr(i as u128)).collect();
495 let dims = [Dimensions {
496 height: 256,
497 width: 1,
498 }];
499 let coset = [4, 5, 6, 7];
501
502 let coset_leaves = GroupedCodewordMmcs::for_folding(mmcs(), &config);
503 assert_eq!(
504 coset_leaves.group_size_at(256),
505 Some(1 << config.log_folding_factor())
506 );
507 let (root, data) = coset_leaves.commit_matrix(RowMajorMatrix::new(values.clone(), 1));
508 let (rows, proof) = coset_leaves.open_multi_batch(&coset, &data);
509 assert!(proof.missing_symbols.is_empty());
510 coset_leaves
511 .verify_multi_batch(&root, &dims, &coset, &rows, &proof)
512 .unwrap();
513
514 let wide_leaves = GroupedCodewordMmcs::with_group_size(mmcs(), &config, 16);
516 let (wide_root, wide_data) = wide_leaves.commit_matrix(RowMajorMatrix::new(values, 1));
517 assert_ne!(wide_root, root);
518 let (wide_rows, wide_proof) = wide_leaves.open_multi_batch(&coset, &wide_data);
519 assert_eq!(wide_proof.missing_symbols.len(), 12);
520 wide_leaves
521 .verify_multi_batch(&wide_root, &dims, &coset, &wide_rows, &wide_proof)
522 .unwrap();
523 }
524
525 #[test]
526 fn malformed_openings_and_unauthenticated_group_members_are_rejected() {
527 let grouped = GroupedCodewordMmcs::new(mmcs(), 4);
528 let values: Vec<_> = (0..32).map(F::from_repr).collect();
529 let (root, data) = grouped.commit_matrix(RowMajorMatrix::new(values, 1));
530 let indices = [0, 1, 0];
531 let dims = [Dimensions {
532 width: 1,
533 height: 32,
534 }];
535 let (rows, proof) = grouped.open_multi_batch(&indices, &data);
536 let verify = |rows: &[Vec<Vec<F>>], proof: &_| {
537 grouped.verify_multi_batch(&root, &dims, &indices, rows, proof)
538 };
539 assert!(verify(&rows, &proof).is_ok());
540 for position in 0..proof.missing_symbols.len() {
541 let mut bad = proof.clone();
542 bad.missing_symbols[position] += F::ONE;
543 assert!(verify(&rows, &bad).is_err());
544 }
545 let mut bad = proof.clone();
546 bad.missing_symbols.pop();
547 assert!(verify(&rows, &bad).is_err());
548 let mut bad = proof.clone();
549 bad.missing_symbols.push(F::ZERO);
550 assert!(verify(&rows, &bad).is_err());
551 let mut bad_rows = rows.clone();
552 bad_rows[2][0][0] += F::ONE;
553 assert!(verify(&bad_rows, &proof).is_err());
554 assert!(verify(&rows[..2], &proof).is_err());
555 let mut bad_rows = rows.clone();
556 bad_rows[0][0].push(F::ZERO);
557 assert!(verify(&bad_rows, &proof).is_err());
558 let mut bad_rows = rows.clone();
559 bad_rows[0].clear();
560 assert!(verify(&bad_rows, &proof).is_err());
561 assert!(
562 grouped
563 .verify_multi_batch(&root, &dims, &[0, 1, 32], &rows, &proof)
564 .is_err()
565 );
566 for dims in [
567 vec![],
568 vec![Dimensions {
569 width: 2,
570 height: 16,
571 }],
572 vec![Dimensions {
573 width: 1,
574 height: 0,
575 }],
576 vec![Dimensions {
577 width: 1,
578 height: 31,
579 }],
580 ] {
581 assert!(
582 grouped
583 .verify_multi_batch(&root, &dims, &indices, &rows, &proof)
584 .is_err()
585 );
586 }
587 assert!(
589 GroupedCodewordMmcs::new(mmcs(), 8)
590 .verify_multi_batch(&root, &dims, &indices, &rows, &proof)
591 .is_err()
592 );
593 assert!(
594 grouped
595 .verify_batch(&root, &dims, 0, BatchOpeningRef::new(&rows[0], &proof))
596 .is_err()
597 );
598 }
599
600 #[test]
601 fn full_pcs_proofs_round_trip_with_grouped_leaves() {
602 use p3_commit::MultilinearPcs;
603 use p3_sumcheck::layout::{Layout, SuffixProver, Table};
604 use p3_sumcheck::{OpeningBatch, OpeningProtocol, TableShape, TableSpec};
605 use rand::SeedableRng;
606 use rand::rngs::SmallRng;
607
608 use crate::test_util::{MyMmcs, challenger};
609 use crate::{BinaryPcs, BinaryPcsConfig, BinaryPcsParams, BinaryPcsProof};
610
611 for num_variables in [1, 6, 10] {
612 for group_size in [1, 2, 4, 8, 16] {
613 let config = BinaryPcsConfig::try_new::<F, F>(
614 num_variables,
615 BinaryPcsParams {
616 log_inv_rate: 2,
617 security_level: 100,
618 pow_bits: 0,
619 },
620 )
621 .unwrap();
622 let pcs = BinaryPcs::new(
623 config,
624 GroupedCodewordMmcs::new(mmcs(), group_size),
625 GroupedCodewordMmcs::new(mmcs(), group_size),
626 )
627 .unwrap();
628 let mut rng = SmallRng::seed_from_u64(0x6710);
629 let table = Table::rand(&mut rng, 1, num_variables);
630 let witness = SuffixProver::<F, F>::new_witness(vec![table], 0);
631 let protocol = OpeningProtocol::new(vec![TableSpec::new(
632 TableShape::new(num_variables, 1),
633 vec![OpeningBatch::new(vec![0], Vec::new())],
634 )]);
635 let mut prover_challenger = challenger();
636 let (commitment, data) = pcs.commit(witness, &mut prover_challenger).unwrap();
637 let proof = pcs
638 .open(data, protocol.clone(), &mut prover_challenger)
639 .unwrap();
640 let bytes = postcard::to_allocvec(&proof).unwrap();
641 let decoded: BinaryPcsProof<
642 F,
643 F,
644 GroupedCodewordMmcs<MyMmcs>,
645 GroupedCodewordMmcs<MyMmcs>,
646 > = postcard::from_bytes(&bytes).unwrap();
647 pcs.verify(&commitment, &decoded, &mut challenger(), protocol)
648 .unwrap();
649 }
650 }
651 }
652
653 #[test]
654 fn vertically_packed_row_matches_the_default_trait_implementation_on_wrap_around() {
655 for (len, group_size) in [(24usize, 3usize), (30, 5), (14, 7), (2, 1), (9, 9)] {
658 let values: Vec<_> = (0..len as u128).map(F::from_repr).collect();
659 let grouped = GroupedCodeword {
660 matrix: RowMajorMatrix::new(values.clone(), 1),
661 group_size,
662 };
663 let reference = DefaultPackedGroupedCodeword {
664 matrix: RowMajorMatrix::new(values, 1),
665 group_size,
666 };
667 let height = grouped.height();
668 assert_eq!(height, reference.height());
669
670 for r in 0..height {
671 let actual: Vec<F> = grouped.vertically_packed_row::<F>(r).collect();
673 let expected: Vec<F> = reference.vertically_packed_row::<F>(r).collect();
674 assert_eq!(
675 actual, expected,
676 "width 1, len {len}, group {group_size}, r {r}"
677 );
678
679 let actual: Vec<Packed> = grouped.vertically_packed_row::<Packed>(r).collect();
682 let expected: Vec<Packed> = reference.vertically_packed_row::<Packed>(r).collect();
683 assert_eq!(
684 actual, expected,
685 "width 3, len {len}, group {group_size}, r {r}"
686 );
687 }
688 }
689 }
690}