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

use std::path::Path;
use std::sync::atomic::{AtomicU8, Ordering};
use std::sync::Arc;

use rayon::prelude::*;

use obicompactvec::{FamilyMask, SiblingAnnexBuilder};
use obikpartitionner::KmerPartition;
use obipipeline::ThrottleGuard;
use obikseq::CanonicalKmer;
use obilayeredmap::MphfLayer;
use obilayeredmap::meta::PartitionMeta;
use obisys::progress_bar;

use obikindex::{OKIError, OKIResult};
use obikindex::KmerIndex;

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

// ── 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 `docmd/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();
        let n_bits = n_parts.trailing_zeros() as usize;

        let partition = KmerPartition::open_with_config(
            self.root_path(),
            self.kmer_size(),
            self.minimizer_size(),
            n_bits,
        )
        .map_err(OKIError::Partition)?;

        tracing::info!("opening {n_parts} partition(s) for the sibling-annex sweep");
        let cache = Arc::new(PartitionCache::build(&partition, 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.partition().part_dir(part).join(INDEX_SUBDIR);
            if !index_dir.exists() {
                pb.inc(1);
                continue;
            }
            let meta = PartitionMeta::load(&index_dir).map_err(olm_to_ok)?;

            let mut part_slots: u64 = 0;
            for l in 0..meta.n_layers {
                let layer_dir = index_dir.join(format!("layer_{l}"));
                part_slots += build_layer_sibling_annex(self, &layer_dir, n_parts, &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,
    n_parts: usize,
    cache: &Arc<PartitionCache>,
) -> OKIResult<u64> {
    let index_dir = layer_dir.parent().expect("layer_dir has a parent index dir");
    let meta = PartitionMeta::load(index_dir).map_err(olm_to_ok)?;
    let mphf = MphfLayer::open(layer_dir, &meta.mode).map_err(olm_to_ok)?;
    let k = index.kmer_size();
    let n = mphf.n();

        // ── 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
        // `docmd/architecture/siblings.md`: the MPHF is not invertible, and
        // evidence answers membership, not identity). The annex is written
        // under this same iteration order end to end, so a reader can later
        // zip `iter_kmers()` with the annex file directly, with no
        // MPHF/slot indirection at read time either.
        //
        // `Arc<Vec<AtomicU8>>`, not `Vec<AtomicU8>` — `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. `AtomicU8`, not `FamilyMask`,
        // because the gather phase below parallelises across destination
        // partitions (independent `query_partition_with` calls, safe to run
        // concurrently) and their `Found` hits can land on arbitrary,
        // possibly-shared source entries — a lock-free `fetch_or` avoids
        // needing any synchronisation beyond that. ─────────────────────────
        let mask: Arc<Vec<AtomicU8>> = Arc::new((0..n).map(|_| AtomicU8::new(0)).collect());

        // ── 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. ──────────────────────────────────────────────
        const BATCH_SIZE: usize = 4096;
        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 into `mask` in this same
        // pass, since this is exactly the iteration order `mask` is keyed
        // on — no separate seeding pass needed.
        let seed_mask = Arc::clone(&mask);
        let batches = mphf.enumerate_kmers_batch(BATCH_SIZE).map(move |(start, kmers)| {
            kmers
                .into_iter()
                .enumerate()
                .map(|(i, kmer)| {
                    seed_mask[start + i].fetch_or(1 << central_base(kmer, k), Ordering::Relaxed);
                    (start + i, kmer)
                })
                .collect::<Vec<_>>()
        });
        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));
            }
        }

        // ── Resolve each partition's batch against the cache in one
        // contiguous pass; parallelised across partitions (independent,
        // read-only) so this keeps using multiple cores without giving up
        // the per-partition locality above. ─────────────────────────────
        outgoing.par_iter().enumerate().filter(|(_, q)| !q.is_empty()).for_each(|(dest, queries)| {
            for &(variant, source_order, base) in queries {
                if cache.find(dest, variant) {
                    mask[source_order].fetch_or(1 << base, Ordering::Relaxed);
                }
            }
        });

        // ── 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).
        let annex_path = layer_dir.join(ANNEX_FILE_NAME);
        let mut builder = SiblingAnnexBuilder::new(n, &annex_path)?;
        for (order, kmer) in mphf.enumerate_kmers() {
            let final_mask = FamilyMask::from_bits(mask[order].load(Ordering::Relaxed));
            let minorant = is_minorant(kmer, final_mask, k);
            builder.set(order, final_mask.with_minorant(minorant));
        }
        builder.close()?;

    Ok(n as u64)
}