p3_matrix/
row_index_mapped.rs

1use alloc::vec::Vec;
2use core::ops::Deref;
3
4use p3_field::PackedValue;
5
6use crate::Matrix;
7use crate::dense::RowMajorMatrix;
8
9/// A trait for remapping row indices of a matrix.
10///
11/// Implementations can change the number of visible rows (`height`)
12/// and define how a given logical row index maps to a physical one.
13pub trait RowIndexMap: Send + Sync {
14    /// Returns the number of rows exposed by the mapping.
15    fn height(&self) -> usize;
16
17    /// Maps a visible row index `r` to the corresponding row index in the underlying matrix.
18    ///
19    /// The input `r` is assumed to lie in the range `0..self.height()` and the output
20    /// will lie in the range `0..self.inner.height()`.
21    ///
22    /// It is considered undefined behaviour to call `map_row_index` with `r >= self.height()`.
23    fn map_row_index(&self, r: usize) -> usize;
24
25    /// Converts the mapped matrix into a dense row-major matrix.
26    ///
27    /// This default implementation iterates over all mapped rows,
28    /// collects them in order, and builds a dense representation.
29    fn to_row_major_matrix<T: Clone + Send + Sync, Inner: Matrix<T>>(
30        &self,
31        inner: Inner,
32    ) -> RowMajorMatrix<T> {
33        RowMajorMatrix::new(
34            unsafe {
35                // Safety: The output of `map_row_index` is less than `inner.height()` for all inputs in the range `0..self.height()`.
36                (0..self.height())
37                    .flat_map(|r| inner.row_unchecked(self.map_row_index(r)))
38                    .collect()
39            },
40            inner.width(),
41        )
42    }
43}
44
45/// A matrix view that applies a row index mapping to an inner matrix.
46///
47/// The mapping changes which rows are visible and in what order.
48/// The width remains unchanged.
49#[derive(Copy, Clone, Debug)]
50pub struct RowIndexMappedView<IndexMap, Inner> {
51    /// A row index mapping that defines the number and order of visible rows.
52    pub index_map: IndexMap,
53    /// The inner matrix that holds actual data.
54    pub inner: Inner,
55}
56
57impl<T: Send + Sync + Clone, IndexMap: RowIndexMap, Inner: Matrix<T>> Matrix<T>
58    for RowIndexMappedView<IndexMap, Inner>
59{
60    fn width(&self) -> usize {
61        self.inner.width()
62    }
63
64    fn height(&self) -> usize {
65        self.index_map.height()
66    }
67
68    unsafe fn get_unchecked(&self, r: usize, c: usize) -> T {
69        unsafe {
70            // Safety: The caller must ensure that r < self.height() and c < self.width().
71            self.inner.get_unchecked(self.index_map.map_row_index(r), c)
72        }
73    }
74
75    unsafe fn row_unchecked(
76        &self,
77        r: usize,
78    ) -> impl IntoIterator<Item = T, IntoIter = impl Iterator<Item = T> + Send + Sync> {
79        unsafe {
80            // Safety: The caller must ensure that r < self.height().
81            self.inner.row_unchecked(self.index_map.map_row_index(r))
82        }
83    }
84
85    unsafe fn row_subseq_unchecked(
86        &self,
87        r: usize,
88        start: usize,
89        end: usize,
90    ) -> impl IntoIterator<Item = T, IntoIter = impl Iterator<Item = T> + Send + Sync> {
91        unsafe {
92            // Safety: The caller must ensure that r < self.height() and start <= end <= self.width().
93            self.inner
94                .row_subseq_unchecked(self.index_map.map_row_index(r), start, end)
95        }
96    }
97
98    unsafe fn row_slice_unchecked(&self, r: usize) -> impl Deref<Target = [T]> {
99        unsafe {
100            // Safety: The caller must ensure that r < self.height().
101            self.inner
102                .row_slice_unchecked(self.index_map.map_row_index(r))
103        }
104    }
105
106    unsafe fn row_subslice_unchecked(
107        &self,
108        r: usize,
109        start: usize,
110        end: usize,
111    ) -> impl Deref<Target = [T]> {
112        unsafe {
113            // Safety: The caller must ensure that r < self.height() and start <= end <= self.width().
114            self.inner
115                .row_subslice_unchecked(self.index_map.map_row_index(r), start, end)
116        }
117    }
118
119    fn to_row_major_matrix(self) -> RowMajorMatrix<T>
120    where
121        Self: Sized,
122        T: Clone,
123    {
124        // Use Perm's optimized permutation routine, if it has one.
125        self.index_map.to_row_major_matrix(self.inner)
126    }
127
128    fn horizontally_packed_row<'a, P>(
129        &'a self,
130        r: usize,
131    ) -> (
132        impl Iterator<Item = P> + Send + Sync,
133        impl Iterator<Item = T> + Send + Sync,
134    )
135    where
136        P: PackedValue<Value = T>,
137        T: Clone + 'a,
138    {
139        self.inner
140            .horizontally_packed_row(self.index_map.map_row_index(r))
141    }
142
143    fn padded_horizontally_packed_row<'a, P>(
144        &'a self,
145        r: usize,
146    ) -> impl Iterator<Item = P> + Send + Sync
147    where
148        P: PackedValue<Value = T>,
149        T: Clone + Default + 'a,
150    {
151        self.inner
152            .padded_horizontally_packed_row(self.index_map.map_row_index(r))
153    }
154
155    #[inline]
156    fn vertically_packed_row<P>(&self, r: usize) -> impl Iterator<Item = P>
157    where
158        T: Copy,
159        P: PackedValue<Value = T>,
160    {
161        let height = self.height();
162        let width = self.width();
163        // The row permutation generally scatters `r..r + P::WIDTH` across non-contiguous
164        // inner rows, so unlike `DenseMatrix` we cannot take a contiguous-slice fast path.
165        // Reading elements directly still avoids the `Vec` of row-slice guards that the
166        // default implementation allocates via `wrapping_row_slices`.
167        let no_wrap = P::WIDTH != 1 && r + P::WIDTH <= height;
168        (0..width).map(move |c| {
169            if no_wrap {
170                // Safety: r + i < height (fast-path guard), and c < width (loop bound).
171                P::from_fn(|i| unsafe { self.get_unchecked(r + i, c) })
172            } else {
173                // Safety: (r + i) % height < height, and c < width (loop bound).
174                P::from_fn(|i| unsafe { self.get_unchecked((r + i) % height, c) })
175            }
176        })
177    }
178
179    #[inline]
180    fn vertically_packed_row_pair<P>(&self, r: usize, step: usize) -> Vec<P>
181    where
182        T: Copy,
183        P: PackedValue<Value = T>,
184    {
185        let height = self.height();
186        let width = self.width();
187        let no_wrap = P::WIDTH != 1 && r + P::WIDTH <= height;
188        let next_no_wrap = P::WIDTH != 1 && r + step + P::WIDTH <= height;
189
190        (0..width)
191            .map(move |c| {
192                if no_wrap {
193                    // Safety: r + i < height (fast-path guard), and c < width (loop bound).
194                    P::from_fn(|i| unsafe { self.get_unchecked(r + i, c) })
195                } else {
196                    // Safety: (r + i) % height < height, and c < width (loop bound).
197                    P::from_fn(|i| unsafe { self.get_unchecked((r + i) % height, c) })
198                }
199            })
200            .chain((0..width).map(move |c| {
201                if next_no_wrap {
202                    // Safety: r + step + i < height (fast-path guard), and c < width (loop bound).
203                    P::from_fn(|i| unsafe { self.get_unchecked(r + step + i, c) })
204                } else {
205                    // Safety: (r + step + i) % height < height, and c < width (loop bound).
206                    P::from_fn(|i| unsafe { self.get_unchecked((r + step + i) % height, c) })
207                }
208            }))
209            .collect()
210    }
211}
212
213#[cfg(test)]
214mod tests {
215    use alloc::vec;
216    use alloc::vec::Vec;
217
218    use itertools::Itertools;
219    use p3_baby_bear::BabyBear;
220    use p3_field::FieldArray;
221
222    use super::*;
223    use crate::dense::RowMajorMatrix;
224
225    /// Mock implementation of RowIndexMap
226    struct IdentityMap(usize);
227
228    impl RowIndexMap for IdentityMap {
229        fn height(&self) -> usize {
230            self.0
231        }
232
233        fn map_row_index(&self, r: usize) -> usize {
234            r
235        }
236    }
237
238    /// Another mock implementation for reversing rows
239    struct ReverseMap(usize);
240
241    impl RowIndexMap for ReverseMap {
242        fn height(&self) -> usize {
243            self.0
244        }
245
246        fn map_row_index(&self, r: usize) -> usize {
247            self.0 - 1 - r
248        }
249    }
250
251    /// A final Mock implementation of RowIndexMap
252    struct ConstantMap;
253
254    impl RowIndexMap for ConstantMap {
255        fn height(&self) -> usize {
256            1
257        }
258
259        fn map_row_index(&self, _r: usize) -> usize {
260            0
261        }
262    }
263
264    #[test]
265    fn test_identity_row_index_map() {
266        // Create an inner matrix.
267        // The matrix will be:
268        // [ 1  2  3 ]
269        // [ 4  5  6 ]
270        let inner = RowMajorMatrix::new(vec![1, 2, 3, 4, 5, 6], 3);
271
272        // Create a mapped view using an `IdentityMap`, which does not alter row indices.
273        let mapped_view = RowIndexMappedView {
274            index_map: IdentityMap(inner.height()),
275            inner,
276        };
277
278        // Check dimensions.
279        assert_eq!(mapped_view.height(), 2);
280        assert_eq!(mapped_view.width(), 3);
281
282        // Check values.
283        assert_eq!(mapped_view.get(0, 0).unwrap(), 1);
284        assert_eq!(mapped_view.get(1, 2).unwrap(), 6);
285
286        unsafe {
287            assert_eq!(mapped_view.get_unchecked(0, 1), 2);
288            assert_eq!(mapped_view.get_unchecked(1, 0), 4);
289        }
290
291        // Check rows.
292        let rows: Vec<Vec<_>> = mapped_view.rows().map(|row| row.collect()).collect();
293        assert_eq!(rows, vec![vec![1, 2, 3], vec![4, 5, 6]]);
294
295        // Check dense matrix.
296        let dense = mapped_view.to_row_major_matrix();
297        assert_eq!(dense.values, vec![1, 2, 3, 4, 5, 6]);
298    }
299
300    #[test]
301    fn test_reverse_row_index_map() {
302        // Create an inner matrix.
303        // The matrix will be:
304        // [ 1  2  3 ]
305        // [ 4  5  6 ]
306        let inner = RowMajorMatrix::new(vec![1, 2, 3, 4, 5, 6], 3);
307
308        // Create a mapped view using a ReverseMap, which reverses row indices.
309        let mapped_view = RowIndexMappedView {
310            index_map: ReverseMap(inner.height()),
311            inner,
312        };
313
314        // Check dimensions.
315        assert_eq!(mapped_view.height(), 2);
316        assert_eq!(mapped_view.width(), 3);
317
318        // Check the first element of the mapped view (originally the second row, first column).
319        assert_eq!(mapped_view.get(0, 0).unwrap(), 4);
320        // Check the last element of the mapped view (originally the first row, last column).
321        assert_eq!(mapped_view.get(1, 2).unwrap(), 3);
322
323        unsafe {
324            assert_eq!(mapped_view.get_unchecked(0, 1), 5);
325            assert_eq!(mapped_view.get_unchecked(1, 0), 1);
326        }
327
328        // Check rows.
329        let rows: Vec<Vec<_>> = mapped_view.rows().map(|row| row.collect()).collect();
330        assert_eq!(rows, vec![vec![4, 5, 6], vec![1, 2, 3]]);
331
332        // Check dense matrix.
333        let dense = mapped_view.to_row_major_matrix();
334        assert_eq!(dense.values, vec![4, 5, 6, 1, 2, 3]);
335    }
336
337    #[test]
338    fn test_horizontally_packed_row() {
339        // Define the packed type with width 2
340        type Packed = FieldArray<BabyBear, 2>;
341
342        // Create an inner matrix of BabyBear elements.
343        // Matrix layout:
344        // [ 1  2 ]
345        // [ 3  4 ]
346        let inner = RowMajorMatrix::new(
347            vec![
348                BabyBear::new(1),
349                BabyBear::new(2),
350                BabyBear::new(3),
351                BabyBear::new(4),
352            ],
353            2,
354        );
355
356        // Apply a reverse row index mapping.
357        let mapped_view = RowIndexMappedView {
358            index_map: ReverseMap(inner.height()),
359            inner,
360        };
361
362        // Extract the packed and suffix iterators from row 0 (which is reversed row 1).
363        let (packed_iter, mut suffix_iter) = mapped_view.horizontally_packed_row::<Packed>(0);
364
365        // Collect iterators to concrete values.
366        let packed: Vec<_> = packed_iter.collect();
367
368        // Check the packed row values match reversed second row.
369        assert_eq!(
370            packed,
371            &[Packed::from([BabyBear::new(3), BabyBear::new(4)])]
372        );
373
374        // Check there are no suffix leftovers.
375        assert!(suffix_iter.next().is_none());
376    }
377
378    #[test]
379    fn test_padded_horizontally_packed_row() {
380        // Define a packed type with width 3
381        type Packed = FieldArray<BabyBear, 3>;
382
383        // Create a 2x2 matrix of BabyBear elements:
384        // [ 1  2 ]
385        // [ 3  4 ]
386        let inner = RowMajorMatrix::new(
387            vec![
388                BabyBear::new(1),
389                BabyBear::new(2),
390                BabyBear::new(3),
391                BabyBear::new(4),
392            ],
393            2,
394        );
395
396        // Use identity mapping (rows remain unchanged).
397        let mapped_view = RowIndexMappedView {
398            index_map: IdentityMap(inner.height()),
399            inner,
400        };
401
402        // Pad the second row (row 1) into chunks of size 3.
403        let packed: Vec<_> = mapped_view
404            .padded_horizontally_packed_row::<Packed>(1)
405            .collect();
406
407        // Verify the packed result includes padding with zero at the end.
408        assert_eq!(
409            packed,
410            vec![Packed::from([
411                BabyBear::new(3),
412                BabyBear::new(4),
413                BabyBear::new(0),
414            ])]
415        );
416    }
417
418    #[test]
419    fn test_vertically_packed_row() {
420        // Define the packed type with width 2
421        type Packed = FieldArray<BabyBear, 2>;
422
423        // Create a 4x2 matrix of BabyBear elements:
424        // [ 1  2 ]
425        // [ 3  4 ]
426        // [ 5  6 ]
427        // [ 7  8 ]
428        let inner = RowMajorMatrix::new((1..=8).map(BabyBear::new).collect::<Vec<_>>(), 2);
429
430        // Reverse row mapping: visible row i maps to inner row (height - 1 - i).
431        // So the view (in inner-row order) is rows [3, 2, 1, 0].
432        let mapped_view = RowIndexMappedView {
433            index_map: ReverseMap(inner.height()),
434            inner,
435        };
436
437        // Non-wrapping gather: visible rows 0 and 1 map to inner rows 3 and 2.
438        let packed: Vec<_> = mapped_view.vertically_packed_row::<Packed>(0).collect();
439        assert_eq!(
440            packed,
441            vec![
442                Packed::from([BabyBear::new(7), BabyBear::new(5)]),
443                Packed::from([BabyBear::new(8), BabyBear::new(6)]),
444            ]
445        );
446
447        // Wrapping gather: visible rows 3 and 0 (wrapping past height=4) map to inner rows 0 and 3.
448        let packed: Vec<_> = mapped_view.vertically_packed_row::<Packed>(3).collect();
449        assert_eq!(
450            packed,
451            vec![
452                Packed::from([BabyBear::new(1), BabyBear::new(7)]),
453                Packed::from([BabyBear::new(2), BabyBear::new(8)]),
454            ]
455        );
456    }
457
458    #[test]
459    fn test_vertically_packed_row_pair() {
460        // Define the packed type with width 2
461        type Packed = FieldArray<BabyBear, 2>;
462
463        // Create a 4x2 matrix of BabyBear elements:
464        // [ 1  2 ]
465        // [ 3  4 ]
466        // [ 5  6 ]
467        // [ 7  8 ]
468        let inner = RowMajorMatrix::new((1..=8).map(BabyBear::new).collect::<Vec<_>>(), 2);
469
470        // Reverse row mapping: visible row i maps to inner row (height - 1 - i).
471        // So the view (in inner-row order) is rows [3, 2, 1, 0].
472        let mapped_view = RowIndexMappedView {
473            index_map: ReverseMap(inner.height()),
474            inner,
475        };
476
477        // Both halves non-wrapping: visible rows [0, 1] and [2, 3] map to inner rows [3, 2] and [1, 0].
478        let packed = mapped_view.vertically_packed_row_pair::<Packed>(0, 2);
479        assert_eq!(
480            packed,
481            vec![
482                Packed::from([BabyBear::new(7), BabyBear::new(5)]),
483                Packed::from([BabyBear::new(8), BabyBear::new(6)]),
484                Packed::from([BabyBear::new(3), BabyBear::new(1)]),
485                Packed::from([BabyBear::new(4), BabyBear::new(2)]),
486            ]
487        );
488
489        // First half non-wrapping, second half wrapping past height=4.
490        let packed = mapped_view.vertically_packed_row_pair::<Packed>(1, 2);
491        assert_eq!(
492            packed,
493            vec![
494                Packed::from([BabyBear::new(5), BabyBear::new(3)]),
495                Packed::from([BabyBear::new(6), BabyBear::new(4)]),
496                Packed::from([BabyBear::new(1), BabyBear::new(7)]),
497                Packed::from([BabyBear::new(2), BabyBear::new(8)]),
498            ]
499        );
500
501        // Both halves wrapping past height=4.
502        let packed = mapped_view.vertically_packed_row_pair::<Packed>(3, 2);
503        assert_eq!(
504            packed,
505            vec![
506                Packed::from([BabyBear::new(1), BabyBear::new(7)]),
507                Packed::from([BabyBear::new(2), BabyBear::new(8)]),
508                Packed::from([BabyBear::new(5), BabyBear::new(3)]),
509                Packed::from([BabyBear::new(6), BabyBear::new(4)]),
510            ]
511        );
512    }
513
514    #[test]
515    fn test_row_and_row_slice_methods() {
516        // Create a 2x3 matrix of integers:
517        // [ 10  20  30 ]
518        // [ 40  50  60 ]
519        let inner = RowMajorMatrix::new(vec![10, 20, 30, 40, 50, 60], 3);
520
521        // Apply reverse row mapping (row 0 becomes 1, row 1 becomes 0).
522        let mapped_view = RowIndexMappedView {
523            index_map: ReverseMap(inner.height()),
524            inner,
525        };
526
527        // Get row slices through dereferencing and verify content.
528        assert_eq!(mapped_view.row_slice(0).unwrap().deref(), &[40, 50, 60]); // was row 1
529        assert_eq!(
530            mapped_view.row(1).unwrap().into_iter().collect_vec(),
531            vec![10, 20, 30]
532        ); // was row 0
533
534        unsafe {
535            // Check unsafe row slices.
536            assert_eq!(
537                mapped_view.row_unchecked(0).into_iter().collect_vec(),
538                vec![40, 50, 60]
539            ); // was row 1
540            assert_eq!(mapped_view.row_slice_unchecked(1).deref(), &[10, 20, 30]); // was row 0
541
542            assert_eq!(
543                mapped_view.row_subslice_unchecked(0, 1, 3).deref(),
544                &[50, 60]
545            ); // was row 1
546            assert_eq!(
547                mapped_view
548                    .row_subseq_unchecked(1, 0, 2)
549                    .into_iter()
550                    .collect_vec(),
551                vec![10, 20]
552            ); // was row 0
553        }
554
555        assert!(mapped_view.row(2).is_none()); // Height out of bounds.
556        assert!(mapped_view.row_slice(2).is_none()); // Height out of bounds.
557    }
558
559    #[test]
560    fn test_out_of_bounds_access() {
561        // Create a 2x2 matrix:
562        // [ 1  2 ]
563        // [ 3  4 ]
564        let inner = RowMajorMatrix::new(vec![1, 2, 3, 4], 2);
565
566        // Use identity mapping.
567        let mapped_view = RowIndexMappedView {
568            index_map: IdentityMap(inner.height()),
569            inner,
570        };
571
572        // Attempt to access out-of-bounds row (index 2). Should panic.
573        assert_eq!(mapped_view.get(2, 1), None);
574        assert!(mapped_view.row(5).is_none());
575        assert!(mapped_view.row_slice(11).is_none());
576        assert_eq!(mapped_view.get(0, 20), None);
577    }
578
579    #[test]
580    fn test_out_of_bounds_access_with_bad_map() {
581        // Create a 2x2 matrix:
582        // [ 1  2 ]
583        // [ 3  4 ]
584        let inner = RowMajorMatrix::new(vec![1, 2, 3, 4], 4);
585
586        // Use identity mapping.
587        let mapped_view = RowIndexMappedView {
588            index_map: ConstantMap,
589            inner,
590        };
591
592        assert_eq!(mapped_view.get(0, 2), Some(3));
593
594        // Attempt to access out-of-bounds row (index 1). Should panic.
595        assert_eq!(mapped_view.get(1, 0), None);
596        assert!(mapped_view.row(1).is_none());
597        assert!(mapped_view.row_slice(1).is_none());
598    }
599}
p3_matrix/row_index_mapped.rs

p3_matrix/
row_index_mapped.rs
