diff --git a/src/obicompactvec/src/bitvec.rs b/src/obicompactvec/src/bitvec.rs index c425a919..a1bd7798 100644 --- a/src/obicompactvec/src/bitvec.rs +++ b/src/obicompactvec/src/bitvec.rs @@ -14,16 +14,20 @@ const MAGIC: [u8; 4] = *b"PBIV"; const HEADER_SIZE: usize = 16; #[inline] -pub(crate) fn n_words(n: usize) -> usize { n.div_ceil(64) } +pub(crate) fn n_words(n: usize) -> usize { + n.div_ceil(64) +} #[inline] -fn n_bytes_for_words(n: usize) -> usize { n_words(n) * 8 } +fn n_bytes_for_words(n: usize) -> usize { + n_words(n) * 8 +} // ── PersistentBitVec ────────────────────────────────────────────────────────── pub struct PersistentBitVec { mmap: Mmap, - n: usize, + n: usize, path: PathBuf, } @@ -31,18 +35,31 @@ impl PersistentBitVec { pub fn open(path: &Path) -> io::Result { let mmap = unsafe { Mmap::map(&File::open(path)?)? }; if mmap.len() < HEADER_SIZE { - return Err(io::Error::new(io::ErrorKind::InvalidData, "PBIV file too short")); + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "PBIV file too short", + )); } if &mmap[0..4] != &MAGIC { return Err(io::Error::new(io::ErrorKind::InvalidData, "bad PBIV magic")); } let n = u64::from_le_bytes(mmap[8..16].try_into().unwrap()) as usize; - Ok(Self { mmap, n, path: path.to_path_buf() }) + Ok(Self { + mmap, + n, + path: path.to_path_buf(), + }) } - pub fn path(&self) -> &Path { &self.path } - pub fn len(&self) -> usize { self.n } - pub fn is_empty(&self) -> bool { self.n == 0 } + pub fn path(&self) -> &Path { + &self.path + } + pub fn len(&self) -> usize { + self.n + } + pub fn is_empty(&self) -> bool { + self.n == 0 + } pub fn get(&self, slot: usize) -> bool { (self.mmap[HEADER_SIZE + (slot >> 3)] >> (slot & 7)) & 1 != 0 @@ -64,7 +81,9 @@ impl PersistentBitVec { pub fn fill_batch(&self, slots: &[usize], out: &mut [bool]) { assert_eq!(slots.len(), out.len()); let n = slots.len(); - if n == 0 { return; } + if n == 0 { + return; + } let mut perm: Vec = (0..n).collect(); perm.sort_by_key(|&i| slots[i]); let sorted: Vec = perm.iter().map(|&i| slots[i]).collect(); @@ -86,7 +105,7 @@ impl PersistentBitVec { // SAFETY: mmap is page-aligned, HEADER_SIZE=16 divisible by 8 → u64-aligned. fn data_words(&self) -> &[u64] { - let nw = n_words(self.n); + let nw = n_words(self.n); let ptr = self.mmap[HEADER_SIZE..].as_ptr() as *const u64; unsafe { std::slice::from_raw_parts(ptr, nw) } } @@ -95,10 +114,16 @@ impl PersistentBitVec { BitSliceView::new(self.data_words(), self.n) } - pub fn words(&self) -> &[u64] { self.data_words() } + pub fn words(&self) -> &[u64] { + self.data_words() + } - pub fn count_ones(&self) -> u64 { self.view().count_ones() } - pub fn count_zeros(&self) -> u64 { self.view().count_zeros() } + pub fn count_ones(&self) -> u64 { + self.view().count_ones() + } + pub fn count_zeros(&self) -> u64 { + self.view().count_zeros() + } pub fn partial_jaccard_dist(&self, other: &PersistentBitVec) -> (u64, u64) { self.view().partial_jaccard_dist(other.view()) @@ -111,22 +136,28 @@ impl PersistentBitVec { } pub fn iter(&self) -> BitIter<'_> { - BitIter { words: self.data_words(), slot: 0, n: self.n } + BitIter { + words: self.data_words(), + slot: 0, + n: self.n, + } } } impl<'a> IntoIterator for &'a PersistentBitVec { type Item = bool; type IntoIter = BitIter<'a>; - fn into_iter(self) -> BitIter<'a> { self.iter() } + fn into_iter(self) -> BitIter<'a> { + self.iter() + } } // ── BitIter ─────────────────────────────────────────────────────────────────── pub struct BitIter<'a> { words: &'a [u64], - slot: usize, - n: usize, + slot: usize, + n: usize, } impl ExactSizeIterator for BitIter<'_> {} @@ -134,7 +165,9 @@ impl ExactSizeIterator for BitIter<'_> {} impl Iterator for BitIter<'_> { type Item = bool; fn next(&mut self) -> Option { - if self.slot >= self.n { return None; } + if self.slot >= self.n { + return None; + } let v = (self.words[self.slot >> 6] >> (self.slot & 63)) & 1 != 0; self.slot += 1; Some(v) @@ -149,7 +182,7 @@ impl Iterator for BitIter<'_> { pub struct PersistentBitVecBuilder { mmap: MmapMut, - n: usize, + n: usize, path: PathBuf, } @@ -157,7 +190,10 @@ impl PersistentBitVecBuilder { pub fn new(n: usize, path: &Path) -> io::Result { let file_size = HEADER_SIZE + n_bytes_for_words(n); let mut file = OpenOptions::new() - .read(true).write(true).create(true).truncate(true) + .read(true) + .write(true) + .create(true) + .truncate(true) .open(path)?; file.write_all(&MAGIC)?; file.write_all(&[0u8; 4])?; @@ -165,20 +201,31 @@ impl PersistentBitVecBuilder { file.seek(SeekFrom::Start(0))?; file.set_len(file_size as u64)?; let mmap = unsafe { MmapMut::map_mut(&file)? }; - Ok(Self { mmap, n, path: path.to_path_buf() }) + Ok(Self { + mmap, + n, + path: path.to_path_buf(), + }) } pub fn from_raw_bytes(bytes: &[u8], n: usize, path: &Path) -> io::Result { let file_size = HEADER_SIZE + n_bytes_for_words(n); let file = OpenOptions::new() - .read(true).write(true).create(true).truncate(true) + .read(true) + .write(true) + .create(true) + .truncate(true) .open(path)?; file.set_len(file_size as u64)?; let mut mmap = unsafe { MmapMut::map_mut(&file)? }; mmap[0..4].copy_from_slice(&MAGIC); mmap[8..16].copy_from_slice(&(n as u64).to_le_bytes()); mmap[HEADER_SIZE..HEADER_SIZE + bytes.len()].copy_from_slice(bytes); - Ok(Self { mmap, n, path: path.to_path_buf() }) + Ok(Self { + mmap, + n, + path: path.to_path_buf(), + }) } /// Create an all-ones bit vector of length `n` at `path`. @@ -186,10 +233,13 @@ impl PersistentBitVecBuilder { /// More efficient than `new(n, path)` + `not()`: the data is written as /// 0xFF bytes in a single sequential pass, with no intermediate all-zeros state. pub fn new_ones(n: usize, path: &Path) -> io::Result { - let nw = n_words(n); + let nw = n_words(n); let file_size = HEADER_SIZE + nw * 8; - let mut file = OpenOptions::new() - .read(true).write(true).create(true).truncate(true) + let mut file = OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(true) .open(path)?; file.write_all(&MAGIC)?; file.write_all(&[0u8; 4])?; @@ -201,11 +251,15 @@ impl PersistentBitVecBuilder { // Clear padding bits in the last word so trailing bits are always 0. let rem = n % 64; if rem != 0 { - let ptr = mmap[HEADER_SIZE..].as_mut_ptr() as *mut u64; + let ptr = mmap[HEADER_SIZE..].as_mut_ptr() as *mut u64; let words = unsafe { std::slice::from_raw_parts_mut(ptr, nw) }; words[nw - 1] &= (1u64 << rem) - 1; } - Ok(Self { mmap, n, path: path.to_path_buf() }) + Ok(Self { + mmap, + n, + path: path.to_path_buf(), + }) } pub fn build_from(source: &PersistentBitVec, path: &Path) -> io::Result { @@ -213,14 +267,25 @@ impl PersistentBitVecBuilder { let file = OpenOptions::new().read(true).write(true).open(path)?; let mmap = unsafe { MmapMut::map_mut(&file)? }; let n = source.len(); - Ok(Self { mmap, n, path: path.to_path_buf() }) + Ok(Self { + mmap, + n, + path: path.to_path_buf(), + }) } - pub fn build_from_counts(source: &PersistentCompactIntVec, threshold: u32, path: &Path) -> io::Result { + pub fn build_from_counts( + source: &PersistentCompactIntVec, + threshold: u32, + path: &Path, + ) -> io::Result { let n = source.len(); let file_size = HEADER_SIZE + n_bytes_for_words(n); let mut file = OpenOptions::new() - .read(true).write(true).create(true).truncate(true) + .read(true) + .write(true) + .create(true) + .truncate(true) .open(path)?; file.write_all(&MAGIC)?; file.write_all(&[0u8; 4])?; @@ -229,22 +294,32 @@ impl PersistentBitVecBuilder { file.set_len(file_size as u64)?; let mut mmap = unsafe { MmapMut::map_mut(&file)? }; { - let nw = n_words(n); + let nw = n_words(n); let ptr = mmap[HEADER_SIZE..].as_mut_ptr() as *mut u64; let words = unsafe { std::slice::from_raw_parts_mut(ptr, nw) }; for (slot, count) in source.iter().enumerate() { - if count >= threshold { words[slot >> 6] |= 1u64 << (slot & 63); } + if count >= threshold { + words[slot >> 6] |= 1u64 << (slot & 63); + } } } - Ok(Self { mmap, n, path: path.to_path_buf() }) + Ok(Self { + mmap, + n, + path: path.to_path_buf(), + }) } pub fn build_from_presence(source: &PersistentCompactIntVec, path: &Path) -> io::Result { Self::build_from_counts(source, 1, path) } - pub fn len(&self) -> usize { self.n } - pub fn is_empty(&self) -> bool { self.n == 0 } + pub fn len(&self) -> usize { + self.n + } + pub fn is_empty(&self) -> bool { + self.n == 0 + } pub fn get(&self, slot: usize) -> bool { (self.mmap[HEADER_SIZE + (slot >> 3)] >> (slot & 7)) & 1 != 0 @@ -252,19 +327,22 @@ impl PersistentBitVecBuilder { pub fn set(&mut self, slot: usize, value: bool) { let bit = 1u64 << (slot & 63); - if value { self.data_words_mut()[slot >> 6] |= bit; } - else { self.data_words_mut()[slot >> 6] &= !bit; } + if value { + self.data_words_mut()[slot >> 6] |= bit; + } else { + self.data_words_mut()[slot >> 6] &= !bit; + } } fn data_words(&self) -> &[u64] { - let nw = n_words(self.n); + let nw = n_words(self.n); let ptr = self.mmap[HEADER_SIZE..].as_ptr() as *const u64; unsafe { std::slice::from_raw_parts(ptr, nw) } } // SAFETY: same alignment argument as PersistentBitVec::data_words. fn data_words_mut(&mut self) -> &mut [u64] { - let nw = n_words(self.n); + let nw = n_words(self.n); let ptr = self.mmap[HEADER_SIZE..].as_mut_ptr() as *mut u64; unsafe { std::slice::from_raw_parts_mut(ptr, nw) } } @@ -273,7 +351,9 @@ impl PersistentBitVecBuilder { BitSliceView::new(self.data_words(), self.n) } - pub fn words(&self) -> &[u64] { self.data_words() } + pub fn words(&self) -> &[u64] { + self.data_words() + } pub fn copy_from(&mut self, src: BitSliceView<'_>) { assert_eq!(self.n, src.len(), "BitSliceView length mismatch"); @@ -282,25 +362,35 @@ impl PersistentBitVecBuilder { pub fn and(&mut self, other: BitSliceView<'_>) { assert_eq!(self.n, other.len(), "BitSliceView length mismatch"); - for (w, &o) in self.data_words_mut().iter_mut().zip(other.words()) { *w &= o; } + for (w, &o) in self.data_words_mut().iter_mut().zip(other.words()) { + *w &= o; + } } pub fn or(&mut self, other: BitSliceView<'_>) { assert_eq!(self.n, other.len(), "BitSliceView length mismatch"); - for (w, &o) in self.data_words_mut().iter_mut().zip(other.words()) { *w |= o; } + for (w, &o) in self.data_words_mut().iter_mut().zip(other.words()) { + *w |= o; + } } pub fn xor(&mut self, other: BitSliceView<'_>) { assert_eq!(self.n, other.len(), "BitSliceView length mismatch"); - for (w, &o) in self.data_words_mut().iter_mut().zip(other.words()) { *w ^= o; } + for (w, &o) in self.data_words_mut().iter_mut().zip(other.words()) { + *w ^= o; + } } pub fn not(&mut self) { - let rem = self.n % 64; + let rem = self.n % 64; let words = self.data_words_mut(); - for w in words.iter_mut() { *w ^= u64::MAX; } + for w in words.iter_mut() { + *w ^= u64::MAX; + } if rem != 0 { - if let Some(last) = words.last_mut() { *last &= (1u64 << rem) - 1; } + if let Some(last) = words.last_mut() { + *last &= (1u64 << rem) - 1; + } } } @@ -312,17 +402,21 @@ impl PersistentBitVecBuilder { let words = self.data_words_mut(); let nw = n_words(n); for wi in 0..nw { - let base = wi * 64; + let base = wi * 64; let limit = (base + 64).min(n); let mut mask = 0u64; for bit in 0..(limit - base) { let b = primary[base + bit]; - if b < 255 && pred(b as u32) { mask |= 1u64 << bit; } + if b < 255 && pred(b as u32) { + mask |= 1u64 << bit; + } } words[wi] |= mask; } for (slot, val) in col.overflow_entries() { - if pred(val) { words[slot >> 6] |= 1u64 << (slot & 63); } + if pred(val) { + words[slot >> 6] |= 1u64 << (slot & 63); + } } } @@ -334,17 +428,21 @@ impl PersistentBitVecBuilder { let words = self.data_words_mut(); let nw = n_words(n); for wi in 0..nw { - let base = wi * 64; + let base = wi * 64; let limit = (base + 64).min(n); let mut mask = 0u64; for bit in 0..(limit - base) { let b = primary[base + bit]; - if b < 255 && !pred(b as u32) { mask |= 1u64 << bit; } + if b < 255 && !pred(b as u32) { + mask |= 1u64 << bit; + } } words[wi] &= !mask; } for (slot, val) in col.overflow_entries() { - if !pred(val) { words[slot >> 6] &= !(1u64 << (slot & 63)); } + if !pred(val) { + words[slot >> 6] &= !(1u64 << (slot & 63)); + } } } @@ -356,17 +454,21 @@ impl PersistentBitVecBuilder { let words = self.data_words_mut(); let nw = n_words(n); for wi in 0..nw { - let base = wi * 64; + let base = wi * 64; let limit = (base + 64).min(n); let mut mask = 0u64; for bit in 0..(limit - base) { let b = primary[base + bit]; - if b < 255 && pred(b as u32) { mask |= 1u64 << bit; } + if b < 255 && pred(b as u32) { + mask |= 1u64 << bit; + } } words[wi] ^= mask; } for (slot, val) in col.overflow_entries() { - if pred(val) { words[slot >> 6] ^= 1u64 << (slot & 63); } + if pred(val) { + words[slot >> 6] ^= 1u64 << (slot & 63); + } } } @@ -374,7 +476,9 @@ impl PersistentBitVecBuilder { self.view().iter() } - pub fn close(self) -> io::Result<()> { self.mmap.flush() } + pub fn close(self) -> io::Result<()> { + self.mmap.flush() + } pub fn finish(self) -> io::Result { let path = self.path.clone(); diff --git a/src/obikphylo/src/siblings/build.rs b/src/obikphylo/src/siblings/build.rs index c2f7a2f5..a57b8977 100644 --- a/src/obikphylo/src/siblings/build.rs +++ b/src/obikphylo/src/siblings/build.rs @@ -132,20 +132,12 @@ fn build_layer_sibling_annex( let n = mphf.n(); // Whether this layer's family-member fields can carry a real layer - // number (fast path, exploited by a future scan-time reader) or must - // stay presence-only (today's behaviour) — a single fact about the - // whole index, decided once here, never stored: `n_layers` is - // guaranteed identical across every partition (a merge adds one layer - // to all of them at once), and `FamilyMask`'s 3-bit field only has room - // for layers `0..=6`. - let fast_mode = meta.n_layers <= 7; - if !fast_mode { - tracing::warn!( - "layer {l} ({layer_dir:?}): index has {} layers (>7) — sibling-annex fast layer \ - lookup disabled for this build; consider compacting this index", - meta.n_layers - ); - } + // number (exploited by `family_scan::scan_layer_families`'s fast path) + // or must stay presence-only — a single fact about the whole index, + // decided once by `cache` (built from the same `PartitionMeta` this + // function would otherwise re-derive) so the writer and every reader + // agree; see `PartitionCache::fast_mode`'s docs. + let fast_mode = cache.fast_mode(); // ── The annex file itself is the reconciliation state, indexed by // this layer's k-mer iteration order (the physical layout of diff --git a/src/obikphylo/src/siblings/cache.rs b/src/obikphylo/src/siblings/cache.rs index 003bd342..99933c6d 100644 --- a/src/obikphylo/src/siblings/cache.rs +++ b/src/obikphylo/src/siblings/cache.rs @@ -117,18 +117,28 @@ pub(super) struct PartitionCache { /// both [`obikindex::KmerIndex::build_sibling_annex`] and /// [`obikindex::KmerIndex::sibling_annex_stats`]. mats: Vec>, + /// Whether every `FamilyMask` field in this index's sibling annexes can + /// be trusted as a real layer index (`n_layers <= 7`, the same decision + /// `build_layer_sibling_annex` makes when writing them) — computed once + /// here, from `PartitionMeta::n_layers` (not `mats[p].len()`, which can + /// undercount if a layer failed to open), reused by both the annex + /// writer (`build.rs`) and reader (`family_scan.rs`) so the two never + /// disagree. `n_layers` is guaranteed identical across every partition + /// (a merge adds one layer to all of them at once), so the first + /// non-empty partition's count speaks for the whole index. + fast_mode: bool, } impl PartitionCache { pub(super) fn build(partition: &KmerPartition, n_parts: usize, with_counts: bool) -> OKIResult { let pb = progress_bar("open_partitions", n_parts as u64, "partitions"); - let built: Vec> = (0..n_parts) + let built: Vec<(Vec, usize)> = (0..n_parts) .into_par_iter() - .map(|part| -> OKIResult> { + .map(|part| -> OKIResult<(Vec, usize)> { let index_dir = partition.part_dir(part).join(INDEX_SUBDIR); if !index_dir.exists() { pb.inc(1); - return Ok(Vec::new()); + return Ok((Vec::new(), 0)); } let meta = PartitionMeta::load(&index_dir).map_err(olm_to_ok)?; let mut mats = Vec::with_capacity(meta.n_layers); @@ -144,11 +154,27 @@ impl PartitionCache { mats.push(mat); } pb.inc(1); - Ok(mats) + Ok((mats, meta.n_layers)) }) .collect::>>()?; pb.finish_and_clear(); - Ok(Self { mats: built }) + + let n_layers = built.iter().map(|&(_, n)| n).find(|&n| n > 0).unwrap_or(0); + let fast_mode = n_layers <= 7; + if !fast_mode { + tracing::warn!( + "index has {n_layers} layers (>7) — sibling-annex fast layer lookup disabled; \ + consider compacting this index" + ); + } + + let mats = built.into_iter().map(|(m, _)| m).collect(); + Ok(Self { mats, fast_mode }) + } + + /// See the [`fast_mode`](Self::fast_mode) field docs. + pub(super) fn fast_mode(&self) -> bool { + self.fast_mode } /// Existence lookup of `variant` in partition `dest_partition`: tries @@ -186,10 +212,16 @@ impl PartitionCache { /// `PersistentBitMatrix::get` dominated wall-clock time, mostly blocked /// on page faults, even after every partition/batch locality fix above /// this in the traversal. + /// `queries`' 4th element (the annex-recorded layer) is ignored here — + /// every layer is probed with evidence regardless, since this is the + /// `!fast_mode` fallback where that field isn't trustworthy (see + /// [`find_presence_batch_fast`](Self::find_presence_batch_fast)). Same + /// tuple shape as the fast path purely so callers don't need two + /// query representations. pub(super) fn find_presence_batch( &self, dest_partition: usize, - queries: &[(CanonicalKmer, usize, u8)], + queries: &[(CanonicalKmer, usize, u8, u8)], n_genomes: usize, mut on_hit: impl FnMut(usize, u8, usize), ) { @@ -198,7 +230,7 @@ impl PartitionCache { // First hit wins, same semantics as the old per-query loop (a // variant present in an earlier layer shadows later ones). let mut by_layer: Vec> = vec![Vec::new(); mats.len()]; - for &(variant, family_idx, base) in queries { + for &(variant, family_idx, base, _layer) in queries { for (li, mat) in mats.iter().enumerate() { if let Some(slot) = mat.find_slot(variant) { by_layer[li].push((slot, family_idx, base)); @@ -211,17 +243,75 @@ impl PartitionCache { if hits.is_empty() { continue; } + resolve_layer_hits(&mats[li], &hits, n_genomes, &mut on_hit); + } + } + + /// Same contract as [`find_presence_batch`](Self::find_presence_batch), + /// but only valid when [`fast_mode`](Self::fast_mode) is true for this + /// whole index: every query's 4th element is trusted outright as its + /// real destination layer (recorded once, evidence-checked, at + /// annex-build time by [`find`](Self::find)) — no per-layer probing, a + /// direct `index_batch` (no evidence) straight to that layer instead. + /// Safe for the same reason `build_layer_sibling_annex`'s own + /// iteration-pipeline lookups are: this exact variant's membership in + /// this exact layer was already positively established once: replaying + /// it via `index`/`index_batch` here isn't a fresh, unverified + /// assumption, it's reusing a fact already paid for. + pub(super) fn find_presence_batch_fast( + &self, + dest_partition: usize, + queries: &[(CanonicalKmer, usize, u8, u8)], + n_genomes: usize, + mut on_hit: impl FnMut(usize, u8, usize), + ) { + let Some(mats) = self.mats.get(dest_partition) else { return }; + + let mut by_layer: Vec> = vec![Vec::new(); mats.len()]; + for &(variant, family_idx, base, layer) in queries { + if let Some(bucket) = by_layer.get_mut(layer as usize) { + bucket.push((variant, family_idx, base)); + } + } + + for (li, entries) in by_layer.into_iter().enumerate() { + if entries.is_empty() { + continue; + } let mat = &mats[li]; - let n_cols = mat.n_cols().min(n_genomes); - let slots: Vec = hits.iter().map(|&(slot, _, _)| slot).collect(); - let mut carries: Vec> = (0..n_cols).map(|_| Vec::new()).collect(); - mat.fill_sub_matrix_carries(&slots, &mut carries); - for (g, col) in carries.iter().enumerate() { - for (&(_, family_idx, base), &carries_it) in hits.iter().zip(col.iter()) { - if carries_it { - on_hit(family_idx, base, g); - } - } + let variants: Vec = entries.iter().map(|&(variant, _, _)| variant).collect(); + let slots = mat.index_batch(&variants); + let hits: Vec<(usize, usize, u8)> = slots + .into_iter() + .zip(entries.iter()) + .map(|(slot, &(_, family_idx, base))| (slot, family_idx, base)) + .collect(); + resolve_layer_hits(mat, &hits, n_genomes, &mut on_hit); + } + } +} + +/// Shared tail of both [`PartitionCache::find_presence_batch`] and +/// [`PartitionCache::find_presence_batch_fast`]: given one layer's already- +/// resolved `(slot, family_idx, base)` hits, sweep genome-major (via +/// `fill_sub_matrix_carries`, buffer-reusing, no per-call allocation) and +/// call `on_hit` for every carry. The two callers differ only in *how* +/// `hits` gets its slots (evidence-probed vs. trusted `index_batch`) — this +/// is everything after that. +fn resolve_layer_hits( + mat: &Mat, + hits: &[(usize, usize, u8)], + n_genomes: usize, + on_hit: &mut impl FnMut(usize, u8, usize), +) { + let n_cols = mat.n_cols().min(n_genomes); + let slots: Vec = hits.iter().map(|&(slot, _, _)| slot).collect(); + let mut carries: Vec> = (0..n_cols).map(|_| Vec::new()).collect(); + mat.fill_sub_matrix_carries(&slots, &mut carries); + for (g, col) in carries.iter().enumerate() { + for (&(_, family_idx, base), &carries_it) in hits.iter().zip(col.iter()) { + if carries_it { + on_hit(family_idx, base, g); } } } diff --git a/src/obikphylo/src/siblings/family_scan.rs b/src/obikphylo/src/siblings/family_scan.rs index fdde2249..46e5f76d 100644 --- a/src/obikphylo/src/siblings/family_scan.rs +++ b/src/obikphylo/src/siblings/family_scan.rs @@ -136,8 +136,11 @@ struct GeneratedBatch { masks: Vec, /// Flat `slots.len() * n_genomes` — `genome_mask[i * n_genomes + g]`. genome_mask: Vec, - /// `outgoing[dest_partition]` = `(variant, family_idx_in_batch, base)`. - outgoing: Vec>, + /// `outgoing[dest_partition]` = `(variant, family_idx_in_batch, base, + /// layer)` — `layer` is the annex-recorded destination layer + /// (`FamilyMask::family_members`'s `layer_value`, `0` if absent), + /// trustworthy only when `PartitionCache::fast_mode()` is true. + outgoing: Vec>, _permit: ThrottleGuard, } @@ -210,27 +213,27 @@ pub(super) fn scan_layer_families( let mut masks = Vec::with_capacity(n); let mut bases = Vec::with_capacity(n); let mut genome_mask = vec![0u8; n * ctx.n_genomes]; - let mut outgoing: Vec> = (0..ctx.n_parts).map(|_| Vec::new()).collect(); + let mut outgoing: Vec> = (0..ctx.n_parts).map(|_| Vec::new()).collect(); // Pass 1: cheap, no matrix access — own base and this // batch's cross-partition queries. Kmer and mask already in // hand from `iter_minorants_batch`, no second annex read, - // no slot lookup needed for this pass. + // no slot lookup needed for this pass. `family_members` + // (not a hand-rolled `central_canonical_neighbors` + + // `mask.has` loop) already filters to present members and + // hands back each one's annex-recorded layer alongside. for (i, entry) in batch.entries.iter().enumerate() { let (kmer, mask) = (entry.kmer, entry.mask); masks.push(mask); let base = central_base(kmer, ctx.k); bases.push(base); - for other in kmer.central_canonical_neighbors() { - if other == kmer { + for (member, layer) in mask.family_members(kmer, ctx.k) { + if member == kmer { continue; // local — resolved below straight from `mat`, no lookup } - let b = central_base(other, ctx.k); - if !mask.has(b) { - continue; - } - let dest = other.partition(ctx.n_parts); - outgoing[dest].push((other, i, b)); + let b = central_base(member, ctx.k); + let dest = member.partition(ctx.n_parts); + outgoing[dest].push((member, i, b, layer.unwrap_or(0))); } } @@ -280,10 +283,16 @@ pub(super) fn scan_layer_families( // module docs), each thread owning one partition's queries // contiguously until this batch is done. let genome_mask: Vec = batch.genome_mask.into_iter().map(AtomicU8::new).collect(); + let fast_mode = cache.fast_mode(); batch.outgoing.par_iter().enumerate().filter(|(_, q)| !q.is_empty()).for_each(|(dest, queries)| { - cache.find_presence_batch(dest, queries, n_genomes, |i, base, g| { + let on_hit = |i: usize, base: u8, g: usize| { genome_mask[i * n_genomes + g].fetch_or(1 << base, Ordering::Relaxed); - }); + }; + if fast_mode { + cache.find_presence_batch_fast(dest, queries, n_genomes, on_hit); + } else { + cache.find_presence_batch(dest, queries, n_genomes, on_hit); + } }); for i in 0..n {