p3_matrix/
row_index_mapped.rs

1use core::ops::Deref;
2
3use p3_field::PackedValue;
4
5use crate::dense::RowMajorMatrix;
6use crate::Matrix;
7
8/// A RowIndexMap remaps row indices, and can change the height.
9pub trait RowIndexMap: Send + Sync {
10    fn height(&self) -> usize;
11    fn map_row_index(&self, r: usize) -> usize;
12
13    /// Permutations can optionally provide an optimized method to
14    /// convert to dense form.
15    fn to_row_major_matrix<T: Clone + Send + Sync, Inner: Matrix<T>>(
16        &self,
17        inner: Inner,
18    ) -> RowMajorMatrix<T> {
19        RowMajorMatrix::new(
20            (0..self.height())
21                .flat_map(|r| inner.row(self.map_row_index(r)))
22                .collect(),
23            inner.width(),
24        )
25    }
26}
27
28#[derive(Debug)]
29pub struct RowIndexMappedView<IndexMap, Inner> {
30    pub index_map: IndexMap,
31    pub inner: Inner,
32}
33
34impl<T: Send + Sync, IndexMap: RowIndexMap, Inner: Matrix<T>> Matrix<T>
35    for RowIndexMappedView<IndexMap, Inner>
36{
37    fn width(&self) -> usize {
38        self.inner.width()
39    }
40    fn height(&self) -> usize {
41        self.index_map.height()
42    }
43
44    type Row<'a>
45        = Inner::Row<'a>
46    where
47        Self: 'a;
48
49    fn row(&self, r: usize) -> Self::Row<'_> {
50        self.inner.row(self.index_map.map_row_index(r))
51    }
52
53    // Override these methods so we use the potentially optimized inner methods instead of defaults.
54
55    fn get(&self, r: usize, c: usize) -> T {
56        self.inner.get(self.index_map.map_row_index(r), c)
57    }
58
59    fn row_slice(&self, r: usize) -> impl Deref<Target = [T]> {
60        self.inner.row_slice(self.index_map.map_row_index(r))
61    }
62
63    fn horizontally_packed_row<'a, P>(
64        &'a self,
65        r: usize,
66    ) -> (impl Iterator<Item = P>, impl Iterator<Item = T>)
67    where
68        P: PackedValue<Value = T>,
69        T: Clone + 'a,
70    {
71        self.inner
72            .horizontally_packed_row(self.index_map.map_row_index(r))
73    }
74
75    fn to_row_major_matrix(self) -> RowMajorMatrix<T>
76    where
77        Self: Sized,
78        T: Clone,
79    {
80        // Use Perm's optimized permutation routine, if it has one.
81        self.index_map.to_row_major_matrix(self.inner)
82    }
83}