diff --git a/src/obikmer2/src/cmd/phylo/args.rs b/src/obikmer2/src/cmd/phylo/args.rs index 6b4b266a..df31e7e0 100644 --- a/src/obikmer2/src/cmd/phylo/args.rs +++ b/src/obikmer2/src/cmd/phylo/args.rs @@ -36,15 +36,21 @@ impl From for DistanceMetric { } /// Partial transfer of `obikmer`'s `phylo` command: the plain distance-metric -/// path only (`--metric`/NJ/UPGMA) — everything sibling-annex-based -/// (`--sibling-annex`, `--snp`, `--sankoff`, `--tnt`/`--phyg`/`--iqtree`, ...) -/// stays in `obikmer` until `obikphylo::siblings` is reconnected (see the -/// project memory on this). +/// path (`--metric`/NJ/UPGMA) plus annex construction (`--sibling-annex`) — +/// everything else sibling-annex-based (`--snp`, `--sankoff`, +/// `--tnt`/`--phyg`/`--iqtree`, stats, ...) stays in `obikmer` until the rest +/// of `obikphylo::siblings` is reconnected (see the project memory on this). #[derive(Args)] pub struct PhyloArgs { /// Index directory pub index: PathBuf, + /// Build (or rebuild) the sibling-count/minorant annex — independent of + /// the distance metric below, meant to be run routinely, ahead of any + /// SNP-family distance computation that will later consume it. + #[arg(long)] + pub sibling_annex: bool, + /// Distance metric to compute #[arg(long, value_enum, default_value = "jaccard")] pub metric: MetricArg, diff --git a/src/obikmer2/src/cmd/phylo/mod.rs b/src/obikmer2/src/cmd/phylo/mod.rs index 49615970..f5661322 100644 --- a/src/obikmer2/src/cmd/phylo/mod.rs +++ b/src/obikmer2/src/cmd/phylo/mod.rs @@ -5,6 +5,7 @@ use std::sync::Arc; use obikidxcache::index_cache::IndexCache; use obikindex::KmerIndex; +use obikphylo::siblings::SiblingAnnexBuildExt; use obikphylo::{Metrics, neighbor_joining, upgma}; use obisys::{Reporter, Stage}; use tracing::info; @@ -31,13 +32,34 @@ pub fn run(args: PhyloArgs) { let mut rep = Reporter::new(); - info!("computing {:?} distances for {} genome(s)", args.metric, n); - // Every partition/layer this needs is opened once, up front, and - // handed to `Metrics::distance` — see `obikquery`'s own use of - // `IndexCache` for the same reasoning (one open, many in-memory reads). + // handed to `Metrics::distance`/`SiblingAnnexBuildExt::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); + // ── Sibling-count/minorant annex (independent of the distance metric) ── + // Meant to be (re)built routinely, ahead of any SNP-family distance + // computation that will later consume it. + if args.sibling_annex { + // Writes into the index directory — hold an exclusive lock for the + // duration so a second, concurrent `--sibling-annex` run on the + // same index can't corrupt these writes (see obisys::DirLock). + let _lock = obisys::DirLock::acquire(&args.index).unwrap_or_else(|e| { + eprintln!("error locking index directory {}: {e}", args.index.display()); + std::process::exit(1); + }); + info!("building sibling-count/minorant annex"); + let t = Stage::start("sibling_annex"); + cache.build_sibling_annex().unwrap_or_else(|e| { + eprintln!("error building sibling annex: {e}"); + std::process::exit(1); + }); + rep.push(t.stop()); + } + + info!("computing {:?} distances for {} genome(s)", args.metric, n); + let need_shared = args.shared_kmers || args.nj || args.upgma; let t = Stage::start("distance"); let result = cache diff --git a/src/obikphylo/src/lib.rs b/src/obikphylo/src/lib.rs index d6611cfc..3720aa6e 100644 --- a/src/obikphylo/src/lib.rs +++ b/src/obikphylo/src/lib.rs @@ -7,16 +7,16 @@ //! cardinality, pseudo-alignment); further phylo-domain functionality //! (currently [`distance`]) moves here incrementally. //! -//! `siblings` is temporarily disconnected (commented out below) while -//! `obikindex`'s stateless audit is being caught up with piece by piece — -//! see the project memory on this. `cardcomp` already moved into -//! `siblings/` (it's sibling-specific) and is unaffected by the -//! disconnection itself compiling; it just isn't reachable from outside -//! this crate until `siblings` comes back. +//! `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 +//! `siblings_old` (SNP distance, cardinality, pseudo-alignment, stats) is +//! still disconnected. mod distance; mod tree; -// pub mod siblings; // temporarily disconnected — see module doc above. +pub mod siblings; pub use distance::{DistanceMetric, DistanceOutput, Metrics}; pub use tree::{Tree, neighbor_joining, upgma}; diff --git a/src/obikphylo/src/siblings/algorithms/annex.rs b/src/obikphylo/src/siblings/algorithms/annex.rs new file mode 100644 index 00000000..cc649149 --- /dev/null +++ b/src/obikphylo/src/siblings/algorithms/annex.rs @@ -0,0 +1,248 @@ +use std::sync::Arc; +use std::sync::atomic::Ordering; + +use rayon::prelude::*; + +use obikidxcache::index_cache::IndexCache; +use obikindex::OKIResult; +use obikindex::layer::KmerLayer; +use obikseq::CanonicalKmer; +use obipipeline::ThrottleGuard; + +use crate::siblings::helpers::{central_base, is_minorant}; +use crate::siblings::{ANNEX_FILE_NAME, FamilyMask, SiblingAnnexBuilder}; + +// ── obipipeline data types ───────────────────────────────────────────────── + +/// A batch of this layer's distinct k-mers (iteration-order index + k-mer), +/// the pipeline's source item — batched, not one k-mer per item, so that +/// pipeline messages and their synchronisation cost stay amortised over +/// thousands of lookups (see `build_layer_sibling_annex`'s comment on +/// `BATCH_SIZE`). Carries the throttle permit for the whole batch, moved +/// (not cloned) into the corresponding `VariantBatch` — a 1-to-1 transform, +/// unlike a fan-out stage, needs no `Arc` sharing. +/// +/// The `usize` is this k-mer's position in `enumerate_kmers()`'s enumeration +/// order, **not** an MPHF slot. It is the same index the annex file is +/// written under. +struct SourceBatch { + items: Vec<(usize, CanonicalKmer)>, + _permit: ThrottleGuard, +} + +/// One batch's worth of central-substitution variants (up to 3 per source +/// k-mer), each already routed to its destination partition and carrying +/// its own central base (0=A/1=C/2=G/3=T) — the mask bit it will set on a +/// hit. `(dest_partition, variant, source_order, base)` per entry, where +/// `source_order` is the source k-mer's iteration-order index (see +/// `SourceBatch`), not an MPHF slot. +struct VariantBatch { + items: Vec<(usize, CanonicalKmer, usize, u8)>, + _permit: ThrottleGuard, +} + +enum SibData { + Batch(SourceBatch), + Variants(VariantBatch), +} + +/// Builds one layer's sibling-count/minorant annex, sweeping every other +/// cached layer of `cache` for each of this layer's k-mers' central-base +/// variants. Returns the number of distinct k-mers (annex entries) +/// processed, for progress reporting. +/// +/// Cross-partition/cross-layer lookups go through `cache` — already opened +/// once for the whole run (see `IndexCache::new`), so no partition/layer is +/// ever reopened or re-mmap'd here, no matter how many source layers this +/// is called for in turn. +pub(crate) fn build_layer_sibling_annex( + cache: &IndexCache, + layer: &KmerLayer, + l: usize, + n_parts: usize, + fast_mode: bool, +) -> OKIResult { + let k = cache.index().kmer_size(); + let n = layer.n(); + + // ── The annex file itself is the reconciliation state, indexed by this + // layer's k-mer iteration order (the physical layout of `unitigs.bin`), + // never by MPHF slot — this layer's own k-mers are known members by + // construction, so no evidence check, no MPHF slot, and no + // slot -> k-mer reconstruction is needed or legitimate here. No + // separate in-memory accumulator: every concurrent write goes straight + // through `SiblingAnnexBuilder::atomic_slot` into the mmap — a layer + // can hold billions of k-mers, so a `Vec` shadowing the whole file in + // RAM just to copy it out again at the end doubles memory for no + // benefit once the builder itself can be written concurrently. + // + // `Arc`, not a bare `SiblingAnnexBuilder` — + // `Pipe::apply` requires its source iterator to be `Send + 'static` + // (its items are dispatched to worker threads that outlive this call), + // so the batch-generating closure below needs an owned handle it can + // move in, not a borrow of a local. `atomic_slot`, not `set`, for every + // concurrent write below: the gather phase parallelises across + // destination partitions (independent `cache.find_in_partition` calls, + // safe to run concurrently) and their hits can land on arbitrary, + // possibly-shared source entries — a lock-free `fetch_or` avoids + // needing any synchronisation beyond that. Reclaimed as a plain owned + // value (`Arc::try_unwrap`) once every concurrent phase below is done, + // for the sequential finalisation pass. ───────────────────────────── + let annex_path = layer.dir().join(ANNEX_FILE_NAME); + let builder = Arc::new(SiblingAnnexBuilder::new(n, &annex_path)?); + + // ── obipipeline: a *batch* transform, not a per-k-mer fan-out one — + // the actual cross-partition lookup reuses `IndexCache::find_in_partition` + // (every partition already open, no reopening/re-mmap'ing per lookup). + // Batching `BATCH_SIZE` source k-mers into one pipeline item — a plain + // 1-to-1 transform, batch in, batch of variants out, one message either + // way — keeps the per-message synchronisation cost amortised over + // thousands of lookups instead of one to three. + const BATCH_SIZE: usize = 32768; + let n_workers = obisys::effective_parallelism(); + let capacity = 256; + + // Throttling limits how many *batches* are in flight at once — the + // permit is acquired per batch (not per k-mer) in the source thread, + // and released once its `VariantBatch` has been read out of the + // pipeline by the accumulation loop below. + // + // Batches stream straight from the layer via `enumerate_kmers_batch` — + // never a full-layer `Vec` collect (a layer can hold billions of + // k-mers). Each k-mer's own base is seeded straight into the annex in + // this same pass, since this is exactly the iteration order the annex + // is keyed on — no separate seeding pass needed. + let seed_builder = Arc::clone(&builder); + let batches = layer + .enumerate_kmers_batch(BATCH_SIZE) + .map(move |(start, kmers)| { + kmers + .into_iter() + .enumerate() + .map(|(i, kmer)| { + let base = central_base(kmer, k); + let bits = if fast_mode { + FamilyMask::layer_bits(base, l) + } else { + FamilyMask::presence_bits(base) + }; + seed_builder + .atomic_slot(start + i) + .fetch_or(bits, Ordering::Relaxed); + (start + i, kmer) + }) + .collect::>() + }); + + // Diagnostic: count still-empty annex slots after seeding + cross-partition + // resolution. A non-zero count here means some k-mer's own base was never + // written (should be impossible after seeding above). + let check_empty = |label: &str| { + let empty = (0..n).filter(|&slot| builder.get(slot).is_none()).count(); + if empty > 0 { + tracing::warn!( + "{label}: {empty} mask slots still empty after resolution ({annex_path:?})" + ); + } + }; + let throttled = obipipeline::throttle(batches, n_workers).map(|t| SourceBatch { + items: t.item, + _permit: t.guard, + }); + + let pipe = obipipeline::make_pipe! { + SibData : SourceBatch => VariantBatch, + | { + move |batch: SourceBatch| -> VariantBatch { + let mut items = Vec::with_capacity(batch.items.len() * 3); + for (order, kmer) in batch.items { + for variant in kmer.central_canonical_neighbors() { + if variant == kmer { + continue; + } + items.push(( + variant.partition(n_parts), + variant, + order, + central_base(variant, k), + )); + } + } + VariantBatch { items, _permit: batch._permit } + } + } : Batch => Variants, + }; + + // ── Group generated variants by destination partition — keeps one + // partition's pages hot while its whole batch is resolved, instead of + // faulting pages in and out as lookups jump between partitions in + // whatever order the pipeline happens to produce them. Each batch's + // throttle permit drops here, once accumulated. + let mut outgoing: Vec> = + (0..n_parts).map(|_| Vec::new()).collect(); + for vb in pipe.apply(throttled, n_workers, capacity) { + for (dest_partition, variant, source_order, base) in vb.items { + outgoing[dest_partition].push((variant, source_order, base)); + } + } + + check_empty("after_seeding"); + + // ── Resolve each partition's batch against the cache, parallelised over + // roughly equal-sized *chunks*, not over partitions — a plain + // `outgoing.par_iter()` over `n_parts` buckets gives one thread the + // whole of one partition's bucket, however large, and a lot of buckets + // are far from equal (see the design doc's discussion of the + // self-partition bucket being disproportionately large). Splitting each + // non-empty bucket into `chunk_size`-sized pieces first keeps this + // pass's per-partition locality while letting Rayon spread a single + // oversized bucket across several threads instead of pinning it to + // one. ───────────────────────────────────────────────────────────── + let chunk_size = ((outgoing.iter().map(Vec::len).sum::() / n_workers).max(1)).min(4096); + let work: Vec<(usize, &[(CanonicalKmer, usize, u8)])> = outgoing + .iter() + .enumerate() + .filter(|(_, q)| !q.is_empty()) + .flat_map(|(dest, q)| q.chunks(chunk_size).map(move |c| (dest, c))) + .collect(); + work.par_iter().for_each(|&(dest, chunk)| { + for &(variant, source_order, base) in chunk { + if let Some((_, dest_layer, _)) = cache.find_in_partition(dest, variant) { + let bits = if fast_mode { + FamilyMask::layer_bits(base, dest_layer) + } else { + FamilyMask::presence_bits(base) + }; + builder + .atomic_slot(source_order) + .fetch_or(bits, Ordering::Relaxed); + } + } + }); + + check_empty("after_resolution"); + + // ── Write the layer's annex file, indexed by iteration order — a + // second streamed pass over the layer's k-mers (via `enumerate_kmers`), + // now that every entry's mask is final; never a `Vec` hold of the whole + // layer. The minorant flag is computed here, not by every later reader: + // this is the one place the whole family's *final* mask and this + // entry's own k-mer (already in hand, no extra lookup) are both + // available together. + // Every concurrent phase above is done and its `Arc` clone (the + // seeding pass's `seed_builder`) has already been dropped along with + // the now-fully-drained `batches` iterator — this is the only + // remaining strong reference, so reclaiming plain ownership for the + // sequential `set` calls below is guaranteed to succeed. + let mut builder = Arc::try_unwrap(builder).unwrap_or_else(|_| { + panic!("sibling-annex builder still shared after every concurrent phase completed") + }); + for (order, kmer) in layer.enumerate_kmers() { + let final_mask = builder.get(order).expect("seeded by construction"); + let minorant = is_minorant(kmer, final_mask, k); + builder.set(order, final_mask.with_minorant(minorant)); + } + builder.close()?; + + Ok(n as u64) +} diff --git a/src/obikphylo/src/siblings/algorithms/mod.rs b/src/obikphylo/src/siblings/algorithms/mod.rs new file mode 100644 index 00000000..3d19732a --- /dev/null +++ b/src/obikphylo/src/siblings/algorithms/mod.rs @@ -0,0 +1,11 @@ +//! 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 +//! 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; + +pub(crate) use annex::build_layer_sibling_annex; diff --git a/src/obikphylo/src/siblings/extensions/annex_build.rs b/src/obikphylo/src/siblings/extensions/annex_build.rs new file mode 100644 index 00000000..5c696297 --- /dev/null +++ b/src/obikphylo/src/siblings/extensions/annex_build.rs @@ -0,0 +1,79 @@ +//! [`SiblingAnnexBuildExt`] — 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`] +//! — same split as `obikphylo::distance::Metrics`, whose `IndexCache` +//! extension impl is itself thin dispatch onto `obicompactvec`'s +//! finalisation methods. + +use obikidxcache::index_cache::IndexCache; +use obikindex::OKIResult; +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 { + /// 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 + /// merge, on a freshly-rebuilt `IndexCache`) — each run simply + /// overwrites the annex files of the layers it covers. + /// + /// Construction only — no statistics gathered here on purpose: this is + /// meant to run routinely (it is the artefact the SNP-family distances + /// consume), while the sibling-count distribution is a separate, + /// occasional diagnostic pass over the result. + /// + /// Cross-partition/cross-layer lookups are required (a k-mer's siblings + /// can live in any partition), but the layer loop itself — and thus the + /// annex file this produces — stays local to one layer at a time. + fn build_sibling_annex(&self) -> OKIResult<()>; +} + +impl SiblingAnnexBuildExt for IndexCache { + fn build_sibling_annex(&self) -> OKIResult<()> { + let n_parts = self.index().n_partitions(); + + // Whether every layer number fits in a `FamilyMask` field (`< 7`) — + // a single fact about the whole index, decided once here so the + // writer (this pass) and every future reader agree. Layer count is + // guaranteed identical across every partition (a merge adds one + // layer to all of them at once), so the max over cached partitions + // speaks for the whole index. + let fast_mode = self + .partitions() + .filter_map(|p| self.n_layer(p)) + .max() + .unwrap_or(0) + <= 7; + if !fast_mode { + tracing::warn!( + "index has more than 7 layers per partition — sibling-annex fast layer lookup \ + disabled; consider compacting this index" + ); + } + + let pb = progress_bar("sibling_annex", self.n_partition() as u64, "partitions"); + let mut total_slots: u64 = 0; + for part in self.partitions() { + let n_layer = self.n_layer(part).unwrap_or(0); + let mut part_slots: u64 = 0; + for l in 0..n_layer { + let layer = self + .get_layer(part, l) + .expect("cache-consistent: n_layer(part) already bounds l"); + part_slots += build_layer_sibling_annex(self, layer, l, n_parts, fast_mode)?; + } + total_slots += part_slots; + pb.inc(1); + pb.set_message(format!( + "partition {part}: {part_slots} kmers ({total_slots} total)" + )); + } + pb.finish_and_clear(); + tracing::info!("sibling annex built — {total_slots} kmers across {n_parts} partitions"); + + Ok(()) + } +} diff --git a/src/obikphylo/src/siblings/extensions/mod.rs b/src/obikphylo/src/siblings/extensions/mod.rs new file mode 100644 index 00000000..d5988cad --- /dev/null +++ b/src/obikphylo/src/siblings/extensions/mod.rs @@ -0,0 +1,9 @@ +//! Extension traits over `obikidxcache::IndexCache` — Rust's orphan rule +//! requires such traits to be defined in the crate that implements them for +//! a foreign type, not in `obikidxcache` itself; same reasoning as every +//! other `obikindex`/`obikidxcache` extension trait in this codebase (see +//! `obikphylo::distance::Metrics`, `obikindexer::extensions`). + +mod annex_build; + +pub use annex_build::SiblingAnnexBuildExt; diff --git a/src/obikphylo/src/siblings/helpers.rs b/src/obikphylo/src/siblings/helpers.rs new file mode 100644 index 00000000..d6e9ab8f --- /dev/null +++ b/src/obikphylo/src/siblings/helpers.rs @@ -0,0 +1,22 @@ +use obikseq::CanonicalKmer; + +use super::FamilyMask; + +/// Central-position base of a canonical k-mer, in the fixed 0=A/1=C/2=G/3=T +/// encoding — the mask's bit index. `k` must be odd (project invariant). +#[inline] +pub(super) fn central_base(kmer: CanonicalKmer, k: usize) -> u8 { + kmer.nucleotide((k - 1) / 2) +} + +/// Is `kmer` the minorant of its family, given the family's presence mask? +/// Regenerates the family's 4 canonical forms from `kmer` itself (cheap, no +/// lookup — see the design doc's "Definitions" section for why this is +/// always safe: the set of 4 forms is invariant regardless of which member +/// you start from), and compares the raw encodings of whichever are marked +/// present in `mask`. +pub(super) fn is_minorant(kmer: CanonicalKmer, mask: FamilyMask, k: usize) -> bool { + kmer.central_canonical_neighbors().into_iter().all(|other| { + other == kmer || !mask.has(central_base(other, k)) || kmer.raw() <= other.raw() + }) +} diff --git a/src/obikphylo/src/siblings/mod.rs b/src/obikphylo/src/siblings/mod.rs new file mode 100644 index 00000000..85a541ac --- /dev/null +++ b/src/obikphylo/src/siblings/mod.rs @@ -0,0 +1,31 @@ +//! Family presence-mask annex. +//! +//! See `DevDocMD/theory/evolutionary_distances.md`, "Definitions: family, and +//! the canonical form of a family" and "Step 2b", for the full design +//! discussion this implements. +//! +//! For each distinct k-mer of each layer of the (already built/merged) +//! index, computes a 4-bit presence mask for its "family" (the up to 4 +//! k-mers sharing its flanks, differing only at the central base — +//! well-defined for odd k): bit `b` set iff the family member whose own +//! canonical central base is `b` (0=A, 1=C, 2=G, 3=T) is observed anywhere +//! in the current multi-genome index. Sibling count and minorant are +//! *derived* from the mask by callers, not stored (see [`FamilyMask`]). +//! +//! Data model ([`FamilyMask`], [`SiblingAnnex`]/[`SiblingAnnexBuilder`] in +//! `siblingannex`, plus the small pure `helpers`) versus construction +//! ([`algorithms`], driven by the [`extensions`] trait it's exposed +//! through) are kept apart the same way `obikindexer` splits its own +//! indexing pipeline: algorithms depend on the data model, never the +//! reverse. + +pub mod algorithms; +pub mod extensions; +mod helpers; +mod siblingannex; + +pub(crate) use siblingannex::{FamilyMask, SiblingAnnex, SiblingAnnexBuilder}; + +pub use extensions::SiblingAnnexBuildExt; + +pub(crate) const ANNEX_FILE_NAME: &str = "siblings.psib"; diff --git a/src/obikphylo/src/siblings/siblingannex.rs b/src/obikphylo/src/siblings/siblingannex.rs new file mode 100644 index 00000000..0ec81d5e --- /dev/null +++ b/src/obikphylo/src/siblings/siblingannex.rs @@ -0,0 +1,433 @@ +//! Family presence-mask annex: a compact, read-only-after-build, per-slot +//! derived value used by the central-position SNP distance estimator (see +//! `DevDocMD/theory/evolutionary_distances.md`, "Step 2b" and "Definitions: +//! family, and the canonical form of a family"). +//! +//! Two bytes are stored per MPHF slot of a partition/layer, packed as four +//! 3-bit fields (one per central base, in the fixed A/C/G/T = 0/1/2/3 +//! encoding) plus a minorant flag: +//! +//! - Each 3-bit field: `0` = the family member whose *own* central base — +//! in its own canonical orientation — is that base is absent from the +//! whole multi-genome index (a property of the whole index, not of any +//! one genome); a non-zero value `v` (`1..=7`) means it's present, and +//! *may* additionally encode `layer = v - 1` (the destination layer +//! within its partition) once a caller populates it via +//! [`with_layer`](FamilyMask::with_layer) — see that method's docs. +//! `FamilyMask` itself is policy-free about whether a given non-zero +//! value is a trustworthy layer index or just "present, no layer +//! recorded": that decision belongs to callers, who already have +//! `PartitionMeta::n_layers` in scope. +//! - The minorant bit: whether *this slot's own k-mer* is its family's +//! minorant (the smallest raw encoding among the family's observed +//! members) — computed once, when the whole family's mask is already +//! final (see `build::build_layer_sibling_annex`), and read back by every +//! consumer that would otherwise have to reconstruct this slot's k-mer +//! from `unitigs.bin` and hash it through the MPHF again just to ask the +//! same question. +//! +//! Both facts an even earlier design stored explicitly are still derived, +//! not stored: +//! - sibling count = `family_size() - 1`; +//! - which of the (up to 4) family members are present, and their canonical +//! form — [`FamilyMask::family_members`], regenerated from the slot's own +//! k-mer (`CanonicalKmer::central_canonical_neighbors`, cheap, no +//! lookup), not stored. +//! +//! Mask value 0 is logically unreachable as a real result (a slot's own +//! base is always present in its own family) and is reused as the "not yet +//! computed" sentinel: annex files are pre-initialised to all-zero, and a +//! real value is only ever written once, by the computation pass. +//! +//! Widened from the original 1-byte/slot design (4 presence bits + 1 +//! minorant bit) to carry a per-member layer number without an extra +//! lookup pass — see the project discussion this implements. An old +//! (1-byte/slot) `.psib` file is *not* silently misread by the new reader: +//! its length no longer matches `HEADER_SIZE + n * 2`, so `SiblingAnnex::open` +//! fails loudly ("PSIB file truncated") rather than producing garbage. +//! Annexes built before this change must be rebuilt (`--sibling-annex`). + +use std::fs::{File, OpenOptions}; +use std::io; +use std::path::{Path, PathBuf}; +use std::sync::atomic::AtomicU16; + +use memmap2::{Mmap, MmapMut}; + +use obikseq::CanonicalKmer; + +use super::helpers::central_base; + +const MAGIC: [u8; 4] = *b"PSIB"; + +// Header: magic(4) + _pad(4) + n(8) = 16 bytes. Data (2 bytes/slot) follows. +const HEADER_SIZE: usize = 16; + +/// Width, in bits, of one base's field. +const FIELD_BITS: u32 = 3; +/// Mask for one base's field once shifted into position. +const FIELD_MASK: u16 = 0b111; +/// Highest raw field value a base can carry (`1..=MAX_FIELD_VALUE`, `0` = absent). +const MAX_FIELD_VALUE: u8 = 7; + +const MINORANT_BIT: u16 = 1 << 12; + +/// A family presence mask — see the module docs for the full bit layout. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct FamilyMask(u16); + +impl FamilyMask { + /// The empty mask — never a valid *computed* result (a slot's own base + /// is always present in its own family) — used only to build up a mask + /// via repeated [`with`](Self::with)/[`with_layer`](Self::with_layer) + /// calls before storing it. + pub const EMPTY: FamilyMask = FamilyMask(0); + + #[inline] + fn field_shift(base: u8) -> u32 { + debug_assert!(base < 4, "base out of range: {base}"); + base as u32 * FIELD_BITS + } + + #[inline] + fn field(self, base: u8) -> u8 { + ((self.0 >> Self::field_shift(base)) & FIELD_MASK) as u8 + } + + #[inline] + fn with_field(self, base: u8, value: u8) -> Self { + debug_assert!(value <= MAX_FIELD_VALUE, "field value out of range: {value}"); + let shift = Self::field_shift(base); + let cleared = self.0 & !(FIELD_MASK << shift); + FamilyMask(cleared | ((value as u16) << shift)) + } + + /// Raw bits to `fetch_or` into a shared `AtomicU16` accumulator to mark + /// `base` present at `layer`, without allocating a whole `FamilyMask` + /// per write — for `build::build_layer_sibling_annex`'s concurrent + /// construction, where each `(word, base)` field is written by exactly + /// one thread (see that module's docs), so a lock-free OR is safe. + #[inline] + pub(crate) fn layer_bits(base: u8, layer: usize) -> u16 { + debug_assert!(layer < MAX_FIELD_VALUE as usize, "layer out of range: {layer}"); + (layer as u16 + 1) << Self::field_shift(base) + } + + /// Same as [`layer_bits`](Self::layer_bits), presence-only (no layer + /// recorded) — the accumulator-side equivalent of [`with`](Self::with). + #[inline] + pub(crate) fn presence_bits(base: u8) -> u16 { + 1u16 << Self::field_shift(base) + } + + /// Mark the member with central base `base` (0=A, 1=C, 2=G, 3=T) as + /// present, without recording a layer (compat path: same observable + /// effect as the original 1-byte design's `with`). Use + /// [`with_layer`](Self::with_layer) instead when the destination layer + /// is already known. + #[inline] + pub fn with(self, base: u8) -> Self { + self.with_field(base, 1) + } + + /// Mark the member with central base `base` as present *and* record its + /// destination layer. `layer` must be `< 7` (`debug_assert`ed) — a + /// caller facing more layers than that has no compact field to record + /// them in and must fall back to [`with`](Self::with) instead (still + /// correct, just without the fast-path payoff); see the module docs on + /// why `FamilyMask` doesn't decide that threshold itself. + #[inline] + pub fn with_layer(self, base: u8, layer: usize) -> Self { + debug_assert!(layer < MAX_FIELD_VALUE as usize, "layer out of range: {layer}"); + self.with_field(base, layer as u8 + 1) + } + + /// Is the member with central base `base` (0..3) present? + #[inline] + pub fn has(self, base: u8) -> bool { + self.field(base) != 0 + } + + /// Raw stored value for `base`'s field, if present: `Some(v - 1)` where + /// `v` is the non-zero field value. This is the *raw* stored value, not + /// a validated layer index — a caller must cross-check it against + /// `PartitionMeta::n_layers` (`<= 7`) before trusting it as a real + /// layer, since a mask written via [`with`](Self::with) (no layer + /// known) also reads back as `Some(0)` here. + #[inline] + pub fn layer_value(self, base: u8) -> Option { + let v = self.field(base); + if v == 0 { None } else { Some(v - 1) } + } + + /// Number of family members observed anywhere in the index (1..=4). + #[inline] + pub fn family_size(self) -> u32 { + (0..4).filter(|&b| self.has(b)).count() as u32 + } + + /// Number of *other* members observed (0..=3) — `family_size() - 1`. + #[inline] + pub fn siblings(self) -> u32 { + self.family_size() - 1 + } + + /// Set or clear the minorant flag (see the module docs). + #[inline] + pub fn with_minorant(self, is_minorant: bool) -> Self { + if is_minorant { + FamilyMask(self.0 | MINORANT_BIT) + } else { + FamilyMask(self.0 & !MINORANT_BIT) + } + } + + /// Is this slot's own k-mer its family's minorant? Only meaningful once + /// `with_minorant` has been called with the family's *final* mask (i.e. + /// after construction) — see the module docs. + #[inline] + pub fn is_minorant(self) -> bool { + self.0 & MINORANT_BIT != 0 + } + + /// Raw presence bitmask (bit `b` = base `b` present, low 4 bits only — + /// never includes the minorant flag or any layer information) — for + /// callers that only need the old 1-bit-per-base view, e.g. to compare + /// against a previously-computed value built via their own bit + /// operations. + #[inline] + pub fn bits(self) -> u8 { + (0..4).fold(0u8, |acc, b| if self.has(b) { acc | (1 << b) } else { acc }) + } + + /// Construct from a raw presence bitmask (only the low 4 bits are kept + /// — the minorant flag is not part of this, use + /// [`with_minorant`](Self::with_minorant) separately). Each set bit is + /// recorded as "present, no layer known" — see [`with`](Self::with). + #[inline] + pub fn from_bits(bits: u8) -> Self { + (0..4).fold(FamilyMask::EMPTY, |mask, b| { + if bits & (1 << b) != 0 { mask.with(b) } else { mask } + }) + } + + /// This family's present members, in canonical form, paired with their + /// raw stored field value (see [`layer_value`](Self::layer_value) for + /// why it's not directly a validated layer index). `owner` is this + /// slot's own k-mer (any member works — the 4 canonical forms are + /// invariant regardless of which member you start from). + pub fn family_members(self, owner: CanonicalKmer, k: usize) -> impl Iterator)> { + owner + .central_canonical_neighbors() + .into_iter() + .filter_map(move |member| { + let base = central_base(member, k); + self.has(base).then(|| (member, self.layer_value(base))) + }) + } + + #[inline] + fn encode(self) -> u16 { + self.0 + } + + #[inline] + fn decode(word: u16) -> Option { + if word == 0 { + // Unreachable for a real result — reserved as the "not yet + // computed" sentinel. A real entry always has at least one + // non-zero base field, so word == 0 always means "nothing + // written yet", never a genuine minorant-only value. + return None; + } + Some(FamilyMask(word)) + } +} + +// ── SiblingAnnex (reader) ─────────────────────────────────────────────────── + +pub struct SiblingAnnex { + mmap: Mmap, + n: usize, + path: PathBuf, +} + +impl SiblingAnnex { + 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, "PSIB file too short")); + } + if mmap[0..4] != MAGIC { + return Err(io::Error::new(io::ErrorKind::InvalidData, "bad PSIB magic")); + } + let n = u64::from_le_bytes(mmap[8..16].try_into().unwrap()) as usize; + if mmap.len() < HEADER_SIZE + n * 2 { + return Err(io::Error::new(io::ErrorKind::InvalidData, "PSIB file truncated")); + } + 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 } + + /// `None` means the slot has not (yet) been computed — see module docs. + pub fn get(&self, slot: usize) -> Option { + let off = HEADER_SIZE + slot * 2; + FamilyMask::decode(u16::from_le_bytes(self.mmap[off..off + 2].try_into().unwrap())) + } +} + +// ── SiblingAnnexBuilder (writer) ──────────────────────────────────────────── + +pub struct SiblingAnnexBuilder { + mmap: MmapMut, +} + +impl SiblingAnnexBuilder { + /// Create a new annex of `n` slots at `path`, pre-initialised to the + /// "not yet computed" sentinel (all-zero). + pub fn new(n: usize, path: &Path) -> io::Result { + let file_size = HEADER_SIZE + n * 2; + let file = OpenOptions::new() + .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[4..8].copy_from_slice(&[0u8; 4]); + mmap[8..16].copy_from_slice(&(n as u64).to_le_bytes()); + // Data region left at 0 by `set_len`/mmap — the sentinel value. + Ok(Self { mmap }) + } + + /// `None` means the slot has not (yet) been computed — see module + /// docs. Used to read back a slot's concurrently-accumulated value + /// (via [`atomic_slot`](Self::atomic_slot)) before finalising it with + /// [`set`](Self::set). + pub fn get(&self, slot: usize) -> Option { + let off = HEADER_SIZE + slot * 2; + FamilyMask::decode(u16::from_le_bytes(self.mmap[off..off + 2].try_into().unwrap())) + } + + /// Atomic handle to `slot`'s raw storage, for lock-free concurrent + /// construction straight into the mmap — no separate in-memory + /// accumulator needed (see `build::build_layer_sibling_annex`, the + /// only caller). Safe: every concurrent write goes through + /// `fetch_or`/atomic RMW, and by that caller's own invariant, a given + /// 3-bit field within a slot's word is written by exactly one thread — + /// only *different* fields of the same word can race, which `fetch_or` + /// already serialises correctly. Alignment holds: `HEADER_SIZE` is + /// even and every slot is 2 bytes, so `HEADER_SIZE + slot * 2` is + /// always even, and the mmap's own base address is page-aligned. + pub(crate) fn atomic_slot(&self, slot: usize) -> &AtomicU16 { + let off = HEADER_SIZE + slot * 2; + debug_assert!(off + 2 <= self.mmap.len()); + let ptr = unsafe { self.mmap.as_ptr().add(off) as *mut u16 }; + unsafe { AtomicU16::from_ptr(ptr) } + } + + /// Plain (non-atomic) store — for the sequential finalisation pass + /// only, once every concurrent write (via + /// [`atomic_slot`](Self::atomic_slot)) is done. + pub fn set(&mut self, slot: usize, mask: FamilyMask) { + let off = HEADER_SIZE + slot * 2; + self.mmap[off..off + 2].copy_from_slice(&mask.encode().to_le_bytes()); + } + + pub fn close(self) -> io::Result<()> { self.mmap.flush() } +} + +#[cfg(test)] +mod tests { + use super::*; + use tempfile::tempdir; + + #[test] + fn sentinel_is_zero_and_unset_slots_read_as_uncomputed() { + let dir = tempdir().unwrap(); + let path = dir.path().join("test.psib"); + let builder = SiblingAnnexBuilder::new(4, &path).unwrap(); + builder.close().unwrap(); + let annex = SiblingAnnex::open(&path).unwrap(); + for slot in 0..4 { + assert_eq!(annex.get(slot), None); + } + } + + #[test] + fn roundtrip_all_valid_masks() { + let dir = tempdir().unwrap(); + let path = dir.path().join("test.psib"); + let mut builder = SiblingAnnexBuilder::new(4, &path).unwrap(); + + let masks = [ + FamilyMask::EMPTY.with(0), // just A: family size 1 + FamilyMask::EMPTY.with(0).with(3), // A + T: size 2 + FamilyMask::EMPTY.with(1).with(2).with(3), // C+G+T: size 3 + FamilyMask::EMPTY.with(0).with(1).with(2).with(3), // all 4 + ]; + for (slot, mask) in masks.iter().enumerate() { + builder.set(slot, *mask); + } + builder.close().unwrap(); + let annex = SiblingAnnex::open(&path).unwrap(); + for (slot, mask) in masks.iter().enumerate() { + assert_eq!(annex.get(slot), Some(*mask)); + } + assert_eq!(annex.get(0).unwrap().siblings(), 0); + assert_eq!(annex.get(1).unwrap().siblings(), 1); + assert_eq!(annex.get(2).unwrap().siblings(), 2); + assert_eq!(annex.get(3).unwrap().siblings(), 3); + assert_eq!(annex.get(3).unwrap().family_size(), 4); + } + + #[test] + fn has_reflects_individual_bits() { + let mask = FamilyMask::EMPTY.with(0).with(2); + assert!(mask.has(0)); + assert!(!mask.has(1)); + assert!(mask.has(2)); + assert!(!mask.has(3)); + } + + #[test] + fn with_layer_roundtrips_and_leaves_absent_bases_none() { + let mask = FamilyMask::EMPTY.with_layer(0, 0).with_layer(2, 5); + assert_eq!(mask.layer_value(0), Some(0)); + assert_eq!(mask.layer_value(2), Some(5)); + assert_eq!(mask.layer_value(1), None); + assert_eq!(mask.layer_value(3), None); + assert!(mask.has(0) && mask.has(2)); + assert!(!mask.has(1) && !mask.has(3)); + assert_eq!(mask.family_size(), 2); + } + + #[test] + fn with_marks_present_without_a_layer() { + // Compat path: `with` still means "present", now reading back as + // `layer_value == Some(0)` (raw field value 1, i.e. "no layer + // recorded") — distinguishable from a real layer 0 only by a + // caller that already knows whether layers were ever wired in. + let mask = FamilyMask::EMPTY.with(1); + assert!(mask.has(1)); + assert_eq!(mask.layer_value(1), Some(0)); + } + + #[test] + fn family_members_reconstructs_present_canonical_forms() { + use obikseq::{Kmer, Sequence}; + + const K: usize = 11; + obikseq::params::set_k(K); + // Centre (index 5) is 'C' (base 1) in this k-mer's own orientation. + let owner = Kmer::from_ascii(b"AACCGCTTAAG").unwrap().canonical(); + let mask = FamilyMask::EMPTY.with_layer(1, 0).with_layer(2, 3); // C (own) + G present + + let members: Vec<_> = mask.family_members(owner, K).collect(); + assert_eq!(members.len(), 2, "only the 2 present members should be yielded"); + assert!(members.iter().any(|(k, layer)| *k == owner && *layer == Some(0))); + assert!(members.iter().any(|(k, layer)| *k != owner && *layer == Some(3))); + } +}