Refactor sibling annex to use mmap-backed concurrent storage

Shift the sibling annex construction pipeline from an in-memory atomic mask to a memory-mapped file backend. This enables lock-free concurrent writes directly into the mapped region, streamlining the two-phase write process to accumulate bits atomically before finalization. Adjusted the cache lookup to return the specific matching layer index rather than a boolean flag, and added tests to verify correct layer tracking and cross-partition resolution in merged indexes.

src/obikphylo/src/siblings/build.rs

@@ -1,5 +1,5 @@
 use std::path::Path;
-use std::sync::atomic::{AtomicU8, Ordering};
+use std::sync::atomic::Ordering;
 use std::sync::Arc;
 
 use rayon::prelude::*;
@@ -102,7 +102,7 @@ impl SiblingAnnexBuildExt for KmerIndex {
             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)?;
+                part_slots += build_layer_sibling_annex(self, &layer_dir, n_parts, l, &cache)?;
             }
             total_slots += part_slots;
             pb.inc(1);
@@ -122,6 +122,7 @@ fn build_layer_sibling_annex(
     index: &KmerIndex,
     layer_dir: &Path,
     n_parts: usize,
+    l: usize,
     cache: &Arc<PartitionCache>,
 ) -> OKIResult<u64> {
     let index_dir = layer_dir.parent().expect("layer_dir has a parent index dir");
@@ -130,28 +131,51 @@ fn build_layer_sibling_annex(
     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.
+    // Whether this layer's family-member fields can carry a real layer
+    // number (fast path, exploited by a future scan-time reader) or must
+    // stay presence-only (today's behaviour) — a single fact about the
+    // whole index, decided once here, never stored: `n_layers` is
+    // guaranteed identical across every partition (a merge adds one layer
+    // to all of them at once), and `FamilyMask`'s 3-bit field only has room
+    // for layers `0..=6`.
+    let fast_mode = meta.n_layers <= 7;
+    if !fast_mode {
+        tracing::warn!(
+            "layer {l} ({layer_dir:?}): index has {} layers (>7) — sibling-annex fast layer \
+             lookup disabled for this build; consider compacting this index",
+            meta.n_layers
+        );
+    }
+
+        // ── 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 `docmd/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<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());
+        // `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
@@ -196,26 +220,28 @@ fn build_layer_sibling_annex(
         // 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);
+        // 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)| {
-                    seed_mask[start + i].fetch_or(1 << central_base(kmer, k), Ordering::Relaxed);
+                    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 mask slots after seeding + cross-partition
+        // 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 = mask.iter().filter(|v| v.load(Ordering::Relaxed) == 0).count();
+            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:?})");
             }
@@ -294,8 +320,9 @@ fn build_layer_sibling_annex(
             .collect();
         work.par_iter().for_each(|&(dest, chunk)| {
             for &(variant, source_order, base) in chunk {
-                if cache.find(dest, variant) {
-                    mask[source_order].fetch_or(1 << base, Ordering::Relaxed);
+                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);
                 }
             }
         });
@@ -315,10 +342,14 @@ fn build_layer_sibling_annex(
         // `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)?;
+        // 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 = FamilyMask::from_bits(mask[order].load(Ordering::Relaxed));
+            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));
         }

src/obikphylo/src/siblings/cache.rs

@@ -151,14 +151,17 @@ impl PartitionCache {
         Ok(Self { mats: built })
     }
 
-    /// Existence-only lookup of `variant` in partition `dest_partition`:
-    /// tries each of the partition's already-open layers in turn, stopping
-    /// at the first hit. `find_slot`, not `sub_matrix`/`carries` — no data
-    /// read needed for a plain existence check.
-    pub(super) fn find(&self, dest_partition: usize, variant: CanonicalKmer) -> bool {
+    /// Existence lookup of `variant` in partition `dest_partition`: tries
+    /// each of the partition's already-open layers in turn, stopping at the
+    /// first hit and reporting which layer it was — `find_slot`, not
+    /// `sub_matrix`/`carries`, no data read needed for a plain existence
+    /// check.
+    pub(super) fn find(&self, dest_partition: usize, variant: CanonicalKmer) -> Option<usize> {
         self.mats
-            .get(dest_partition)
-            .is_some_and(|mats| mats.iter().any(|mat| mat.find_slot(variant).is_some()))
+            .get(dest_partition)?
+            .iter()
+            .enumerate()
+            .find_map(|(li, mat)| mat.find_slot(variant).map(|_| li))
     }
 
     /// Resolve many `(variant, family_idx, base)` queries against one

src/obikphylo/src/siblings/siblingannex.rs

@@ -50,6 +50,7 @@
 use std::fs::{File, OpenOptions};
 use std::io;
 use std::path::{Path, PathBuf};
+use std::sync::atomic::AtomicU16;
 
 use memmap2::{Mmap, MmapMut};
 
@@ -101,6 +102,24 @@ impl FamilyMask {
         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
@@ -283,12 +302,36 @@ impl SiblingAnnexBuilder {
         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) {
-        // Redundant concurrent writes from independent recomputation paths
-        // converge to the same encoded value for a given slot, so a plain
-        // store here is safe even without external synchronisation, as long
-        // as the write itself doesn't tear — true for a 2-byte-aligned
-        // `u16` store on every platform this project targets.
         let off = HEADER_SIZE + slot * 2;
         self.mmap[off..off + 2].copy_from_slice(&mask.encode().to_le_bytes());
     }

src/obikphylo/src/siblings/tests.rs

@@ -297,3 +297,37 @@ fn sibling_histogram_does_not_panic_on_partial_last_batch() {
     let total: u64 = hist.iter().sum();
     assert!(total > 0, "histogram should contain at least one family");
 }
+
+/// `sibling_annex_one_sibling_each`'s fixture happens to put g2's k-mer in
+/// layer 0 and g1's in layer 1 (a fresh single-genome index, then one
+/// merge that adds exactly one new layer for content absent from the
+/// first — confirmed empirically, not assumed) — a real two-layer index,
+/// not a contrived one, so it exercises both write paths
+/// `build_layer_sibling_annex` now has to get right: the seeding pass (a
+/// slot's own base gets *this* layer's index directly) and the
+/// cross-partition/cross-layer resolution pass (`PartitionCache::find`
+/// reporting which layer the other family member actually lives in, not
+/// just that it exists).
+#[test]
+fn sibling_annex_records_the_real_layer_of_each_family_member() {
+    let dir = tempdir().unwrap();
+    let g1 = build_single_genome_index(dir.path(), "g1", b"AACCGCTTAAG");
+    let g2 = build_single_genome_index(dir.path(), "g2", b"AACCGGTTAAG");
+    let merged = merge_two(dir.path(), &g1, &g2);
+    merged.build_sibling_annex().expect("build_sibling_annex");
+
+    let index_dir = merged.partition().part_dir(0).join(INDEX_SUBDIR);
+    let meta = PartitionMeta::load(&index_dir).unwrap();
+    assert_eq!(meta.n_layers, 2, "fixture assumption: one merge, one new layer");
+
+    let g1_kmer = canonical(b"AACCGCTTAAG"); // own base C (1), sibling base G (2) — lives in layer 1
+    let g2_kmer = canonical(b"AACCGGTTAAG"); // own base G (2), sibling base C (1) — lives in layer 0
+
+    let a = annex_info_for(&merged, g1_kmer);
+    assert_eq!(a.layer_value(1), Some(1), "g1's own base: seeded with its own layer (1)");
+    assert_eq!(a.layer_value(2), Some(0), "g1's sibling (g2's form): resolved to its real layer (0)");
+
+    let b = annex_info_for(&merged, g2_kmer);
+    assert_eq!(b.layer_value(2), Some(0), "g2's own base: seeded with its own layer (0)");
+    assert_eq!(b.layer_value(1), Some(1), "g2's sibling (g1's form): resolved to its real layer (1)");
+}