obikmer/src/obikphylo/src/siblings/build.rs

use std::path::Path;
use std::sync::Arc;
use std::sync::atomic::Ordering;

use rayon::prelude::*;

use obikseq::CanonicalKmer;
use obikindex::layer::MphfLayer;
use obikindex::layer::meta::IndexMode;
use obipipeline::ThrottleGuard;
use obisys::progress_bar;

use obikindex::KmerIndex;
use obikindex::OKIResult;

use super::cache::PartitionCache;
use super::helpers::{central_base, is_minorant};
use super::{ANNEX_FILE_NAME, FamilyMask, SiblingAnnexBuilder, olm_to_ok};

// ── 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 the fan-out `Flat` stage this replaced, needs no `Arc` sharing.
///
/// The `usize` is this k-mer's position in `iter_kmers()`'s enumeration
/// order, **not** an MPHF slot — see `DevDocMD/architecture/siblings.md`. 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),
}

/// Adds [`build_sibling_annex`](Self::build_sibling_annex) to `KmerIndex`.
pub trait SiblingAnnexBuildExt {
    /// Build the sibling-count/minorant annex for every layer of every
    /// partition of this (already built) index, writing one annex file per
    /// layer alongside its existing index files. Safe to call again later
    /// (e.g. after a fresh `merge`) — each run simply overwrites the annex
    /// files of the index it is called on.
    ///
    /// Construction only — no statistics gathered here on purpose: this is
    /// meant to run routinely (it is the artefact the SNP-family distances
    /// will consume), while the sibling-count distribution
    /// ([`sibling_annex_stats`](super::stats::SiblingStatsExt::sibling_annex_stats))
    /// is a separate, occasional diagnostic pass over the result, not run
    /// every time.
    ///
    /// 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 KmerIndex {
    fn build_sibling_annex(&self) -> OKIResult<()> {
        let n_parts = self.n_partitions();

        tracing::info!("opening {n_parts} partition(s) for the sibling-annex sweep");
        let cache = Arc::new(PartitionCache::build(
            self,
            n_parts,
            self.meta().config.with_counts,
        )?);

        let pb = progress_bar("sibling_annex", n_parts as u64, "partitions");
        let mut total_slots: u64 = 0;
        for part in 0..n_parts {
            let index_dir = self.index_dir(part);
            if !index_dir.exists() {
                pb.inc(1);
                continue;
            }
            let meta = self.partition_meta(part)?;

            let mut part_slots: u64 = 0;
            for l in 0..meta.n_layers {
                part_slots += build_layer_sibling_annex(
                    self,
                    &self.layer_dir(part, l),
                    &meta.mode,
                    n_parts,
                    l,
                    &cache,
                )?;
            }
            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(())
    }
}

/// Returns the number of distinct k-mers (annex entries) processed, for
/// progress reporting. A free function, not a `KmerIndex` method — called
/// only from `build_sibling_annex` above, in the same file.
fn build_layer_sibling_annex(
    index: &KmerIndex,
    layer_dir: &Path,
    mode: &IndexMode,
    n_parts: usize,
    l: usize,
    cache: &Arc<PartitionCache>,
) -> OKIResult<u64> {
    let mphf = MphfLayer::open(layer_dir, mode).map_err(olm_to_ok)?;
    let k = index.kmer_size();
    let n = mphf.n();

    // Whether this layer's family-member fields can carry a real layer
    // number (exploited by `family_scan::scan_layer_families`'s fast path)
    // or must stay presence-only — a single fact about the whole index,
    // decided once by `cache` (built from the same `PartitionMeta` this
    // function would otherwise re-derive) so the writer and every reader
    // agree; see `PartitionCache::fast_mode`'s docs.
    let fast_mode = cache.fast_mode();

    // ── The annex file itself is the reconciliation state, indexed by
    // this layer's k-mer iteration order (the physical layout of
    // `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 (see `DevDocMD/architecture/siblings.md`: the MPHF is not
    // invertible, and evidence answers membership, not identity). 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 (the
    // previous design) 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` 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 `Flat` one —
    // the actual cross-partition lookup reuses
    // `KmerPartition::query_partition_with` (the same batching mechanism
    // `obikmer query` already uses: open a partition's files once,
    // answer a whole batch of queries against it) instead of one lookup
    // per pipeline item. A per-item lookup (tried first) reopened/
    // re-mmap'd every target partition's files on every single variant
    // — fine at toy scale, but ~90% system time against a real index,
    // observed in practice. A *later* attempt still pushed one pipeline
    // message per generated variant (a `Flat` stage, `SourceItem` =>
    // `VariantQuery`, one k-mer in => up to 3 variants out as separate
    // messages) — cheaper than reopening files, but sampling a real run
    // showed most wall-clock time going into per-message channel
    // send/notify syscalls instead of the lookup itself: the pipeline's
    // whole point is amortising synchronisation over a batch, and a
    // single k-mer's ≤3 variants is far too fine a granularity for
    // that. Batching `BATCH_SIZE` source k-mers into one pipeline item
    // — a plain 1-to-1 (`|`, not `||`) 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. `BATCH_SIZE` was originally tuned back when the
    // per-k-mer cost inside this stage (minimiser via `RollingStat`)
    // was ~3x higher than it is now (see `CanonicalKmerOf::minimizer`,
    // a direct O(k) bit-arithmetic replacement) — at the old per-k-mer
    // cost, dispatch overhead (the scheduler's single-threaded `Select`
    // loop: one channel round-trip per batch) was negligible next to
    // the work each batch represented; now that the work itself is
    // cheaper, that fixed per-batch overhead is proportionally larger,
    // so a bigger batch amortises it over more k-mers again. ─────────
    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. See
    // `obipipeline::throttle`'s docs for why this is required, not
    // optional, once a `Flat`-style stage sits in the pipeline.
    //
    // Batches stream straight from `unitigs.bin` via
    // `enumerate_kmers_batch` — never a full-layer `Vec` collect (a
    // layer can hold billions of k-mers; see the "no full collect"
    // rule). 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 = mphf
        .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 the batch_offset fix 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 (layer {layer_dir:?})"
            );
        }
    };
    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. `cache`
    // holds every partition already mmap'd (no more `open()` cost), but
    // `mmap` pages are still loaded on demand and can be evicted — a
    // lookup is not free just because the file isn't reopened. Grouping
    // 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: a central-base substitution changes
    // the minimiser (and thus routes to a different partition) only
    // when the winning minimiser window overlaps the central position.
    // For k=31/m=11 that is ~11 of the 21 possible windows — so *most*
    // of the remaining ~10/21 send the variant right back to the
    // partition already being built. That self-partition bucket ends up
    // far larger than any other, so the naive per-partition split
    // leaves one thread grinding through it alone long after every
    // other partition's (tiny) bucket is done — confirmed by sampling a
    // real run: one thread solely in `MphfLayer::find` while the rest
    // of the pool sits idle. Splitting each non-empty bucket into
    // `chunk_size`-sized pieces first keeps this pass's per-partition
    // locality (each chunk is still contiguous within one partition,
    // still resolved via `cache.find` grouped by that partition) 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(layer) = cache.find(dest, variant) {
                let bits = if fast_mode {
                    FamilyMask::layer_bits(base, 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 `unitigs.bin` (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 consumer that only
    // needs to know "is this k-mer the family's minorant" — the common
    // case, since a family is tallied once, at its minorant — reads
    // the bit straight back instead of re-deriving it (re-scanning
    // `unitigs.bin` and re-hashing through the MPHF, the cost
    // `is_minorant` was cheap to *compute* but expensive to *get the
    // inputs for* every time).
    // 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 mphf.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)
}