p3_matrix/
horizontally_truncated.rs

1use alloc::vec::Vec;
2use core::marker::PhantomData;
3use core::ops::Range;
4
5use p3_field::PackedValue;
6
7use crate::Matrix;
8use crate::bitrev::BitReversibleMatrix;
9
10/// A matrix wrapper that exposes a contiguous range of columns from an inner matrix.
11///
12/// This struct:
13/// - wraps another matrix,
14/// - restricts access to only the columns within the specified `column_range`.
15#[derive(Clone)]
16pub struct HorizontallyTruncated<T, Inner> {
17    /// The underlying full matrix being wrapped.
18    inner: Inner,
19    /// The range of columns to expose from the inner matrix.
20    column_range: Range<usize>,
21    /// Marker for the element type `T`, not used at runtime.
22    _phantom: PhantomData<T>,
23}
24
25impl<T, Inner: Matrix<T>> HorizontallyTruncated<T, Inner>
26where
27    T: Send + Sync + Clone,
28{
29    /// Construct a new horizontally truncated view of a matrix.
30    ///
31    /// # Arguments
32    /// - `inner`: The full inner matrix to be wrapped.
33    /// - `truncated_width`: The number of columns to expose from the start (must be ≤ `inner.width()`).
34    ///
35    /// This is equivalent to `new_with_range(inner, 0..truncated_width)`.
36    ///
37    /// Returns `None` if `truncated_width` is greater than the width of the inner matrix.
38    pub fn new(inner: Inner, truncated_width: usize) -> Option<Self> {
39        Self::new_with_range(inner, 0..truncated_width)
40    }
41
42    /// Construct a new view exposing a specific column range of a matrix.
43    ///
44    /// # Arguments
45    /// - `inner`: The full inner matrix to be wrapped.
46    /// - `column_range`: The columns to expose, satisfying `start <= end <= inner.width()`.
47    ///
48    /// Returns `None` if the range is inverted (`start > end`) or extends beyond the inner width.
49    pub fn new_with_range(inner: Inner, column_range: Range<usize>) -> Option<Self> {
50        // Accept the range only when:
51        // - `start <= end`: the range is not inverted.
52        // - `end <= inner.width()`: it stays within the inner matrix.
53        let valid = column_range.start <= column_range.end && column_range.end <= inner.width();
54        valid.then(|| Self {
55            inner,
56            column_range,
57            _phantom: PhantomData,
58        })
59    }
60}
61
62impl<T, Inner> Matrix<T> for HorizontallyTruncated<T, Inner>
63where
64    T: Send + Sync + Clone,
65    Inner: Matrix<T>,
66{
67    /// Returns the number of columns exposed by the truncated matrix.
68    #[inline(always)]
69    fn width(&self) -> usize {
70        self.column_range.len()
71    }
72
73    /// Returns the number of rows in the matrix (same as the inner matrix).
74    #[inline(always)]
75    fn height(&self) -> usize {
76        self.inner.height()
77    }
78
79    #[inline(always)]
80    unsafe fn get_unchecked(&self, r: usize, c: usize) -> T {
81        unsafe {
82            // Safety: The caller must ensure that `c < self.width()` and `r < self.height()`.
83            //
84            // We translate the column index by adding `column_range.start`.
85            self.inner.get_unchecked(r, self.column_range.start + c)
86        }
87    }
88
89    unsafe fn row_unchecked(
90        &self,
91        r: usize,
92    ) -> impl IntoIterator<Item = T, IntoIter = impl Iterator<Item = T> + Send + Sync> {
93        unsafe {
94            // Safety: The caller must ensure that `r < self.height()`.
95            self.inner
96                .row_subseq_unchecked(r, self.column_range.start, self.column_range.end)
97        }
98    }
99
100    unsafe fn row_subseq_unchecked(
101        &self,
102        r: usize,
103        start: usize,
104        end: usize,
105    ) -> impl IntoIterator<Item = T, IntoIter = impl Iterator<Item = T> + Send + Sync> {
106        unsafe {
107            // Safety: The caller must ensure that r < self.height() and start <= end <= self.width().
108            //
109            // We translate the column indices by adding `column_range.start`.
110            self.inner.row_subseq_unchecked(
111                r,
112                self.column_range.start + start,
113                self.column_range.start + end,
114            )
115        }
116    }
117
118    unsafe fn row_subslice_unchecked(
119        &self,
120        r: usize,
121        start: usize,
122        end: usize,
123    ) -> impl core::ops::Deref<Target = [T]> {
124        unsafe {
125            // Safety: The caller must ensure that `r < self.height()` and `start <= end <= self.width()`.
126            //
127            // We translate the column indices by adding `column_range.start`.
128            self.inner.row_subslice_unchecked(
129                r,
130                self.column_range.start + start,
131                self.column_range.start + end,
132            )
133        }
134    }
135
136    #[inline]
137    fn vertically_packed_row<P>(&self, r: usize) -> impl Iterator<Item = P>
138    where
139        T: Copy,
140        P: PackedValue<Value = T>,
141    {
142        self.inner
143            .vertically_packed_row::<P>(r)
144            .skip(self.column_range.start)
145            .take(self.width())
146    }
147
148    #[inline]
149    fn vertically_packed_row_pair<P>(&self, r: usize, step: usize) -> Vec<P>
150    where
151        T: Copy,
152        P: PackedValue<Value = T>,
153    {
154        self.vertically_packed_row::<P>(r)
155            .chain(self.vertically_packed_row::<P>(r + step))
156            .collect()
157    }
158}
159
160impl<T: Clone + Send + Sync, Inner: BitReversibleMatrix<T>> BitReversibleMatrix<T>
161    for HorizontallyTruncated<T, Inner>
162{
163    type BitRev = HorizontallyTruncated<T, Inner::BitRev>;
164
165    fn bit_reverse_rows(self) -> Self::BitRev {
166        HorizontallyTruncated {
167            inner: self.inner.bit_reverse_rows(),
168            column_range: self.column_range,
169            _phantom: PhantomData,
170        }
171    }
172}
173
174#[cfg(test)]
175mod tests {
176    use alloc::vec;
177    use alloc::vec::Vec;
178
179    use super::*;
180    use crate::dense::RowMajorMatrix;
181
182    #[test]
183    fn test_truncate_width_by_one() {
184        // Create a 3x4 matrix:
185        // [ 1  2  3  4]
186        // [ 5  6  7  8]
187        // [ 9 10 11 12]
188        let inner = RowMajorMatrix::new(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12], 4);
189
190        // Truncate to width 3.
191        let truncated = HorizontallyTruncated::new(inner, 3).unwrap();
192
193        // Width should be 3.
194        assert_eq!(truncated.width(), 3);
195
196        // Height remains unchanged.
197        assert_eq!(truncated.height(), 3);
198
199        // Check individual elements.
200        assert_eq!(truncated.get(0, 0), Some(1)); // row 0, col 0
201        assert_eq!(truncated.get(1, 1), Some(6)); // row 1, col 1
202        unsafe {
203            assert_eq!(truncated.get_unchecked(0, 1), 2); // row 0, col 1
204            assert_eq!(truncated.get_unchecked(2, 2), 11); // row 1, col 0
205        }
206
207        // Row 0: should return [1, 2, 3]
208        let row0: Vec<_> = truncated.row(0).unwrap().into_iter().collect();
209        assert_eq!(row0, vec![1, 2, 3]);
210        unsafe {
211            // Row 2: should return [5, 6, 7]
212            let row1: Vec<_> = truncated.row_unchecked(1).into_iter().collect();
213            assert_eq!(row1, vec![5, 6, 7]);
214
215            // Row 3: is equal to return [9, 10, 11]
216            let row3_subset: Vec<_> = truncated
217                .row_subseq_unchecked(2, 1, 2)
218                .into_iter()
219                .collect();
220            assert_eq!(row3_subset, vec![10]);
221        }
222
223        unsafe {
224            let row1 = truncated.row_slice(1).unwrap();
225            assert_eq!(&*row1, &[5, 6, 7]);
226
227            let row2 = truncated.row_slice_unchecked(2);
228            assert_eq!(&*row2, &[9, 10, 11]);
229
230            let row0_subslice = truncated.row_subslice_unchecked(0, 0, 2);
231            assert_eq!(&*row0_subslice, &[1, 2]);
232        }
233
234        assert!(truncated.get(0, 3).is_none()); // Width out of bounds
235        assert!(truncated.get(3, 0).is_none()); // Height out of bounds
236        assert!(truncated.row(3).is_none()); // Height out of bounds
237        assert!(truncated.row_slice(3).is_none()); // Height out of bounds
238
239        // Convert the truncated view to a RowMajorMatrix and check contents.
240        let as_matrix = truncated.to_row_major_matrix();
241
242        // The expected matrix after truncation:
243        // [1  2  3]
244        // [5  6  7]
245        // [9 10 11]
246        let expected = RowMajorMatrix::new(vec![1, 2, 3, 5, 6, 7, 9, 10, 11], 3);
247
248        assert_eq!(as_matrix, expected);
249    }
250
251    #[test]
252    fn test_no_truncation() {
253        // 2x2 matrix:
254        // [ 7  8 ]
255        // [ 9 10 ]
256        let inner = RowMajorMatrix::new(vec![7, 8, 9, 10], 2);
257
258        // Truncate to full width (no change).
259        let truncated = HorizontallyTruncated::new(inner, 2).unwrap();
260
261        assert_eq!(truncated.width(), 2);
262        assert_eq!(truncated.height(), 2);
263        assert_eq!(truncated.get(0, 1).unwrap(), 8);
264        assert_eq!(truncated.get(1, 0).unwrap(), 9);
265
266        unsafe {
267            assert_eq!(truncated.get_unchecked(0, 0), 7);
268            assert_eq!(truncated.get_unchecked(1, 1), 10);
269        }
270
271        let row0: Vec<_> = truncated.row(0).unwrap().into_iter().collect();
272        assert_eq!(row0, vec![7, 8]);
273
274        let row1: Vec<_> = unsafe { truncated.row_unchecked(1).into_iter().collect() };
275        assert_eq!(row1, vec![9, 10]);
276
277        assert!(truncated.get(0, 2).is_none()); // Width out of bounds
278        assert!(truncated.get(2, 0).is_none()); // Height out of bounds
279        assert!(truncated.row(2).is_none()); // Height out of bounds
280        assert!(truncated.row_slice(2).is_none()); // Height out of bounds
281    }
282
283    #[test]
284    fn test_truncate_to_zero_width() {
285        // 1x3 matrix: [11 12 13]
286        let inner = RowMajorMatrix::new(vec![11, 12, 13], 3);
287
288        // Truncate to width 0.
289        let truncated = HorizontallyTruncated::new(inner, 0).unwrap();
290
291        assert_eq!(truncated.width(), 0);
292        assert_eq!(truncated.height(), 1);
293
294        // Row should be empty.
295        assert!(truncated.row(0).unwrap().into_iter().next().is_none());
296
297        assert!(truncated.get(0, 0).is_none()); // Width out of bounds
298        assert!(truncated.get(1, 0).is_none()); // Height out of bounds
299        assert!(truncated.row(1).is_none()); // Height out of bounds
300        assert!(truncated.row_slice(1).is_none()); // Height out of bounds
301    }
302
303    #[test]
304    fn test_invalid_truncation_width() {
305        // 2x2 matrix:
306        // [1 2]
307        // [3 4]
308        let inner = RowMajorMatrix::new(vec![1, 2, 3, 4], 2);
309
310        // Attempt to truncate beyond inner width (invalid).
311        assert!(HorizontallyTruncated::new(inner, 5).is_none());
312    }
313
314    #[test]
315    fn test_column_range_middle() {
316        // Create a 3x5 matrix:
317        // [ 1  2  3  4  5]
318        // [ 6  7  8  9 10]
319        // [11 12 13 14 15]
320        let inner = RowMajorMatrix::new(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15], 5);
321
322        // Select columns 1..4 (columns 1, 2, 3).
323        let view = HorizontallyTruncated::new_with_range(inner, 1..4).unwrap();
324
325        // Width should be 3 (columns 1, 2, 3).
326        assert_eq!(view.width(), 3);
327
328        // Height remains unchanged.
329        assert_eq!(view.height(), 3);
330
331        // Check individual elements (column indices are relative to the view).
332        assert_eq!(view.get(0, 0), Some(2)); // row 0, col 0 -> inner col 1
333        assert_eq!(view.get(0, 1), Some(3)); // row 0, col 1 -> inner col 2
334        assert_eq!(view.get(0, 2), Some(4)); // row 0, col 2 -> inner col 3
335        assert_eq!(view.get(1, 0), Some(7)); // row 1, col 0 -> inner col 1
336        assert_eq!(view.get(2, 2), Some(14)); // row 2, col 2 -> inner col 3
337
338        unsafe {
339            assert_eq!(view.get_unchecked(1, 1), 8); // row 1, col 1 -> inner col 2
340            assert_eq!(view.get_unchecked(2, 0), 12); // row 2, col 0 -> inner col 1
341        }
342
343        // Row 0: should return [2, 3, 4]
344        let row0: Vec<_> = view.row(0).unwrap().into_iter().collect();
345        assert_eq!(row0, vec![2, 3, 4]);
346
347        // Row 1: should return [7, 8, 9]
348        let row1: Vec<_> = view.row(1).unwrap().into_iter().collect();
349        assert_eq!(row1, vec![7, 8, 9]);
350
351        unsafe {
352            // Row 2: should return [12, 13, 14]
353            let row2: Vec<_> = view.row_unchecked(2).into_iter().collect();
354            assert_eq!(row2, vec![12, 13, 14]);
355
356            // Subsequence of row 1, cols 1..3 (view indices) -> [8, 9]
357            let row1_subseq: Vec<_> = view.row_subseq_unchecked(1, 1, 3).into_iter().collect();
358            assert_eq!(row1_subseq, vec![8, 9]);
359        }
360
361        // Out of bounds checks.
362        assert!(view.get(0, 3).is_none()); // Width out of bounds
363        assert!(view.get(3, 0).is_none()); // Height out of bounds
364
365        // Convert the view to a RowMajorMatrix and check contents.
366        let as_matrix = view.to_row_major_matrix();
367
368        // The expected matrix after selecting columns 1..4:
369        // [2  3  4]
370        // [7  8  9]
371        // [12 13 14]
372        let expected = RowMajorMatrix::new(vec![2, 3, 4, 7, 8, 9, 12, 13, 14], 3);
373
374        assert_eq!(as_matrix, expected);
375    }
376
377    #[test]
378    fn test_column_range_end() {
379        // Create a 2x4 matrix:
380        // [1 2 3 4]
381        // [5 6 7 8]
382        let inner = RowMajorMatrix::new(vec![1, 2, 3, 4, 5, 6, 7, 8], 4);
383
384        // Select columns 2..4 (columns 2, 3).
385        let view = HorizontallyTruncated::new_with_range(inner, 2..4).unwrap();
386
387        assert_eq!(view.width(), 2);
388        assert_eq!(view.height(), 2);
389
390        // Row 0: should return [3, 4]
391        let row0: Vec<_> = view.row(0).unwrap().into_iter().collect();
392        assert_eq!(row0, vec![3, 4]);
393
394        // Row 1: should return [7, 8]
395        let row1: Vec<_> = view.row(1).unwrap().into_iter().collect();
396        assert_eq!(row1, vec![7, 8]);
397
398        assert_eq!(view.get(0, 0), Some(3));
399        assert_eq!(view.get(1, 1), Some(8));
400    }
401
402    #[test]
403    fn test_column_range_single_column() {
404        // Create a 3x4 matrix:
405        // [1 2 3 4]
406        // [5 6 7 8]
407        // [9 10 11 12]
408        let inner = RowMajorMatrix::new(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12], 4);
409
410        // Select only column 2.
411        let view = HorizontallyTruncated::new_with_range(inner, 2..3).unwrap();
412
413        assert_eq!(view.width(), 1);
414        assert_eq!(view.height(), 3);
415
416        assert_eq!(view.get(0, 0), Some(3));
417        assert_eq!(view.get(1, 0), Some(7));
418        assert_eq!(view.get(2, 0), Some(11));
419
420        // Row 0: should return [3]
421        let row0: Vec<_> = view.row(0).unwrap().into_iter().collect();
422        assert_eq!(row0, vec![3]);
423    }
424
425    #[test]
426    fn test_column_range_empty() {
427        // Create a 2x3 matrix:
428        // [1 2 3]
429        // [4 5 6]
430        let inner = RowMajorMatrix::new(vec![1, 2, 3, 4, 5, 6], 3);
431
432        // Select empty range (2..2).
433        let view = HorizontallyTruncated::new_with_range(inner, 2..2).unwrap();
434
435        assert_eq!(view.width(), 0);
436        assert_eq!(view.height(), 2);
437
438        // Row should be empty.
439        assert!(view.row(0).unwrap().into_iter().next().is_none());
440    }
441
442    #[test]
443    fn test_invalid_column_range() {
444        // Create a 2x3 matrix:
445        // [1 2 3]
446        // [4 5 6]
447        let inner = RowMajorMatrix::new(vec![1, 2, 3, 4, 5, 6], 3);
448
449        // Attempt to select columns 1..5 (extends beyond width).
450        assert!(HorizontallyTruncated::new_with_range(inner, 1..5).is_none());
451    }
452
453    #[test]
454    fn test_inverted_column_range_is_rejected() {
455        // 2x3 matrix:
456        // [1 2 3]
457        // [4 5 6]
458        let inner = RowMajorMatrix::new(vec![1, 2, 3, 4, 5, 6], 3);
459
460        // An inverted range has `start > end` while `end` is still inside the inner width.
461        // Built explicitly so the literal does not trip the empty-range lint.
462        //
463        //     2..1: start = 2 > end = 1, yet end = 1 <= width = 3
464        //
465        // The constructor must reject it rather than return a view.
466        let inverted = Range { start: 2, end: 1 };
467        assert!(HorizontallyTruncated::new_with_range(inner, inverted).is_none());
468    }
469
470    #[test]
471    fn test_vertically_packed_row_scalar_width_1() {
472        use p3_baby_bear::BabyBear;
473
474        type Packed = BabyBear;
475
476        // Full matrix (4x4):
477        // [  1   2   3   4  ]  <-- Row 0
478        // [  5   6   7   8  ]  <-- Row 1
479        // [  9  10  11  12  ]  <-- Row 2
480        // [ 13  14  15  16  ]  <-- Row 3
481        let inner = RowMajorMatrix::new((1..17).map(BabyBear::new).collect::<Vec<_>>(), 4);
482
483        // Truncate to the first two columns.
484        let truncated = HorizontallyTruncated::new(inner, 2).unwrap();
485
486        let packed = truncated
487            .vertically_packed_row::<Packed>(2)
488            .collect::<Vec<_>>();
489
490        assert_eq!(packed, vec![BabyBear::new(9), BabyBear::new(10)]);
491    }
492
493    #[test]
494    fn test_vertically_packed_row_pair_middle_column_range() {
495        use p3_baby_bear::BabyBear;
496        use p3_field::FieldArray;
497
498        type Packed = FieldArray<BabyBear, 2>;
499
500        let inner = RowMajorMatrix::new((1..17).map(BabyBear::new).collect::<Vec<_>>(), 4);
501
502        // Select the middle two columns (1, 2), so the packed-row skip/take logic is exercised
503        // with a non-zero `column_range.start`.
504        let view = HorizontallyTruncated::new_with_range(inner, 1..3).unwrap();
505
506        let packed = view.vertically_packed_row_pair::<Packed>(0, 2);
507
508        assert_eq!(
509            packed,
510            vec![
511                [BabyBear::new(2), BabyBear::new(6)].into(),
512                [BabyBear::new(3), BabyBear::new(7)].into(),
513                [BabyBear::new(10), BabyBear::new(14)].into(),
514                [BabyBear::new(11), BabyBear::new(15)].into(),
515            ]
516        );
517    }
518}
p3_matrix/horizontally_truncated.rs

p3_matrix/
horizontally_truncated.rs
