Improve M4riTable ergonomics

SpectralSequences · Feb 12, 2024 · 3f0914c · 3f0914c
1 parent 55244e5
commit 3f0914c
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 11 deletions.
diff --git a/ext/crates/fp/src/limb.rs b/ext/crates/fp/src/limb.rs
@@ -2,7 +2,7 @@ pub(crate) use crate::constants::Limb;
 use crate::constants::BITS_PER_LIMB;
 
 /// A struct containing the information required to access a specific entry in an array of `Limb`s.
-#[derive(Copy, Clone)]
+#[derive(Debug, Copy, Clone)]
 pub(crate) struct LimbBitIndexPair {
     pub(crate) limb: usize,
     pub(crate) bit_index: usize,

diff --git a/ext/crates/fp/src/matrix/m4ri.rs b/ext/crates/fp/src/matrix/m4ri.rs
@@ -1,9 +1,8 @@
 use itertools::Itertools;
 
 use crate::{
-    constants::BITS_PER_LIMB,
     field::{limb::LimbMethods, Fp},
-    limb::Limb,
+    limb::{Limb, LimbBitIndexPair},
     matrix::Matrix,
     prime::P2,
     simd,
@@ -27,7 +26,7 @@ pub(crate) struct M4riTable {
     /// The indices of new rows in the table
     rows: Vec<usize>,
     /// The list of pivot columns of the rows, in the format (limb, bit)
-    columns: Vec<(usize, usize)>,
+    columns: Vec<LimbBitIndexPair>,
     /// The 2^k linear combinations of the k rows, apart from the first one which is identically
     /// zero.
     data: Vec<Limb>,
@@ -70,8 +69,7 @@ impl M4riTable {
     ///  - `column`: pivot column of the row
     ///  - `row`: index of the row
     pub fn add(&mut self, column: usize, row: usize) {
-        self.columns
-            .push((column / BITS_PER_LIMB, column % BITS_PER_LIMB));
+        self.columns.push(Fp(P2).limb_bit_index_pair(column));
         self.rows.push(row);
     }
 
@@ -95,21 +93,21 @@ impl M4riTable {
                 simd::add_simd(
                     &mut self.data[i * num_limbs..(i + 1) * num_limbs],
                     matrix[r].limbs(),
-                    c.0,
+                    c.limb,
                 );
             }
-            self.min_limb = std::cmp::min(self.min_limb, c.0);
+            self.min_limb = std::cmp::min(self.min_limb, c.limb);
         }
     }
 
     pub fn reduce_naive(&self, matrix: &mut Matrix, target: usize) {
         for (&row, col) in self.rows.iter().zip_eq(&self.columns) {
             assert!(target != row);
             unsafe {
-                let coef = (matrix[target].limbs()[col.0] >> col.1) & 1;
+                let coef = (matrix[target].limbs()[col.limb] >> col.bit_index) & 1;
                 if coef != 0 {
                     let (target, source) = matrix.split_borrow(target, row);
-                    simd::add_simd(target.limbs_mut(), source.limbs(), col.0)
+                    simd::add_simd(target.limbs_mut(), source.limbs(), col.limb)
                 }
             }
         }
@@ -120,7 +118,7 @@ impl M4riTable {
         let mut index: usize = 0;
         for &col in self.columns.iter().rev() {
             index <<= 1;
-            index += ((v[col.0] >> col.1) & 1) as usize;
+            index += ((v[col.limb] >> col.bit_index) & 1) as usize;
         }
         if index != 0 {
             simd::add_simd(v, &self.data[(index - 1) * num_limbs..], self.min_limb);