Skip to main content

p3_matrix/
dense.rs

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/// A dense matrix in row-major format, with customizable backing storage.
20///
21/// The data is stored as a flat buffer, where rows are laid out consecutively.
22#[derive(Copy, Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
23pub struct DenseMatrix<T, V = Vec<T>> {
24    /// Flat buffer of matrix values in row-major order.
25    pub values: V,
26    /// Number of columns in the matrix.
27    ///
28    /// The number of rows is implicitly determined as `values.len() / width`.
29    pub width: usize,
30    /// Marker for the element type `T`, unused directly.
31    ///
32    /// Required to retain type information when `V` does not own or contain `T`.
33    _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
45// Cow doesn't impl IntoOwned so we can't blanket it
46impl<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    /// Create a new dense matrix of the given dimensions, backed by a `Vec`, and filled with
72    /// default values.
73    #[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    /// Create a new dense matrix of the given dimensions, backed by the given storage.
81    ///
82    /// # Panics
83    /// Panics in debug builds if `values.len() % width != 0`. Release builds silently
84    /// construct a matrix whose row/column indexing is inconsistent with the storage —
85    /// callers must validate dimensions before calling.
86    #[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    /// Create a new RowMajorMatrix containing a single row.
97    #[must_use]
98    pub fn new_row(values: S) -> Self {
99        let width = values.borrow().len();
100        Self::new(values, width)
101    }
102
103    /// Create a new RowMajorMatrix containing a single column.
104    #[must_use]
105    pub fn new_col(values: S) -> Self {
106        Self::new(values, 1)
107    }
108
109    /// Get a view of the matrix, i.e. a reference to the underlying data.
110    pub fn as_view(&self) -> RowMajorMatrixView<'_, T> {
111        RowMajorMatrixView::new(self.values.borrow(), self.width)
112    }
113
114    /// Get a mutable view of the matrix, i.e. a mutable reference to the underlying data.
115    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    /// Copy the values from the given matrix into this matrix.
123    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        // Equivalent to:
131        // self.values.borrow_mut().copy_from_slice(source.values.borrow());
132        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    /// Flatten an extension field matrix to a base field matrix.
140    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    /// Get an iterator over the rows of the matrix.
150    pub fn row_slices(&self) -> impl DoubleEndedIterator<Item = &[T]> {
151        self.values.borrow().chunks_exact(self.width)
152    }
153
154    /// Get a parallel iterator over the rows of the matrix.
155    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    /// Returns a slice of the given row.
163    ///
164    /// # Panics
165    /// Panics if `r` larger than self.height().
166    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    /// Get a mutable iterator over the rows of the matrix.
174    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    /// Get a mutable parallel iterator over the rows of the matrix.
182    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    /// Get a mutable iterator over the rows of the matrix which packs the rows into packed values.
191    ///
192    /// If `P::WIDTH` does not divide `self.width`, the remainder of the row will be returned as a
193    /// base slice.
194    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    /// Scale the given row by the given value.
203    ///
204    /// # Panics
205    /// Panics if `r` larger than `self.height()`.
206    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    /// Scale the given row by the given value.
215    ///
216    /// # Performance
217    /// This function is parallelized, which may introduce some overhead compared to
218    /// [`Self::scale_row`] when the width is small.
219    ///
220    /// # Panics
221    /// Panics if `r` larger than `self.height()`.
222    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    /// Scale the entire matrix by the given value.
231    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    /// Split the matrix into two matrix views, one with the first `r` rows and one with the remaining rows.
240    ///
241    /// # Panics
242    /// Panics if `r` larger than `self.height()`.
243    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    /// Split the matrix into two mutable matrix views, one with the first `r` rows and one with the remaining rows.
252    ///
253    /// # Panics
254    /// Panics if `r` larger than `self.height()`.
255    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    /// Get an iterator over the rows of the matrix which takes `chunk_rows` rows at a time.
270    ///
271    /// If `chunk_rows` does not divide the height of the matrix, the last chunk will be smaller.
272    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    /// Get a parallel iterator over the rows of the matrix which takes `chunk_rows` rows at a time.
286    ///
287    /// If `chunk_rows` does not divide the height of the matrix, the last chunk will be smaller.
288    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    /// Get a mutable iterator over the rows of the matrix which takes `chunk_rows` rows at a time.
302    ///
303    /// If `chunk_rows` does not divide the height of the matrix, the last chunk will be smaller.
304    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    /// Get a mutable iterator over the rows of the matrix which takes `chunk_rows` rows at a time.
319    ///
320    /// If `chunk_rows` does not divide the height of the matrix, the last up to `chunk_rows - 1` rows
321    /// of the matrix will be omitted.
322    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    /// Get a parallel mutable iterator over the rows of the matrix which takes `chunk_rows` rows at a time.
337    ///
338    /// If `chunk_rows` does not divide the height of the matrix, the last up to `chunk_rows - 1` rows
339    /// of the matrix will be omitted.
340    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    /// Get a pair of mutable slices of the given rows.
355    ///
356    /// # Panics
357    /// Panics if `row_1` or `row_2` are out of bounds or if `row_1 >= row_2`.
358    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    /// Get a pair of mutable slices of the given rows, both packed into packed field elements.
370    ///
371    /// If `P:WIDTH` does not divide `self.width`, the remainder of the row will be returned as a base slice.
372    ///
373    /// # Panics
374    /// Panics if `row_1` or `row_2` are out of bounds or if `row_1 >= row_2`.
375    #[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    /// Append zeros to the "end" of the given matrix, except that the matrix is in bit-reversed order,
393    /// so in actuality we're interleaving zero rows.
394    #[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        // This is equivalent to:
404        //     reverse_matrix_index_bits(mat);
405        //     mat
406        //         .values
407        //         .resize(mat.values.len() << added_bits, F::ZERO);
408        //     reverse_matrix_index_bits(mat);
409        // But rather than implement it with bit reversals, we directly construct the resulting matrix,
410        // whose rows are zero except for rows whose low `added_bits` bits are zero.
411
412        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            // Safety: The caller must ensure that r < self.height() and c < self.width().
443            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            // Safety: The caller must ensure that r < self.height() and start <= end <= self.width().
459            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            // Safety: The caller must ensure that r < self.height()
476            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    /// Return a copy of this matrix with additional columns filled with random
606    /// values appended on the right.
607    ///
608    /// The original columns are preserved unchanged and the new trailing
609    /// columns in each row are populated independently from the provided
610    /// random number generator.
611    ///
612    /// # Memory Layout
613    ///
614    /// ```text
615    ///     Original (h × w):          Result (h × (w + num_cols)):
616    ///     [ a00  a01  …  a0w ]  →    [ a00  a01  …  a0w | r0  r1  …  rN ]
617    ///     [ a10  a11  …  a1w ]       [ a10  a11  …  a1w | r0  r1  …  rN ]
618    ///     …                          …
619    /// ```
620    ///
621    /// # Arguments
622    ///
623    /// - `num_cols`: number of random columns to append.
624    /// - `rng`: random number generator used to sample each new element.
625    ///
626    /// # Returns
627    ///
628    /// A new matrix with width equal to `self.width() + num_cols`.
629    #[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        // Record the original width so we know where to split each row.
637        let old_w = self.width();
638        let new_w = old_w + num_cols;
639
640        // Allocate a zero-initialized buffer for the widened matrix.
641        let new_values = T::zero_vec(new_w * self.height());
642        let mut result = Self::new(new_values, new_w);
643
644        // - Copy original data into the left portion of each row,
645        // - Then fill the right portion with independent random samples.
646        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    /// Return a copy of this matrix with additional zero-filled columns
657    /// appended on the right.
658    ///
659    /// Delegates to cloning the matrix and calling the in-place widening
660    /// method with a zero fill value.
661    ///
662    /// # Memory Layout
663    ///
664    /// ```text
665    ///     Original (h × w):          Result (h × (w + num_cols)):
666    ///     [ a00  a01  …  a0w ]  →    [ a00  a01  …  a0w | 0  0  …  0 ]
667    ///     [ a10  a11  …  a1w ]       [ a10  a11  …  a1w | 0  0  …  0 ]
668    ///     …                          …
669    /// ```
670    ///
671    /// # Arguments
672    ///
673    /// - `num_cols`: number of zero columns to append.
674    ///
675    /// # Returns
676    ///
677    /// A new matrix with width equal to `self.width() + num_cols`.
678    #[instrument(level = "debug", skip_all)]
679    pub fn with_zero_cols(&self, num_cols: usize) -> Self
680    where
681        T: Field,
682    {
683        // Clone the original matrix and widen it in-place with zero fill.
684        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    /// Pad the matrix height to the next power of two by appending rows filled with `fill`.
695    ///
696    /// This is commonly used in proof systems where trace matrices must have power-of-two heights.
697    ///
698    /// # Behavior
699    ///
700    /// - If the matrix is empty (height = 0), it is padded to have exactly one row of `fill` values.
701    /// - If the height is already a power of two, the matrix is unchanged.
702    /// - Otherwise, the matrix is padded to the next power of two height.
703    pub fn pad_to_power_of_two_height(&mut self, fill: T) {
704        // Compute the target height as the next power of two.
705        let target_height = self.height().next_power_of_two();
706
707        // If target_height == height, resize will have no effect.
708        // Otherwise we pad the matrix to a power of two height by filling with the supplied value.
709        self.values.resize(self.width * target_height, fill);
710    }
711
712    /// Pad the matrix height to at least a given minimum, rounded up to the next power of two.
713    ///
714    /// Appends rows filled with the provided fill value.
715    /// Useful in batch proving where multiple trace matrices must share a minimum height
716    /// while still satisfying the power-of-two requirement.
717    ///
718    /// # Logic
719    ///
720    /// - Round both the current height and the minimum up to the next power of two.
721    /// - Take the maximum of those two values as the target.
722    /// - Append fill rows until the target is reached.
723    ///
724    /// # Behavior
725    ///
726    /// - If the matrix already meets or exceeds the target, it is unchanged.
727    /// - If the minimum is 0, this reduces to padding to the next power of two.
728    /// - If the matrix is empty (height = 0), it is padded entirely with fill values.
729    pub fn pad_to_min_power_of_two_height(&mut self, min_height: usize, fill: T) {
730        // Compute the target as the larger of the two power-of-two ceilings.
731        let target_height = self
732            .height()
733            .next_power_of_two()
734            .max(min_height.next_power_of_two());
735
736        // Extend with fill values to reach the target height. No-op if already there.
737        self.values.resize(self.width * target_height, fill);
738    }
739
740    /// Build a matrix from a flat buffer whose length may not be a multiple of the
741    /// requested width.
742    ///
743    /// Useful when constructing trace matrices from a stream of values where the
744    /// final row may be incomplete.
745    ///
746    /// # Arguments
747    ///
748    /// - `values`: flat row-major data, ownership transferred to avoid a copy.
749    /// - `width`: number of columns (must be > 0).
750    /// - `fill`: value used to complete the last row or to create an empty row.
751    ///
752    /// # Returns
753    ///
754    /// A dense matrix with `ceil(values.len() / width)` rows (at least 1).
755    ///
756    /// # Panics
757    ///
758    /// Panics if `width` is zero.
759    #[must_use]
760    pub fn from_flat_padded(mut values: Vec<T>, width: usize, fill: T) -> Self {
761        // Zero width would cause a division-by-zero when computing height.
762        assert!(width > 0, "width must be positive");
763
764        // How many elements the last row is missing.
765        let len = values.len();
766        let rem = len % width;
767
768        // Complete the partial trailing row.
769        //
770        // `resize` is a single capacity check + contiguous fill,
771        // faster than `extend(repeat_n(..))` which uses iterator machinery.
772        if rem != 0 {
773            values.resize(len + (width - rem), fill.clone());
774        }
775
776        // Guarantee at least one row so callers never get a zero-height matrix.
777        if values.is_empty() {
778            values.resize(width, fill);
779        }
780
781        Self::new(values, width)
782    }
783
784    /// Return a new matrix with additional columns appended to the right of
785    /// every row, filled with a constant value.
786    ///
787    /// Useful when a trace matrix needs extra selector or flag columns that are
788    /// initialised to a default.
789    ///
790    /// # Memory Layout
791    ///
792    /// ```text
793    ///  Before (width = W):          After (width = W + extra):
794    ///  [ d0  d1 ... d_{W-1} ]       [ d0  d1 ... d_{W-1}  fill ... fill ]
795    ///  [ ..                 ]       [ ..                                ]
796    /// ```
797    ///
798    /// # Algorithm
799    ///
800    /// Rows are relocated back-to-front so that earlier (lower-address) source
801    /// data is never overwritten before it is read.
802    ///
803    /// - Grow the backing buffer to `height * new_width`.
804    /// - Walk rows from the last to the first.
805    /// - For each row, move its elements to the new position and fill the
806    ///   trailing gap with the provided value.
807    ///
808    /// # Arguments
809    ///
810    /// - `extra_cols`: number of columns to append (0 is a no-op).
811    /// - `fill`: value written into every new column.
812    pub fn widen_right(&mut self, extra_cols: usize, fill: T)
813    where
814        T: Copy,
815    {
816        // No columns to add.
817        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        // Grow the buffer to the widened size.
826        //
827        // The last row's trailing columns are filled for free by `resize`.
828        // Interior gaps still contain stale data — fixed up below.
829        self.values.resize(h * new_w, fill);
830
831        // Reverse iteration prevents clobbering: each row moves to a
832        // higher offset than its source.
833        //
834        // After relocating row r, fill the trailing columns of row r-1.
835        for r in (1..h).rev() {
836            // Source offset in the old (compact) layout.
837            let src_start = r * old_w;
838
839            // Destination offset in the new (widened) layout.
840            let dst_start = r * new_w;
841
842            // Move the row data. Compiles to a single `memmove`.
843            self.values
844                .copy_within(src_start..src_start + old_w, dst_start);
845
846            // Fill row (r-1)'s trailing columns, right before this row.
847            self.values[dst_start - extra_cols..dst_start].fill(fill);
848        }
849
850        // - h >= 2: the r == 1 iteration already filled row 0's gap.
851        // - h == 1: the loop never ran, so row 0's trailing columns are stale.
852        // - h == 0: the buffer is empty — nothing to do.
853        if h == 1 {
854            self.values[old_w..new_w].fill(fill);
855        }
856
857        // Commit the new width so subsequent accesses use the widened stride.
858        self.width = new_w;
859    }
860}
861
862impl<T: Copy + Default + Send + Sync, V: DenseStorage<T>> DenseMatrix<T, V> {
863    /// Return the transpose of this matrix.
864    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    /// Transpose the matrix returning the result in `other` without intermediate allocation.
877    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); // Height = 3, Width = 2
938        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); // Height out of bounds
947        assert_eq!(matrix.get(0, 2), None); // Width out of bounds
948    }
949
950    #[test]
951    fn test_row_methods() {
952        let matrix = RowMajorMatrix::new(vec![1, 2, 3, 4, 5, 6, 7, 8], 4); // Height = 2, Width = 4
953        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()); // Height out of bounds
966    }
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); // Height = 3, Width = 3
971        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()); // Height out of bounds
983    }
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        // The converted matrix should have the same values and width
1145        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        // Width = 4 is exactly divisible by P::WIDTH = 2.
1213        // The iterator must return exactly div_ceil(4, 2) = 2 packed elements,
1214        // with no extra zero-filled element appended.
1215        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        // Original matrix:
1246        // [ 1  2  3 ]
1247        // [ 4  5  6 ] (height = 2)
1248
1249        matrix.pad_to_height(4, 9);
1250
1251        // Expected matrix after padding:
1252        // [ 1  2  3 ]
1253        // [ 4  5  6 ]
1254        // [ 9  9  9 ]  <-- Newly added row
1255        // [ 9  9  9 ]  <-- Newly added row
1256
1257        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        // Test 1: Non-power-of-two height (3 rows -> 4 rows) with fill value 0.
1264        //
1265        // - Original matrix has 3 rows, which is not a power of two.
1266        // - After padding, it should have 4 rows (next power of two).
1267        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        // Original values preserved, new row filled with 0.
1272        assert_eq!(matrix.values, vec![1, 2, 3, 4, 5, 6, 0, 0]);
1273
1274        // Test 2: Already power-of-two height (4 rows -> 4 rows, unchanged).
1275        //
1276        // Matrix height is already a power of two, so no padding occurs.
1277        // Fill value is ignored when no padding is needed.
1278        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        // Values unchanged (fill value not used).
1283        assert_eq!(matrix.values, vec![1, 2, 3, 4, 5, 6, 7, 8]);
1284
1285        // Test 3: Single row matrix (1 row -> 1 row, unchanged).
1286        //
1287        // Height of 1 is a power of two (2^0 = 1).
1288        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        // Test 4: 5 rows -> 8 rows with custom fill value (-1).
1295        //
1296        // Demonstrates padding across a larger gap with a non-zero fill value.
1297        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        // Original 10 values plus 6 fill values (3 new rows * 2 width).
1302        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        // Empty matrix (0 rows) should be padded to 1 row of fill values.
1310        // This ensures the matrix is valid for downstream operations.
1311        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        // After padding: 1 row with 3 columns, all filled with 7.
1316        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        // Test 1: min_height dominates (3 rows, min_height = 5 -> 8 rows).
1323        //
1324        // - Current height 3 rounds to 4.
1325        // - min_height 5 rounds to 8.
1326        // - Target is max(4, 8) = 8.
1327        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        // Test 2: Current height dominates (5 rows, min_height = 2 -> 8 rows).
1335        //
1336        // - Current height 5 rounds to 8.
1337        // - min_height 2 rounds to 2.
1338        // - Target is max(8, 2) = 8.
1339        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        // Test 3: Already at target (4 rows, min_height = 3 -> 4 rows, unchanged).
1347        //
1348        // - Current height 4 is already a power of two.
1349        // - min_height 3 rounds to 4.
1350        // - Target is max(4, 4) = 4, no padding needed.
1351        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        // Test 4: min_height = 0 behaves like pad_to_power_of_two_height.
1358        //
1359        // - min_height 0 rounds to 1.
1360        // - Current height 3 rounds to 4.
1361        // - Target is max(4, 1) = 4.
1362        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        // Test 5: min_height is already a power of two (2 rows, min_height = 8 -> 8 rows).
1369        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); // 8 rows * 3 width
1376    }
1377
1378    #[test]
1379    fn test_pad_to_min_power_of_two_height_empty_matrix() {
1380        // Empty matrix (0 rows) with min_height = 5 should pad to 8 rows.
1381        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        // Test 1: Buffer length is an exact multiple of width (no padding needed).
1392        //
1393        // 6 values with width 3 -> 2 complete rows, no fill appended.
1394        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        // Test 2: Partial last row is padded with fill.
1400        //
1401        // 5 values with width 3 -> 1 complete row + 1 partial row (2 values + 1 fill).
1402        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        // Test 3: Single value with width 3 -> padded to one full row.
1408        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        // Test 4: Empty buffer -> one row filled entirely with fill.
1413        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        // Test 5: Width of 1 never needs padding.
1419        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        // Test 1: Widen a 2x2 matrix by 1 column.
1433        //
1434        // Original:        Widened:
1435        // [ 1  2 ]    ->   [ 1  2  0 ]
1436        // [ 3  4 ]         [ 3  4  0 ]
1437        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        // Test 2: Widen by 3 columns with a non-zero fill.
1444        //
1445        // Original:             Widened:
1446        // [ 1  2 ]    ->       [ 1  2  -1  -1  -1 ]
1447        // [ 3  4 ]             [ 3  4  -1  -1  -1 ]
1448        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        // Test 3: extra_cols = 0 leaves the matrix unchanged.
1455        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        // Test 4: Single-row matrix.
1461        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        // Test 5: Single-column matrix widened to 3 columns.
1468        //
1469        // Original:    Widened:
1470        // [ 1 ]   ->   [ 1  0  0 ]
1471        // [ 2 ]        [ 2  0  0 ]
1472        // [ 3 ]        [ 3  0  0 ]
1473        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        // Empty matrix (0 rows) widened should remain empty with updated width.
1483        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        // Original matrix:
1495        // [ 1  2  3 ]
1496        // [ 4  5  6 ]
1497
1498        let mut transposed = RowMajorMatrix::new(vec![0; 6], 2);
1499
1500        matrix.transpose_into(&mut transposed);
1501
1502        // Expected transposed matrix:
1503        // [ 1  4 ]
1504        // [ 2  5 ]
1505        // [ 3  6 ]
1506
1507        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; // 512 * 256
1724        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; // 512 * 256
1748        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        // Calling the function with r = 0 and step = 2
1795        let packed = matrix.vertically_packed_row_pair::<Packed>(0, 2);
1796
1797        // Matrix visualization:
1798        //
1799        // [  1   2   3   4  ]  <-- Row 0
1800        // [  5   6   7   8  ]  <-- Row 1
1801        // [  9  10  11  12  ]  <-- Row 2
1802        // [ 13  14  15  16  ]  <-- Row 3
1803        //
1804        // Packing rows 0-1 together, then rows 2-3 together:
1805        //
1806        // Packed result:
1807        // [
1808        //   (1, 5), (2, 6), (3, 7), (4, 8),   // First packed row (Row 0 & Row 1)
1809        //   (9, 13), (10, 14), (11, 15), (12, 16),   // Second packed row (Row 2 & Row 3)
1810        // ]
1811
1812        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        // Original matrix visualization:
1850        //
1851        // [  1   2   3   4  ]  <-- Row 0
1852        // [  5   6   7   8  ]  <-- Row 1
1853        // [  9  10  11  12  ]  <-- Row 2
1854        // [ 13  14  15  16  ]  <-- Row 3
1855        //
1856        // Packing rows 0-1 together, then rows 1-2 together:
1857        //
1858        // Expected packed result:
1859        // [
1860        //   (1, 5), (2, 6), (3, 7), (4, 8),   // First packed row (Row 0 & Row 1)
1861        //   (5, 9), (6, 10), (7, 11), (8, 12) // Second packed row (Row 1 & Row 2)
1862        // ]
1863
1864        // Calling the function with overlapping rows (r = 0 and step = 1)
1865        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        // Original matrix visualization:
1886        //
1887        // [  1   2   3   4  ]  <-- Row 0
1888        // [  5   6   7   8  ]  <-- Row 1
1889        // [  9  10  11  12  ]  <-- Row 2
1890        // [ 13  14  15  16  ]  <-- Row 3
1891        //
1892        // Packing starts from row 1, skipping 2 rows (step = 2):
1893        // - The first packed row should contain row 1 & row 2.
1894        // - The second packed row should contain row 3 & row 1 (wraparound case).
1895        //
1896        // Expected packed result:
1897        // [
1898        //   (5, 9), (6, 10), (7, 11), (8, 12),   // Packed row (Row 1 & Row 2)
1899        //   (13, 1), (14, 2), (15, 3), (16, 4)    // Packed row (Row 3 & Row 1)
1900        // ]
1901
1902        // Calling the function with wraparound scenario (starting at r = 1 with step = 2)
1903        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        // Test 1: Append 2 zero columns to a 2×3 matrix.
1923        //
1924        //     Original:             Result:
1925        //     [ 1  2  3 ]    →      [ 1  2  3  0  0 ]
1926        //     [ 4  5  6 ]           [ 4  5  6  0  0 ]
1927        let mat: RowMajorMatrix<BabyBear> =
1928            RowMajorMatrix::new((1..=6).map(BabyBear::new).collect(), 3);
1929        let widened = mat.with_zero_cols(2);
1930
1931        // Verify the new dimensions: width grows by 2, height stays the same.
1932        assert_eq!(widened.width(), 5);
1933        assert_eq!(widened.height(), 2);
1934
1935        // Row 0: original values followed by zeros.
1936        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        // Row 1: original values followed by zeros.
1948        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        // Test 2: Appending 0 columns returns an identical copy.
1960        let same = mat.with_zero_cols(0);
1961        assert_eq!(same.width(), mat.width());
1962        assert_eq!(same.values, mat.values);
1963
1964        // Test 3: Single-row matrix.
1965        //
1966        //     [ 7  8 ]  →  [ 7  8  0  0  0 ]
1967        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        // Test 4: Single-column matrix widened to 3 columns.
1984        //
1985        //     [ 1 ]        [ 1  0  0 ]
1986        //     [ 2 ]   →    [ 2  0  0 ]
1987        //     [ 3 ]        [ 3  0  0 ]
1988        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            // Each row has the original value followed by two zeros.
1997            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        // Test 5: Empty matrix stays empty with updated width.
2003        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        // Both paths must produce identical results: cloning + in-place widen
2013        // with zero fill versus the dedicated method.
2014        //
2015        //     Original (3×4):
2016        //     [  1   2   3   4 ]
2017        //     [  5   6   7   8 ]
2018        //     [  9  10  11  12 ]
2019        let mat: RowMajorMatrix<BabyBear> =
2020            RowMajorMatrix::new((1..=12).map(BabyBear::new).collect(), 4);
2021
2022        // Produce the result via the functional method.
2023        let via_method = mat.with_zero_cols(3);
2024
2025        // Produce the result via move + in-place widen.
2026        let mut via_widen = mat;
2027        via_widen.widen_right(3, BabyBear::ZERO);
2028
2029        // Both matrices must be identical in dimensions and content.
2030        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        // Append 3 random columns to a 2×2 matrix using a seeded RNG.
2038        // We replay the same seed independently to build the exact expected
2039        // matrix, so the assertion is fully deterministic.
2040        //
2041        //     Original:          Result:
2042        //     [ 1  2 ]    →      [ 1  2  r00  r01  r02 ]
2043        //     [ 3  4 ]           [ 3  4  r10  r11  r12 ]
2044        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        // Verify dimensions: width grows by 3, height unchanged.
2051        assert_eq!(widened.width(), 5);
2052        assert_eq!(widened.height(), 2);
2053
2054        // Replay the same seed to produce the expected random values in the
2055        // exact same order the method consumes them: row-by-row, left to right
2056        // within the appended portion.
2057        let mut reference_rng = SmallRng::seed_from_u64(seed);
2058        for (new_row, old_row) in widened.row_slices().zip(mat.row_slices()) {
2059            // Left portion must be the original data, unchanged.
2060            assert_eq!(&new_row[..2], old_row);
2061
2062            // Right portion must match the reference RNG output exactly.
2063            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        // Appending 0 random columns returns an exact copy of the original.
2073        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        // An empty matrix (0 rows) remains empty with updated width.
2083        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        // Verify that two different seeds each produce the correct output by
2093        // replaying both seeds independently. This is fully deterministic.
2094        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        // Replay each seed and verify every element exactly.
2105        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                // Left portion must be the original data, unchanged.
2109                assert_eq!(&new_row[..2], old_row);
2110
2111                // Right portion must match the reference RNG output exactly.
2112                for val in &new_row[2..] {
2113                    let expected: BabyBear = reference_rng.random();
2114                    assert_eq!(*val, expected);
2115                }
2116            }
2117        }
2118    }
2119}