1use alloc::borrow::Cow;
2use alloc::vec;
3use alloc::vec::Vec;
4use core::borrow::{Borrow, BorrowMut};
5use core::marker::PhantomData;
6use core::ops::Deref;
7
8use p3_field::{
9 ExtensionField, Field, PackedValue, par_scale_slice_in_place, scale_slice_in_place_single_core,
10};
11use p3_maybe_rayon::prelude::*;
12use rand::distr::{Distribution, StandardUniform};
13use rand::{Rng, RngExt};
14use serde::{Deserialize, Serialize};
15use tracing::instrument;
16
17use crate::Matrix;
18
19#[derive(Copy, Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
23pub struct DenseMatrix<T, V = Vec<T>> {
24 pub values: V,
26 pub width: usize,
30 _phantom: PhantomData<T>,
34}
35
36pub type RowMajorMatrix<T> = DenseMatrix<T>;
37pub type RowMajorMatrixView<'a, T> = DenseMatrix<T, &'a [T]>;
38pub type RowMajorMatrixViewMut<'a, T> = DenseMatrix<T, &'a mut [T]>;
39pub type RowMajorMatrixCow<'a, T> = DenseMatrix<T, Cow<'a, [T]>>;
40
41pub trait DenseStorage<T>: Borrow<[T]> + Send + Sync {
42 fn to_vec(self) -> Vec<T>;
43}
44
45impl<T: Clone + Send + Sync> DenseStorage<T> for Vec<T> {
47 fn to_vec(self) -> Self {
48 self
49 }
50}
51
52impl<T: Clone + Send + Sync> DenseStorage<T> for &[T] {
53 fn to_vec(self) -> Vec<T> {
54 <[T]>::to_vec(self)
55 }
56}
57
58impl<T: Clone + Send + Sync> DenseStorage<T> for &mut [T] {
59 fn to_vec(self) -> Vec<T> {
60 <[T]>::to_vec(self)
61 }
62}
63
64impl<T: Clone + Send + Sync> DenseStorage<T> for Cow<'_, [T]> {
65 fn to_vec(self) -> Vec<T> {
66 self.into_owned()
67 }
68}
69
70impl<T: Clone + Send + Sync + Default> DenseMatrix<T> {
71 #[must_use]
74 pub fn default(width: usize, height: usize) -> Self {
75 Self::new(vec![T::default(); width * height], width)
76 }
77}
78
79impl<T: Clone + Send + Sync, S: DenseStorage<T>> DenseMatrix<T, S> {
80 #[must_use]
87 pub fn new(values: S, width: usize) -> Self {
88 debug_assert!(values.borrow().len().is_multiple_of(width));
89 Self {
90 values,
91 width,
92 _phantom: PhantomData,
93 }
94 }
95
96 #[must_use]
98 pub fn new_row(values: S) -> Self {
99 let width = values.borrow().len();
100 Self::new(values, width)
101 }
102
103 #[must_use]
105 pub fn new_col(values: S) -> Self {
106 Self::new(values, 1)
107 }
108
109 pub fn as_view(&self) -> RowMajorMatrixView<'_, T> {
111 RowMajorMatrixView::new(self.values.borrow(), self.width)
112 }
113
114 pub fn as_view_mut(&mut self) -> RowMajorMatrixViewMut<'_, T>
116 where
117 S: BorrowMut<[T]>,
118 {
119 RowMajorMatrixViewMut::new(self.values.borrow_mut(), self.width)
120 }
121
122 pub fn copy_from<S2>(&mut self, source: &DenseMatrix<T, S2>)
124 where
125 T: Copy,
126 S: BorrowMut<[T]>,
127 S2: DenseStorage<T>,
128 {
129 assert_eq!(self.dimensions(), source.dimensions());
130 self.par_rows_mut()
133 .zip(source.par_row_slices())
134 .for_each(|(dst, src)| {
135 dst.copy_from_slice(src);
136 });
137 }
138
139 pub fn flatten_to_base<F: Field>(self) -> RowMajorMatrix<F>
141 where
142 T: ExtensionField<F>,
143 {
144 let width = self.width * T::DIMENSION;
145 let values = T::flatten_to_base(self.values.to_vec());
146 RowMajorMatrix::new(values, width)
147 }
148
149 pub fn row_slices(&self) -> impl DoubleEndedIterator<Item = &[T]> {
151 self.values.borrow().chunks_exact(self.width)
152 }
153
154 pub fn par_row_slices(&self) -> impl IndexedParallelIterator<Item = &[T]>
156 where
157 T: Sync,
158 {
159 self.values.borrow().par_chunks_exact(self.width)
160 }
161
162 pub fn row_mut(&mut self, r: usize) -> &mut [T]
167 where
168 S: BorrowMut<[T]>,
169 {
170 &mut self.values.borrow_mut()[r * self.width..(r + 1) * self.width]
171 }
172
173 pub fn rows_mut(&mut self) -> impl Iterator<Item = &mut [T]>
175 where
176 S: BorrowMut<[T]>,
177 {
178 self.values.borrow_mut().chunks_exact_mut(self.width)
179 }
180
181 pub fn par_rows_mut<'a>(&'a mut self) -> impl IndexedParallelIterator<Item = &'a mut [T]>
183 where
184 T: 'a + Send,
185 S: BorrowMut<[T]>,
186 {
187 self.values.borrow_mut().par_chunks_exact_mut(self.width)
188 }
189
190 pub fn horizontally_packed_row_mut<P>(&mut self, r: usize) -> (&mut [P], &mut [T])
195 where
196 P: PackedValue<Value = T>,
197 S: BorrowMut<[T]>,
198 {
199 P::pack_slice_with_suffix_mut(self.row_mut(r))
200 }
201
202 pub fn scale_row(&mut self, r: usize, scale: T)
207 where
208 T: Field,
209 S: BorrowMut<[T]>,
210 {
211 scale_slice_in_place_single_core(self.row_mut(r), scale);
212 }
213
214 pub fn par_scale_row(&mut self, r: usize, scale: T)
223 where
224 T: Field,
225 S: BorrowMut<[T]>,
226 {
227 par_scale_slice_in_place(self.row_mut(r), scale);
228 }
229
230 pub fn scale(&mut self, scale: T)
232 where
233 T: Field,
234 S: BorrowMut<[T]>,
235 {
236 par_scale_slice_in_place(self.values.borrow_mut(), scale);
237 }
238
239 pub fn split_rows(&self, r: usize) -> (RowMajorMatrixView<'_, T>, RowMajorMatrixView<'_, T>) {
244 let (lo, hi) = self.values.borrow().split_at(r * self.width);
245 (
246 DenseMatrix::new(lo, self.width),
247 DenseMatrix::new(hi, self.width),
248 )
249 }
250
251 pub fn split_rows_mut(
256 &mut self,
257 r: usize,
258 ) -> (RowMajorMatrixViewMut<'_, T>, RowMajorMatrixViewMut<'_, T>)
259 where
260 S: BorrowMut<[T]>,
261 {
262 let (lo, hi) = self.values.borrow_mut().split_at_mut(r * self.width);
263 (
264 DenseMatrix::new(lo, self.width),
265 DenseMatrix::new(hi, self.width),
266 )
267 }
268
269 pub fn par_row_chunks(
273 &self,
274 chunk_rows: usize,
275 ) -> impl IndexedParallelIterator<Item = RowMajorMatrixView<'_, T>>
276 where
277 T: Send,
278 {
279 self.values
280 .borrow()
281 .par_chunks(self.width * chunk_rows)
282 .map(|slice| RowMajorMatrixView::new(slice, self.width))
283 }
284
285 pub fn par_row_chunks_exact(
289 &self,
290 chunk_rows: usize,
291 ) -> impl IndexedParallelIterator<Item = RowMajorMatrixView<'_, T>>
292 where
293 T: Send,
294 {
295 self.values
296 .borrow()
297 .par_chunks_exact(self.width * chunk_rows)
298 .map(|slice| RowMajorMatrixView::new(slice, self.width))
299 }
300
301 pub fn par_row_chunks_mut(
305 &mut self,
306 chunk_rows: usize,
307 ) -> impl IndexedParallelIterator<Item = RowMajorMatrixViewMut<'_, T>>
308 where
309 T: Send,
310 S: BorrowMut<[T]>,
311 {
312 self.values
313 .borrow_mut()
314 .par_chunks_mut(self.width * chunk_rows)
315 .map(|slice| RowMajorMatrixViewMut::new(slice, self.width))
316 }
317
318 pub fn row_chunks_exact_mut(
323 &mut self,
324 chunk_rows: usize,
325 ) -> impl Iterator<Item = RowMajorMatrixViewMut<'_, T>>
326 where
327 T: Send,
328 S: BorrowMut<[T]>,
329 {
330 self.values
331 .borrow_mut()
332 .chunks_exact_mut(self.width * chunk_rows)
333 .map(|slice| RowMajorMatrixViewMut::new(slice, self.width))
334 }
335
336 pub fn par_row_chunks_exact_mut(
341 &mut self,
342 chunk_rows: usize,
343 ) -> impl IndexedParallelIterator<Item = RowMajorMatrixViewMut<'_, T>>
344 where
345 T: Send,
346 S: BorrowMut<[T]>,
347 {
348 self.values
349 .borrow_mut()
350 .par_chunks_exact_mut(self.width * chunk_rows)
351 .map(|slice| RowMajorMatrixViewMut::new(slice, self.width))
352 }
353
354 pub fn row_pair_mut(&mut self, row_1: usize, row_2: usize) -> (&mut [T], &mut [T])
359 where
360 S: BorrowMut<[T]>,
361 {
362 debug_assert_ne!(row_1, row_2);
363 let start_1 = row_1 * self.width;
364 let start_2 = row_2 * self.width;
365 let (lo, hi) = self.values.borrow_mut().split_at_mut(start_2);
366 (&mut lo[start_1..][..self.width], &mut hi[..self.width])
367 }
368
369 #[allow(clippy::type_complexity)]
376 pub fn packed_row_pair_mut<P>(
377 &mut self,
378 row_1: usize,
379 row_2: usize,
380 ) -> ((&mut [P], &mut [T]), (&mut [P], &mut [T]))
381 where
382 S: BorrowMut<[T]>,
383 P: PackedValue<Value = T>,
384 {
385 let (slice_1, slice_2) = self.row_pair_mut(row_1, row_2);
386 (
387 P::pack_slice_with_suffix_mut(slice_1),
388 P::pack_slice_with_suffix_mut(slice_2),
389 )
390 }
391
392 #[instrument(level = "debug", skip_all)]
395 pub fn bit_reversed_zero_pad(self, added_bits: usize) -> RowMajorMatrix<T>
396 where
397 T: Field,
398 {
399 if added_bits == 0 {
400 return self.to_row_major_matrix();
401 }
402
403 let w = self.width;
413 let mut padded =
414 RowMajorMatrix::new(T::zero_vec(self.values.borrow().len() << added_bits), w);
415 padded
416 .par_row_chunks_exact_mut(1 << added_bits)
417 .zip(self.par_row_slices())
418 .for_each(|(mut ch, r)| ch.row_mut(0).copy_from_slice(r));
419
420 padded
421 }
422}
423
424impl<T: Clone + Send + Sync, S: DenseStorage<T>> Matrix<T> for DenseMatrix<T, S> {
425 #[inline]
426 fn width(&self) -> usize {
427 self.width
428 }
429
430 #[inline]
431 fn height(&self) -> usize {
432 self.values
433 .borrow()
434 .len()
435 .checked_div(self.width)
436 .unwrap_or(0)
437 }
438
439 #[inline]
440 unsafe fn get_unchecked(&self, r: usize, c: usize) -> T {
441 unsafe {
442 self.values
444 .borrow()
445 .get_unchecked(r * self.width + c)
446 .clone()
447 }
448 }
449
450 #[inline]
451 unsafe fn row_subseq_unchecked(
452 &self,
453 r: usize,
454 start: usize,
455 end: usize,
456 ) -> impl IntoIterator<Item = T, IntoIter = impl Iterator<Item = T> + Send + Sync> {
457 unsafe {
458 self.values
460 .borrow()
461 .get_unchecked(r * self.width + start..r * self.width + end)
462 .iter()
463 .cloned()
464 }
465 }
466
467 #[inline]
468 unsafe fn row_subslice_unchecked(
469 &self,
470 r: usize,
471 start: usize,
472 end: usize,
473 ) -> impl Deref<Target = [T]> {
474 unsafe {
475 self.values
477 .borrow()
478 .get_unchecked(r * self.width + start..r * self.width + end)
479 }
480 }
481
482 fn to_row_major_matrix(self) -> RowMajorMatrix<T>
483 where
484 Self: Sized,
485 T: Clone,
486 {
487 RowMajorMatrix::new(self.values.to_vec(), self.width)
488 }
489
490 #[inline]
491 fn horizontally_packed_row<'a, P>(
492 &'a self,
493 r: usize,
494 ) -> (
495 impl Iterator<Item = P> + Send + Sync,
496 impl Iterator<Item = T> + Send + Sync,
497 )
498 where
499 P: PackedValue<Value = T>,
500 T: Clone + 'a,
501 {
502 let buf = &self.values.borrow()[r * self.width..(r + 1) * self.width];
503 let (packed, sfx) = P::pack_slice_with_suffix(buf);
504 (packed.iter().copied(), sfx.iter().cloned())
505 }
506
507 #[inline]
508 fn padded_horizontally_packed_row<'a, P>(
509 &'a self,
510 r: usize,
511 ) -> impl Iterator<Item = P> + Send + Sync
512 where
513 P: PackedValue<Value = T>,
514 T: Clone + Default + 'a,
515 {
516 let buf = &self.values.borrow()[r * self.width..(r + 1) * self.width];
517 let (packed, sfx) = P::pack_slice_with_suffix(buf);
518 packed.iter().copied().chain(
519 (!sfx.is_empty()).then(|| P::from_fn(|i| sfx.get(i).cloned().unwrap_or_default())),
520 )
521 }
522
523 #[inline]
524 fn vertically_packed_row<P>(&self, r: usize) -> impl Iterator<Item = P>
525 where
526 T: Copy,
527 P: PackedValue<Value = T>,
528 {
529 let values = self.values.borrow();
530 let width = self.width;
531 let height = self.height();
532 let row = r % height;
533 let rows = (P::WIDTH != 1).then(|| self.wrapping_row_slices(r, P::WIDTH));
534
535 (0..width).map(move |c| {
536 if P::WIDTH == 1 {
537 unsafe { P::broadcast(*values.get_unchecked(row * width + c)) }
538 } else {
539 let rows = rows.as_ref().unwrap();
540 P::from_fn(|i| rows[i][c])
541 }
542 })
543 }
544
545 #[inline]
546 fn vertically_packed_row_pair<P>(&self, r: usize, step: usize) -> Vec<P>
547 where
548 T: Copy,
549 P: PackedValue<Value = T>,
550 {
551 if P::WIDTH == 1 {
552 let values = self.values.borrow();
553 let width = self.width;
554 let height = self.height();
555 let row = r % height;
556 let next_row = (r + step) % height;
557 let mut out = Vec::with_capacity(width * 2);
558
559 out.extend(
560 (0..width).map(|c| unsafe { P::broadcast(*values.get_unchecked(row * width + c)) }),
561 );
562 out.extend(
563 (0..width)
564 .map(|c| unsafe { P::broadcast(*values.get_unchecked(next_row * width + c)) }),
565 );
566 out
567 } else {
568 let rows = self.wrapping_row_slices(r, P::WIDTH);
569 let next_rows = self.wrapping_row_slices(r + step, P::WIDTH);
570
571 (0..self.width())
572 .map(|c| P::from_fn(|i| rows[i][c]))
573 .chain((0..self.width()).map(|c| P::from_fn(|i| next_rows[i][c])))
574 .collect::<Vec<_>>()
575 }
576 }
577}
578
579impl<T: Clone + Default + Send + Sync> DenseMatrix<T> {
580 pub fn as_cow<'a>(self) -> RowMajorMatrixCow<'a, T> {
581 RowMajorMatrixCow::new(Cow::Owned(self.values), self.width)
582 }
583
584 pub fn rand<R: Rng>(rng: &mut R, rows: usize, cols: usize) -> Self
585 where
586 StandardUniform: Distribution<T>,
587 {
588 let values = rng.sample_iter(StandardUniform).take(rows * cols).collect();
589 Self::new(values, cols)
590 }
591
592 pub fn rand_nonzero<R: Rng>(rng: &mut R, rows: usize, cols: usize) -> Self
593 where
594 T: Field,
595 StandardUniform: Distribution<T>,
596 {
597 let values = rng
598 .sample_iter(StandardUniform)
599 .filter(|x| !x.is_zero())
600 .take(rows * cols)
601 .collect();
602 Self::new(values, cols)
603 }
604
605 #[instrument(level = "debug", skip_all)]
630 pub fn with_random_cols<R>(&self, num_cols: usize, mut rng: R) -> Self
631 where
632 T: Field,
633 R: Rng + Send + Sync,
634 StandardUniform: Distribution<T>,
635 {
636 let old_w = self.width();
638 let new_w = old_w + num_cols;
639
640 let new_values = T::zero_vec(new_w * self.height());
642 let mut result = Self::new(new_values, new_w);
643
644 result
647 .rows_mut()
648 .zip(self.row_slices())
649 .for_each(|(new_row, old_row)| {
650 new_row[..old_w].copy_from_slice(old_row);
651 new_row[old_w..].iter_mut().for_each(|v| *v = rng.random());
652 });
653 result
654 }
655
656 #[instrument(level = "debug", skip_all)]
679 pub fn with_zero_cols(&self, num_cols: usize) -> Self
680 where
681 T: Field,
682 {
683 let mut result = self.clone();
685 result.widen_right(num_cols, T::ZERO);
686 result
687 }
688
689 pub fn pad_to_height(&mut self, new_height: usize, fill: T) {
690 assert!(new_height >= self.height());
691 self.values.resize(self.width * new_height, fill);
692 }
693
694 pub fn pad_to_power_of_two_height(&mut self, fill: T) {
704 let target_height = self.height().next_power_of_two();
706
707 self.values.resize(self.width * target_height, fill);
710 }
711
712 pub fn pad_to_min_power_of_two_height(&mut self, min_height: usize, fill: T) {
730 let target_height = self
732 .height()
733 .next_power_of_two()
734 .max(min_height.next_power_of_two());
735
736 self.values.resize(self.width * target_height, fill);
738 }
739
740 #[must_use]
760 pub fn from_flat_padded(mut values: Vec<T>, width: usize, fill: T) -> Self {
761 assert!(width > 0, "width must be positive");
763
764 let len = values.len();
766 let rem = len % width;
767
768 if rem != 0 {
773 values.resize(len + (width - rem), fill.clone());
774 }
775
776 if values.is_empty() {
778 values.resize(width, fill);
779 }
780
781 Self::new(values, width)
782 }
783
784 pub fn widen_right(&mut self, extra_cols: usize, fill: T)
813 where
814 T: Copy,
815 {
816 if extra_cols == 0 {
818 return;
819 }
820
821 let old_w = self.width;
822 let new_w = old_w + extra_cols;
823 let h = self.height();
824
825 self.values.resize(h * new_w, fill);
830
831 for r in (1..h).rev() {
836 let src_start = r * old_w;
838
839 let dst_start = r * new_w;
841
842 self.values
844 .copy_within(src_start..src_start + old_w, dst_start);
845
846 self.values[dst_start - extra_cols..dst_start].fill(fill);
848 }
849
850 if h == 1 {
854 self.values[old_w..new_w].fill(fill);
855 }
856
857 self.width = new_w;
859 }
860}
861
862impl<T: Copy + Default + Send + Sync, V: DenseStorage<T>> DenseMatrix<T, V> {
863 pub fn transpose(&self) -> RowMajorMatrix<T> {
865 let nelts = self.height() * self.width();
866 let mut values = vec![T::default(); nelts];
867 p3_util::transpose::transpose(
868 self.values.borrow(),
869 &mut values,
870 self.width(),
871 self.height(),
872 );
873 RowMajorMatrix::new(values, self.height())
874 }
875
876 pub fn transpose_into<W: DenseStorage<T> + BorrowMut<[T]>>(
878 &self,
879 other: &mut DenseMatrix<T, W>,
880 ) {
881 assert_eq!(self.height(), other.width());
882 assert_eq!(other.height(), self.width());
883 p3_util::transpose::transpose(
884 self.values.borrow(),
885 other.values.borrow_mut(),
886 self.width(),
887 self.height(),
888 );
889 }
890}
891
892impl<'a, T: Clone + Default + Send + Sync> RowMajorMatrixView<'a, T> {
893 pub fn as_cow(self) -> RowMajorMatrixCow<'a, T> {
894 RowMajorMatrixCow::new(Cow::Borrowed(self.values), self.width)
895 }
896}
897
898#[cfg(test)]
899mod tests {
900 use p3_baby_bear::BabyBear;
901 use p3_field::{FieldArray, PrimeCharacteristicRing};
902 use rand::SeedableRng;
903 use rand::rngs::SmallRng;
904
905 use super::*;
906
907 #[test]
908 fn test_new() {
909 let matrix = RowMajorMatrix::new(vec![1, 2, 3, 4, 5, 6], 2);
910 assert_eq!(matrix.width, 2);
911 assert_eq!(matrix.height(), 3);
912 assert_eq!(matrix.values, vec![1, 2, 3, 4, 5, 6]);
913 }
914
915 #[test]
916 fn test_new_row() {
917 let matrix = RowMajorMatrix::new_row(vec![1, 2, 3]);
918 assert_eq!(matrix.width, 3);
919 assert_eq!(matrix.height(), 1);
920 }
921
922 #[test]
923 fn test_new_col() {
924 let matrix = RowMajorMatrix::new_col(vec![1, 2, 3]);
925 assert_eq!(matrix.width, 1);
926 assert_eq!(matrix.height(), 3);
927 }
928
929 #[test]
930 fn test_height_with_zero_width() {
931 let matrix: DenseMatrix<i32> = RowMajorMatrix::new(vec![], 0);
932 assert_eq!(matrix.height(), 0);
933 }
934
935 #[test]
936 fn test_get_methods() {
937 let matrix = RowMajorMatrix::new(vec![1, 2, 3, 4, 5, 6], 2); assert_eq!(matrix.get(0, 0), Some(1));
939 assert_eq!(matrix.get(1, 1), Some(4));
940 assert_eq!(matrix.get(2, 0), Some(5));
941 unsafe {
942 assert_eq!(matrix.get_unchecked(0, 1), 2);
943 assert_eq!(matrix.get_unchecked(1, 0), 3);
944 assert_eq!(matrix.get_unchecked(2, 1), 6);
945 }
946 assert_eq!(matrix.get(3, 0), None); assert_eq!(matrix.get(0, 2), None); }
949
950 #[test]
951 fn test_row_methods() {
952 let matrix = RowMajorMatrix::new(vec![1, 2, 3, 4, 5, 6, 7, 8], 4); let row: Vec<_> = matrix.row(1).unwrap().into_iter().collect();
954 assert_eq!(row, vec![5, 6, 7, 8]);
955 unsafe {
956 let row: Vec<_> = matrix.row_unchecked(0).into_iter().collect();
957 assert_eq!(row, vec![1, 2, 3, 4]);
958 let row: Vec<_> = matrix.row_subseq_unchecked(0, 0, 3).into_iter().collect();
959 assert_eq!(row, vec![1, 2, 3]);
960 let row: Vec<_> = matrix.row_subseq_unchecked(0, 1, 3).into_iter().collect();
961 assert_eq!(row, vec![2, 3]);
962 let row: Vec<_> = matrix.row_subseq_unchecked(0, 2, 4).into_iter().collect();
963 assert_eq!(row, vec![3, 4]);
964 }
965 assert!(matrix.row(2).is_none()); }
967
968 #[test]
969 fn test_row_slice_methods() {
970 let matrix = RowMajorMatrix::new(vec![1, 2, 3, 4, 5, 6, 7, 8, 9], 3); let slice0 = matrix.row_slice(0);
972 let slice2 = matrix.row_slice(2);
973 assert_eq!(slice0.unwrap().deref(), &[1, 2, 3]);
974 assert_eq!(slice2.unwrap().deref(), &[7, 8, 9]);
975 unsafe {
976 assert_eq!(&[1, 2, 3], matrix.row_slice_unchecked(0).deref());
977 assert_eq!(&[7, 8, 9], matrix.row_slice_unchecked(2).deref());
978
979 assert_eq!(&[1, 2, 3], matrix.row_subslice_unchecked(0, 0, 3).deref());
980 assert_eq!(&[8], matrix.row_subslice_unchecked(2, 1, 2).deref());
981 }
982 assert!(matrix.row_slice(3).is_none()); }
984
985 #[test]
986 fn test_as_view() {
987 let matrix = RowMajorMatrix::new(vec![1, 2, 3, 4], 2);
988 let view = matrix.as_view();
989 assert_eq!(view.values, &[1, 2, 3, 4]);
990 assert_eq!(view.width, 2);
991 }
992
993 #[test]
994 fn test_as_view_mut() {
995 let mut matrix = RowMajorMatrix::new(vec![1, 2, 3, 4], 2);
996 let view = matrix.as_view_mut();
997 view.values[0] = 10;
998 assert_eq!(matrix.values, vec![10, 2, 3, 4]);
999 }
1000
1001 #[test]
1002 fn test_copy_from() {
1003 let mut matrix1 = RowMajorMatrix::new(vec![0, 0, 0, 0], 2);
1004 let matrix2 = RowMajorMatrix::new(vec![1, 2, 3, 4], 2);
1005 matrix1.copy_from(&matrix2);
1006 assert_eq!(matrix1.values, vec![1, 2, 3, 4]);
1007 }
1008
1009 #[test]
1010 fn test_split_rows() {
1011 let matrix = RowMajorMatrix::new(vec![1, 2, 3, 4, 5, 6], 2);
1012 let (top, bottom) = matrix.split_rows(1);
1013 assert_eq!(top.values, vec![1, 2]);
1014 assert_eq!(bottom.values, vec![3, 4, 5, 6]);
1015 }
1016
1017 #[test]
1018 fn test_split_rows_mut() {
1019 let mut matrix = RowMajorMatrix::new(vec![1, 2, 3, 4, 5, 6], 2);
1020 let (top, bottom) = matrix.split_rows_mut(1);
1021 assert_eq!(top.values, vec![1, 2]);
1022 assert_eq!(bottom.values, vec![3, 4, 5, 6]);
1023 }
1024
1025 #[test]
1026 fn test_row_mut() {
1027 let mut matrix = RowMajorMatrix::new(vec![1, 2, 3, 4, 5, 6], 2);
1028 matrix.row_mut(1)[0] = 10;
1029 assert_eq!(matrix.values, vec![1, 2, 10, 4, 5, 6]);
1030 }
1031
1032 #[test]
1033 fn test_bit_reversed_zero_pad() {
1034 let matrix = RowMajorMatrix::new(
1035 vec![
1036 BabyBear::new(1),
1037 BabyBear::new(2),
1038 BabyBear::new(3),
1039 BabyBear::new(4),
1040 ],
1041 2,
1042 );
1043 let padded = matrix.bit_reversed_zero_pad(1);
1044 assert_eq!(padded.width, 2);
1045 assert_eq!(
1046 padded.values,
1047 vec![
1048 BabyBear::new(1),
1049 BabyBear::new(2),
1050 BabyBear::new(0),
1051 BabyBear::new(0),
1052 BabyBear::new(3),
1053 BabyBear::new(4),
1054 BabyBear::new(0),
1055 BabyBear::new(0)
1056 ]
1057 );
1058 }
1059
1060 #[test]
1061 fn test_bit_reversed_zero_pad_no_change() {
1062 let matrix = RowMajorMatrix::new(
1063 vec![
1064 BabyBear::new(1),
1065 BabyBear::new(2),
1066 BabyBear::new(3),
1067 BabyBear::new(4),
1068 ],
1069 2,
1070 );
1071 let padded = matrix.bit_reversed_zero_pad(0);
1072
1073 assert_eq!(padded.width, 2);
1074 assert_eq!(
1075 padded.values,
1076 vec![
1077 BabyBear::new(1),
1078 BabyBear::new(2),
1079 BabyBear::new(3),
1080 BabyBear::new(4),
1081 ]
1082 );
1083 }
1084
1085 #[test]
1086 fn test_scale() {
1087 let mut matrix = RowMajorMatrix::new(
1088 vec![
1089 BabyBear::new(1),
1090 BabyBear::new(2),
1091 BabyBear::new(3),
1092 BabyBear::new(4),
1093 BabyBear::new(5),
1094 BabyBear::new(6),
1095 ],
1096 2,
1097 );
1098 matrix.scale(BabyBear::new(2));
1099 assert_eq!(
1100 matrix.values,
1101 vec![
1102 BabyBear::new(2),
1103 BabyBear::new(4),
1104 BabyBear::new(6),
1105 BabyBear::new(8),
1106 BabyBear::new(10),
1107 BabyBear::new(12)
1108 ]
1109 );
1110 }
1111
1112 #[test]
1113 fn test_scale_row() {
1114 let mut matrix = RowMajorMatrix::new(
1115 vec![
1116 BabyBear::new(1),
1117 BabyBear::new(2),
1118 BabyBear::new(3),
1119 BabyBear::new(4),
1120 BabyBear::new(5),
1121 BabyBear::new(6),
1122 ],
1123 2,
1124 );
1125 matrix.scale_row(1, BabyBear::new(3));
1126 assert_eq!(
1127 matrix.values,
1128 vec![
1129 BabyBear::new(1),
1130 BabyBear::new(2),
1131 BabyBear::new(9),
1132 BabyBear::new(12),
1133 BabyBear::new(5),
1134 BabyBear::new(6),
1135 ]
1136 );
1137 }
1138
1139 #[test]
1140 fn test_to_row_major_matrix() {
1141 let matrix = RowMajorMatrix::new(vec![1, 2, 3, 4, 5, 6], 2);
1142 let converted = matrix.to_row_major_matrix();
1143
1144 assert_eq!(converted.width, 2);
1146 assert_eq!(converted.height(), 3);
1147 assert_eq!(converted.values, vec![1, 2, 3, 4, 5, 6]);
1148 }
1149
1150 #[test]
1151 fn test_horizontally_packed_row() {
1152 type Packed = FieldArray<BabyBear, 2>;
1153
1154 let matrix = RowMajorMatrix::new(
1155 vec![
1156 BabyBear::new(1),
1157 BabyBear::new(2),
1158 BabyBear::new(3),
1159 BabyBear::new(4),
1160 BabyBear::new(5),
1161 BabyBear::new(6),
1162 ],
1163 3,
1164 );
1165
1166 let (packed_iter, suffix_iter) = matrix.horizontally_packed_row::<Packed>(1);
1167
1168 let packed: Vec<_> = packed_iter.collect();
1169 let suffix: Vec<_> = suffix_iter.collect();
1170
1171 assert_eq!(
1172 packed,
1173 vec![Packed::from([BabyBear::new(4), BabyBear::new(5)])]
1174 );
1175 assert_eq!(suffix, vec![BabyBear::new(6)]);
1176 }
1177
1178 #[test]
1179 fn test_padded_horizontally_packed_row() {
1180 use p3_baby_bear::BabyBear;
1181
1182 type Packed = FieldArray<BabyBear, 2>;
1183
1184 let matrix = RowMajorMatrix::new(
1185 vec![
1186 BabyBear::new(1),
1187 BabyBear::new(2),
1188 BabyBear::new(3),
1189 BabyBear::new(4),
1190 BabyBear::new(5),
1191 BabyBear::new(6),
1192 ],
1193 3,
1194 );
1195
1196 let packed_iter = matrix.padded_horizontally_packed_row::<Packed>(1);
1197 let packed: Vec<_> = packed_iter.collect();
1198
1199 assert_eq!(
1200 packed,
1201 vec![
1202 Packed::from([BabyBear::new(4), BabyBear::new(5)]),
1203 Packed::from([BabyBear::new(6), BabyBear::new(0)])
1204 ]
1205 );
1206 }
1207
1208 #[test]
1209 fn test_padded_horizontally_packed_row_exact_width() {
1210 type Packed = FieldArray<BabyBear, 2>;
1211
1212 let matrix = RowMajorMatrix::new(
1216 vec![
1217 BabyBear::new(1),
1218 BabyBear::new(2),
1219 BabyBear::new(3),
1220 BabyBear::new(4),
1221 BabyBear::new(5),
1222 BabyBear::new(6),
1223 BabyBear::new(7),
1224 BabyBear::new(8),
1225 ],
1226 4,
1227 );
1228
1229 let packed: Vec<_> = matrix.padded_horizontally_packed_row::<Packed>(1).collect();
1230
1231 assert_eq!(packed.len(), 2);
1232 assert_eq!(
1233 packed,
1234 vec![
1235 Packed::from([BabyBear::new(5), BabyBear::new(6)]),
1236 Packed::from([BabyBear::new(7), BabyBear::new(8)]),
1237 ]
1238 );
1239 }
1240
1241 #[test]
1242 fn test_pad_to_height() {
1243 let mut matrix = RowMajorMatrix::new(vec![1, 2, 3, 4, 5, 6], 3);
1244
1245 matrix.pad_to_height(4, 9);
1250
1251 assert_eq!(matrix.height(), 4);
1258 assert_eq!(matrix.values, vec![1, 2, 3, 4, 5, 6, 9, 9, 9, 9, 9, 9]);
1259 }
1260
1261 #[test]
1262 fn test_pad_to_power_of_two_height() {
1263 let mut matrix = RowMajorMatrix::new(vec![1, 2, 3, 4, 5, 6], 2);
1268 assert_eq!(matrix.height(), 3);
1269 matrix.pad_to_power_of_two_height(0);
1270 assert_eq!(matrix.height(), 4);
1271 assert_eq!(matrix.values, vec![1, 2, 3, 4, 5, 6, 0, 0]);
1273
1274 let mut matrix = RowMajorMatrix::new(vec![1, 2, 3, 4, 5, 6, 7, 8], 2);
1279 assert_eq!(matrix.height(), 4);
1280 matrix.pad_to_power_of_two_height(99);
1281 assert_eq!(matrix.height(), 4);
1282 assert_eq!(matrix.values, vec![1, 2, 3, 4, 5, 6, 7, 8]);
1284
1285 let mut matrix = RowMajorMatrix::new(vec![1, 2, 3], 3);
1289 assert_eq!(matrix.height(), 1);
1290 matrix.pad_to_power_of_two_height(42);
1291 assert_eq!(matrix.height(), 1);
1292 assert_eq!(matrix.values, vec![1, 2, 3]);
1293
1294 let mut matrix = RowMajorMatrix::new(vec![1; 10], 2);
1298 assert_eq!(matrix.height(), 5);
1299 matrix.pad_to_power_of_two_height(-1);
1300 assert_eq!(matrix.height(), 8);
1301 assert_eq!(matrix.values.len(), 16);
1303 assert!(matrix.values[..10].iter().all(|&v| v == 1));
1304 assert!(matrix.values[10..].iter().all(|&v| v == -1));
1305 }
1306
1307 #[test]
1308 fn test_pad_to_power_of_two_height_empty_matrix() {
1309 let mut matrix: RowMajorMatrix<i32> = RowMajorMatrix::new(vec![], 3);
1312 assert_eq!(matrix.height(), 0);
1313 assert_eq!(matrix.width, 3);
1314 matrix.pad_to_power_of_two_height(7);
1315 assert_eq!(matrix.height(), 1);
1317 assert_eq!(matrix.values, vec![7, 7, 7]);
1318 }
1319
1320 #[test]
1321 fn test_pad_to_min_power_of_two_height() {
1322 let mut matrix = RowMajorMatrix::new(vec![1, 2, 3, 4, 5, 6], 2);
1328 assert_eq!(matrix.height(), 3);
1329 matrix.pad_to_min_power_of_two_height(5, 0);
1330 assert_eq!(matrix.height(), 8);
1331 assert_eq!(matrix.values[..6], [1, 2, 3, 4, 5, 6]);
1332 assert!(matrix.values[6..].iter().all(|&v| v == 0));
1333
1334 let mut matrix = RowMajorMatrix::new(vec![1; 10], 2);
1340 assert_eq!(matrix.height(), 5);
1341 matrix.pad_to_min_power_of_two_height(2, -1);
1342 assert_eq!(matrix.height(), 8);
1343 assert!(matrix.values[..10].iter().all(|&v| v == 1));
1344 assert!(matrix.values[10..].iter().all(|&v| v == -1));
1345
1346 let mut matrix = RowMajorMatrix::new(vec![1, 2, 3, 4, 5, 6, 7, 8], 2);
1352 assert_eq!(matrix.height(), 4);
1353 matrix.pad_to_min_power_of_two_height(3, 99);
1354 assert_eq!(matrix.height(), 4);
1355 assert_eq!(matrix.values, vec![1, 2, 3, 4, 5, 6, 7, 8]);
1356
1357 let mut matrix = RowMajorMatrix::new(vec![1, 2, 3, 4, 5, 6], 2);
1363 assert_eq!(matrix.height(), 3);
1364 matrix.pad_to_min_power_of_two_height(0, 0);
1365 assert_eq!(matrix.height(), 4);
1366 assert_eq!(matrix.values, vec![1, 2, 3, 4, 5, 6, 0, 0]);
1367
1368 let mut matrix = RowMajorMatrix::new(vec![1, 2, 3, 4, 5, 6], 3);
1370 assert_eq!(matrix.height(), 2);
1371 matrix.pad_to_min_power_of_two_height(8, 7);
1372 assert_eq!(matrix.height(), 8);
1373 assert_eq!(matrix.values[..6], [1, 2, 3, 4, 5, 6]);
1374 assert!(matrix.values[6..].iter().all(|&v| v == 7));
1375 assert_eq!(matrix.values.len(), 24); }
1377
1378 #[test]
1379 fn test_pad_to_min_power_of_two_height_empty_matrix() {
1380 let mut matrix: RowMajorMatrix<i32> = RowMajorMatrix::new(vec![], 3);
1382 assert_eq!(matrix.height(), 0);
1383 matrix.pad_to_min_power_of_two_height(5, 7);
1384 assert_eq!(matrix.height(), 8);
1385 assert_eq!(matrix.values.len(), 24);
1386 assert!(matrix.values.iter().all(|&v| v == 7));
1387 }
1388
1389 #[test]
1390 fn test_from_flat_padded() {
1391 let matrix = RowMajorMatrix::from_flat_padded(vec![1, 2, 3, 4, 5, 6], 3, 0);
1395 assert_eq!(matrix.height(), 2);
1396 assert_eq!(matrix.width, 3);
1397 assert_eq!(matrix.values, vec![1, 2, 3, 4, 5, 6]);
1398
1399 let matrix = RowMajorMatrix::from_flat_padded(vec![1, 2, 3, 4, 5], 3, 99);
1403 assert_eq!(matrix.height(), 2);
1404 assert_eq!(matrix.width, 3);
1405 assert_eq!(matrix.values, vec![1, 2, 3, 4, 5, 99]);
1406
1407 let matrix = RowMajorMatrix::from_flat_padded(vec![42], 3, 0);
1409 assert_eq!(matrix.height(), 1);
1410 assert_eq!(matrix.values, vec![42, 0, 0]);
1411
1412 let matrix = RowMajorMatrix::from_flat_padded(vec![], 4, 7);
1414 assert_eq!(matrix.height(), 1);
1415 assert_eq!(matrix.width, 4);
1416 assert_eq!(matrix.values, vec![7, 7, 7, 7]);
1417
1418 let matrix = RowMajorMatrix::from_flat_padded(vec![10, 20, 30], 1, 0);
1420 assert_eq!(matrix.height(), 3);
1421 assert_eq!(matrix.values, vec![10, 20, 30]);
1422 }
1423
1424 #[test]
1425 #[should_panic(expected = "width must be positive")]
1426 fn test_from_flat_padded_zero_width_panics() {
1427 let _ = RowMajorMatrix::from_flat_padded(vec![1, 2, 3], 0, 0);
1428 }
1429
1430 #[test]
1431 fn test_widen_right() {
1432 let mut matrix = RowMajorMatrix::new(vec![1, 2, 3, 4], 2);
1438 matrix.widen_right(1, 0);
1439 assert_eq!(matrix.width, 3);
1440 assert_eq!(matrix.height(), 2);
1441 assert_eq!(matrix.values, vec![1, 2, 0, 3, 4, 0]);
1442
1443 let mut matrix = RowMajorMatrix::new(vec![1, 2, 3, 4], 2);
1449 matrix.widen_right(3, -1);
1450 assert_eq!(matrix.width, 5);
1451 assert_eq!(matrix.height(), 2);
1452 assert_eq!(matrix.values, vec![1, 2, -1, -1, -1, 3, 4, -1, -1, -1]);
1453
1454 let mut matrix = RowMajorMatrix::new(vec![1, 2, 3, 4], 2);
1456 matrix.widen_right(0, 99);
1457 assert_eq!(matrix.width, 2);
1458 assert_eq!(matrix.values, vec![1, 2, 3, 4]);
1459
1460 let mut matrix = RowMajorMatrix::new(vec![10, 20, 30], 3);
1462 matrix.widen_right(2, 0);
1463 assert_eq!(matrix.width, 5);
1464 assert_eq!(matrix.height(), 1);
1465 assert_eq!(matrix.values, vec![10, 20, 30, 0, 0]);
1466
1467 let mut matrix = RowMajorMatrix::new(vec![1, 2, 3], 1);
1474 matrix.widen_right(2, 0);
1475 assert_eq!(matrix.width, 3);
1476 assert_eq!(matrix.height(), 3);
1477 assert_eq!(matrix.values, vec![1, 0, 0, 2, 0, 0, 3, 0, 0]);
1478 }
1479
1480 #[test]
1481 fn test_widen_right_empty_matrix() {
1482 let mut matrix: RowMajorMatrix<i32> = RowMajorMatrix::new(vec![], 3);
1484 matrix.widen_right(2, 0);
1485 assert_eq!(matrix.width, 5);
1486 assert_eq!(matrix.height(), 0);
1487 assert!(matrix.values.is_empty());
1488 }
1489
1490 #[test]
1491 fn test_transpose_into() {
1492 let matrix = RowMajorMatrix::new(vec![1, 2, 3, 4, 5, 6], 3);
1493
1494 let mut transposed = RowMajorMatrix::new(vec![0; 6], 2);
1499
1500 matrix.transpose_into(&mut transposed);
1501
1502 assert_eq!(transposed.width, 2);
1508 assert_eq!(transposed.height(), 3);
1509 assert_eq!(transposed.values, vec![1, 4, 2, 5, 3, 6]);
1510 }
1511
1512 #[test]
1513 fn test_flatten_to_base() {
1514 let matrix = RowMajorMatrix::new(
1515 vec![
1516 BabyBear::new(2),
1517 BabyBear::new(3),
1518 BabyBear::new(4),
1519 BabyBear::new(5),
1520 ],
1521 2,
1522 );
1523
1524 let flattened: RowMajorMatrix<BabyBear> = matrix.flatten_to_base();
1525
1526 assert_eq!(flattened.width, 2);
1527 assert_eq!(
1528 flattened.values,
1529 vec![
1530 BabyBear::new(2),
1531 BabyBear::new(3),
1532 BabyBear::new(4),
1533 BabyBear::new(5),
1534 ]
1535 );
1536 }
1537
1538 #[test]
1539 fn test_horizontally_packed_row_mut() {
1540 type Packed = FieldArray<BabyBear, 2>;
1541
1542 let mut matrix = RowMajorMatrix::new(
1543 vec![
1544 BabyBear::new(1),
1545 BabyBear::new(2),
1546 BabyBear::new(3),
1547 BabyBear::new(4),
1548 BabyBear::new(5),
1549 BabyBear::new(6),
1550 ],
1551 3,
1552 );
1553
1554 let (packed, suffix) = matrix.horizontally_packed_row_mut::<Packed>(1);
1555 packed[0] = Packed::from([BabyBear::new(9), BabyBear::new(10)]);
1556 suffix[0] = BabyBear::new(11);
1557
1558 assert_eq!(
1559 matrix.values,
1560 vec![
1561 BabyBear::new(1),
1562 BabyBear::new(2),
1563 BabyBear::new(3),
1564 BabyBear::new(9),
1565 BabyBear::new(10),
1566 BabyBear::new(11),
1567 ]
1568 );
1569 }
1570
1571 #[test]
1572 fn test_par_row_chunks() {
1573 let matrix = RowMajorMatrix::new(vec![1, 2, 3, 4, 5, 6, 7, 8], 2);
1574
1575 let chunks: Vec<_> = matrix.par_row_chunks(2).collect();
1576
1577 assert_eq!(chunks.len(), 2);
1578 assert_eq!(chunks[0].values, vec![1, 2, 3, 4]);
1579 assert_eq!(chunks[1].values, vec![5, 6, 7, 8]);
1580 }
1581
1582 #[test]
1583 fn test_par_row_chunks_exact() {
1584 let matrix = RowMajorMatrix::new(vec![1, 2, 3, 4, 5, 6], 2);
1585
1586 let chunks: Vec<_> = matrix.par_row_chunks_exact(1).collect();
1587
1588 assert_eq!(chunks.len(), 3);
1589 assert_eq!(chunks[0].values, vec![1, 2]);
1590 assert_eq!(chunks[1].values, vec![3, 4]);
1591 assert_eq!(chunks[2].values, vec![5, 6]);
1592 }
1593
1594 #[test]
1595 fn test_par_row_chunks_mut() {
1596 let mut matrix = RowMajorMatrix::new(vec![1, 2, 3, 4, 5, 6, 7, 8], 2);
1597
1598 matrix
1599 .par_row_chunks_mut(2)
1600 .for_each(|chunk| chunk.values.iter_mut().for_each(|x| *x += 10));
1601
1602 assert_eq!(matrix.values, vec![11, 12, 13, 14, 15, 16, 17, 18]);
1603 }
1604
1605 #[test]
1606 fn test_row_chunks_exact_mut() {
1607 let mut matrix = RowMajorMatrix::new(vec![1, 2, 3, 4, 5, 6], 2);
1608
1609 for chunk in matrix.row_chunks_exact_mut(1) {
1610 chunk.values.iter_mut().for_each(|x| *x *= 2);
1611 }
1612
1613 assert_eq!(matrix.values, vec![2, 4, 6, 8, 10, 12]);
1614 }
1615
1616 #[test]
1617 fn test_par_row_chunks_exact_mut() {
1618 let mut matrix = RowMajorMatrix::new(vec![1, 2, 3, 4, 5, 6], 2);
1619
1620 matrix
1621 .par_row_chunks_exact_mut(1)
1622 .for_each(|chunk| chunk.values.iter_mut().for_each(|x| *x += 5));
1623
1624 assert_eq!(matrix.values, vec![6, 7, 8, 9, 10, 11]);
1625 }
1626
1627 #[test]
1628 fn test_row_pair_mut() {
1629 let mut matrix = RowMajorMatrix::new(vec![1, 2, 3, 4, 5, 6], 2);
1630
1631 let (row1, row2) = matrix.row_pair_mut(0, 2);
1632 row1[0] = 9;
1633 row2[1] = 10;
1634
1635 assert_eq!(matrix.values, vec![9, 2, 3, 4, 5, 10]);
1636 }
1637
1638 #[test]
1639 fn test_packed_row_pair_mut() {
1640 type Packed = FieldArray<BabyBear, 2>;
1641
1642 let mut matrix = RowMajorMatrix::new(
1643 vec![
1644 BabyBear::new(1),
1645 BabyBear::new(2),
1646 BabyBear::new(3),
1647 BabyBear::new(4),
1648 BabyBear::new(5),
1649 BabyBear::new(6),
1650 ],
1651 3,
1652 );
1653
1654 let ((packed1, sfx1), (packed2, sfx2)) = matrix.packed_row_pair_mut::<Packed>(0, 1);
1655 packed1[0] = Packed::from([BabyBear::new(7), BabyBear::new(8)]);
1656 packed2[0] = Packed::from([BabyBear::new(33), BabyBear::new(44)]);
1657 sfx1[0] = BabyBear::new(99);
1658 sfx2[0] = BabyBear::new(9);
1659
1660 assert_eq!(
1661 matrix.values,
1662 vec![
1663 BabyBear::new(7),
1664 BabyBear::new(8),
1665 BabyBear::new(99),
1666 BabyBear::new(33),
1667 BabyBear::new(44),
1668 BabyBear::new(9),
1669 ]
1670 );
1671 }
1672
1673 #[test]
1674 fn test_transpose_square_matrix() {
1675 const START_INDEX: usize = 1;
1676 const VALUE_LEN: usize = 9;
1677 const WIDTH: usize = 3;
1678 const HEIGHT: usize = 3;
1679
1680 let matrix_values = (START_INDEX..=VALUE_LEN).collect::<Vec<_>>();
1681 let matrix = RowMajorMatrix::new(matrix_values, WIDTH);
1682 let transposed = matrix.transpose();
1683 let should_be_transposed_values = vec![1, 4, 7, 2, 5, 8, 3, 6, 9];
1684 let should_be_transposed = RowMajorMatrix::new(should_be_transposed_values, HEIGHT);
1685 assert_eq!(transposed, should_be_transposed);
1686 }
1687
1688 #[test]
1689 fn test_transpose_row_matrix() {
1690 const START_INDEX: usize = 1;
1691 const VALUE_LEN: usize = 30;
1692 const WIDTH: usize = 1;
1693 const HEIGHT: usize = 30;
1694
1695 let matrix_values = (START_INDEX..=VALUE_LEN).collect::<Vec<_>>();
1696 let matrix = RowMajorMatrix::new(matrix_values.clone(), WIDTH);
1697 let transposed = matrix.transpose();
1698 let should_be_transposed = RowMajorMatrix::new(matrix_values, HEIGHT);
1699 assert_eq!(transposed, should_be_transposed);
1700 }
1701
1702 #[test]
1703 fn test_transpose_rectangular_matrix() {
1704 const START_INDEX: usize = 1;
1705 const VALUE_LEN: usize = 30;
1706 const WIDTH: usize = 5;
1707 const HEIGHT: usize = 6;
1708
1709 let matrix_values = (START_INDEX..=VALUE_LEN).collect::<Vec<_>>();
1710 let matrix = RowMajorMatrix::new(matrix_values, WIDTH);
1711 let transposed = matrix.transpose();
1712 let should_be_transposed_values = vec![
1713 1, 6, 11, 16, 21, 26, 2, 7, 12, 17, 22, 27, 3, 8, 13, 18, 23, 28, 4, 9, 14, 19, 24, 29,
1714 5, 10, 15, 20, 25, 30,
1715 ];
1716 let should_be_transposed = RowMajorMatrix::new(should_be_transposed_values, HEIGHT);
1717 assert_eq!(transposed, should_be_transposed);
1718 }
1719
1720 #[test]
1721 fn test_transpose_larger_rectangular_matrix() {
1722 const START_INDEX: usize = 1;
1723 const VALUE_LEN: usize = 131072; const WIDTH: usize = 256;
1725 const HEIGHT: usize = 512;
1726
1727 let matrix_values = (START_INDEX..=VALUE_LEN).collect::<Vec<_>>();
1728 let matrix = RowMajorMatrix::new(matrix_values, WIDTH);
1729 let transposed = matrix.transpose();
1730
1731 assert_eq!(transposed.width(), HEIGHT);
1732 assert_eq!(transposed.height(), WIDTH);
1733
1734 for col_index in 0..WIDTH {
1735 for row_index in 0..HEIGHT {
1736 assert_eq!(
1737 matrix.values[row_index * WIDTH + col_index],
1738 transposed.values[col_index * HEIGHT + row_index]
1739 );
1740 }
1741 }
1742 }
1743
1744 #[test]
1745 fn test_transpose_very_large_rectangular_matrix() {
1746 const START_INDEX: usize = 1;
1747 const VALUE_LEN: usize = 1048576; const WIDTH: usize = 1024;
1749 const HEIGHT: usize = 1024;
1750
1751 let matrix_values = (START_INDEX..=VALUE_LEN).collect::<Vec<_>>();
1752 let matrix = RowMajorMatrix::new(matrix_values, WIDTH);
1753 let transposed = matrix.transpose();
1754
1755 assert_eq!(transposed.width(), HEIGHT);
1756 assert_eq!(transposed.height(), WIDTH);
1757
1758 for col_index in 0..WIDTH {
1759 for row_index in 0..HEIGHT {
1760 assert_eq!(
1761 matrix.values[row_index * WIDTH + col_index],
1762 transposed.values[col_index * HEIGHT + row_index]
1763 );
1764 }
1765 }
1766 }
1767
1768 #[test]
1769 fn test_vertically_packed_row_scalar_width_1() {
1770 type Packed = BabyBear;
1771
1772 let matrix = RowMajorMatrix::new((1..17).map(BabyBear::new).collect::<Vec<_>>(), 4);
1773 let packed = matrix
1774 .vertically_packed_row::<Packed>(2)
1775 .collect::<Vec<_>>();
1776
1777 assert_eq!(
1778 packed,
1779 vec![
1780 BabyBear::new(9),
1781 BabyBear::new(10),
1782 BabyBear::new(11),
1783 BabyBear::new(12),
1784 ]
1785 );
1786 }
1787
1788 #[test]
1789 fn test_vertically_packed_row_pair() {
1790 type Packed = FieldArray<BabyBear, 2>;
1791
1792 let matrix = RowMajorMatrix::new((1..17).map(BabyBear::new).collect::<Vec<_>>(), 4);
1793
1794 let packed = matrix.vertically_packed_row_pair::<Packed>(0, 2);
1796
1797 assert_eq!(
1813 packed,
1814 (1..5)
1815 .chain(9..13)
1816 .map(|i| [BabyBear::new(i), BabyBear::new(i + 4)].into())
1817 .collect::<Vec<_>>(),
1818 );
1819 }
1820
1821 #[test]
1822 fn test_vertically_packed_row_pair_scalar_width_1() {
1823 type Packed = BabyBear;
1824
1825 let matrix = RowMajorMatrix::new((1..17).map(BabyBear::new).collect::<Vec<_>>(), 4);
1826 let packed = matrix.vertically_packed_row_pair::<Packed>(1, 2);
1827
1828 assert_eq!(
1829 packed,
1830 vec![
1831 BabyBear::new(5),
1832 BabyBear::new(6),
1833 BabyBear::new(7),
1834 BabyBear::new(8),
1835 BabyBear::new(13),
1836 BabyBear::new(14),
1837 BabyBear::new(15),
1838 BabyBear::new(16),
1839 ]
1840 );
1841 }
1842
1843 #[test]
1844 fn test_vertically_packed_row_pair_overlap() {
1845 type Packed = FieldArray<BabyBear, 2>;
1846
1847 let matrix = RowMajorMatrix::new((1..17).map(BabyBear::new).collect::<Vec<_>>(), 4);
1848
1849 let packed = matrix.vertically_packed_row_pair::<Packed>(0, 1);
1866
1867 assert_eq!(
1868 packed,
1869 (1..5)
1870 .chain(5..9)
1871 .map(|i| [BabyBear::new(i), BabyBear::new(i + 4)].into())
1872 .collect::<Vec<_>>(),
1873 );
1874 }
1875
1876 #[test]
1877 fn test_vertically_packed_row_pair_wraparound_start_1() {
1878 use p3_baby_bear::BabyBear;
1879 use p3_field::FieldArray;
1880
1881 type Packed = FieldArray<BabyBear, 2>;
1882
1883 let matrix = RowMajorMatrix::new((1..17).map(BabyBear::new).collect::<Vec<_>>(), 4);
1884
1885 let packed = matrix.vertically_packed_row_pair::<Packed>(1, 2);
1904
1905 assert_eq!(
1906 packed,
1907 vec![
1908 Packed::from([BabyBear::new(5), BabyBear::new(9)]),
1909 Packed::from([BabyBear::new(6), BabyBear::new(10)]),
1910 Packed::from([BabyBear::new(7), BabyBear::new(11)]),
1911 Packed::from([BabyBear::new(8), BabyBear::new(12)]),
1912 Packed::from([BabyBear::new(13), BabyBear::new(1)]),
1913 Packed::from([BabyBear::new(14), BabyBear::new(2)]),
1914 Packed::from([BabyBear::new(15), BabyBear::new(3)]),
1915 Packed::from([BabyBear::new(16), BabyBear::new(4)]),
1916 ]
1917 );
1918 }
1919
1920 #[test]
1921 fn test_with_zero_cols() {
1922 let mat: RowMajorMatrix<BabyBear> =
1928 RowMajorMatrix::new((1..=6).map(BabyBear::new).collect(), 3);
1929 let widened = mat.with_zero_cols(2);
1930
1931 assert_eq!(widened.width(), 5);
1933 assert_eq!(widened.height(), 2);
1934
1935 assert_eq!(
1937 widened.row_slices().next().unwrap(),
1938 &[
1939 BabyBear::new(1),
1940 BabyBear::new(2),
1941 BabyBear::new(3),
1942 BabyBear::ZERO,
1943 BabyBear::ZERO,
1944 ]
1945 );
1946
1947 assert_eq!(
1949 widened.row_slices().nth(1).unwrap(),
1950 &[
1951 BabyBear::new(4),
1952 BabyBear::new(5),
1953 BabyBear::new(6),
1954 BabyBear::ZERO,
1955 BabyBear::ZERO,
1956 ]
1957 );
1958
1959 let same = mat.with_zero_cols(0);
1961 assert_eq!(same.width(), mat.width());
1962 assert_eq!(same.values, mat.values);
1963
1964 let single_row: RowMajorMatrix<BabyBear> =
1968 RowMajorMatrix::new(vec![BabyBear::new(7), BabyBear::new(8)], 2);
1969 let widened = single_row.with_zero_cols(3);
1970 assert_eq!(widened.width(), 5);
1971 assert_eq!(widened.height(), 1);
1972 assert_eq!(
1973 widened.row_slices().next().unwrap(),
1974 &[
1975 BabyBear::new(7),
1976 BabyBear::new(8),
1977 BabyBear::ZERO,
1978 BabyBear::ZERO,
1979 BabyBear::ZERO,
1980 ]
1981 );
1982
1983 let single_col: RowMajorMatrix<BabyBear> = RowMajorMatrix::new(
1989 vec![BabyBear::new(1), BabyBear::new(2), BabyBear::new(3)],
1990 1,
1991 );
1992 let widened = single_col.with_zero_cols(2);
1993 assert_eq!(widened.width(), 3);
1994 assert_eq!(widened.height(), 3);
1995 for (i, row) in widened.row_slices().enumerate() {
1996 assert_eq!(row[0], BabyBear::new((i + 1) as u32));
1998 assert_eq!(row[1], BabyBear::ZERO);
1999 assert_eq!(row[2], BabyBear::ZERO);
2000 }
2001
2002 let empty: RowMajorMatrix<BabyBear> = RowMajorMatrix::new(vec![], 3);
2004 let widened = empty.with_zero_cols(2);
2005 assert_eq!(widened.width(), 5);
2006 assert_eq!(widened.height(), 0);
2007 assert!(widened.values.is_empty());
2008 }
2009
2010 #[test]
2011 fn test_with_zero_cols_matches_widen_right() {
2012 let mat: RowMajorMatrix<BabyBear> =
2020 RowMajorMatrix::new((1..=12).map(BabyBear::new).collect(), 4);
2021
2022 let via_method = mat.with_zero_cols(3);
2024
2025 let mut via_widen = mat;
2027 via_widen.widen_right(3, BabyBear::ZERO);
2028
2029 assert_eq!(via_method.width(), via_widen.width());
2031 assert_eq!(via_method.height(), via_widen.height());
2032 assert_eq!(via_method.values, via_widen.values);
2033 }
2034
2035 #[test]
2036 fn test_with_random_cols() {
2037 let mat: RowMajorMatrix<BabyBear> =
2045 RowMajorMatrix::new((1..=4).map(BabyBear::new).collect(), 2);
2046
2047 let seed = 42u64;
2048 let widened = mat.with_random_cols(3, SmallRng::seed_from_u64(seed));
2049
2050 assert_eq!(widened.width(), 5);
2052 assert_eq!(widened.height(), 2);
2053
2054 let mut reference_rng = SmallRng::seed_from_u64(seed);
2058 for (new_row, old_row) in widened.row_slices().zip(mat.row_slices()) {
2059 assert_eq!(&new_row[..2], old_row);
2061
2062 for val in &new_row[2..] {
2064 let expected: BabyBear = reference_rng.random();
2065 assert_eq!(*val, expected);
2066 }
2067 }
2068 }
2069
2070 #[test]
2071 fn test_with_random_cols_zero_extra() {
2072 let mat: RowMajorMatrix<BabyBear> =
2074 RowMajorMatrix::new((1..=6).map(BabyBear::new).collect(), 3);
2075 let same = mat.with_random_cols(0, SmallRng::seed_from_u64(0));
2076 assert_eq!(same.width(), mat.width());
2077 assert_eq!(same.values, mat.values);
2078 }
2079
2080 #[test]
2081 fn test_with_random_cols_empty_matrix() {
2082 let empty: RowMajorMatrix<BabyBear> = RowMajorMatrix::new(vec![], 3);
2084 let widened = empty.with_random_cols(2, SmallRng::seed_from_u64(0));
2085 assert_eq!(widened.width(), 5);
2086 assert_eq!(widened.height(), 0);
2087 assert!(widened.values.is_empty());
2088 }
2089
2090 #[test]
2091 fn test_with_random_cols_different_seeds() {
2092 let mat: RowMajorMatrix<BabyBear> =
2095 RowMajorMatrix::new((1..=4).map(BabyBear::new).collect(), 2);
2096
2097 let num_random = 4;
2098 let seed_a = 1u64;
2099 let seed_b = 2u64;
2100
2101 let result_a = mat.with_random_cols(num_random, SmallRng::seed_from_u64(seed_a));
2102 let result_b = mat.with_random_cols(num_random, SmallRng::seed_from_u64(seed_b));
2103
2104 for (seed, result) in [(seed_a, &result_a), (seed_b, &result_b)] {
2106 let mut reference_rng = SmallRng::seed_from_u64(seed);
2107 for (new_row, old_row) in result.row_slices().zip(mat.row_slices()) {
2108 assert_eq!(&new_row[..2], old_row);
2110
2111 for val in &new_row[2..] {
2113 let expected: BabyBear = reference_rng.random();
2114 assert_eq!(*val, expected);
2115 }
2116 }
2117 }
2118 }
2119}