refactor: centralize nonzero slot iteration logic across views

Extract duplicated traversal logic into a shared `NonzeroSlotsView` trait and `nonzero_triples` helper. Update slice view iterators to consistently yield `(position, value)` tuples and use a temporary buffer with explicit permutation mapping for sorted access. Delegate manual slot sorting and column iteration to the new shared helper, eliminating per-format duplication and eager collection while preserving existing public API signatures.
This commit is contained in:
Eric Coissac
2026-08-20 14:20:24 +02:00
parent 82374deca5
commit 19f9954050
3 changed files with 77 additions and 65 deletions
+5 -19
View File
@@ -7,7 +7,7 @@ use crate::bitvec::{PersistentBitVec, PersistentBitVecBuilder};
use crate::layer_meta::LayerMeta; use crate::layer_meta::LayerMeta;
use crate::meta::MatrixMeta; use crate::meta::MatrixMeta;
use crate::traits::{BitPartials, ColumnWeights}; use crate::traits::{BitPartials, ColumnWeights};
use crate::views::BitSliceView; use crate::views::{nonzero_triples, BitSliceView};
use super::columnar::ColumnarBitMatrix; use super::columnar::ColumnarBitMatrix;
use super::packed::PackedBitMatrix; use super::packed::PackedBitMatrix;
@@ -199,10 +199,9 @@ impl PersistentBitMatrix {
/// `Sparse` delegates to its native row-major decode (no `n_cols`-wide /// `Sparse` delegates to its native row-major decode (no `n_cols`-wide
/// buffer, ever — the whole point, see /// buffer, ever — the whole point, see
/// `PersistentSparseBitMatrix::nonzero_iter`); `Columnar`/`Packed` /// `PersistentSparseBitMatrix::nonzero_iter`); `Columnar`/`Packed`
/// reuse the same sorted-slot batching `fill_sub_matrix` always used /// reuse the shared sorted-slot batching in `views::nonzero_triples`
/// (`BitSliceView::enumerate_nonzero_slots`), just eagerly collected here /// (same as `IntMatrix::nonzero_iter`) — same `get` calls, same mmap
/// rather than filled into a dense buffer — same `get` calls, same /// locality, no per-format duplication of that logic elsewhere.
/// mmap locality, no per-format duplication of that logic elsewhere.
/// `Implicit` is trivial (one column, always present). /// `Implicit` is trivial (one column, always present).
/// ///
/// Boxed (not `impl Iterator`) because the match arms are genuinely /// Boxed (not `impl Iterator`) because the match arms are genuinely
@@ -219,20 +218,7 @@ impl PersistentBitMatrix {
Box::new(slots.iter().enumerate().map(|(i, _)| (i, 0usize, 1u32))) Box::new(slots.iter().enumerate().map(|(i, _)| (i, 0usize, 1u32)))
} }
Self::Columnar(_) | Self::Packed(_) => { Self::Columnar(_) | Self::Packed(_) => {
let n = slots.len(); Box::new(nonzero_triples(slots, self.n_cols(), |c| self.col_view(c)))
let mut perm: Vec<usize> = (0..n).collect();
perm.sort_by_key(|&i| slots[i]);
let sorted_slots: Vec<usize> = perm.iter().map(|&i| slots[i]).collect();
let n_cols = self.n_cols();
let mut hits: Vec<(usize, usize, u32)> = Vec::new();
for c in 0..n_cols {
let view = self.col_view(c);
hits.extend(
view.enumerate_nonzero_slots(&sorted_slots)
.map(|pos| (perm[pos], c, 1u32)),
);
}
Box::new(hits.into_iter())
} }
} }
} }
+4 -17
View File
@@ -14,7 +14,7 @@ use crate::meta::MatrixMeta;
use crate::reader::PersistentCompactIntVec; use crate::reader::PersistentCompactIntVec;
use crate::tempbitvec::{TempBitVec, TempBitVecBuilder}; use crate::tempbitvec::{TempBitVec, TempBitVecBuilder};
use crate::tempintvec::{TempCompactIntVec, TempCompactIntVecBuilder}; use crate::tempintvec::{TempCompactIntVec, TempCompactIntVecBuilder};
use crate::views::IntSliceView; use crate::views::{nonzero_triples, IntSliceView};
fn col_path(dir: &Path, col: usize) -> PathBuf { fn col_path(dir: &Path, col: usize) -> PathBuf {
dir.join(format!("col_{col:06}.pciv")) dir.join(format!("col_{col:06}.pciv"))
@@ -403,23 +403,10 @@ impl PersistentCompactIntMatrix {
/// sparse row-major access" — counts explicitly not excluded from that /// sparse row-major access" — counts explicitly not excluded from that
/// design even though no sparse count format exists yet). No native /// design even though no sparse count format exists yet). No native
/// low-effort case here the way `PersistentSparseBitMatrix` has one — /// low-effort case here the way `PersistentSparseBitMatrix` has one —
/// both variants reuse the same sorted-slot batching `fill_sub_matrix` /// both variants reuse the shared sorted-slot batching in
/// already used, via `IntSliceView::enumerate_nonzero_slots`. /// `views::nonzero_triples`.
pub fn nonzero_iter<'a>(&'a self, slots: &'a [usize]) -> impl Iterator<Item = (usize, usize, u32)> + 'a { pub fn nonzero_iter<'a>(&'a self, slots: &'a [usize]) -> impl Iterator<Item = (usize, usize, u32)> + 'a {
let n = slots.len(); nonzero_triples(slots, self.n_cols(), |c| self.col_view(c))
let mut perm: Vec<usize> = (0..n).collect();
perm.sort_by_key(|&i| slots[i]);
let sorted_slots: Vec<usize> = perm.iter().map(|&i| slots[i]).collect();
let n_cols = self.n_cols();
let mut hits: Vec<(usize, usize, u32)> = Vec::new();
for c in 0..n_cols {
hits.extend(
self.col_view(c)
.enumerate_nonzero_slots(&sorted_slots)
.map(|(pos, v)| (perm[pos], c, v)),
);
}
hits.into_iter()
} }
#[inline] #[inline]
+68 -29
View File
@@ -43,10 +43,8 @@ impl<'a> BitSliceView<'a> {
let mut perm: Vec<usize> = (0..n).collect(); let mut perm: Vec<usize> = (0..n).collect();
perm.sort_by_key(|&i| slots[i]); perm.sort_by_key(|&i| slots[i]);
let sorted: Vec<usize> = perm.iter().map(|&i| slots[i]).collect(); let sorted: Vec<usize> = perm.iter().map(|&i| slots[i]).collect();
let mut tmp = vec![false; n]; for (i, v) in self.enumerate_slots_values(&sorted) {
self.fill_slots_values_sorted(&sorted, &mut tmp); out[perm[i]] = v;
for (i, &orig_idx) in perm.iter().enumerate() {
out[orig_idx] = tmp[i];
} }
} }
@@ -62,25 +60,20 @@ impl<'a> BitSliceView<'a> {
slots.iter().enumerate().map(move |(pos, &slot)| (pos, self.get(slot))) slots.iter().enumerate().map(move |(pos, &slot)| (pos, self.get(slot)))
} }
/// Fill `out` assuming `sorted_slots` is already in ascending order. /// `(position in slots, 1)` for every slot whose bit is set — the value
/// Results are written in `sorted_slots` order (no reordering). /// is always `1u32` (a set bit has no other value to carry), kept in
pub(crate) fn fill_slots_values_sorted(&self, sorted_slots: &[usize], out: &mut [bool]) { /// the tuple only so this has the same shape as
assert_eq!(sorted_slots.len(), out.len()); /// [`IntSliceView::enumerate_nonzero_slots`] and both can feed the same
for (pos, v) in self.enumerate_slots_values(sorted_slots) { /// generic triple-collecting code (see `views::NonzeroSlotsView`).
out[pos] = v;
}
}
/// Positions in `slots` (not the slots themselves) whose bit is set.
/// Order of `slots` never matters for correctness; callers batching /// Order of `slots` never matters for correctness; callers batching
/// several columns over the same slot set sort it once beforehand /// several columns over the same slot set sort it once beforehand
/// purely for cache-friendlier `get` access (see /// purely for cache-friendlier `get` access (see
/// `PersistentBitMatrix::nonzero_iter`). /// `PersistentBitMatrix::nonzero_iter`).
pub fn enumerate_nonzero_slots<'b>(self, slots: &'b [usize]) -> impl Iterator<Item = usize> + 'b pub fn enumerate_nonzero_slots<'b>(self, slots: &'b [usize]) -> impl Iterator<Item = (usize, u32)> + 'b
where where
Self: 'b, Self: 'b,
{ {
self.enumerate_slots_values(slots).filter_map(|(pos, v)| v.then_some(pos)) self.enumerate_slots_values(slots).filter_map(|(pos, v)| v.then_some((pos, 1u32)))
} }
pub fn count_ones(&self) -> u64 { pub fn count_ones(&self) -> u64 {
@@ -209,10 +202,8 @@ impl<'a> IntSliceView<'a> {
let mut perm: Vec<usize> = (0..n).collect(); let mut perm: Vec<usize> = (0..n).collect();
perm.sort_by_key(|&i| slots[i]); perm.sort_by_key(|&i| slots[i]);
let sorted: Vec<usize> = perm.iter().map(|&i| slots[i]).collect(); let sorted: Vec<usize> = perm.iter().map(|&i| slots[i]).collect();
let mut tmp = vec![0u32; n]; for (i, v) in self.enumerate_slots_values(&sorted) {
self.fill_slots_values_sorted(&sorted, &mut tmp); out[perm[i]] = v;
for (i, &orig_idx) in perm.iter().enumerate() {
out[orig_idx] = tmp[i];
} }
} }
@@ -227,15 +218,6 @@ impl<'a> IntSliceView<'a> {
slots.iter().enumerate().map(move |(pos, &slot)| (pos, self.get(slot))) slots.iter().enumerate().map(move |(pos, &slot)| (pos, self.get(slot)))
} }
/// Fill `out` assuming `sorted_slots` is already in ascending order.
/// Results are written in `sorted_slots` order (no reordering).
pub(crate) fn fill_slots_values_sorted(&self, sorted_slots: &[usize], out: &mut [u32]) {
assert_eq!(sorted_slots.len(), out.len());
for (pos, v) in self.enumerate_slots_values(sorted_slots) {
out[pos] = v;
}
}
/// Like [`BitSliceView::enumerate_nonzero_slots`]: `(position in slots, /// Like [`BitSliceView::enumerate_nonzero_slots`]: `(position in slots,
/// value)` for every nonzero value — a `filter_map` over `get`, not a /// value)` for every nonzero value — a `filter_map` over `get`, not a
/// new traversal. /// new traversal.
@@ -370,6 +352,63 @@ impl<'a> IntSliceView<'a> {
} }
} }
// ── NonzeroSlotsView / nonzero_triples ──────────────────────────────────────
/// Column views (`BitSliceView`, `IntSliceView`) that can enumerate their
/// nonzero slots as `(position in slots, value)`. Mirrors each type's own
/// inherent `enumerate_nonzero_slots` — exists only so [`nonzero_triples`]
/// can be written once and reused by every matrix type instead of
/// reimplementing the same sort-then-scan-columns loop per format.
pub(crate) trait NonzeroSlotsView: Sized {
fn enumerate_nonzero_slots<'b>(self, slots: &'b [usize]) -> impl Iterator<Item = (usize, u32)> + 'b
where
Self: 'b;
}
impl<'a> NonzeroSlotsView for BitSliceView<'a> {
#[inline]
fn enumerate_nonzero_slots<'b>(self, slots: &'b [usize]) -> impl Iterator<Item = (usize, u32)> + 'b
where
Self: 'b,
{
BitSliceView::enumerate_nonzero_slots(self, slots)
}
}
impl<'a> NonzeroSlotsView for IntSliceView<'a> {
#[inline]
fn enumerate_nonzero_slots<'b>(self, slots: &'b [usize]) -> impl Iterator<Item = (usize, u32)> + 'b
where
Self: 'b,
{
IntSliceView::enumerate_nonzero_slots(self, slots)
}
}
/// Shared body of every matrix `nonzero_iter`: sort `slots` once, scan each
/// column with the sorted slice for mmap locality, remap positions back to
/// the caller's original order. Generic over the column view type so
/// `IntMatrix` and `PersistentBitMatrix` (Columnar/Packed) don't each carry
/// their own copy of this traversal.
pub(crate) fn nonzero_triples<V: NonzeroSlotsView>(
slots: &[usize],
n_cols: usize,
col_view: impl Fn(usize) -> V,
) -> impl Iterator<Item = (usize, usize, u32)> {
let n = slots.len();
let mut perm: Vec<usize> = (0..n).collect();
perm.sort_by_key(|&i| slots[i]);
let sorted_slots: Vec<usize> = perm.iter().map(|&i| slots[i]).collect();
let mut hits: Vec<(usize, usize, u32)> = Vec::new();
for c in 0..n_cols {
hits.extend(
col_view(c).enumerate_nonzero_slots(&sorted_slots)
.map(|(pos, v)| (perm[pos], c, v)),
);
}
hits.into_iter()
}
// ── IntSliceViewIter ────────────────────────────────────────────────────────── // ── IntSliceViewIter ──────────────────────────────────────────────────────────
pub struct IntSliceViewIter<'a> { pub struct IntSliceViewIter<'a> {