p3_matrix/
bitrev.rs

1use p3_util::{log2_strict_usize, reverse_bits_len};
2
3use crate::dense::{DenseMatrix, DenseStorage, RowMajorMatrix};
4use crate::row_index_mapped::{RowIndexMap, RowIndexMappedView};
5use crate::util::reverse_matrix_index_bits;
6use crate::Matrix;
7
8/// A matrix whose row indices are possibly bit-reversed, enabling easily switching
9/// between orderings. Pretty much just either `RowMajorMatrix` or
10/// `BitReversedMatrixView<RowMajorMatrix>`.
11pub trait BitReversableMatrix<T: Send + Sync>: Matrix<T> {
12    type BitRev: BitReversableMatrix<T>;
13    fn bit_reverse_rows(self) -> Self::BitRev;
14}
15
16#[derive(Debug)]
17pub struct BitReversalPerm {
18    log_height: usize,
19}
20
21impl BitReversalPerm {
22    /// Assumes the inner matrix height is a power of two; panics otherwise.
23    pub fn new_view<T: Send + Sync, Inner: Matrix<T>>(
24        inner: Inner,
25    ) -> BitReversedMatrixView<Inner> {
26        RowIndexMappedView {
27            index_map: Self {
28                log_height: log2_strict_usize(inner.height()),
29            },
30            inner,
31        }
32    }
33}
34
35impl RowIndexMap for BitReversalPerm {
36    fn height(&self) -> usize {
37        1 << self.log_height
38    }
39    fn map_row_index(&self, r: usize) -> usize {
40        reverse_bits_len(r, self.log_height)
41    }
42    // This might not be more efficient than the lazy generic impl
43    // if we have a nested view.
44    fn to_row_major_matrix<T: Clone + Send + Sync, Inner: Matrix<T>>(
45        &self,
46        inner: Inner,
47    ) -> RowMajorMatrix<T> {
48        let mut inner = inner.to_row_major_matrix();
49        reverse_matrix_index_bits(&mut inner);
50        inner
51    }
52}
53
54pub type BitReversedMatrixView<Inner> = RowIndexMappedView<BitReversalPerm, Inner>;
55
56impl<T: Clone + Send + Sync, S: DenseStorage<T>> BitReversableMatrix<T>
57    for BitReversedMatrixView<DenseMatrix<T, S>>
58{
59    type BitRev = DenseMatrix<T, S>;
60    fn bit_reverse_rows(self) -> Self::BitRev {
61        self.inner
62    }
63}
64
65impl<T: Clone + Send + Sync, S: DenseStorage<T>> BitReversableMatrix<T> for DenseMatrix<T, S> {
66    type BitRev = BitReversedMatrixView<DenseMatrix<T, S>>;
67    fn bit_reverse_rows(self) -> Self::BitRev {
68        BitReversalPerm::new_view(self)
69    }
70}