src/obikphylo/src/siblings/algorithms/alignment.rs

@@ -7,7 +7,7 @@ use obikidxcache::index_cache::IndexCache;
 use obikindex::OKIResult;
 
 use super::masking::{iupac_code, masked_state};
-use super::subsample::{EntropyBias, sample_index};
+use super::subsample::{EntropyBias, SurvivingFamily, sample_index};
 
 /// `sequences[i]` is `genome_indices[i]`'s IUPAC-coded row — every row the
 /// same length (one byte per sampled family, in the order [`sample_index`]
@@ -27,7 +27,12 @@ pub struct SnpAlignment {
 }
 
 /// See [`crate::siblings::extensions::SiblingExt::snp_pseudo_alignment`]
-/// for the public-facing docs — this is its implementation.
+/// for the public-facing docs — this is its implementation. A single
+/// "reduce" ([`reduce_alignment`]) over [`sample_index`]'s per-layer
+/// batches — the only consumer today, but structured so a future sibling
+/// reduce (e.g. the Sankoff cost-matrix calibration's cardinality/
+/// composition tallies) can run over the exact same batches, each cloning
+/// its own `Arc`, without this function changing at all.
 pub(crate) fn snp_pseudo_alignment(
     cache: &IndexCache,
     n: usize,
@@ -49,15 +54,31 @@ pub(crate) fn snp_pseudo_alignment(
         no_ambiguity,
         excluded,
         entropy_bias,
-        |_partition, _layer, _family_idx, _mask, genome_mask| {
-            for (row, &g) in genome_indices.iter().enumerate() {
-                let ch = masked_state(genome_mask[g], free_loss, no_ambiguity)
-                    .map(iupac_code)
-                    .unwrap_or(b'?');
-                sequences[row].push(ch);
-            }
+        |_partition, _layer, survivors| {
+            reduce_alignment(&survivors, &genome_indices, free_loss, no_ambiguity, &mut sequences);
         },
     )?;
 
     Ok(SnpAlignment { sequences, genome_indices })
 }
+
+/// Appends one layer's worth of columns onto `sequences` — a "reduce" over
+/// [`sample_index`]'s per-layer batch, self-contained (no sampling/masking
+/// logic of its own, just character encoding) so it can be called
+/// independently of, and alongside, any other reduce over the same batch.
+fn reduce_alignment(
+    survivors: &[SurvivingFamily],
+    genome_indices: &[usize],
+    free_loss: bool,
+    no_ambiguity: bool,
+    sequences: &mut [Vec<u8>],
+) {
+    for family in survivors {
+        for (row, &g) in genome_indices.iter().enumerate() {
+            let ch = masked_state(family.genome_mask[g], free_loss, no_ambiguity)
+                .map(iupac_code)
+                .unwrap_or(b'?');
+            sequences[row].push(ch);
+        }
+    }
+}

src/obikphylo/src/siblings/algorithms/subsample.rs

@@ -44,6 +44,7 @@
 
 use std::collections::HashSet;
 use std::path::Path;
+use std::sync::Arc;
 
 use obicompactvec::TempBitVecBuilder;
 use obikidxcache::index_cache::IndexCache;
@@ -63,6 +64,19 @@ use crate::siblings::extensions::SiblingBuilder;
 /// accepted families rather than resolving them one at a time.
 const MAX_TOPUP_ROUNDS: usize = 5;
 
+/// One surviving family — already drawn (Bernoulli/entropy-bias), already
+/// resolved (`scan_layer_families`), already passed post-hoc masking
+/// (`super::masking::survives_masking`). `genome_mask` is an owned copy: the
+/// borrowed slice `scan_layer_families` hands its own callback doesn't
+/// outlive that callback, but a `SurvivingFamily` is meant to be handed to
+/// [`sample_index`]'s caller as part of a whole layer's batch, well after
+/// that callback returns.
+pub(crate) struct SurvivingFamily {
+    pub family_idx: usize,
+    pub mask: FamilyMask,
+    pub genome_mask: Vec<u8>,
+}
+
 /// Gaussian-kernel entropy bias parameters — `mu`/`sigma` are the
 /// `--entropy`/`--entropy-sd` CLI values (defaulted by the caller, not
 /// here). `pub`, not `pub(crate)`: part of the public signature of
@@ -89,17 +103,29 @@ fn circular_entropy_values(
 }
 
 /// Samples up to `n` non-monomorphic families index-wide, proportionally
-/// per layer, calling `on_family(partition, layer, family_idx, mask,
-/// genome_mask)` once for every family that was both drawn *and* still
-/// carries real polymorphism after `free_loss`/`no_ambiguity`/`excluded`
-/// masking (`algorithms::masking::survives_masking`) — i.e. every family
-/// this produces is already a usable alignment site, no second resolution
-/// pass needed by the caller. `excluded[g]` (see `survives_masking`'s own
-/// docs) never counts genome `g` toward that check — `genome_mask` itself
-/// still carries every genome's own resolved state, excluded or not, so a
-/// caller building per-genome output (e.g. an alignment row) must apply the
-/// same exclusion itself when consuming it. Returns the actual number of
-/// sites kept, which may be less than `n` (see the module docs).
+/// per layer, calling `on_layer(partition, layer, survivors)` once per
+/// cached layer that contributed anything, with every family from that
+/// layer that was both drawn *and* still carries real polymorphism after
+/// `free_loss`/`no_ambiguity`/`excluded` masking
+/// (`algorithms::masking::survives_masking`) — i.e. every
+/// [`SurvivingFamily`] this produces is already a usable alignment site, no
+/// second resolution pass needed by the caller. `excluded[g]` (see
+/// `survives_masking`'s own docs) never counts genome `g` toward that check
+/// — `genome_mask` itself still carries every genome's own resolved state,
+/// excluded or not, so a caller building per-genome output (e.g. an
+/// alignment row) must apply the same exclusion itself when consuming it.
+///
+/// `survivors` is a single `Arc<Vec<SurvivingFamily>>` — one allocation per
+/// layer (bounded by that layer's own quota, never the whole index), handed
+/// out already shared: a caller that wants several independent reduces over
+/// the same batch (an alignment builder, a tally builder, ...) just clones
+/// the `Arc` once per reduce (an `O(1)` refcount bump, no data copy) instead
+/// of re-deriving or re-resolving anything — including handing a clone to
+/// its own `rayon::join`/`scope` worker, since `Arc` (unlike a plain
+/// borrowed slice) satisfies the `Send + 'static` those need.
+///
+/// Returns the actual number of sites kept, which may be less than `n` (see
+/// the module docs).
 pub(crate) fn sample_index(
     cache: &IndexCache,
     n: usize,
@@ -107,7 +133,7 @@ pub(crate) fn sample_index(
     no_ambiguity: bool,
     excluded: &[bool],
     entropy_bias: Option<EntropyBias>,
-    mut on_family: impl FnMut(usize, usize, usize, FamilyMask, &[u8]),
+    mut on_layer: impl FnMut(usize, usize, Arc<Vec<SurvivingFamily>>),
 ) -> OKIResult<usize> {
     let n_genomes = cache.meta().genomes().len();
     let fast_mode = is_fast_mode(cache);
@@ -164,7 +190,7 @@ pub(crate) fn sample_index(
             no_ambiguity,
             excluded,
             entropy_bias,
-            |family_idx, mask, genome_mask| on_family(part, l, family_idx, mask, genome_mask),
+            &mut on_layer,
         )?;
     }
 
@@ -187,7 +213,7 @@ fn sample_layer(
     no_ambiguity: bool,
     excluded: &[bool],
     entropy_bias: Option<EntropyBias>,
-    mut on_family: impl FnMut(usize, FamilyMask, &[u8]),
+    on_layer: &mut impl FnMut(usize, usize, Arc<Vec<SurvivingFamily>>),
 ) -> OKIResult<usize> {
     if n_minorants == 0 {
         return Ok(0);
@@ -208,6 +234,7 @@ fn sample_layer(
     let mut visited = 0usize;
     let mut kept = 0usize;
     let mut resolved = 0usize; // total accepted (and resolved) so far, across every round
+    let mut survivors: Vec<SurvivingFamily> = Vec::new();
 
     for _round in 0..MAX_TOPUP_ROUNDS {
         // Stop on success (quota reached) or true exhaustion (every
@@ -286,12 +313,20 @@ fn sample_layer(
                     "eligible bitset must only accept non-monomorphic minorants"
                 );
                 if survives_masking(genome_mask, excluded, free_loss, no_ambiguity) {
-                    on_family(family_idx, mask, genome_mask);
+                    survivors.push(SurvivingFamily {
+                        family_idx,
+                        mask,
+                        genome_mask: genome_mask.to_vec(),
+                    });
                     kept += 1;
                 }
             },
         )?;
     }
 
+    if !survivors.is_empty() {
+        on_layer(partition, layer_idx, Arc::new(survivors));
+    }
+
     Ok(kept)
 }