p3_matrix/
stack.rs

1use alloc::vec::Vec;
2use core::ops::Deref;
3
4use p3_field::PackedValue;
5
6use crate::Matrix;
7use crate::bitrev::BitReversibleMatrix;
8use crate::dense::RowMajorMatrixView;
9
10/// A type alias representing a vertical composition of two row-major matrix views.
11///
12/// `ViewPair` combines two [`RowMajorMatrixView`]'s with the same element type `T`
13/// and lifetime `'a` into a single virtual matrix stacked vertically.
14///
15/// Both views must have the same width; the resulting view has a height equal
16/// to the sum of the two original heights.
17pub type ViewPair<'a, T> = VerticalPair<RowMajorMatrixView<'a, T>, RowMajorMatrixView<'a, T>>;
18
19/// A matrix composed by stacking two matrices vertically, one on top of the other.
20///
21/// Both matrices must have the same `width`.
22/// The resulting matrix has dimensions:
23/// - `width`: The same as the inputs.
24/// - `height`: The sum of the `heights` of the input matrices.
25///
26/// Element access and iteration will first access the rows of the top matrix,
27/// followed by the rows of the bottom matrix.
28#[derive(Copy, Clone, Debug)]
29pub struct VerticalPair<Top, Bottom> {
30    /// The top matrix in the vertical composition.
31    pub top: Top,
32    /// The bottom matrix in the vertical composition.
33    pub bottom: Bottom,
34}
35
36/// A matrix composed by placing two matrices side-by-side horizontally.
37///
38/// Both matrices must have the same `height`.
39/// The resulting matrix has dimensions:
40/// - `width`: The sum of the `widths` of the input matrices.
41/// - `height`: The same as the inputs.
42///
43/// Element access and iteration for a given row `i` will first access the elements in the `i`'th row of the left matrix,
44/// followed by elements in the `i'`th row of the right matrix.
45#[derive(Copy, Clone, Debug)]
46pub struct HorizontalPair<Left, Right> {
47    /// The left matrix in the horizontal composition.
48    pub left: Left,
49    /// The right matrix in the horizontal composition.
50    pub right: Right,
51}
52
53impl<Top, Bottom> VerticalPair<Top, Bottom> {
54    /// Create a new `VerticalPair` by stacking two matrices vertically.
55    ///
56    /// # Panics
57    /// Panics if the two matrices do not have the same width (i.e., number of columns),
58    /// since vertical composition requires column alignment.
59    ///
60    /// # Returns
61    /// A `VerticalPair` that represents the combined matrix.
62    pub fn new<T>(top: Top, bottom: Bottom) -> Self
63    where
64        T: Send + Sync + Clone,
65        Top: Matrix<T>,
66        Bottom: Matrix<T>,
67    {
68        assert_eq!(top.width(), bottom.width());
69        Self { top, bottom }
70    }
71}
72
73impl<Left, Right> HorizontalPair<Left, Right> {
74    /// Create a new `HorizontalPair` by joining two matrices side by side.
75    ///
76    /// # Panics
77    /// Panics if the two matrices do not have the same height (i.e., number of rows),
78    /// since horizontal composition requires row alignment.
79    ///
80    /// # Returns
81    /// A `HorizontalPair` that represents the combined matrix.
82    pub fn new<T>(left: Left, right: Right) -> Self
83    where
84        T: Send + Sync + Clone,
85        Left: Matrix<T>,
86        Right: Matrix<T>,
87    {
88        assert_eq!(left.height(), right.height());
89        Self { left, right }
90    }
91}
92
93impl<T: Send + Sync + Clone, Top: Matrix<T>, Bottom: Matrix<T>> Matrix<T>
94    for VerticalPair<Top, Bottom>
95{
96    fn width(&self) -> usize {
97        self.top.width()
98    }
99
100    fn height(&self) -> usize {
101        self.top.height() + self.bottom.height()
102    }
103
104    unsafe fn get_unchecked(&self, r: usize, c: usize) -> T {
105        unsafe {
106            // Safety: The caller must ensure that r < self.height() and c < self.width()
107            if r < self.top.height() {
108                self.top.get_unchecked(r, c)
109            } else {
110                self.bottom.get_unchecked(r - self.top.height(), c)
111            }
112        }
113    }
114
115    unsafe fn row_unchecked(
116        &self,
117        r: usize,
118    ) -> impl IntoIterator<Item = T, IntoIter = impl Iterator<Item = T> + Send + Sync> {
119        unsafe {
120            // Safety: The caller must ensure that r < self.height()
121            if r < self.top.height() {
122                EitherRow::Left(self.top.row_unchecked(r).into_iter())
123            } else {
124                EitherRow::Right(self.bottom.row_unchecked(r - self.top.height()).into_iter())
125            }
126        }
127    }
128
129    unsafe fn row_subseq_unchecked(
130        &self,
131        r: usize,
132        start: usize,
133        end: usize,
134    ) -> impl IntoIterator<Item = T, IntoIter = impl Iterator<Item = T> + Send + Sync> {
135        unsafe {
136            // Safety: The caller must ensure that r < self.height() and start <= end <= self.width()
137            if r < self.top.height() {
138                EitherRow::Left(self.top.row_subseq_unchecked(r, start, end).into_iter())
139            } else {
140                EitherRow::Right(
141                    self.bottom
142                        .row_subseq_unchecked(r - self.top.height(), start, end)
143                        .into_iter(),
144                )
145            }
146        }
147    }
148
149    unsafe fn row_slice_unchecked(&self, r: usize) -> impl Deref<Target = [T]> {
150        unsafe {
151            // Safety: The caller must ensure that r < self.height()
152            if r < self.top.height() {
153                EitherRow::Left(self.top.row_slice_unchecked(r))
154            } else {
155                EitherRow::Right(self.bottom.row_slice_unchecked(r - self.top.height()))
156            }
157        }
158    }
159
160    unsafe fn row_subslice_unchecked(
161        &self,
162        r: usize,
163        start: usize,
164        end: usize,
165    ) -> impl Deref<Target = [T]> {
166        unsafe {
167            // Safety: The caller must ensure that r < self.height() and start <= end <= self.width()
168            if r < self.top.height() {
169                EitherRow::Left(self.top.row_subslice_unchecked(r, start, end))
170            } else {
171                EitherRow::Right(self.bottom.row_subslice_unchecked(
172                    r - self.top.height(),
173                    start,
174                    end,
175                ))
176            }
177        }
178    }
179}
180
181impl<T: Send + Sync + Clone, Left: Matrix<T>, Right: Matrix<T>> Matrix<T>
182    for HorizontalPair<Left, Right>
183{
184    fn width(&self) -> usize {
185        self.left.width() + self.right.width()
186    }
187
188    fn height(&self) -> usize {
189        self.left.height()
190    }
191
192    unsafe fn get_unchecked(&self, r: usize, c: usize) -> T {
193        unsafe {
194            // Safety: The caller must ensure that r < self.height() and c < self.width()
195            if c < self.left.width() {
196                self.left.get_unchecked(r, c)
197            } else {
198                self.right.get_unchecked(r, c - self.left.width())
199            }
200        }
201    }
202
203    unsafe fn row_unchecked(
204        &self,
205        r: usize,
206    ) -> impl IntoIterator<Item = T, IntoIter = impl Iterator<Item = T> + Send + Sync> {
207        unsafe {
208            // Safety: The caller must ensure that r < self.height()
209            self.left
210                .row_unchecked(r)
211                .into_iter()
212                .chain(self.right.row_unchecked(r))
213        }
214    }
215
216    #[inline]
217    fn vertically_packed_row<P>(&self, r: usize) -> impl Iterator<Item = P>
218    where
219        T: Copy,
220        P: PackedValue<Value = T>,
221    {
222        self.left
223            .vertically_packed_row::<P>(r)
224            .chain(self.right.vertically_packed_row::<P>(r))
225    }
226
227    #[inline]
228    fn vertically_packed_row_pair<P>(&self, r: usize, step: usize) -> Vec<P>
229    where
230        T: Copy,
231        P: PackedValue<Value = T>,
232    {
233        self.vertically_packed_row::<P>(r)
234            .chain(self.vertically_packed_row::<P>(r + step))
235            .collect()
236    }
237}
238
239/// We use this to wrap both the row iterator and the row slice.
240#[derive(Debug)]
241pub enum EitherRow<L, R> {
242    Left(L),
243    Right(R),
244}
245
246impl<T, L, R> Iterator for EitherRow<L, R>
247where
248    L: Iterator<Item = T>,
249    R: Iterator<Item = T>,
250{
251    type Item = T;
252
253    fn next(&mut self) -> Option<Self::Item> {
254        match self {
255            Self::Left(l) => l.next(),
256            Self::Right(r) => r.next(),
257        }
258    }
259}
260
261impl<T, L, R> Deref for EitherRow<L, R>
262where
263    L: Deref<Target = [T]>,
264    R: Deref<Target = [T]>,
265{
266    type Target = [T];
267    fn deref(&self) -> &Self::Target {
268        match self {
269            Self::Left(l) => l,
270            Self::Right(r) => r,
271        }
272    }
273}
274
275impl<T: Clone + Send + Sync, Left: BitReversibleMatrix<T>, Right: BitReversibleMatrix<T>>
276    BitReversibleMatrix<T> for HorizontalPair<Left, Right>
277{
278    type BitRev = HorizontalPair<Left::BitRev, Right::BitRev>;
279
280    fn bit_reverse_rows(self) -> Self::BitRev {
281        HorizontalPair {
282            left: self.left.bit_reverse_rows(),
283            right: self.right.bit_reverse_rows(),
284        }
285    }
286}
287
288#[cfg(test)]
289mod tests {
290    use alloc::vec;
291    use alloc::vec::Vec;
292
293    use itertools::Itertools;
294
295    use super::*;
296    use crate::RowMajorMatrix;
297
298    #[test]
299    fn test_vertical_pair_empty_top() {
300        let top = RowMajorMatrix::new(vec![], 2); // 0x2
301        let bottom = RowMajorMatrix::new(vec![1, 2, 3, 4], 2); // 2x2
302        let vpair = VerticalPair::new::<i32>(top, bottom);
303        assert_eq!(vpair.height(), 2);
304        assert_eq!(vpair.get(1, 1), Some(4));
305        unsafe {
306            assert_eq!(vpair.get_unchecked(0, 0), 1);
307        }
308    }
309
310    #[test]
311    fn test_vertical_pair_composition() {
312        let top = RowMajorMatrix::new(vec![1, 2, 3, 4], 2); // 2x2
313        let bottom = RowMajorMatrix::new(vec![5, 6, 7, 8], 2); // 2x2
314        let vertical = VerticalPair::new::<i32>(top, bottom);
315
316        // Dimensions
317        assert_eq!(vertical.width(), 2);
318        assert_eq!(vertical.height(), 4);
319
320        // Values from top
321        assert_eq!(vertical.get(0, 0), Some(1));
322        assert_eq!(vertical.get(1, 1), Some(4));
323
324        // Values from bottom
325        unsafe {
326            assert_eq!(vertical.get_unchecked(2, 0), 5);
327            assert_eq!(vertical.get_unchecked(3, 1), 8);
328        }
329
330        // Row iter from bottom
331        let row = vertical.row(3).unwrap().into_iter().collect_vec();
332        assert_eq!(row, vec![7, 8]);
333
334        unsafe {
335            // Row iter from top
336            let row = vertical.row_unchecked(1).into_iter().collect_vec();
337            assert_eq!(row, vec![3, 4]);
338
339            let row = vertical
340                .row_subseq_unchecked(0, 0, 1)
341                .into_iter()
342                .collect_vec();
343            assert_eq!(row, vec![1]);
344        }
345
346        // Row slice
347        assert_eq!(vertical.row_slice(2).unwrap().deref(), &[5, 6]);
348
349        unsafe {
350            // Row slice unchecked
351            assert_eq!(vertical.row_slice_unchecked(3).deref(), &[7, 8]);
352            assert_eq!(vertical.row_subslice_unchecked(1, 1, 2).deref(), &[4]);
353        }
354
355        assert_eq!(vertical.get(0, 2), None); // Width out of bounds
356        assert_eq!(vertical.get(4, 0), None); // Height out of bounds
357        assert!(vertical.row(4).is_none()); // Height out of bounds
358        assert!(vertical.row_slice(4).is_none()); // Height out of bounds
359    }
360
361    #[test]
362    fn test_horizontal_pair_composition() {
363        let left = RowMajorMatrix::new(vec![1, 2, 3, 4], 2); // 2x2
364        let right = RowMajorMatrix::new(vec![5, 6, 7, 8], 2); // 2x2
365        let horizontal = HorizontalPair::new::<i32>(left, right);
366
367        // Dimensions
368        assert_eq!(horizontal.height(), 2);
369        assert_eq!(horizontal.width(), 4);
370
371        // Left values
372        assert_eq!(horizontal.get(0, 0), Some(1));
373        assert_eq!(horizontal.get(1, 1), Some(4));
374
375        // Right values
376        unsafe {
377            assert_eq!(horizontal.get_unchecked(0, 2), 5);
378            assert_eq!(horizontal.get_unchecked(1, 3), 8);
379        }
380
381        // Row iter
382        let row = horizontal.row(0).unwrap().into_iter().collect_vec();
383        assert_eq!(row, vec![1, 2, 5, 6]);
384
385        unsafe {
386            let row = horizontal.row_unchecked(1).into_iter().collect_vec();
387            assert_eq!(row, vec![3, 4, 7, 8]);
388        }
389
390        assert_eq!(horizontal.get(0, 4), None); // Width out of bounds
391        assert_eq!(horizontal.get(2, 0), None); // Height out of bounds
392        assert!(horizontal.row(2).is_none()); // Height out of bounds
393    }
394
395    #[test]
396    fn test_either_row_iterator_behavior() {
397        type Iter = alloc::vec::IntoIter<i32>;
398
399        // Left variant
400        let left: EitherRow<Iter, Iter> = EitherRow::Left(vec![10, 20].into_iter());
401        assert_eq!(left.collect::<Vec<_>>(), vec![10, 20]);
402
403        // Right variant
404        let right: EitherRow<Iter, Iter> = EitherRow::Right(vec![30, 40].into_iter());
405        assert_eq!(right.collect::<Vec<_>>(), vec![30, 40]);
406    }
407
408    #[test]
409    fn test_either_row_deref_behavior() {
410        let left: EitherRow<&[i32], &[i32]> = EitherRow::Left(&[1, 2, 3]);
411        let right: EitherRow<&[i32], &[i32]> = EitherRow::Right(&[4, 5]);
412
413        assert_eq!(&*left, &[1, 2, 3]);
414        assert_eq!(&*right, &[4, 5]);
415    }
416
417    #[test]
418    #[should_panic]
419    fn test_vertical_pair_width_mismatch_should_panic() {
420        let a = RowMajorMatrix::new(vec![1, 2, 3], 1); // 3x1
421        let b = RowMajorMatrix::new(vec![4, 5], 2); // 1x2
422        let _ = VerticalPair::new::<i32>(a, b);
423    }
424
425    #[test]
426    #[should_panic]
427    fn test_horizontal_pair_height_mismatch_should_panic() {
428        let a = RowMajorMatrix::new(vec![1, 2, 3], 3); // 1x3
429        let b = RowMajorMatrix::new(vec![4, 5], 1); // 2x1
430        let _ = HorizontalPair::new::<i32>(a, b);
431    }
432
433    #[test]
434    fn test_horizontal_pair_vertically_packed_row_scalar_width_1() {
435        use p3_baby_bear::BabyBear;
436
437        type Packed = BabyBear;
438
439        // Full matrix (4x4), split into two 4x2 halves:
440        // [  1   2 |  3   4  ]  <-- Row 0
441        // [  5   6 |  7   8  ]  <-- Row 1
442        // [  9  10 | 11  12  ]  <-- Row 2
443        // [ 13  14 | 15  16  ]  <-- Row 3
444        let left = RowMajorMatrix::new([1, 2, 5, 6, 9, 10, 13, 14].map(BabyBear::new).to_vec(), 2);
445        let right =
446            RowMajorMatrix::new([3, 4, 7, 8, 11, 12, 15, 16].map(BabyBear::new).to_vec(), 2);
447        let horizontal = HorizontalPair::new::<BabyBear>(left, right);
448
449        let packed = horizontal
450            .vertically_packed_row::<Packed>(2)
451            .collect::<Vec<_>>();
452
453        assert_eq!(
454            packed,
455            vec![
456                BabyBear::new(9),
457                BabyBear::new(10),
458                BabyBear::new(11),
459                BabyBear::new(12),
460            ]
461        );
462    }
463
464    #[test]
465    fn test_horizontal_pair_vertically_packed_row_pair() {
466        use p3_baby_bear::BabyBear;
467        use p3_field::FieldArray;
468
469        type Packed = FieldArray<BabyBear, 2>;
470
471        let left = RowMajorMatrix::new([1, 2, 5, 6, 9, 10, 13, 14].map(BabyBear::new).to_vec(), 2);
472        let right =
473            RowMajorMatrix::new([3, 4, 7, 8, 11, 12, 15, 16].map(BabyBear::new).to_vec(), 2);
474        let horizontal = HorizontalPair::new::<BabyBear>(left, right);
475
476        let packed = horizontal.vertically_packed_row_pair::<Packed>(0, 2);
477
478        assert_eq!(
479            packed,
480            (1..5)
481                .chain(9..13)
482                .map(|i| [BabyBear::new(i), BabyBear::new(i + 4)].into())
483                .collect::<Vec<_>>(),
484        );
485    }
486}
p3_matrix/stack.rs

p3_matrix/
stack.rs
