1use alloc::vec::Vec;
37use core::marker::PhantomData;
38use core::{ptr, slice};
39
40use p3_binary_field::{
41 BinaryField8, BinaryField16, BinaryField32, BinaryField64, BinaryField128, Gf2, PackedGf2,
42 TowerLevel, Underlier,
43};
44use p3_field::PrimeCharacteristicRing;
45use p3_matrix::dense::RowMajorMatrix;
46use p3_maybe_rayon::prelude::*;
47use p3_multilinear_util::poly::Poly;
48use p3_sumcheck::layout::Table;
49use p3_util::log2_strict_usize;
50use thiserror::Error;
51
52pub unsafe trait Coordinates: Copy + Send + Sync + 'static {
68 const COORDINATES: usize;
70}
71
72const fn check_coordinates<A: Coordinates>() {
74 assert!(
75 A::COORDINATES == 8 * size_of::<A>(),
76 "a coordinate run must be exactly the bits of its bytes"
77 );
78}
79
80unsafe impl Coordinates for BinaryField8 {
87 const COORDINATES: usize = 8;
88}
89
90unsafe impl Coordinates for BinaryField16 {
91 const COORDINATES: usize = 16;
92}
93
94unsafe impl Coordinates for BinaryField32 {
95 const COORDINATES: usize = 32;
96}
97
98unsafe impl Coordinates for BinaryField64 {
99 const COORDINATES: usize = 64;
100}
101
102unsafe impl Coordinates for BinaryField128 {
103 const COORDINATES: usize = 128;
104}
105
106unsafe impl<U: Underlier> Coordinates for PackedGf2<U> {
109 const COORDINATES: usize = U::BITS;
110}
111
112const PACK_CHUNK_BYTES: usize = 1 << 20;
116
117#[must_use]
125pub fn pack<A, EF>(cells: &[A]) -> Vec<EF>
126where
127 A: Coordinates,
128 EF: Coordinates,
129{
130 const {
131 check_coordinates::<A>();
132 check_coordinates::<EF>();
133 assert!(
134 cfg!(target_endian = "little"),
135 "packing a coordinate run needs a little-endian target"
136 );
137 }
138 let coordinates = cells.len() * A::COORDINATES;
139 assert_eq!(
140 coordinates % EF::COORDINATES,
141 0,
142 "a packed run must fill whole elements"
143 );
144 let len = coordinates / EF::COORDINATES;
145 let mut packed = Vec::<EF>::with_capacity(len);
146
147 let elements_per_chunk = (PACK_CHUNK_BYTES / size_of::<EF>()).max(1);
152 packed.spare_capacity_mut()[..len]
153 .par_chunks_mut(elements_per_chunk)
154 .zip(coordinate_bytes(cells).par_chunks(elements_per_chunk * size_of::<EF>()))
155 .for_each(|(destination, source)| {
156 debug_assert_eq!(source.len(), size_of_val(destination));
157 unsafe {
166 ptr::copy_nonoverlapping(
167 source.as_ptr(),
168 destination.as_mut_ptr().cast::<u8>(),
169 source.len(),
170 );
171 }
172 });
173 unsafe {
175 packed.set_len(len);
176 }
177 packed
178}
179
180#[must_use]
184pub fn unpack<A, EF>(elements: &[EF]) -> Vec<A>
185where
186 A: Coordinates,
187 EF: Coordinates,
188{
189 const {
190 check_coordinates::<A>();
191 check_coordinates::<EF>();
192 assert!(
193 cfg!(target_endian = "little"),
194 "unpacking a coordinate run needs a little-endian target"
195 );
196 }
197 let coordinates = elements.len() * EF::COORDINATES;
198 assert_eq!(
199 coordinates % A::COORDINATES,
200 0,
201 "an unpacked run must fill whole cells"
202 );
203 let len = coordinates / A::COORDINATES;
204 let mut cells = Vec::<A>::with_capacity(len);
205
206 unsafe {
208 ptr::copy_nonoverlapping(
209 elements.as_ptr().cast::<u8>(),
210 cells.as_mut_ptr().cast::<u8>(),
211 len * size_of::<A>(),
212 );
213 cells.set_len(len);
214 }
215 cells
216}
217
218#[must_use]
220pub const fn coordinate_bytes<A: Coordinates>(cells: &[A]) -> &[u8] {
221 const {
222 check_coordinates::<A>();
223 assert!(
224 cfg!(target_endian = "little"),
225 "viewing coordinate bytes needs a little-endian target"
226 );
227 }
228
229 unsafe { slice::from_raw_parts(cells.as_ptr().cast::<u8>(), size_of_val(cells)) }
234}
235
236#[derive(Clone, Debug, PartialEq, Eq)]
250pub struct PackedStack<A, EF> {
251 elements: Vec<EF>,
253 column_len: usize,
255 _alphabet: PhantomData<A>,
257}
258
259impl<A, EF> PackedStack<A, EF>
260where
261 A: Coordinates,
262 EF: Coordinates + TowerLevel,
263{
264 pub const CELLS_PER_ELEMENT: usize = EF::COORDINATES / A::COORDINATES;
266
267 pub fn from_columns(columns: &[&[A]]) -> Result<Self, PackError> {
277 if A::COORDINATES > EF::COORDINATES {
278 return Err(PackError::CellWiderThanElement {
279 cell: A::COORDINATES,
280 element: EF::COORDINATES,
281 });
282 }
283 let Some((&first, rest)) = columns.split_first() else {
284 return Err(PackError::NoColumns);
285 };
286 if let Some(other) = rest.iter().find(|column| column.len() != first.len()) {
287 return Err(PackError::RaggedColumns {
288 expected: first.len(),
289 actual: other.len(),
290 });
291 }
292 if !first.len().is_multiple_of(Self::CELLS_PER_ELEMENT) {
293 return Err(PackError::PartialElement {
294 cells: first.len(),
295 per_element: Self::CELLS_PER_ELEMENT,
296 });
297 }
298
299 let column_len = first.len() / Self::CELLS_PER_ELEMENT;
300 if column_len == 0 || !column_len.is_power_of_two() {
301 return Err(PackError::NotAHypercube {
302 elements: column_len,
303 });
304 }
305
306 let elements = if let [column] = columns {
310 pack::<A, EF>(column)
311 } else {
312 let mut elements = Vec::with_capacity(column_len * columns.len());
313 for column in columns {
314 elements.extend(pack::<A, EF>(column));
315 }
316 elements
317 };
318
319 Ok(Self {
320 elements,
321 column_len,
322 _alphabet: PhantomData,
323 })
324 }
325
326 #[must_use]
328 pub fn elements(&self) -> &[EF] {
329 &self.elements
330 }
331
332 #[must_use]
334 pub const fn num_columns(&self) -> usize {
335 self.elements.len() / self.column_len
336 }
337
338 #[must_use]
340 pub const fn column_num_variables(&self) -> usize {
341 log2_strict_usize(self.column_len)
342 }
343
344 #[must_use]
346 pub const fn column_num_cells(&self) -> usize {
347 self.column_len * Self::CELLS_PER_ELEMENT
348 }
349
350 #[must_use]
352 pub fn column(&self, index: usize) -> &[EF] {
353 assert!(index < self.num_columns(), "column index out of range");
354 &self.elements[index * self.column_len..][..self.column_len]
355 }
356
357 #[must_use]
365 pub fn locate(&self, column: usize, cell: usize) -> (usize, usize) {
366 assert!(column < self.num_columns(), "column index out of range");
367 assert!(cell < self.column_num_cells(), "cell index out of range");
368 let element = column * self.column_len + cell / Self::CELLS_PER_ELEMENT;
369 (element, cell % Self::CELLS_PER_ELEMENT)
370 }
371
372 #[must_use]
374 pub fn cell(&self, column: usize, cell: usize) -> A {
375 let (element, inside) = self.locate(column, cell);
376 unpack::<A, EF>(&self.elements[element..=element])[inside]
377 }
378
379 #[must_use]
381 pub fn unpack_column(&self, index: usize) -> Vec<A> {
382 unpack::<A, EF>(self.column(index))
383 }
384
385 pub fn into_poly(self) -> Poly<EF> {
387 Poly::new(self.elements)
388 }
389
390 pub fn into_table(self) -> Table<EF> {
394 Table::new(RowMajorMatrix::new(self.elements, self.column_len))
395 }
396}
397
398impl<U, EF> PackedStack<PackedGf2<U>, EF>
399where
400 U: Underlier,
401 EF: Coordinates + TowerLevel,
402{
403 pub fn from_bit_rows(rows: &[PackedGf2<U>], num_columns: usize) -> Result<Self, PackError> {
422 let width = PackedGf2::<U>::WIDTH;
423 if !rows.len().is_multiple_of(width) {
424 return Err(PackError::PartialBitBlock {
425 rows: rows.len(),
426 width,
427 });
428 }
429 if num_columns > width {
430 return Err(PackError::TooManyColumns {
431 requested: num_columns,
432 width,
433 });
434 }
435
436 let blocks = rows.len() / width;
438 let mut columns = alloc::vec![Vec::with_capacity(blocks); num_columns];
439 let mut block = alloc::vec![PackedGf2::<U>::from(Gf2::ZERO); width];
440
441 for source in rows.chunks_exact(width) {
442 block.copy_from_slice(source);
444 PackedGf2::transpose(&mut block);
445
446 for (column, &transposed) in columns.iter_mut().zip(&block) {
447 column.push(transposed);
448 }
449 }
450
451 let views: Vec<&[PackedGf2<U>]> = columns.iter().map(alloc::vec::Vec::as_slice).collect();
452 Self::from_columns(&views)
453 }
454}
455
456#[derive(Clone, Copy, Debug, Error, PartialEq, Eq)]
458#[non_exhaustive]
459pub enum PackError {
460 #[error("a {cell}-coordinate cell does not fit a {element}-coordinate element")]
462 CellWiderThanElement {
463 cell: usize,
465 element: usize,
467 },
468
469 #[error("a packed stack needs at least one column")]
471 NoColumns,
472
473 #[error("a column holds {actual} cells, expected {expected}")]
475 RaggedColumns {
476 expected: usize,
478 actual: usize,
480 },
481
482 #[error("{cells} cells do not fill whole elements of {per_element} cells each")]
484 PartialElement {
485 cells: usize,
487 per_element: usize,
489 },
490
491 #[error("a packed column holds {elements} elements, which is no hypercube")]
493 NotAHypercube {
494 elements: usize,
496 },
497
498 #[error("{rows} bit-sliced rows are no whole number of {width}-row blocks")]
500 PartialBitBlock {
501 rows: usize,
503 width: usize,
505 },
506
507 #[error("{requested} columns exceed the {width} lanes a block holds")]
509 TooManyColumns {
510 requested: usize,
512 width: usize,
514 },
515}
516
517#[cfg(test)]
518mod tests {
519 use p3_binary_field::{PackedGf2x8, PackedGf2x64, PackedGf2x128};
520 use p3_field::{BasedVectorSpace, PrimeCharacteristicRing};
521 use proptest::prelude::*;
522 use rand::rngs::SmallRng;
523 use rand::{RngExt, SeedableRng};
524
525 use super::*;
526
527 type EF = BinaryField128;
528
529 fn coordinate<A: Coordinates>(cells: &[A], index: usize) -> bool {
531 let bytes = coordinate_bytes(cells);
532 (bytes[index / 8] >> (index % 8)) & 1 == 1
533 }
534
535 #[test]
536 fn packing_preserves_every_coordinate_in_order() {
537 let mut rng = SmallRng::seed_from_u64(0xC0DE);
540 let cells: Vec<BinaryField32> = (0..32).map(|_| rng.random()).collect();
541 let packed = pack::<BinaryField32, EF>(&cells);
542
543 assert_eq!(packed.len(), 8);
544 for index in 0..32 * 32 {
545 assert_eq!(
546 coordinate(&packed, index),
547 coordinate(&cells, index),
548 "coordinate {index}"
549 );
550 }
551 }
552
553 #[test]
554 fn the_reinterpretation_agrees_with_the_field_basis_accessor() {
555 let mut rng = SmallRng::seed_from_u64(0xBA515);
562
563 let words: Vec<BinaryField32> = (0..64).map(|_| rng.random()).collect();
564 let by_basis: Vec<EF> = words
565 .as_chunks::<4>()
566 .0
567 .iter()
568 .map(|chunk| EF::from_basis_coefficients_fn(|index| chunk[index]))
569 .collect();
570 assert_eq!(pack::<BinaryField32, EF>(&words), by_basis);
571
572 let halves: Vec<BinaryField64> = (0..64).map(|_| rng.random()).collect();
573 let by_basis: Vec<EF> = halves
574 .as_chunks::<2>()
575 .0
576 .iter()
577 .map(|chunk| EF::from_basis_coefficients_fn(|index| chunk[index]))
578 .collect();
579 assert_eq!(pack::<BinaryField64, EF>(&halves), by_basis);
580 }
581
582 #[test]
583 fn a_round_trip_through_the_packing_is_the_identity() {
584 let mut rng = SmallRng::seed_from_u64(0xB1A5);
586
587 let words: Vec<BinaryField64> = (0..64).map(|_| rng.random()).collect();
588 assert_eq!(unpack::<BinaryField64, EF>(&pack::<_, EF>(&words)), words);
589
590 let bytes: Vec<BinaryField8> = (0..256).map(|_| rng.random()).collect();
591 assert_eq!(unpack::<BinaryField8, EF>(&pack::<_, EF>(&bytes)), bytes);
592
593 let blocks: Vec<PackedGf2x64> = (0..16)
594 .map(|_| PackedGf2x64::new(rng.random::<u64>()))
595 .collect();
596 assert_eq!(unpack::<PackedGf2x64, EF>(&pack::<_, EF>(&blocks)), blocks);
597 }
598
599 #[test]
600 fn a_packed_bit_block_carries_the_lanes_it_was_given() {
601 let mut rng = SmallRng::seed_from_u64(0xB175);
605 let blocks: Vec<PackedGf2x64> = (0..8)
606 .map(|_| PackedGf2x64::new(rng.random::<u64>()))
607 .collect();
608 let packed = pack::<PackedGf2x64, EF>(&blocks);
609
610 assert_eq!(packed.len(), 4);
611 for (b, block) in blocks.iter().enumerate() {
612 for lane in 0..PackedGf2x64::WIDTH {
613 let set = block.get(lane) == Gf2::ONE;
614 assert_eq!(coordinate(&packed, 64 * b + lane), set, "b={b} lane={lane}");
615 }
616 }
617 }
618
619 #[test]
620 fn the_stack_lays_columns_out_back_to_back() {
621 let mut rng = SmallRng::seed_from_u64(0x57AC);
624 let columns: Vec<Vec<BinaryField32>> = (0..3)
625 .map(|_| (0..8).map(|_| rng.random()).collect())
626 .collect();
627 let views: Vec<&[BinaryField32]> = columns.iter().map(Vec::as_slice).collect();
628
629 let stack = PackedStack::<BinaryField32, EF>::from_columns(&views).unwrap();
630
631 assert_eq!(stack.num_columns(), 3);
632 assert_eq!(stack.column_num_variables(), 1);
633 assert_eq!(stack.column_num_cells(), 8);
634 assert_eq!(stack.elements().len(), 6);
635
636 for (index, column) in columns.iter().enumerate() {
638 assert_eq!(&stack.unpack_column(index), column);
639 for (cell, &value) in column.iter().enumerate() {
640 assert_eq!(stack.cell(index, cell), value, "column={index} cell={cell}");
641 }
642 }
643
644 let (element, inside) = stack.locate(2, 5);
646 assert_eq!(element, 2 * 2 + 1);
647 assert_eq!(inside, 1);
648 }
649
650 #[test]
651 fn the_stacked_table_keeps_one_row_per_packed_column() {
652 let mut rng = SmallRng::seed_from_u64(0x7AB1);
654 let columns: Vec<Vec<BinaryField64>> = (0..2)
655 .map(|_| (0..8).map(|_| rng.random()).collect())
656 .collect();
657 let views: Vec<&[BinaryField64]> = columns.iter().map(Vec::as_slice).collect();
658 let stack = PackedStack::<BinaryField64, EF>::from_columns(&views).unwrap();
659
660 let table = stack.clone().into_table();
661 assert_eq!(table.num_polys(), 2);
662 assert_eq!(table.num_variables(), 2);
663 for index in 0..2 {
664 assert_eq!(table.poly(index).as_slice(), stack.column(index));
665 }
666 }
667
668 fn transpose_reference<U: Underlier>(
670 rows: &[PackedGf2<U>],
671 num_columns: usize,
672 ) -> Vec<Vec<PackedGf2<U>>> {
673 let width = PackedGf2::<U>::WIDTH;
674 (0..num_columns)
675 .map(|column| {
676 rows.chunks_exact(width)
677 .map(|block| PackedGf2::<U>::from_fn(|lane| block[lane].get(column)))
678 .collect()
679 })
680 .collect()
681 }
682
683 #[test]
684 fn a_row_major_bit_block_packs_as_its_own_transpose() {
685 let mut rng = SmallRng::seed_from_u64(0x7A05);
692 let rows: Vec<PackedGf2x64> = (0..128)
693 .map(|_| PackedGf2x64::new(rng.random::<u64>()))
694 .collect();
695
696 let stack = PackedStack::<PackedGf2x64, EF>::from_bit_rows(&rows, 40).unwrap();
697 let expected = transpose_reference(&rows, 40);
698
699 assert_eq!(stack.num_columns(), 40);
700 for (column, blocks) in expected.iter().enumerate() {
701 assert_eq!(&stack.unpack_column(column), blocks, "column={column}");
702 }
703 }
704
705 #[test]
706 fn a_row_major_block_is_refused_unless_it_is_square_and_wide_enough() {
707 let rows = alloc::vec![PackedGf2x8::from(Gf2::ZERO); 12];
709 assert_eq!(
710 PackedStack::<PackedGf2x8, EF>::from_bit_rows(&rows, 8).unwrap_err(),
711 PackError::PartialBitBlock { rows: 12, width: 8 }
712 );
713
714 let rows = alloc::vec![PackedGf2x8::from(Gf2::ZERO); 128];
715 assert_eq!(
716 PackedStack::<PackedGf2x8, EF>::from_bit_rows(&rows, 9).unwrap_err(),
717 PackError::TooManyColumns {
718 requested: 9,
719 width: 8,
720 }
721 );
722 }
723
724 #[test]
725 fn a_stack_that_describes_no_commitment_is_refused() {
726 assert_eq!(
728 PackedStack::<BinaryField32, EF>::from_columns(&[]).unwrap_err(),
729 PackError::NoColumns
730 );
731
732 let long = alloc::vec![BinaryField32::ZERO; 8];
734 let short = alloc::vec![BinaryField32::ZERO; 4];
735 assert_eq!(
736 PackedStack::<BinaryField32, EF>::from_columns(&[&long, &short]).unwrap_err(),
737 PackError::RaggedColumns {
738 expected: 8,
739 actual: 4,
740 }
741 );
742
743 let ragged = alloc::vec![BinaryField32::ZERO; 6];
745 assert_eq!(
746 PackedStack::<BinaryField32, EF>::from_columns(&[&ragged]).unwrap_err(),
747 PackError::PartialElement {
748 cells: 6,
749 per_element: 4,
750 }
751 );
752
753 let odd = alloc::vec![BinaryField32::ZERO; 12];
755 assert_eq!(
756 PackedStack::<BinaryField32, EF>::from_columns(&[&odd]).unwrap_err(),
757 PackError::NotAHypercube { elements: 3 }
758 );
759
760 let wide = alloc::vec![PackedGf2x128::from(Gf2::ZERO); 4];
762 assert_eq!(
763 PackedStack::<PackedGf2x128, BinaryField64>::from_columns(&[&wide]).unwrap_err(),
764 PackError::CellWiderThanElement {
765 cell: 128,
766 element: 64,
767 }
768 );
769 }
770
771 proptest! {
772 #[test]
776 fn packing_is_a_bijection_on_coordinates(
777 log_cells in 4usize..=9,
778 seed: u64,
779 ) {
780 let mut rng = SmallRng::seed_from_u64(seed);
781 let cells = 1usize << log_cells;
782
783 let bits: Vec<PackedGf2x8> = (0..cells)
785 .map(|_| PackedGf2x8::new(rng.random::<u8>()))
786 .collect();
787 prop_assert_eq!(unpack::<PackedGf2x8, EF>(&pack::<_, EF>(&bits)), bits);
788
789 let bytes: Vec<BinaryField8> = (0..cells).map(|_| rng.random()).collect();
791 prop_assert_eq!(unpack::<BinaryField8, EF>(&pack::<_, EF>(&bytes)), bytes);
792
793 let words: Vec<BinaryField32> = (0..cells).map(|_| rng.random()).collect();
795 prop_assert_eq!(&unpack::<BinaryField32, EF>(&pack::<_, EF>(&words)), &words);
796
797 let packed = pack::<BinaryField32, EF>(&words);
799 for index in 0..cells * 32 {
800 prop_assert_eq!(coordinate(&packed, index), coordinate(&words, index));
801 }
802 }
803 }
804}