Implement --sibling-annex flag and k-mer family annex builder

Introduces a new `--sibling-annex` CLI flag that triggers the construction of a sibling-count/minorant annex for multi-genome indices. The implementation adds a `FamilyMask` data structure with memory-mapped I/O, enabling lock-free concurrent updates via atomic bitwise operations. Batched processing improves cache locality and parallelism, while helper functions derive family presence and minorant flags dynamically. The feature is exposed through an `IndexCache` extension trait, protected by an exclusive directory lock to prevent index corruption during construction.
This commit is contained in:
Eric Coissac
2026-08-28 20:29:01 +02:00
parent d084396aba
commit 920120d20e
10 changed files with 876 additions and 15 deletions
+10 -4
View File
@@ -36,15 +36,21 @@ impl From<MetricArg> 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,
+26 -4
View File
@@ -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
+7 -7
View File
@@ -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};
@@ -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<u64> {
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<SiblingAnnexBuilder>`, 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::<Vec<_>>()
});
// 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<Vec<(CanonicalKmer, usize, u8)>> =
(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::<usize>() / 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)
}
@@ -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;
@@ -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(())
}
}
@@ -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;
+22
View File
@@ -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()
})
}
+31
View File
@@ -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";
+433
View File
@@ -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<u8> {
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<Item = (CanonicalKmer, Option<u8>)> {
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<Self> {
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<Self> {
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<FamilyMask> {
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<Self> {
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<FamilyMask> {
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)));
}
}