Add FactorEncoding abstraction

Author: leostimpfle

crates/within/src/domain.rs

@@ -78,10 +78,29 @@ impl<T> Loading<T> {
     }
 }
 
+/// Mapping between caller-visible factor labels and numerical level positions.
+///
+/// This is an identity mapping for the existing dense `u32` API. Later
+/// encodings can store an explicit mapping without changing `TermMeta`.
+#[derive(Debug, Clone)]
+pub(crate) struct FactorEncoding {
+    n_levels: usize,
+}
+
+impl FactorEncoding {
+    fn identity(n_levels: usize) -> Self {
+        Self { n_levels }
+    }
+
+    pub(crate) fn n_levels(&self) -> usize {
+        self.n_levels
+    }
+}
+
 /// Per-term metadata; coefficient `c` of `level` lives at `offset + c · n_levels + level`.
 #[derive(Debug, Clone)]
 pub(crate) struct TermMeta {
-    pub n_levels: usize,
+    pub(crate) encoding: FactorEncoding,
     pub offset: usize,
     /// Non-decreasing in the design's internal row order (fixed at construction).
     pub sorted: bool,
@@ -90,17 +109,21 @@ pub(crate) struct TermMeta {
 }
 
 impl TermMeta {
+    pub fn n_levels(&self) -> usize {
+        self.encoding.n_levels()
+    }
+
     pub fn n_columns(&self) -> usize {
         self.columns.len()
     }
 
     pub fn n_dofs(&self) -> usize {
-        self.n_columns() * self.n_levels
+        self.n_columns() * self.n_levels()
     }
 
     /// Global DOF base of coefficient column `column`.
     pub fn column_base(&self, column: usize) -> usize {
-        self.offset + column * self.n_levels
+        self.offset + column * self.n_levels()
     }
 }
 
@@ -209,8 +232,9 @@ impl<'a> Design<'a> {
                 sorted &= v >= prev;
                 prev = v;
             }
+            let encoding = FactorEncoding::identity(max as usize + 1);
             let meta = TermMeta {
-                n_levels: max as usize + 1,
+                encoding,
                 offset,
                 sorted,
                 columns,
@@ -228,7 +252,7 @@ impl<'a> Design<'a> {
         let dominant = (0..terms.len()).max_by_key(|&q| terms[q].n_dofs());
         let (frame, obs_perm) = match dominant {
             Some(d) if locality_sort && !terms[d].sorted && u32::try_from(n_obs).is_ok() => {
-                let perm = stable_argsort(frame.level_column(d), terms[d].n_levels);
+                let perm = stable_argsort(frame.level_column(d), terms[d].n_levels());
                 let sorted_frame = frame.permuted(&perm);
                 // Factors nested in the dominant one come out sorted, keeping coalesced scatter.
                 for (q, meta) in terms.iter_mut().enumerate() {

crates/within/src/domain/cross_tab.rs

@@ -26,7 +26,7 @@ pub(crate) fn find_all_active_levels(design: &Design<'_>) -> Vec<Vec<bool>> {
     let mut active: Vec<Vec<bool>> = design
         .terms
         .iter()
-        .map(|f| vec![false; f.n_levels])
+        .map(|f| vec![false; f.n_levels()])
         .collect();
     // Factor-outer/obs-inner: all writes for a factor land in one `active[f]` buffer.
     for (f, col) in active.iter_mut().enumerate() {

crates/within/src/domain/level_moments.rs

@@ -70,9 +70,9 @@ impl LevelMoments {
         let v = covariates.len();
         let mut moments = Self {
             intercept: v < meta.columns.len(),
-            w_sum: vec![0.0; meta.n_levels],
-            mean: vec![0.0; meta.n_levels * v],
-            comoment: vec![0.0; meta.n_levels * tri_len(v)],
+            w_sum: vec![0.0; meta.n_levels()],
+            mean: vec![0.0; meta.n_levels() * v],
+            comoment: vec![0.0; meta.n_levels() * tri_len(v)],
             covariates,
         };
         let zs: Vec<&[f64]> = moments

crates/within/src/operator/design/gather.rs

@@ -21,7 +21,7 @@ pub(crate) fn gather_apply(
 
     for_each_chunk(dst, |chunk, row_start| {
         for (q, t) in design.terms.iter().enumerate() {
-            let (offset, n_levels) = (t.offset, t.n_levels);
+            let (offset, n_levels) = (t.offset, t.n_levels());
             let levels = design.level_column(q);
             let col = |c: usize| &src[offset + c * n_levels..offset + (c + 1) * n_levels];
             match &*t.columns {

crates/within/src/operator/design/scatter.rs

@@ -52,8 +52,8 @@ pub(super) fn scatter_apply(
             }
             columns => {
                 for (c, loading) in columns.iter().enumerate() {
-                    let start = c * t.n_levels;
-                    let slot = &mut block[start..start + t.n_levels];
+                    let start = c * t.n_levels();
+                    let slot = &mut block[start..start + t.n_levels()];
                     match loading {
                         Loading::Constant => {
                             scatter_term::<1>(slot, t, levels, parallel, scratch, |i| [base(i)]);
@@ -80,7 +80,7 @@ fn scatter_term<const C: usize>(
     scratch: &[AtomicF64],
     values: impl Fn(usize) -> [f64; C] + Sync,
 ) {
-    let n_levels = meta.n_levels;
+    let n_levels = meta.n_levels();
     debug_assert_eq!(block.len(), C * n_levels);
     match ScatterStrategy::pick(parallel, C * n_levels, meta.sorted) {
         ScatterStrategy::Sequential => scatter_sequential::<C>(block, n_levels, levels, &values),

crates/within/src/operator/design/tests.rs

@@ -307,7 +307,7 @@ mod slope_design_tests {
         for (q, t) in design.terms.iter().enumerate() {
             let levels = design.level_column(q);
             for (c, loading) in t.columns.iter().enumerate() {
-                let base = t.offset + c * t.n_levels;
+                let base = t.offset + c * t.n_levels();
                 for (i, &lev) in levels.iter().enumerate() {
                     d[i][base + lev as usize] = match loading {
                         Loading::Constant => 1.0,

crates/within/src/operator/schwarz.rs

@@ -221,7 +221,7 @@ fn build_diagonal(
         let w = |uid: usize| weights.map_or(1.0, |ws| ws[uid]);
         for (column, loading) in term.columns.iter().enumerate() {
             let base = term.column_base(column);
-            let slice = &mut diag[base..base + term.n_levels];
+            let slice = &mut diag[base..base + term.n_levels()];
             match loading {
                 Loading::Constant => {
                     for (uid, &level) in levels.iter().enumerate() {

crates/within/src/solver.rs

@@ -143,7 +143,7 @@ impl CoefficientLayout {
             .iter()
             .map(|t| TermLayout {
                 offset: t.offset,
-                n_levels: t.n_levels,
+                n_levels: t.n_levels(),
                 n_columns: t.n_columns(),
             })
             .collect();

crates/within/src/solver/reparam.rs

@@ -80,7 +80,7 @@ impl TermReparam {
         unidentified: &mut Vec<CoefficientAddress>,
     ) -> Self {
         let meta = &design.terms[term];
-        let (offset, n_levels) = (meta.offset, meta.n_levels);
+        let (offset, n_levels) = (meta.offset, meta.n_levels());
         let mut intercept_column = None;
         let mut slope_columns = Vec::new();
         for (column, loading) in meta.columns.iter().enumerate() {

crates/within/src/solver/reparam/tests.rs

@@ -40,7 +40,7 @@ fn build_whitens_each_slope_bearing_term() {
             .filter_map(|c| c.covariate())
             .map(|&k| design.loading_column(k as usize))
             .collect();
-        for level in 0..meta.n_levels {
+        for level in 0..meta.n_levels() {
             let obs: Vec<usize> = (0..levels.len())
                 .filter(|&i| levels[i] as usize == level)
                 .collect();