diff --git a/src/obikmer/src/cmd/phylo/mod.rs b/src/obikmer/src/cmd/phylo/mod.rs index 1b822efd..8e3e31d8 100644 --- a/src/obikmer/src/cmd/phylo/mod.rs +++ b/src/obikmer/src/cmd/phylo/mod.rs @@ -14,7 +14,7 @@ use obikphylo::{ cardinality_transition_probs, composition_transition_probs, pairwise_cost_matrix, siblings::{ DistanceExt, EntropyBias, RawSnpDistanceOutput, SankoffBundleExt, ShannonEntropyExt, - SiblingAnnexBuildExt, SiblingStatsExt, SnpAlignment, SnpAlignmentExt, + SiblingExt, SiblingStatsExt, SnpAlignment, SnpAlignmentExt, }, }; use obisys::{Reporter, Stage}; diff --git a/src/obikmer2/src/cmd/phylo/mod.rs b/src/obikmer2/src/cmd/phylo/mod.rs index f5661322..e7d13270 100644 --- a/src/obikmer2/src/cmd/phylo/mod.rs +++ b/src/obikmer2/src/cmd/phylo/mod.rs @@ -5,7 +5,7 @@ use std::sync::Arc; use obikidxcache::index_cache::IndexCache; use obikindex::KmerIndex; -use obikphylo::siblings::SiblingAnnexBuildExt; +use obikphylo::siblings::SiblingExt; use obikphylo::{Metrics, neighbor_joining, upgma}; use obisys::{Reporter, Stage}; use tracing::info; @@ -33,7 +33,7 @@ pub fn run(args: PhyloArgs) { let mut rep = Reporter::new(); // Every partition/layer this needs is opened once, up front, and - // handed to `Metrics::distance`/`SiblingAnnexBuildExt::build_sibling_annex` + // handed to `Metrics::distance`/`SiblingExt::build_sibling_annex` // — see `obikquery`'s own use of `IndexCache` for the same reasoning // (one open, many in-memory reads). let cache = IndexCache::new(Arc::clone(&idx), None); diff --git a/src/obikphylo/src/lib.rs b/src/obikphylo/src/lib.rs index 3720aa6e..b422a7ab 100644 --- a/src/obikphylo/src/lib.rs +++ b/src/obikphylo/src/lib.rs @@ -10,7 +10,7 @@ //! `siblings` is being reconnected piece by piece while `obikindex`'s //! stateless audit is caught up with — see the project memory on this. So //! far it only carries the annex-construction path -//! (`SiblingAnnexBuildExt::build_sibling_annex`); the rest of the old +//! (`SiblingExt::build_sibling_annex`); the rest of the old //! `siblings_old` (SNP distance, cardinality, pseudo-alignment, stats) is //! still disconnected. diff --git a/src/obikphylo/src/siblings/algorithms/family_scan.rs b/src/obikphylo/src/siblings/algorithms/family_scan.rs new file mode 100644 index 00000000..53f426d0 --- /dev/null +++ b/src/obikphylo/src/siblings/algorithms/family_scan.rs @@ -0,0 +1,420 @@ +//! Shared per-layer family traversal — resolves each minorant family's +//! per-genome base-presence, with the same locality discipline as +//! [`super::annex::build_layer_sibling_annex`]: cross-partition lookups are +//! grouped by destination partition and resolved in one contiguous sweep +//! per partition, instead of one lookup jumping between partitions in +//! family order. +//! +//! Two concerns, kept on two different mechanisms because they need +//! opposite things: +//! +//! - **Generating** each batch's cross-partition queries is pure CPU (bit +//! tests, one hash per variant) — cheap, and safe to run for several +//! batches at once. It runs on an `obipipeline::throttle` + `make_pipe!` +//! pipeline, the same mechanism `build_layer_sibling_annex` uses, for the +//! same reason: it overlaps this CPU work with the *previous* batch's +//! resolution instead of leaving cores idle between batches. +//! - **Resolving** those queries against `cache` is I/O (mmap page faults on +//! a real index far larger than RAM) — one batch at a time, `rayon`- +//! parallel *across partitions* (like `build_sibling_annex`'s own +//! `outgoing.par_iter()`), never several batches concurrently: 16 +//! concurrent sweeps across the same partition space scatters exactly what +//! the grouping is meant to prevent, just spread over threads instead of +//! over layers. Resolving one batch's worth at a time, with each of +//! `rayon`'s threads owning one partition contiguously until that batch is +//! done, keeps only one partition set "hot" at once. +//! +//! `obipipeline` does not guarantee output order, so generated batches carry +//! their own starting family index and are replayed through a small reorder +//! buffer before resolution — bounded by the throttle's own concurrency +//! (`n_workers` batches), not by the layer's size. +//! +//! The pipeline's own working context (`LayerCtx`) needs an owned, +//! `Send + 'static` `KmerLayer` — `IndexCache`'s own copy of this same layer +//! is only ever borrowed (`&'_ KmerLayer`, tied to the caller's stack frame), +//! so it can't be moved into the worker threads `Pipe::apply` spawns; this +//! reopens the layer once, independently, for that purpose. Cross-partition +//! resolution, by contrast, runs synchronously on the calling thread (via +//! `rayon::scope`-free `par_iter`, no `thread::spawn`), so it can — and +//! does — borrow `cache` directly for the whole call. + +use std::collections::HashMap; +use std::path::Path; +use std::sync::Arc; +use std::sync::atomic::{AtomicU8, Ordering}; + +use rayon::prelude::*; + +use obikidxcache::index_cache::IndexCache; +use obikindex::layer::KmerLayer; +use obikindex::{OKIError, OKIResult}; +use obikseq::CanonicalKmer; +use obipipeline::ThrottleGuard; + +use crate::siblings::helpers::central_base; +use crate::siblings::iter::{SiblingEntry, SiblingLayerExt}; +use crate::siblings::{ANNEX_FILE_NAME, FamilyMask, SiblingAnnex}; + +/// Families per batch — trades off per-batch memory (`FAMILY_BATCH * +/// genome_count` bytes) against giving each partition a decent number of +/// queries to resolve in one go; too small a batch starves that density +/// regardless of how the resolution step is parallelised. +const FAMILY_BATCH: usize = 65536; + +/// Restricts [`scan_layer_families`] to a subset of a layer's families, +/// identified by their iteration-order index. `All` (the only variant any +/// current caller uses) keeps a full, unsampled scan; `Some` is exercised +/// once sampling (`--subsample`/`--entropy-bias`) is reconnected. +pub enum Selection<'a> { + All, + Some(&'a std::collections::HashSet), +} + +fn is_selected(selected: &Option>, family_idx: usize) -> bool { + match selected { + None => true, + Some(set) => set.contains(&family_idx), + } +} + +/// Read-only state shared (via `Arc`) across every pipeline worker +/// generating this layer's batches — opened once, not per batch. No +/// `cache` here: generation never touches the cross-partition cache, only +/// this layer's own already-open matrix. +struct LayerCtx { + mat: KmerLayer, + n_parts: usize, + n_genomes: usize, + n_cols: usize, + k: usize, +} + +struct SourceBatch { + start_family_idx: usize, + /// One entry per minorant family in this batch, straight from + /// `iter_minorants_batch` — `order` (iteration-order index, not an MPHF + /// slot), `kmer`, and `mask` already carried through, no second annex + /// read. + entries: Vec, + _permit: ThrottleGuard, +} + +/// One batch's generated work: local presence already resolved (straight +/// from this layer's own matrix, no lookup needed), cross-partition queries +/// collected but not yet resolved against the cache. +struct GeneratedBatch { + start_family_idx: usize, + 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, + /// layer)` — `layer` is the annex-recorded destination layer + /// (`FamilyMask::family_members`'s `layer_value`, `0` if absent), + /// trustworthy only when `fast_mode` is true. + outgoing: Vec>, + _permit: ThrottleGuard, +} + +enum FamData { + Batch(SourceBatch), + Generated(GeneratedBatch), +} + +/// Visits every minorant family of one layer, in iteration order, batched +/// [`FAMILY_BATCH`] at a time — see the module docs for why generation and +/// resolution use different concurrency. `on_family` is called once per +/// family, in iteration order, with that family's iteration-order index +/// (the same numbering [`Selection`] indices are drawn from), its own annex +/// mask, and its per-genome base-presence (`genome_mask[g]`: bit `b` set +/// iff genome `g` carries the member whose own canonical central base is +/// `b`), backed by a scratch buffer reused across every call — callers that +/// need to keep data past the call must copy it themselves. +/// +/// `source_partition`/`source_layer` identify which of `cache`'s layers to +/// scan; `fast_mode` must be the same value `build_layer_sibling_annex` +/// used when writing this layer's annex (whether every family member's +/// recorded field is a trustworthy layer index — see +/// `crate::siblings::FamilyMask`'s docs). +pub(crate) fn scan_layer_families( + cache: &IndexCache, + source_partition: usize, + source_layer: usize, + n_genomes: usize, + fast_mode: bool, + selection: &Selection, + mut on_family: impl FnMut(usize, FamilyMask, &[u8]), +) -> OKIResult<()> { + let n_parts = cache.index().n_partitions(); + let k = cache.index().kmer_size(); + + // An owned, independently-opened copy of this same layer — the shared + // `cache`'s own copy is only ever borrowed, so it can't be moved into + // the `Send + 'static` pipeline workers below (see the module docs). + let partition = cache + .index() + .partition(source_partition) + .map_err(|e| OKIError::InvalidInput(format!("partition {source_partition}: {e}")))?; + let mat = partition + .layer(source_layer) + .and_then(KmerLayer::open) + .map_err(|e| OKIError::InvalidInput(format!("layer {source_layer}: {e}")))?; + let n_cols = mat.n_cols().min(n_genomes); + let layer_dir: &Path = mat.dir(); + let annex = Arc::new(SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME)).map_err(OKIError::Io)?); + + let ctx = Arc::new(LayerCtx { mat, n_parts, n_genomes, n_cols, k }); + + // Streamed straight from `iter_minorants_batch` (zips this layer's own + // `iter_kmers()` with the annex, both in iteration order — never an + // MPHF slot) — never collected into a `Vec` first: a layer can hold + // billions of k-mers. + let batches = ctx + .mat + .iter_minorants_batch(annex, FAMILY_BATCH) + .scan(0usize, |offset, batch| { + let start = *offset; + *offset += batch.len(); + Some((start, batch)) + }); + + let n_workers = obisys::effective_parallelism(); + let capacity = 4; + let throttled = obipipeline::throttle(batches, n_workers).map(|t| SourceBatch { + start_family_idx: t.item.0, + entries: t.item.1, + _permit: t.guard, + }); + + // `Selection::Some` borrows a `HashSet` whose lifetime doesn't span the + // pipeline's worker threads — cloned once into an `Arc` so every worker + // can share it cheaply instead of requiring `selection` itself to be + // `'static`. `None` means [`Selection::All`], checked with a plain + // `map_or`-style match at each use instead of allocating an + // always-true set. + let selected: Arc>> = Arc::new(match selection { + Selection::All => None, + Selection::Some(set) => Some((*set).clone()), + }); + + let worker_ctx = Arc::clone(&ctx); + let worker_selected = Arc::clone(&selected); + let pipe = obipipeline::make_pipe! { + FamData : SourceBatch => GeneratedBatch, + | { + move |batch: SourceBatch| -> GeneratedBatch { + let ctx = &worker_ctx; + let n = batch.entries.len(); + 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(); + + // 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. `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. + // Families outside `selected` (sampling only) still get a + // mask/base entry — pass 2 indexes uniformly by `i` — but + // never an outgoing query: that's the expensive part a + // sampled run exists to avoid paying for every family. + 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); + let family_idx = batch.start_family_idx + i; + if !is_selected(&worker_selected, family_idx) { + continue; + } + 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(member, ctx.k); + let dest = member.partition(ctx.n_parts); + outgoing[dest].push((member, i, b, layer.unwrap_or(0))); + } + } + + // Pass 2: genome-major, not family-major — the matrix is + // stored one contiguous block per genome (column), slot as + // the offset within it. `fill_sub_matrix_carries` sorts the + // slots internally for a sequential mmap sweep per column, + // then restores this batch's order. The presence/count + // matrix is still MPHF-slot-indexed (unlike the annex), so + // each entry's kmer is mapped to its slot via `hash_batch` + // — a pure MPHF lookup, no evidence check, since these are + // this layer's own kmers, known members by construction. + let kmers: Vec = batch.entries.iter().map(|e| e.kmer).collect(); + let slots = ctx.mat.hash_batch(&kmers); + let mut carries: Vec> = (0..ctx.n_cols).map(|_| Vec::new()).collect(); + ctx.mat.fill_sub_matrix_carries(&slots, &mut carries); + for (g, col) in carries.iter().enumerate() { + for (i, &carries_it) in col.iter().enumerate() { + if carries_it { + genome_mask[i * ctx.n_genomes + g] |= 1 << bases[i]; + } + } + } + + GeneratedBatch { + start_family_idx: batch.start_family_idx, + masks, + genome_mask, + outgoing, + _permit: batch._permit, + } + } + } : Batch => Generated, + }; + + // Batches finish generation in whatever order their worker completes + // them, not submission order — replay in order through a buffer bounded + // by the throttle's own concurrency (`n_workers` batches can be in + // flight at once, so at most that many can be waiting here). + let mut pending: HashMap = HashMap::new(); + let mut next_expected = 0usize; + let mut scratch = vec![0u8; n_genomes]; + for generated in pipe.apply(throttled, n_workers, capacity) { + pending.insert(generated.start_family_idx, generated); + while let Some(batch) = pending.remove(&next_expected) { + let n = batch.masks.len(); + + // ── Resolve this one batch's cross-partition queries — the + // only place that touches `cache`. Parallel across partitions + // (one batch at a time, never several concurrently — see the + // 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(); + batch + .outgoing + .par_iter() + .enumerate() + .filter(|(_, q)| !q.is_empty()) + .for_each(|(dest, queries)| { + let on_hit = |i: usize, base: u8, g: usize| { + genome_mask[i * n_genomes + g].fetch_or(1 << base, Ordering::Relaxed); + }; + resolve_partition_batch(cache, dest, queries, n_genomes, fast_mode, on_hit); + }); + + for i in 0..n { + let family_idx = batch.start_family_idx + i; + if !is_selected(&selected, family_idx) { + continue; // not in the sample — never resolved above, nothing to report + } + for (g, dst) in scratch.iter_mut().enumerate() { + *dst = genome_mask[i * n_genomes + g].load(Ordering::Relaxed); + } + on_family(family_idx, batch.masks[i], &scratch); + } + next_expected += n; + } + } + debug_assert!( + pending.is_empty(), + "every generated batch must have been replayed" + ); + + Ok(()) +} + +/// Resolve one batch's `(variant, family_idx, base, layer_hint)` queries +/// against `dest_partition`'s already-cached layers, calling +/// `on_hit(family_idx, base, g)` for every genome `g` that carries the +/// resolved variant. `fast_mode` selects between trusting `layer_hint` +/// outright (a direct `hash_batch`, no evidence check — safe because that +/// exact variant's membership in that exact layer was already positively +/// established once, at annex-build time) and probing every layer of the +/// partition in turn with `find_slot` (first hit wins), for indexes with +/// too many layers per partition for `layer_hint` to be trustworthy (see +/// `crate::siblings::FamilyMask`'s docs). +fn resolve_partition_batch( + cache: &IndexCache, + dest_partition: usize, + queries: &[(CanonicalKmer, usize, u8, u8)], + n_genomes: usize, + fast_mode: bool, + mut on_hit: impl FnMut(usize, u8, usize), +) { + let Some(n_layer) = cache.n_layer(dest_partition) else { + return; + }; + + if fast_mode { + let mut by_layer: Vec> = vec![Vec::new(); n_layer]; + for &(variant, family_idx, base, layer_hint) in queries { + if let Some(bucket) = by_layer.get_mut(layer_hint as usize) { + bucket.push((variant, family_idx, base)); + } + } + for (li, entries) in by_layer.into_iter().enumerate() { + if entries.is_empty() { + continue; + } + let layer = cache + .get_layer(dest_partition, li) + .expect("cache-consistent: n_layer(dest_partition) already bounds li"); + let variants: Vec = entries.iter().map(|&(v, _, _)| v).collect(); + let slots = layer.hash_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(layer, &hits, n_genomes, &mut on_hit); + } + } else { + // First hit wins, same semantics as the fast path — a variant + // present in an earlier layer shadows later ones. + let mut by_layer: Vec> = vec![Vec::new(); n_layer]; + for &(variant, family_idx, base, _layer_hint) in queries { + for li in 0..n_layer { + let layer = cache + .get_layer(dest_partition, li) + .expect("cache-consistent: n_layer(dest_partition) already bounds li"); + if let Some(slot) = layer.find_slot(variant) { + by_layer[li].push((slot, family_idx, base)); + break; + } + } + } + for (li, hits) in by_layer.into_iter().enumerate() { + if hits.is_empty() { + continue; + } + let layer = cache + .get_layer(dest_partition, li) + .expect("cache-consistent: n_layer(dest_partition) already bounds li"); + resolve_layer_hits(layer, &hits, n_genomes, &mut on_hit); + } + } +} + +/// Shared tail of both branches of [`resolve_partition_batch`]: 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. +fn resolve_layer_hits( + layer: &KmerLayer, + hits: &[(usize, usize, u8)], + n_genomes: usize, + on_hit: &mut impl FnMut(usize, u8, usize), +) { + let n_cols = layer.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(); + layer.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/algorithms/mod.rs b/src/obikphylo/src/siblings/algorithms/mod.rs index 3d19732a..37595e25 100644 --- a/src/obikphylo/src/siblings/algorithms/mod.rs +++ b/src/obikphylo/src/siblings/algorithms/mod.rs @@ -1,11 +1,13 @@ //! Sibling-annex construction algorithm: the cross-partition sweep that //! fills in one layer's [`crate::siblings::FamilyMask`] slots — driven, one //! layer at a time, by -//! [`crate::siblings::extensions::SiblingAnnexBuildExt`]. Kept apart from +//! [`crate::siblings::extensions::SiblingExt`]. Kept apart from //! that trait (and from the `siblingannex`/`helpers` data model) the same //! way `obikindexer` splits `algorithms` from `extensions`: this module //! depends on the data model, never the reverse. mod annex; +mod family_scan; pub(crate) use annex::build_layer_sibling_annex; +pub(crate) use family_scan::{Selection, scan_layer_families}; diff --git a/src/obikphylo/src/siblings/extensions/annex_build.rs b/src/obikphylo/src/siblings/extensions/annex_build.rs index 5c696297..acd2519b 100644 --- a/src/obikphylo/src/siblings/extensions/annex_build.rs +++ b/src/obikphylo/src/siblings/extensions/annex_build.rs @@ -1,4 +1,4 @@ -//! [`SiblingAnnexBuildExt`] — the public entry point for building the +//! [`SiblingExt`] — the public entry point for building the //! sibling-count/minorant annex of an already-cached index. Thin: it only //! walks `cache`'s partitions/layers and delegates each layer's actual //! cross-partition sweep to [`crate::siblings::algorithms::build_layer_sibling_annex`] @@ -13,7 +13,7 @@ use obisys::progress_bar; use crate::siblings::algorithms::build_layer_sibling_annex; /// Adds [`build_sibling_annex`](Self::build_sibling_annex) to `IndexCache`. -pub trait SiblingAnnexBuildExt { +pub trait SiblingExt { /// Build the sibling-count/minorant annex for every layer of every /// cached partition, writing one annex file per layer alongside its /// existing index files. Safe to call again later (e.g. after a fresh @@ -31,7 +31,7 @@ pub trait SiblingAnnexBuildExt { fn build_sibling_annex(&self) -> OKIResult<()>; } -impl SiblingAnnexBuildExt for IndexCache { +impl SiblingExt for IndexCache { fn build_sibling_annex(&self) -> OKIResult<()> { let n_parts = self.index().n_partitions(); diff --git a/src/obikphylo/src/siblings/extensions/mod.rs b/src/obikphylo/src/siblings/extensions/mod.rs index d5988cad..d5be0833 100644 --- a/src/obikphylo/src/siblings/extensions/mod.rs +++ b/src/obikphylo/src/siblings/extensions/mod.rs @@ -6,4 +6,4 @@ mod annex_build; -pub use annex_build::SiblingAnnexBuildExt; +pub use annex_build::SiblingExt; diff --git a/src/obikphylo/src/siblings/iter.rs b/src/obikphylo/src/siblings/iter.rs new file mode 100644 index 00000000..2c00d2ec --- /dev/null +++ b/src/obikphylo/src/siblings/iter.rs @@ -0,0 +1,250 @@ +//! Phylo/sibling-domain iteration over a layer — an extension trait, not a +//! new field on `TypedLayer`/`KmerLayer`: "family"/"minorant" are phylo +//! concepts, `obikindex::layer` stays kmer/slot-mapping only. +//! +//! The sibling annex is persisted in the same order as `iter_kmers()` (see +//! `crate::siblings::algorithms::build_layer_sibling_annex`), so pairing +//! them is a plain zip — no MPHF, no slot. Both sides are already `Send + +//! 'static` (`KmerIter` owns an `Arc`-backed reader clone; `SiblingAnnex` is +//! mmap-backed and handed in as an `Arc` by the caller), so `SiblingIter` +//! streams straight from disk and can be fed to `obipipeline` batch by +//! batch — never collected whole into memory. +//! +//! Four iterator types, deliberately mirroring `obikindex::layer`'s own +//! `KmerIter`/`KmerBatchIter` pair (single item vs. `Vec` batch) — plus the +//! minorant-filtered variant of each, since "all siblings" and "one row per +//! family" are both common cases: +//! +//! | | all entries | minorants only | +//! |------------|-----------------------|------------------------| +//! | single | [`SiblingIter`] | [`MinorantIter`] | +//! | batch | [`SiblingBatchIter`] | [`MinorantBatchIter`] | +//! +//! No separate "enumerate" variant (unlike `KmerIter`/`enumerate_kmers`): +//! [`SiblingEntry`] already carries `order` for free, since pairing with +//! the annex requires it anyway. + +use std::sync::Arc; + +use obikindex::layer::{KmerIter, LayerData, TypedLayer}; +use obikseq::CanonicalKmer; + +use super::{FamilyMask, SiblingAnnex}; + +/// One layer entry: a k-mer's position in the layer's iteration order (the +/// same index the sibling annex is keyed on — not an MPHF slot), the k-mer +/// itself, and its family mask. +#[derive(Debug, Clone, Copy)] +pub(crate) struct SiblingEntry { + pub order: usize, + pub kmer: CanonicalKmer, + pub mask: FamilyMask, +} + +/// Streams `(order, kmer, mask)` triples for one layer, in iteration order. +/// Produced by [`SiblingLayerExt::iter_siblings`]. +pub(crate) struct SiblingIter { + kmers: KmerIter, + annex: Arc, + order: usize, +} + +impl Iterator for SiblingIter { + type Item = SiblingEntry; + + fn next(&mut self) -> Option { + loop { + let kmer = self.kmers.next()?; + let order = self.order; + self.order += 1; + // `None` means "not yet computed" (see `SiblingAnnex` module + // docs) — shouldn't happen against a fully-built annex, but + // skip rather than misalign the two streams if it does. + if let Some(mask) = self.annex.get(order) { + return Some(SiblingEntry { order, kmer, mask }); + } + } + } +} + +/// Batches of [`SiblingIter`]'s entries, `batch_size` at a time — the last +/// batch may be shorter. Produced by [`SiblingLayerExt::iter_siblings_batch`]. +pub(crate) struct SiblingBatchIter { + inner: SiblingIter, + batch_size: usize, +} + +impl Iterator for SiblingBatchIter { + type Item = Vec; + + fn next(&mut self) -> Option { + collect_batch(&mut self.inner, self.batch_size) + } +} + +/// Like [`SiblingIter`], filtered to the minorant of each family — the +/// common case, since a family is tallied once, at its minorant. Produced +/// by [`SiblingLayerExt::iter_minorants`]. +pub(crate) struct MinorantIter { + inner: SiblingIter, +} + +impl Iterator for MinorantIter { + type Item = SiblingEntry; + + fn next(&mut self) -> Option { + self.inner.by_ref().find(|e| e.mask.is_minorant()) + } +} + +/// Batches of [`MinorantIter`]'s entries, `batch_size` at a time — the last +/// batch may be shorter. Produced by +/// [`SiblingLayerExt::iter_minorants_batch`]. +pub(crate) struct MinorantBatchIter { + inner: MinorantIter, + batch_size: usize, +} + +impl Iterator for MinorantBatchIter { + type Item = Vec; + + fn next(&mut self) -> Option { + collect_batch(&mut self.inner, self.batch_size) + } +} + +/// Shared by [`SiblingBatchIter`] and [`MinorantBatchIter`] — pull up to +/// `batch_size` items, `None` once the source is exhausted with nothing left. +fn collect_batch(inner: &mut I, batch_size: usize) -> Option> { + let mut batch = Vec::with_capacity(batch_size); + for _ in 0..batch_size { + match inner.next() { + Some(item) => batch.push(item), + None => break, + } + } + if batch.is_empty() { None } else { Some(batch) } +} + +/// Adds phylo/sibling iteration to any `TypedLayer` — the extension that +/// turns a plain layer into a "sibling layer". Generic over `D` +/// (`LayerData`) rather than implemented once per matrix kind: kmer +/// iteration doesn't depend on the data payload, and `TypedLayer` already +/// carries `iter_kmers`/`hash_batch`/`fill_sub_matrix_carries` for every `D`. +pub(crate) trait SiblingLayerExt { + /// Zip this layer's k-mers with their sibling-annex entry, in iteration + /// order. `annex` must have been built from this same layer (its length + /// must match the layer's k-mer count). + fn iter_siblings(&self, annex: Arc) -> SiblingIter; + + /// Like [`iter_siblings`](Self::iter_siblings), yielding `batch_size` + /// entries at a time. + fn iter_siblings_batch(&self, annex: Arc, batch_size: usize) -> SiblingBatchIter; + + /// Like [`iter_siblings`](Self::iter_siblings), filtered to the + /// minorant of each family — the common case, since a family is + /// tallied once, at its minorant. + fn iter_minorants(&self, annex: Arc) -> MinorantIter; + + /// Like [`iter_minorants`](Self::iter_minorants), yielding `batch_size` + /// minorants at a time. + fn iter_minorants_batch( + &self, + annex: Arc, + batch_size: usize, + ) -> MinorantBatchIter; +} + +impl SiblingLayerExt for TypedLayer { + fn iter_siblings(&self, annex: Arc) -> SiblingIter { + SiblingIter { + kmers: self.iter_kmers(), + annex, + order: 0, + } + } + + fn iter_siblings_batch(&self, annex: Arc, batch_size: usize) -> SiblingBatchIter { + SiblingBatchIter { + inner: self.iter_siblings(annex), + batch_size, + } + } + + fn iter_minorants(&self, annex: Arc) -> MinorantIter { + MinorantIter { + inner: self.iter_siblings(annex), + } + } + + fn iter_minorants_batch( + &self, + annex: Arc, + batch_size: usize, + ) -> MinorantBatchIter { + MinorantBatchIter { + inner: self.iter_minorants(annex), + batch_size, + } + } +} + +/// Same extension, over `obikindex::layer::KmerLayer` (the format-erased +/// `Count`/`Presence` handle `IndexCache` actually holds) — dispatch only, +/// both arms return the same concrete iterator types (they don't depend on +/// which `D` is inside), so no boxing is needed. `obikindex::layer` itself +/// can't implement this: "family"/"minorant" are phylo concepts, it stays +/// kmer/slot-mapping only (see the module docs above). +impl SiblingLayerExt for obikindex::layer::KmerLayer { + fn iter_siblings(&self, annex: Arc) -> SiblingIter { + match self { + obikindex::layer::KmerLayer::Count { layer, .. } => layer.iter_siblings(annex), + obikindex::layer::KmerLayer::Presence { layer, .. } => layer.iter_siblings(annex), + obikindex::layer::KmerLayer::Empty { .. } => { + panic!("iter_siblings() called on an Empty layer") + } + } + } + + fn iter_siblings_batch(&self, annex: Arc, batch_size: usize) -> SiblingBatchIter { + match self { + obikindex::layer::KmerLayer::Count { layer, .. } => { + layer.iter_siblings_batch(annex, batch_size) + } + obikindex::layer::KmerLayer::Presence { layer, .. } => { + layer.iter_siblings_batch(annex, batch_size) + } + obikindex::layer::KmerLayer::Empty { .. } => { + panic!("iter_siblings_batch() called on an Empty layer") + } + } + } + + fn iter_minorants(&self, annex: Arc) -> MinorantIter { + match self { + obikindex::layer::KmerLayer::Count { layer, .. } => layer.iter_minorants(annex), + obikindex::layer::KmerLayer::Presence { layer, .. } => layer.iter_minorants(annex), + obikindex::layer::KmerLayer::Empty { .. } => { + panic!("iter_minorants() called on an Empty layer") + } + } + } + + fn iter_minorants_batch( + &self, + annex: Arc, + batch_size: usize, + ) -> MinorantBatchIter { + match self { + obikindex::layer::KmerLayer::Count { layer, .. } => { + layer.iter_minorants_batch(annex, batch_size) + } + obikindex::layer::KmerLayer::Presence { layer, .. } => { + layer.iter_minorants_batch(annex, batch_size) + } + obikindex::layer::KmerLayer::Empty { .. } => { + panic!("iter_minorants_batch() called on an Empty layer") + } + } + } +} diff --git a/src/obikphylo/src/siblings/mod.rs b/src/obikphylo/src/siblings/mod.rs index 85a541ac..5d3377a1 100644 --- a/src/obikphylo/src/siblings/mod.rs +++ b/src/obikphylo/src/siblings/mod.rs @@ -22,10 +22,11 @@ pub mod algorithms; pub mod extensions; mod helpers; +mod iter; mod siblingannex; pub(crate) use siblingannex::{FamilyMask, SiblingAnnex, SiblingAnnexBuilder}; -pub use extensions::SiblingAnnexBuildExt; +pub use extensions::SiblingExt; pub(crate) const ANNEX_FILE_NAME: &str = "siblings.psib";