caps_sa/

ext_mem.rs

//! External-memory suffix array construction.
//!
//! Implements upstream CaPS-SA's *sample-sort + LCP-enhanced merge*
//! external-memory algorithm:
//!
//! 1. **Presample pivots.** Sort a small deterministic sample and pick
//!    `p − 1` pivots at evenly-spaced ranks, splitting the suffix order into
//!    `p` partitions.
//! 2. **Sort + distribute.** Split the selected positions into `p` subarrays.
//!    In parallel, sort each with a task-local LCP merge-sort, split it by the
//!    pivots, and write each sorted slice directly to its final partition
//!    bucket. The complete intermediate subarray spill is avoided.
//! 3. **Per-partition merge.** Load each partition's bucket into RAM
//!    (≈ `n / p` records); cascade 2-way LCP-enhanced merges across the
//!    `p` sub-subarrays to produce that partition's globally-sorted slice.
//! 4. **Stream output.** Iterate partitions in order and emit each
//!    position via the caller's closure.
//!
//! Peak RAM is the text plus bounded per-worker phase-1 scratch and in-flight
//! `O(n / p)` partition merge buffers. The full suffix array is never
//! materialized in RAM.

use std::cmp::Ordering;
use std::io;
use std::num::NonZeroUsize;
use std::path::{Path, PathBuf};
use std::sync::Mutex;
use std::time::Instant;

use rayon::prelude::*;

use crate::Index;
use crate::ext_bucket::{BucketPool, BucketRecord, InMemBucket, SaLcp, SaLcpBucketStore};
use crate::lcp::{LcpDispatch, Symbol};
use crate::lcp_memo::{
    GeometricMemo, GeometricMemoizationConfig, LcpMemoizationPolicy, MemoConfig, MemoStats,
};
use crate::limits::{LimitProvider, PlainText};
use crate::sample_sort;

/// Emit a phase-timing line to stderr if `CAPS_SA_PROFILE` is set in
/// the environment. Used to localise where the ext-mem path spends its
/// time without paying the cost of always logging — see
/// `bench/README.md` "Where AVX-512 helps and where it doesn't" for
/// how this is used.
fn profile_log(message: &str) {
    if std::env::var_os("CAPS_SA_PROFILE").is_some() {
        eprintln!("caps-sa profile  {message}");
    }
}

/// Tunable options for [`build_ext_mem`].
///
/// This structure is non-exhaustive so future releases can add options
/// without repeatedly breaking callers that use struct literals. Start from
/// [`Self::default`] or [`Self::from_env`] and use the builder methods.
#[non_exhaustive]
#[derive(Clone, Debug)]
pub struct ExtMemOpts {
    /// Bound on suffix-comparison context. Comparisons equal through this many
    /// symbols are resolved by [`LimitProvider::boundary_order`], so
    /// `usize::MAX` (the default) is required for ordinary full lexicographic
    /// ordering unless the caller deliberately wants truncated contexts.
    pub max_context: usize,
    /// Number of subarrays (`p` in upstream CaPS-SA). `0` (default) targets
    /// roughly 65,536 selected positions per subarray, clamped to at least one
    /// subarray per Rayon worker and at most 8,192 (and never above `n`).
    pub subproblem_count: usize,
    /// Directory for temp files. Defaults to [`std::env::temp_dir`].
    pub work_dir: PathBuf,
    /// Number of physical files backing the final partition buckets. `0`
    /// (default) picks `rayon::current_num_threads()` —
    /// the right answer in practice: one writable inode per worker
    /// keeps kernel-level write contention bounded.
    ///
    /// The `p` logical buckets (typically thousands at genome scale) collapse
    /// onto this pool of anonymous tempfiles via
    /// `bucket_id % physical_file_count`. Larger values lower kernel
    /// write contention; smaller values are kinder to networked
    /// filesystems with high metadata cost. The `CAPS_SA_N_PHYS` env
    /// var overrides this for one-off benches.
    pub physical_file_count: usize,
    /// Use the bounded ordered phase-4 emitter instead of the default
    /// chunk collect-then-emit path. This can reduce transient merged
    /// partition residency on skewed workloads, but is slower on the
    /// GRCh38 32-thread benchmark because of channel coordination and
    /// backpressure, so it is opt-in.
    pub ordered_phase4_emit: bool,
    /// Policy for reusing exact long-LCP intervals during phase-4 partition
    /// merges. Disabled by default; callers with long repeated contexts can
    /// opt into [`LcpMemoizationPolicy::Geometric`].
    pub lcp_memoization: LcpMemoizationPolicy,
    // Diagnostic-only implementation detail selected by `from_env()`. It is
    // intentionally not public API: counters are emitted to the profiling log
    // rather than returned to callers, and their shape may evolve freely.
    collect_lcp_memoization_stats: bool,
}

impl Default for ExtMemOpts {
    fn default() -> Self {
        Self {
            max_context: usize::MAX,
            subproblem_count: 0,
            work_dir: std::env::temp_dir(),
            physical_file_count: 0,
            ordered_phase4_emit: false,
            lcp_memoization: LcpMemoizationPolicy::Disabled,
            collect_lcp_memoization_stats: false,
        }
    }
}

impl ExtMemOpts {
    /// Convenience constructor with the supplied `work_dir` and defaults
    /// for everything else.
    pub fn with_work_dir(work_dir: impl AsRef<Path>) -> Self {
        Self {
            work_dir: work_dir.as_ref().to_path_buf(),
            ..Self::default()
        }
    }

    /// Defaults plus caps-sa environment overrides.
    ///
    /// Recognised variables:
    /// - `CAPS_SA_WORK_DIR` / `CAPS_SA_TMPDIR`: temp-file directory
    /// - `CAPS_SA_SUBPROBLEMS`: subarray count (`p`)
    /// - `CAPS_SA_N_PHYS`: physical backing-file count
    /// - `CAPS_SA_MAX_CONTEXT`: LCP comparison cap
    /// - `CAPS_SA_ORDERED_PHASE4=1|true|yes|on`: bounded ordered phase-4 emit
    /// - `CAPS_SA_GEOMETRIC_MEMO=1|true|yes|on`: geometric LCP memoization
    /// - `CAPS_SA_MEMO_PROBE`: ordinary symbols compared before table lookup
    /// - `CAPS_SA_MEMO_MIN_LCP`: minimum exact LCP admitted to the table
    /// - `CAPS_SA_MEMO_ACTIVATE_ENTRIES`: entries learned before table lookup
    /// - `CAPS_SA_MEMO_CAPACITY`: maximum entries per partition table
    /// - `CAPS_SA_MEMO_STATS=1|true|yes|on`: detailed memoization counters
    ///
    /// Invalid and zero-valued numeric overrides are ignored, preserving the
    /// corresponding defaults. Memoization tuning variables take effect only
    /// when `CAPS_SA_GEOMETRIC_MEMO` enables the policy.
    pub fn from_env() -> Self {
        let mut opts = Self::default();
        if let Some(dir) =
            std::env::var_os("CAPS_SA_WORK_DIR").or_else(|| std::env::var_os("CAPS_SA_TMPDIR"))
        {
            opts.work_dir = PathBuf::from(dir);
        }
        if let Some(v) = read_env_usize("CAPS_SA_SUBPROBLEMS") {
            opts.subproblem_count = v;
        }
        if let Some(v) = read_env_usize("CAPS_SA_N_PHYS") {
            opts.physical_file_count = v;
        }
        if let Some(v) = read_env_usize("CAPS_SA_MAX_CONTEXT") {
            opts.max_context = v;
        }
        if read_env_bool("CAPS_SA_ORDERED_PHASE4") {
            opts.ordered_phase4_emit = true;
        }
        if read_env_bool("CAPS_SA_GEOMETRIC_MEMO") {
            let mut config = GeometricMemoizationConfig::default();
            if let Some(v) = read_env_nonzero_usize("CAPS_SA_MEMO_PROBE") {
                config = config.with_probe_symbols(v);
            }
            if let Some(v) = read_env_nonzero_usize("CAPS_SA_MEMO_MIN_LCP") {
                config = config.with_min_lcp_symbols(v);
            }
            if let Some(v) = read_env_nonzero_usize("CAPS_SA_MEMO_ACTIVATE_ENTRIES") {
                config = config.with_activate_after_entries(v);
            }
            if let Some(v) = read_env_nonzero_usize("CAPS_SA_MEMO_CAPACITY") {
                config = config.with_max_entries_per_partition(v);
            }
            opts.lcp_memoization = LcpMemoizationPolicy::Geometric(config);
        }
        opts.collect_lcp_memoization_stats = read_env_bool("CAPS_SA_MEMO_STATS");
        opts
    }

    /// Builder-style setter for [`Self::max_context`].
    pub fn max_context(mut self, max_context: usize) -> Self {
        self.max_context = max_context;
        self
    }

    /// Builder-style setter for [`Self::subproblem_count`].
    pub fn subproblem_count(mut self, subproblem_count: usize) -> Self {
        self.subproblem_count = subproblem_count;
        self
    }

    /// Builder-style setter for [`Self::work_dir`].
    pub fn work_dir(mut self, work_dir: impl AsRef<Path>) -> Self {
        self.work_dir = work_dir.as_ref().to_path_buf();
        self
    }

    /// Builder-style setter for [`Self::physical_file_count`].
    pub fn physical_file_count(mut self, physical_file_count: usize) -> Self {
        self.physical_file_count = physical_file_count;
        self
    }

    /// Builder-style setter for [`Self::ordered_phase4_emit`].
    pub fn ordered_phase4_emit(mut self, ordered_phase4_emit: bool) -> Self {
        self.ordered_phase4_emit = ordered_phase4_emit;
        self
    }

    /// Builder-style setter for [`Self::lcp_memoization`].
    pub fn lcp_memoization(mut self, lcp_memoization: impl Into<LcpMemoizationPolicy>) -> Self {
        self.lcp_memoization = lcp_memoization.into();
        self
    }
}

fn read_env_usize(name: &str) -> Option<usize> {
    std::env::var(name).ok()?.parse().ok()
}

fn read_env_nonzero_usize(name: &str) -> Option<NonZeroUsize> {
    NonZeroUsize::new(read_env_usize(name)?)
}

fn read_env_bool(name: &str) -> bool {
    std::env::var(name)
        .ok()
        .is_some_and(|v| matches!(v.as_str(), "1" | "true" | "yes" | "on"))
}

/// Error type for `try_*` builders whose output callback can return
/// an application-specific error.
#[derive(Debug)]
pub enum BuildError<E> {
    /// Temporary-file or bucket I/O failed inside caps-sa.
    Io(io::Error),
    /// The caller's output callback rejected an emitted suffix-array
    /// position.
    Emit(E),
}

impl<E> From<io::Error> for BuildError<E> {
    fn from(err: io::Error) -> Self {
        Self::Io(err)
    }
}

fn into_io_result(result: Result<(), BuildError<io::Error>>) -> io::Result<()> {
    match result {
        Ok(()) => Ok(()),
        Err(BuildError::Io(err) | BuildError::Emit(err)) => Err(err),
    }
}

/// Build the suffix array of `text` with bounded RAM, streaming each
/// output position to `emit` in lexicographic order.
///
/// Returns an [`io::Error`] if temp-file I/O fails. The callback may also
/// return an error to abort construction; partial work is discarded and
/// temp files are cleaned up when their bucket drops.
///
/// Equivalent in semantics to [`crate::build_in_memory`]: produces a
/// standard lexicographic suffix array with the "shorter suffix is
/// smaller when one runs off the end" tie-break.
pub fn build_ext_mem<S, F>(text: &[S], opts: &ExtMemOpts, emit: F) -> io::Result<()>
where
    S: Symbol,
    F: FnMut(u64) -> io::Result<()>,
{
    into_io_result(try_build_ext_mem(text, opts, emit))
}

/// Generic-error variant of [`build_ext_mem`].
pub fn try_build_ext_mem<S, E, F>(
    text: &[S],
    opts: &ExtMemOpts,
    emit: F,
) -> Result<(), BuildError<E>>
where
    S: Symbol,
    F: FnMut(u64) -> Result<(), E>,
{
    try_build_ext_mem_with(text, &PlainText::new(text.len()), opts, emit)
}

/// Variant of [`build_ext_mem`] that accepts a [`LimitProvider`]. With
/// [`PlainText`] this matches [`build_ext_mem`] exactly (and
/// monomorphizes to identical assembly); with
/// [`SegmentedText`][crate::limits::SegmentedText] the LCP scans stop
/// at segment boundaries.
pub fn build_ext_mem_with<S, L, F>(text: &[S], lp: &L, opts: &ExtMemOpts, emit: F) -> io::Result<()>
where
    S: Symbol,
    L: LimitProvider,
    F: FnMut(u64) -> io::Result<()>,
{
    into_io_result(try_build_ext_mem_with(text, lp, opts, emit))
}

/// Generic-error variant of [`build_ext_mem_with`].
pub fn try_build_ext_mem_with<S, L, E, F>(
    text: &[S],
    lp: &L,
    opts: &ExtMemOpts,
    emit: F,
) -> Result<(), BuildError<E>>
where
    S: Symbol,
    L: LimitProvider,
    F: FnMut(u64) -> Result<(), E>,
{
    // Dispatch on text size: when every suffix position fits in `u32`
    // (n ≤ 2^32), use the narrow record type to halve all the SaLcp
    // bytes — bucket disk I/O, phase-1 records, and the per-partition
    // load in phase 4.
    if text.len() <= u32::MAX as usize + 1 {
        build_ext_mem_inner::<S, u32, L, E, F>(
            text,
            PositionSource::Identity(text.len()),
            lp,
            opts,
            emit,
        )
    } else {
        build_ext_mem_inner::<S, u64, L, E, F>(
            text,
            PositionSource::Identity(text.len()),
            lp,
            opts,
            emit,
        )
    }
}

/// Like [`build_ext_mem`] but sorts only the caller-supplied
/// `positions` by the lexicographic order of their suffixes in `text`.
/// Suffix content is always `text[position..]`; no positions are
/// dropped from the input. To filter, the caller constructs
/// `positions` with only the indices they want.
///
/// This lets STAR-style genome indexing skip the bin-padding
/// pathology: pass only the ACGT-starting positions and the
/// spacer-suffix work disappears.
///
/// `positions` does not need to be pre-sorted in any way.
pub fn build_ext_mem_for_positions<S, F>(
    text: &[S],
    positions: Vec<u64>,
    opts: &ExtMemOpts,
    emit: F,
) -> io::Result<()>
where
    S: Symbol,
    F: FnMut(u64) -> io::Result<()>,
{
    into_io_result(try_build_ext_mem_for_positions(text, positions, opts, emit))
}

/// Generic-error variant of [`build_ext_mem_for_positions`].
pub fn try_build_ext_mem_for_positions<S, E, F>(
    text: &[S],
    positions: Vec<u64>,
    opts: &ExtMemOpts,
    emit: F,
) -> Result<(), BuildError<E>>
where
    S: Symbol,
    F: FnMut(u64) -> Result<(), E>,
{
    try_build_ext_mem_for_positions_with(text, positions, &PlainText::new(text.len()), opts, emit)
}

/// Variant of [`build_ext_mem_for_positions`] that accepts a
/// [`LimitProvider`]. See [`build_ext_mem_with`] for the semantics.
pub fn build_ext_mem_for_positions_with<S, L, F>(
    text: &[S],
    positions: Vec<u64>,
    lp: &L,
    opts: &ExtMemOpts,
    emit: F,
) -> io::Result<()>
where
    S: Symbol,
    L: LimitProvider,
    F: FnMut(u64) -> io::Result<()>,
{
    into_io_result(try_build_ext_mem_for_positions_with(
        text, positions, lp, opts, emit,
    ))
}

/// Generic-error variant of [`build_ext_mem_for_positions_with`].
pub fn try_build_ext_mem_for_positions_with<S, L, E, F>(
    text: &[S],
    positions: Vec<u64>,
    lp: &L,
    opts: &ExtMemOpts,
    emit: F,
) -> Result<(), BuildError<E>>
where
    S: Symbol,
    L: LimitProvider,
    F: FnMut(u64) -> Result<(), E>,
{
    // We hold a reference to `positions` for the duration of the build;
    // `phase1_sort_sample_spill` copies each chunk out. The Vec is
    // dropped once phase 1 returns.
    if text.len() <= u32::MAX as usize + 1 {
        build_ext_mem_inner::<S, u32, L, E, F>(
            text,
            PositionSource::Subset(&positions),
            lp,
            opts,
            emit,
        )
    } else {
        build_ext_mem_inner::<S, u64, L, E, F>(
            text,
            PositionSource::Subset(&positions),
            lp,
            opts,
            emit,
        )
    }
}

/// Like [`build_ext_mem_for_positions`] but takes a **predicate** over
/// text positions instead of a pre-materialised `Vec<u64>` of kept
/// positions.
///
/// caps-sa walks the predicate **once** to build a bitmap of kept
/// positions + a tiny per-block popcount prefix-sum (together ~`n / 8`
/// bytes — ~770 MB on the human genome, vs the ~50 GB the equivalent
/// `Vec<u64>` would take). Phase 1's per-subarray fill is then driven
/// by popcount-walking the bitmap; the predicate is **never invoked
/// again** after the initial build. See the crate-internal
/// `FilteredSource` for the memory accounting and the inner loop.
///
/// Use this entry when the caller already has the text in RAM and
/// the kept positions are described by a cheap per-position
/// predicate (e.g. STAR's `text[p] < 4` for ACGT-only suffix
/// sampling). It is **the right entry for genome-scale inputs** —
/// the `Vec<u64>` path can dominate peak RSS otherwise.
///
/// `keep` is invoked from rayon worker threads in parallel during
/// the bitmap build; it must be `Send + Sync` (typically a plain
/// closure capturing only `&[u8]` references is fine).
pub fn build_ext_mem_for_filter<S, F, Pred>(
    text: &[S],
    keep: Pred,
    opts: &ExtMemOpts,
    emit: F,
) -> io::Result<()>
where
    S: Symbol,
    F: FnMut(u64) -> io::Result<()>,
    Pred: Fn(u64) -> bool + Send + Sync,
{
    into_io_result(try_build_ext_mem_for_filter(text, keep, opts, emit))
}

/// Generic-error variant of [`build_ext_mem_for_filter`].
pub fn try_build_ext_mem_for_filter<S, E, F, Pred>(
    text: &[S],
    keep: Pred,
    opts: &ExtMemOpts,
    emit: F,
) -> Result<(), BuildError<E>>
where
    S: Symbol,
    F: FnMut(u64) -> Result<(), E>,
    Pred: Fn(u64) -> bool + Send + Sync,
{
    try_build_ext_mem_for_filter_with(text, keep, &PlainText::new(text.len()), opts, emit)
}

/// Variant of [`build_ext_mem_for_filter`] that accepts a
/// [`LimitProvider`]. See [`build_ext_mem_with`] for the semantics.
pub fn build_ext_mem_for_filter_with<S, L, F, Pred>(
    text: &[S],
    keep: Pred,
    lp: &L,
    opts: &ExtMemOpts,
    emit: F,
) -> io::Result<()>
where
    S: Symbol,
    L: LimitProvider,
    F: FnMut(u64) -> io::Result<()>,
    Pred: Fn(u64) -> bool + Send + Sync,
{
    into_io_result(try_build_ext_mem_for_filter_with(
        text, keep, lp, opts, emit,
    ))
}

/// Generic-error variant of [`build_ext_mem_for_filter_with`].
pub fn try_build_ext_mem_for_filter_with<S, L, E, F, Pred>(
    text: &[S],
    keep: Pred,
    lp: &L,
    opts: &ExtMemOpts,
    emit: F,
) -> Result<(), BuildError<E>>
where
    S: Symbol,
    L: LimitProvider,
    F: FnMut(u64) -> Result<(), E>,
    Pred: Fn(u64) -> bool + Send + Sync,
{
    let filtered = FilteredSource::new(text.len(), keep);
    if text.len() <= u32::MAX as usize + 1 {
        build_ext_mem_inner::<S, u32, L, E, F>(
            text,
            PositionSource::Filtered(filtered),
            lp,
            opts,
            emit,
        )
    } else {
        build_ext_mem_inner::<S, u64, L, E, F>(
            text,
            PositionSource::Filtered(filtered),
            lp,
            opts,
            emit,
        )
    }
}

fn build_ext_mem_inner<S, I, L, E, F>(
    text: &[S],
    source: PositionSource<'_>,
    lp: &L,
    opts: &ExtMemOpts,
    mut emit: F,
) -> Result<(), BuildError<E>>
where
    S: Symbol,
    I: Index,
    L: LimitProvider,
    SaLcp<I>: BucketRecord,
    F: FnMut(u64) -> Result<(), E>,
{
    let n = source.len();
    if n == 0 {
        return Ok(());
    }
    let p = effective_subproblem_count(n, opts.subproblem_count);
    let dispatch = LcpDispatch::detect();
    let work_dir = opts.work_dir.clone();

    // The fused phase 1 writes directly to partition buckets, so this path
    // needs only one file pool. The previous subarray pool held a complete
    // spilled copy that phase 3 read and wrote again before the final merge.
    let n_phys = effective_physical_file_count(opts.physical_file_count);
    let phase3_pool = BucketPool::new(n_phys, &work_dir)?;

    profile_log(&format!(
        "build_ext_mem n={n} p={p} index_width={}b n_phys={n_phys}",
        std::mem::size_of::<I>() * 8
    ));

    let part_factory = |j: usize| phase3_pool.new_bucket::<SaLcp<I>>(j);

    let t = Instant::now();
    let pivots = phase0_presample_pivots::<S, I, L>(text, lp, &source, p, opts, dispatch);
    profile_log(&format!(
        "phase0 (presample pivots)  {:.3}s",
        t.elapsed().as_secs_f64()
    ));

    let t = Instant::now();
    let mut partition_buckets = phase1_sort_and_distribute::<S, I, L, _, _>(
        text,
        lp,
        &source,
        &pivots,
        p,
        opts,
        dispatch,
        part_factory,
    )?;
    profile_log(&format!(
        "phase1 (sort+distribute)   {:.3}s",
        t.elapsed().as_secs_f64()
    ));

    // Drop the position source as soon as phase 1 returns. For
    // `PositionSource::Subset` this frees
    // the caller's `Vec<u64>` (e.g. ~47 GB on a human-scale
    // _for_positions build); for `PositionSource::Filtered` it
    // frees the bitmap + cumsum (~770 MB); for `Identity` it's a
    // no-op. Phase 4 needs only the text and partition buckets.
    drop(source);

    let t = Instant::now();
    let result = phase4_merge_and_emit::<S, I, L, _, E, F>(
        text,
        lp,
        &mut partition_buckets,
        opts.max_context,
        opts.ordered_phase4_emit,
        memo_config(opts.lcp_memoization),
        opts.collect_lcp_memoization_stats,
        &mut emit,
        dispatch,
    );
    profile_log(&format!(
        "phase4 (merge+emit)          {:.3}s",
        t.elapsed().as_secs_f64()
    ));
    result
}

/// Same algorithm as [`build_ext_mem_inner`] but with the disk-backed
/// [`ExtMemBucket`] replaced by [`InMemBucket`] throughout — phase 1
/// sorts each subarray and keeps the result in a `Vec<SaLcp<I>>`,
/// phase 3 distributes into in-RAM partition Vecs, phase 4 cascade-
/// merges the in-RAM partitions. No disk I/O.
///
/// Trades RAM for wall time: peak memory is ~`n × sizeof(SaLcp<I>)`
/// (the post-phase-1 records sitting around until phase 3 consumes
/// them), so ~25 GB on the human genome with `I = u32`. In exchange,
/// the disk-spill / distribute-write / partition-load round-trip is
/// gone — useful on machines with enough RAM to hold the working set.
fn build_in_memory_ss_inner<S, I, L, E, F>(
    text: &[S],
    source: PositionSource<'_>,
    lp: &L,
    opts: &ExtMemOpts,
    mut emit: F,
) -> Result<(), BuildError<E>>
where
    S: Symbol,
    I: Index,
    L: LimitProvider,
    SaLcp<I>: BucketRecord,
    F: FnMut(u64) -> Result<(), E>,
{
    let n = source.len();
    if n == 0 {
        return Ok(());
    }
    let p = effective_subproblem_count(n, opts.subproblem_count);
    let dispatch = LcpDispatch::detect();

    let factory = |_i: usize| InMemBucket::<SaLcp<I>>::new();

    let (mut subarray_buckets, samples) =
        phase1_sort_sample_spill::<S, I, L, _, _>(text, lp, &source, p, opts, dispatch, factory)?;
    // Same rationale as in `build_ext_mem_inner` — drop the source
    // as soon as phase 1's `fill_chunk` calls have stopped.
    drop(source);
    let pivots = phase2_select_pivots::<S, I, L>(text, lp, samples, p, opts.max_context, dispatch);
    let mut partition_buckets = phase3_distribute::<S, I, L, _, _>(
        text,
        lp,
        &mut subarray_buckets,
        &pivots,
        p,
        opts,
        dispatch,
        factory,
    )?;
    drop(subarray_buckets);
    phase4_merge_and_emit::<S, I, L, _, E, F>(
        text,
        lp,
        &mut partition_buckets,
        opts.max_context,
        opts.ordered_phase4_emit,
        memo_config(opts.lcp_memoization),
        opts.collect_lcp_memoization_stats,
        &mut emit,
        dispatch,
    )
}

fn memo_config(policy: LcpMemoizationPolicy) -> Option<MemoConfig> {
    match policy {
        LcpMemoizationPolicy::Disabled => None,
        LcpMemoizationPolicy::Geometric(config) => Some(config.into()),
    }
}

/// In-memory variant of the sample-sort algorithm used by
/// [`build_ext_mem`]. Skips all disk I/O at the cost of holding the
/// (`pos`, `lcp`) records in RAM throughout. Picks `u32` records when
/// `n ≤ 2³²`, falls back to `u64` otherwise. The caller's `emit`
/// closure is called once per output position in lex order, just like
/// in the ext-mem path.
pub fn build_in_memory_sample_sort<S, F>(text: &[S], opts: &ExtMemOpts, emit: F) -> io::Result<()>
where
    S: Symbol,
    F: FnMut(u64) -> io::Result<()>,
{
    into_io_result(try_build_in_memory_sample_sort(text, opts, emit))
}

/// Generic-error variant of [`build_in_memory_sample_sort`].
pub fn try_build_in_memory_sample_sort<S, E, F>(
    text: &[S],
    opts: &ExtMemOpts,
    emit: F,
) -> Result<(), BuildError<E>>
where
    S: Symbol,
    F: FnMut(u64) -> Result<(), E>,
{
    try_build_in_memory_sample_sort_with(text, &PlainText::new(text.len()), opts, emit)
}

/// Variant of [`build_in_memory_sample_sort`] that accepts a
/// [`LimitProvider`].
pub fn build_in_memory_sample_sort_with<S, L, F>(
    text: &[S],
    lp: &L,
    opts: &ExtMemOpts,
    emit: F,
) -> io::Result<()>
where
    S: Symbol,
    L: LimitProvider,
    F: FnMut(u64) -> io::Result<()>,
{
    into_io_result(try_build_in_memory_sample_sort_with(text, lp, opts, emit))
}

/// Generic-error variant of [`build_in_memory_sample_sort_with`].
pub fn try_build_in_memory_sample_sort_with<S, L, E, F>(
    text: &[S],
    lp: &L,
    opts: &ExtMemOpts,
    emit: F,
) -> Result<(), BuildError<E>>
where
    S: Symbol,
    L: LimitProvider,
    F: FnMut(u64) -> Result<(), E>,
{
    if text.len() <= u32::MAX as usize + 1 {
        build_in_memory_ss_inner::<S, u32, L, E, F>(
            text,
            PositionSource::Identity(text.len()),
            lp,
            opts,
            emit,
        )
    } else {
        build_in_memory_ss_inner::<S, u64, L, E, F>(
            text,
            PositionSource::Identity(text.len()),
            lp,
            opts,
            emit,
        )
    }
}

/// Subset-positions variant of [`build_in_memory_sample_sort`]. Same
/// shape as [`build_ext_mem_for_positions`].
pub fn build_in_memory_sample_sort_for_positions<S, F>(
    text: &[S],
    positions: Vec<u64>,
    opts: &ExtMemOpts,
    emit: F,
) -> io::Result<()>
where
    S: Symbol,
    F: FnMut(u64) -> io::Result<()>,
{
    into_io_result(try_build_in_memory_sample_sort_for_positions(
        text, positions, opts, emit,
    ))
}

/// Generic-error variant of [`build_in_memory_sample_sort_for_positions`].
pub fn try_build_in_memory_sample_sort_for_positions<S, E, F>(
    text: &[S],
    positions: Vec<u64>,
    opts: &ExtMemOpts,
    emit: F,
) -> Result<(), BuildError<E>>
where
    S: Symbol,
    F: FnMut(u64) -> Result<(), E>,
{
    try_build_in_memory_sample_sort_for_positions_with(
        text,
        positions,
        &PlainText::new(text.len()),
        opts,
        emit,
    )
}

/// Variant of [`build_in_memory_sample_sort_for_positions`] that
/// accepts a [`LimitProvider`].
pub fn build_in_memory_sample_sort_for_positions_with<S, L, F>(
    text: &[S],
    positions: Vec<u64>,
    lp: &L,
    opts: &ExtMemOpts,
    emit: F,
) -> io::Result<()>
where
    S: Symbol,
    L: LimitProvider,
    F: FnMut(u64) -> io::Result<()>,
{
    into_io_result(try_build_in_memory_sample_sort_for_positions_with(
        text, positions, lp, opts, emit,
    ))
}

/// Generic-error variant of [`build_in_memory_sample_sort_for_positions_with`].
pub fn try_build_in_memory_sample_sort_for_positions_with<S, L, E, F>(
    text: &[S],
    positions: Vec<u64>,
    lp: &L,
    opts: &ExtMemOpts,
    emit: F,
) -> Result<(), BuildError<E>>
where
    S: Symbol,
    L: LimitProvider,
    F: FnMut(u64) -> Result<(), E>,
{
    if text.len() <= u32::MAX as usize + 1 {
        build_in_memory_ss_inner::<S, u32, L, E, F>(
            text,
            PositionSource::Subset(&positions),
            lp,
            opts,
            emit,
        )
    } else {
        build_in_memory_ss_inner::<S, u64, L, E, F>(
            text,
            PositionSource::Subset(&positions),
            lp,
            opts,
            emit,
        )
    }
}

/// Source of the positions to sort.
///
/// - [`PositionSource::Identity`] is the all-suffixes case `0..n`;
///   avoids materialising a `Vec<u64>` of length `n`, which on the
///   human genome would itself be ~50 GB.
/// - [`PositionSource::Subset`] holds a caller-supplied `&[u64]` of
///   the positions to sort, in any order. Random-access by index.
/// - [`PositionSource::Filtered`] is the streaming-predicate variant:
///   the caller hands in a `Fn(u64) -> bool` over text positions and
///   caps-sa walks the filter forward through the text on demand. A
///   tiny prefix-sum (`text.len() / BLOCK_SIZE × 8` B ≈ 760 KB on
///   the human genome at the 64 KB block size) lets `fill_chunk`
///   locate the i-th kept position in `O(log n_blocks + BLOCK_SIZE +
///   chunk_size)`. **No kept-positions list is materialised** — the
///   memory saving over `Subset` is `≈ 8 × n_kept` bytes, dominant
///   on genome-scale inputs.
enum PositionSource<'a> {
    Identity(usize),
    Subset(&'a [u64]),
    Filtered(FilteredSource),
}

/// Block size in **u64 words** for [`FilteredSource`]'s popcount
/// prefix-sum. With 1024 words per block (64 K text bits) the
/// prefix-sum stores one `u64` per block — ~760 KB for the human
/// genome — and each `fill_chunk` walks at most one block (≈ 8 KB)
/// in cache before emitting the first kept position.
const FILTERED_WORDS_PER_BLOCK: usize = 1024;

/// Streaming position source backed by a **bitmap** of kept positions
/// plus a per-block popcount prefix-sum. No `Vec<u64>` of kept positions
/// is ever materialised.
///
/// Memory: `(n + 7) / 8` bytes for the bitmap (~770 MB on the human
/// genome at `n ≈ 6.2 B`) + `8 × n_blocks` for the prefix sum
/// (~760 KB). The 47 GB `Vec<u64>` the [`Subset`][PositionSource::Subset]
/// variant requires goes away entirely.
///
/// Lookup path inside [`fill_chunk`]:
/// 1. `partition_point` the prefix-sum to find the block containing
///    the `start`-th kept position (`O(log n_blocks)` ≈ 17 ops).
/// 2. Walk that block's bitmap words `popcount`-by-`popcount` until
///    we've skipped the right number of set bits to reach `start`.
/// 3. Walk forward through bitmap words using `trailing_zeros`-style
///    iteration, emitting each set bit's position to `dst`. Inner
///    loop is `O(chunk_size / 64)` u64 ops — no per-text-position
///    closure calls, branch-light, cache-resident.
///
/// A truly `O(1)` select-1 structure (darray / Elias-Fano) on top
/// of the bitmap would shrink step (1)+(2) further; with
/// `chunk_size` ≈ 750 K and only `p` ≈ 8192 fill_chunk calls per
/// build, the `O(log n_blocks)` + at-most-one-block-walk cost is
/// already a rounding error. See `bench/README.md` for the
/// follow-up note if that ever changes.
struct FilteredSource {
    text_len: usize,
    total_kept: usize,
    /// Bitmap: `(bitmap[w] >> b) & 1 == 1` iff position `64 * w + b`
    /// is kept. Length = `text_len.div_ceil(64)`.
    bitmap: Vec<u64>,
    /// `cumsum[i] = sum of set bits in
    /// `bitmap[0 .. i * FILTERED_WORDS_PER_BLOCK]`. Length =
    /// `n_blocks + 1`; the last entry equals `total_kept`.
    cumsum: Vec<u64>,
}

impl FilteredSource {
    /// Build a [`FilteredSource`] by walking the predicate once over
    /// `0..text_len` to fill the bitmap, then accumulating per-block
    /// popcounts.
    ///
    /// The predicate is invoked exactly `text_len` times here; once
    /// the bitmap is built, `fill_chunk` never calls it again. This
    /// trades one full predicate pass for zero per-position calls
    /// during all subsequent random-access `fill_chunk`s — a clear
    /// win when (as in caps-sa's phase 1) every position is read at
    /// least once.
    fn new<Pred>(text_len: usize, keep: Pred) -> Self
    where
        Pred: Fn(u64) -> bool + Send + Sync,
    {
        let n_words = text_len.div_ceil(64);
        // Parallel per-word bitmap build. Each word reads 64 text
        // positions (clamped at `text_len`), packs them into a u64.
        let bitmap: Vec<u64> = (0..n_words)
            .into_par_iter()
            .map(|w| {
                let mut word: u64 = 0;
                let base = (w as u64) * 64;
                let limit = ((w + 1) * 64).min(text_len) - w * 64;
                for b in 0..limit {
                    if keep(base + b as u64) {
                        word |= 1u64 << b;
                    }
                }
                word
            })
            .collect();

        // Per-block popcount cumsum. Each block covers
        // `FILTERED_WORDS_PER_BLOCK` words = `FILTERED_BITS_PER_BLOCK`
        // text positions.
        let n_blocks = n_words.div_ceil(FILTERED_WORDS_PER_BLOCK);
        let per_block: Vec<u64> = (0..n_blocks)
            .into_par_iter()
            .map(|i| {
                let start = i * FILTERED_WORDS_PER_BLOCK;
                let end = ((i + 1) * FILTERED_WORDS_PER_BLOCK).min(n_words);
                let mut c: u64 = 0;
                for &word in &bitmap[start..end] {
                    c += word.count_ones() as u64;
                }
                c
            })
            .collect();
        let mut cumsum = Vec::with_capacity(n_blocks + 1);
        let mut s: u64 = 0;
        cumsum.push(0);
        for &k in &per_block {
            s += k;
            cumsum.push(s);
        }
        let total_kept = s as usize;
        Self {
            text_len,
            total_kept,
            bitmap,
            cumsum,
        }
    }

    /// Number of kept positions.
    #[inline]
    fn len(&self) -> usize {
        self.total_kept
    }

    /// Fill `dst` with the next `dst.len()` kept positions starting
    /// from the `start`-th (0-based) kept position. See type-level
    /// doc for the algorithm; this is the hot path during phase 1
    /// fill_chunk.
    fn fill_chunk<I: Index>(&self, start: usize, dst: &mut [I]) {
        debug_assert!(start + dst.len() <= self.total_kept);
        if dst.is_empty() {
            return;
        }

        // (1) Locate the block containing the `start`-th set bit.
        // `partition_point(|c| c <= start)` gives the first cumsum
        // entry strictly greater than `start`; previous index is the
        // containing block.
        let pp = self.cumsum.partition_point(|&c| c <= start as u64);
        debug_assert!(pp > 0);
        let block_idx = pp - 1;
        let mut word_idx = block_idx * FILTERED_WORDS_PER_BLOCK;
        let mut skip = start as u64 - self.cumsum[block_idx];

        // (2) Skip the first `skip` set bits — possibly spanning
        // several bitmap words. Whole words with `popcount ≤ skip`
        // are consumed wholesale; the final partial word has its
        // lowest `skip` set bits cleared so the emit loop sees only
        // un-skipped 1s.
        //
        // Note: a naive `while skip >= 64 { … }` is wrong because a
        // word's popcount can be far less than 64; we must subtract
        // the actual popcount each iteration, not 64. This matters
        // any time the bitmap is sparser than ~50%.
        let n_words = self.bitmap.len();
        let mut word: u64 = if word_idx < n_words {
            self.bitmap[word_idx]
        } else {
            0
        };
        while skip > 0 {
            let pc = word.count_ones() as u64;
            if skip < pc {
                // Consume `skip` lowest set bits inside the current
                // word; emit loop continues from the remaining ones.
                for _ in 0..skip {
                    word &= word - 1;
                }
                break;
            }
            // Skip ≥ pc: consume the whole word and advance.
            skip -= pc;
            word_idx += 1;
            word = if word_idx < n_words {
                self.bitmap[word_idx]
            } else {
                0
            };
        }

        // (3) Walk `word`+subsequent words, emitting one position per
        // set bit. Uses `trailing_zeros` to jump straight to the next
        // 1 inside a word, then clears it via `word &= word - 1`.
        let mut written = 0usize;
        let need = dst.len();
        loop {
            while word != 0 && written < need {
                let bit = word.trailing_zeros() as u64;
                let pos = (word_idx as u64) * 64 + bit;
                debug_assert!((pos as usize) < self.text_len);
                dst[written] = I::from_usize(pos as usize);
                written += 1;
                word &= word - 1;
            }
            if written == need {
                break;
            }
            word_idx += 1;
            debug_assert!(
                word_idx < n_words,
                "FilteredSource::fill_chunk: walked past bitmap end \
                 ({written}/{need} emitted, word_idx={word_idx}, n_words={n_words})"
            );
            word = self.bitmap[word_idx];
        }
    }
}

impl<'a> PositionSource<'a> {
    fn len(&self) -> usize {
        match self {
            Self::Identity(n) => *n,
            Self::Subset(p) => p.len(),
            Self::Filtered(f) => f.len(),
        }
    }

    /// Fill `dst` with positions for the half-open subarray range
    /// `[start, start + dst.len())`, narrowing the caller's `u64`
    /// positions into `I` via [`Index::from_usize`]. For
    /// [`PositionSource::Identity`] this generates the contiguous
    /// integer range on the fly; for [`PositionSource::Subset`] it
    /// reads from the caller's slice; for [`PositionSource::Filtered`]
    /// it walks the predicate forward from the right text block.
    fn fill_chunk<I: Index>(&self, start: usize, dst: &mut [I]) {
        match self {
            Self::Identity(_) => {
                for (i, slot) in dst.iter_mut().enumerate() {
                    *slot = I::from_usize(start + i);
                }
            }
            Self::Subset(p) => {
                let end = start + dst.len();
                for (slot, &v) in dst.iter_mut().zip(p[start..end].iter()) {
                    *slot = I::from_usize(v as usize);
                }
            }
            Self::Filtered(f) => f.fill_chunk(start, dst),
        }
    }
}

/// Target subarray size used by [`effective_subproblem_count`] when
/// auto-picking `p`. Smaller means more (smaller) subarrays — lower
/// per-task phase-1 scratch, at the cost of more pivot splits and logical
/// partition pieces.
const PHASE1_TARGET_CHUNK: usize = 65_536;
/// Hard cap on the number of subarrays. Matches upstream CaPS-SA's
/// default of 8192. The cap keeps pivot splitting, logical partition metadata,
/// and the number of lock acquisitions bounded.
const PHASE1_MAX_PARTITIONS: usize = 8192;

/// Resolve [`ExtMemOpts::physical_file_count`] for the current build.
/// `0` (the default) means "let the runtime decide"; we pick
/// `rayon::current_num_threads()` so the pool has one inode per
/// concurrent writer, which empirically matches per-bucket-file wall
/// time while collapsing thousands of small files into dozens of
/// large ones. The `CAPS_SA_N_PHYS` env var overrides at the call
/// site for benchmarks.
fn effective_physical_file_count(requested: usize) -> usize {
    if let Some(v) = std::env::var("CAPS_SA_N_PHYS")
        .ok()
        .and_then(|s| s.parse::<usize>().ok())
        .filter(|&v| v >= 1)
    {
        return v;
    }
    if requested >= 1 {
        return requested;
    }
    rayon::current_num_threads().max(1)
}

fn effective_subproblem_count(n: usize, requested: usize) -> usize {
    if n == 0 {
        return 0;
    }
    let raw = if requested == 0 {
        let nthreads = rayon::current_num_threads().max(1);
        let p_from_size = n.div_ceil(PHASE1_TARGET_CHUNK);
        // At least one chunk per thread (otherwise we leave cores idle),
        // at most `PHASE1_MAX_PARTITIONS` (so pivot splitting and partition
        // metadata stay manageable). For
        // small inputs `p_from_size` is well below the cap, so the
        // formula degrades gracefully to roughly "one chunk per thread";
        // for human-scale inputs the cap binds and per-task scratch
        // stays in the tens-of-MB range.
        p_from_size.clamp(nthreads, PHASE1_MAX_PARTITIONS)
    } else {
        requested
    };
    raw.clamp(1, n)
}

/// Phase 1: sort each subarray in parallel, sample from it, and spill
/// `(position, lcp)` records to its own [`ExtMemBucket`].
///
/// One Rayon task per subarray. When `p` provides at least one outer task per
/// worker (including every auto-picked build), the sort recursion stays local
/// to that task, avoiding nested scheduler work and preventing a worker from
/// retaining one subarray's scratch while stealing another outer task. An
/// explicit smaller `p` retains recursive parallelism. With auto `p` (target
/// chunk ~64 K records until the 8,192 cap binds), peak scratch is bounded by
/// the number of workers rather than the number of subarrays.
#[allow(clippy::too_many_arguments)]
fn phase1_sort_sample_spill<S, I, L, B, MkB>(
    text: &[S],
    lp: &L,
    source: &PositionSource<'_>,
    p: usize,
    opts: &ExtMemOpts,
    dispatch: LcpDispatch,
    mk_bucket: MkB,
) -> io::Result<(Vec<B>, Vec<I>)>
where
    S: Symbol,
    I: Index,
    L: LimitProvider,
    SaLcp<I>: BucketRecord,
    B: SaLcpBucketStore<I> + Send,
    MkB: Fn(usize) -> B + Send + Sync,
{
    let n = source.len();
    let chunk_size = n.div_ceil(p);
    let samples_target_total = sample_target_total(n, p);
    let task_local_sort = p >= rayon::current_num_threads().max(1);

    let per_subarray: Vec<(B, Vec<I>)> = (0..p)
        .into_par_iter()
        .map(|i| {
            let start = (i * chunk_size).min(n);
            let end = ((i + 1) * chunk_size).min(n);
            let len = end - start;

            let mut bucket = mk_bucket(i);
            if len == 0 {
                return Ok::<_, io::Error>((bucket, Vec::new()));
            }

            // In-memory sort of this subarray with LCP maintenance.
            let mut sa: Vec<I> = vec![I::zero(); len];
            source.fill_chunk(start, &mut sa);
            let mut sa_w = vec![I::zero(); len];
            let mut lcp_arr = vec![I::zero(); len];
            let mut lcp_w = vec![I::zero(); len];
            if task_local_sort {
                sample_sort::merge_sort_task_local(
                    text,
                    lp,
                    &mut sa,
                    &mut sa_w,
                    &mut lcp_arr,
                    &mut lcp_w,
                    opts.max_context,
                    dispatch,
                );
            } else {
                sample_sort::merge_sort(
                    text,
                    lp,
                    &mut sa,
                    &mut sa_w,
                    &mut lcp_arr,
                    &mut lcp_w,
                    opts.max_context,
                    dispatch,
                );
            }

            // Pull `samples_per_subarray` evenly-spaced positions out of
            // the now-sorted subarray. Deterministic — no RNG needed for
            // pivot selection to be globally well-distributed.
            let samples_per_subarray = samples_target_total.div_ceil(p).min(len);
            let samples = evenly_spaced(&sa, samples_per_subarray);

            // Spill (position, lcp) records to the bucket. `lcp[0]`
            // remains 0 (set by the merge-sort base case), making each
            // subarray its own well-formed LCP-annotated sorted run.
            bucket.add_soa(&sa, &lcp_arr)?;

            Ok((bucket, samples))
        })
        .collect::<Result<Vec<_>, _>>()?;

    let mut buckets = Vec::with_capacity(p);
    let mut all_samples = Vec::with_capacity(samples_target_total);
    for (bucket, samples) in per_subarray {
        buckets.push(bucket);
        all_samples.extend(samples);
    }
    Ok((buckets, all_samples))
}

/// Target sample count *across all subarrays*. Matches upstream CaPS-SA's
/// "`c · ln n`" rule per subarray with `c = 4`, so the global pool is
/// `p · 4 · ln n` samples.
fn sample_target_total(n: usize, p: usize) -> usize {
    let ln_n = (n as f64).ln().max(1.0);
    let per = (4.0 * ln_n).ceil() as usize;
    // At least p (so we have enough to pick p-1 pivots) and at most n.
    p.saturating_mul(per).clamp(p, n)
}

/// Pick `count` evenly-spaced elements from a slice. Deterministic, which
/// keeps the algorithm reproducible without an RNG dependency.
fn evenly_spaced<T: Copy>(xs: &[T], count: usize) -> Vec<T> {
    let n = xs.len();
    if count == 0 || n == 0 {
        return Vec::new();
    }
    if count >= n {
        return xs.to_vec();
    }
    // Pick indices at positions (i + 0.5) · n / count for i in 0..count,
    // i.e. evenly-spaced midpoints. Avoids both endpoints — keeps pivots
    // away from extreme corners of the order.
    (0..count)
        .map(|i| xs[(2 * i + 1) * n / (2 * count)])
        .collect()
}

/// Phase 2: globally sort the pooled samples and pick `p − 1` pivots at
/// evenly-spaced ranks.
fn phase2_select_pivots<S, I, L>(
    text: &[S],
    lp: &L,
    mut samples: Vec<I>,
    p: usize,
    max_ctx: usize,
    dispatch: LcpDispatch,
) -> Vec<I>
where
    S: Symbol,
    I: Index,
    L: LimitProvider,
{
    if p <= 1 || samples.is_empty() {
        return Vec::new();
    }
    let n_samples = samples.len();
    let mut sa_w = vec![I::zero(); n_samples];
    let mut lcp = vec![I::zero(); n_samples];
    let mut lcp_w = vec![I::zero(); n_samples];
    sample_sort::merge_sort(
        text,
        lp,
        &mut samples,
        &mut sa_w,
        &mut lcp,
        &mut lcp_w,
        max_ctx,
        dispatch,
    );

    // p-1 pivots at evenly-spaced ranks across the sorted sample pool.
    (1..p).map(|j| samples[(j * n_samples) / p]).collect()
}

/// Phase 3: walk each subarray *in parallel*, load it into RAM,
/// binary-search the pivots to find its `p` sub-subarray boundaries,
/// and append each sub-subarray to the corresponding partition bucket.
///
/// Partition buckets are wrapped in a [`Mutex`] each so multiple
/// threads can write to different partitions concurrently without
/// shard-merging afterwards. With `p` in the thousands and `T` in the
/// tens, lock contention is negligible (probability that two threads
/// want the same partition at the same instant is `~T/p`); the lock
/// scope per acquisition is one `add_slice` + `mark_boundary` of a
/// few-KB sub-subarray.
///
/// Phase 4 doesn't care about the relative order of sub-subarrays
/// within a partition — only that each one between consecutive
/// boundaries is internally sorted. Both properties hold under
/// arbitrary thread interleaving.
#[allow(clippy::too_many_arguments)]
fn phase3_distribute<S, I, L, B, MkB>(
    text: &[S],
    lp: &L,
    subarray_buckets: &mut [B],
    pivots: &[I],
    p: usize,
    opts: &ExtMemOpts,
    dispatch: LcpDispatch,
    mk_bucket: MkB,
) -> io::Result<Vec<B>>
where
    S: Symbol,
    I: Index,
    L: LimitProvider,
    SaLcp<I>: BucketRecord,
    B: SaLcpBucketStore<I> + Send,
    MkB: Fn(usize) -> B + Send + Sync,
{
    let _ = opts; // work_dir is used only by the ext-mem factory closure now
    let partition_buckets: Vec<Mutex<B>> = (0..p).map(|j| Mutex::new(mk_bucket(j))).collect();

    subarray_buckets
        .par_iter_mut()
        .try_for_each(|sub_bucket| -> io::Result<()> {
            if sub_bucket.total_records() == 0 {
                return Ok(());
            }
            let records = sub_bucket.load_all()?;

            // Find p-1 split points by binary-searching each pivot's
            // *upper bound* in the sorted subarray.
            let mut splits = Vec::with_capacity(p + 1);
            splits.push(0usize);
            for &pivot in pivots {
                splits.push(upper_bound_by_pivot(
                    &records,
                    pivot,
                    text,
                    lp,
                    opts.max_context,
                    dispatch,
                ));
            }
            splits.push(records.len());

            // Distribute each sub-subarray. Reset the first record's
            // `lcp` to 0 so the per-partition merge sees a well-formed
            // boundary.
            for j in 0..p {
                let lo = splits[j];
                let hi = splits[j + 1];
                if lo >= hi {
                    continue;
                }
                let mut bucket = partition_buckets[j].lock().unwrap();
                bucket.add_slice_reset_first_lcp(&records[lo..hi])?;
                bucket.mark_boundary();
            }
            Ok(())
        })?;

    // Unwrap the Mutexes — at this point only this thread holds
    // references, so the locks are uncontended.
    Ok(partition_buckets
        .into_iter()
        .map(|m| m.into_inner().expect("partition mutex poisoned"))
        .collect())
}

/// Choose `p - 1` pivots before sorting the phase-1 subarrays.
///
/// Any sorted splitters produce a correct sample sort; their quality affects
/// partition balance, not ordering. Selecting them from a cheap block-strided
/// pre-sample lets phase 1 write sorted pieces directly to final partition
/// buckets instead of spilling and re-reading every record in a separate
/// distribution phase.
fn phase0_presample_pivots<S, I, L>(
    text: &[S],
    lp: &L,
    source: &PositionSource<'_>,
    p: usize,
    opts: &ExtMemOpts,
    dispatch: LcpDispatch,
) -> Vec<I>
where
    S: Symbol,
    I: Index,
    L: LimitProvider,
{
    let n = source.len();
    if p <= 1 || n == 0 {
        return Vec::new();
    }

    // PositionSource::fill_chunk is substantially cheaper for a contiguous
    // block than for the same number of singleton calls, particularly for the
    // filtered bitmap source used by ruSTAR.
    const BLOCK: usize = 64;
    let target = sample_target_total(n, p).min(n);
    let n_blocks = target.div_ceil(BLOCK).max(1);
    let stride = (n / n_blocks).max(1);

    let mut sample: Vec<I> = Vec::with_capacity(n_blocks * BLOCK);
    let mut start = 0usize;
    while start < n && sample.len() < target {
        let len = BLOCK.min(n - start);
        let base = sample.len();
        sample.resize(base + len, I::zero());
        source.fill_chunk(start, &mut sample[base..]);
        start += stride;
    }
    if sample.is_empty() {
        return Vec::new();
    }

    let m = sample.len();
    let mut sa_w = vec![I::zero(); m];
    let mut lcp = vec![I::zero(); m];
    let mut lcp_w = vec![I::zero(); m];
    sample_sort::merge_sort(
        text,
        lp,
        &mut sample,
        &mut sa_w,
        &mut lcp,
        &mut lcp_w,
        opts.max_context,
        dispatch,
    );

    (1..p).map(|i| sample[(i * m / p).min(m - 1)]).collect()
}

/// Sort each phase-1 subarray and distribute its sorted pieces directly to
/// the final partition buckets.
///
/// Pivots have already been selected by [`phase0_presample_pivots`], so the
/// subarrays never need to be written to one bucket pool and read back into a
/// second one. The partition pieces remain individually sorted and retain an
/// LCP reset at every boundary, exactly as the cascade merge requires.
#[allow(clippy::too_many_arguments)]
fn phase1_sort_and_distribute<S, I, L, B, MkB>(
    text: &[S],
    lp: &L,
    source: &PositionSource<'_>,
    pivots: &[I],
    p: usize,
    opts: &ExtMemOpts,
    dispatch: LcpDispatch,
    mk_bucket: MkB,
) -> io::Result<Vec<B>>
where
    S: Symbol,
    I: Index,
    L: LimitProvider,
    SaLcp<I>: BucketRecord,
    B: SaLcpBucketStore<I> + Send,
    MkB: Fn(usize) -> B + Send + Sync,
{
    let n = source.len();
    let chunk_size = n.div_ceil(p);
    let partition_buckets: Vec<Mutex<B>> = (0..p).map(|j| Mutex::new(mk_bucket(j))).collect();
    let task_local_sort = p >= rayon::current_num_threads().max(1);

    (0..p).into_par_iter().try_for_each(|i| -> io::Result<()> {
        let start = (i * chunk_size).min(n);
        let end = ((i + 1) * chunk_size).min(n);
        let len = end - start;
        if len == 0 {
            return Ok(());
        }

        let mut sa: Vec<I> = vec![I::zero(); len];
        source.fill_chunk(start, &mut sa);
        let mut sa_w = vec![I::zero(); len];
        let mut lcp_arr = vec![I::zero(); len];
        let mut lcp_w = vec![I::zero(); len];
        if task_local_sort {
            sample_sort::merge_sort_task_local(
                text,
                lp,
                &mut sa,
                &mut sa_w,
                &mut lcp_arr,
                &mut lcp_w,
                opts.max_context,
                dispatch,
            );
        } else {
            sample_sort::merge_sort(
                text,
                lp,
                &mut sa,
                &mut sa_w,
                &mut lcp_arr,
                &mut lcp_w,
                opts.max_context,
                dispatch,
            );
        }
        drop(sa_w);
        drop(lcp_w);

        // Consecutive pivots have non-decreasing upper bounds. Galloping from
        // the previous split avoids restarting a full binary search p times.
        let mut splits = Vec::with_capacity(p + 1);
        splits.push(0usize);
        let mut from = 0usize;
        for &pivot in pivots {
            from =
                upper_bound_positions_from(&sa, from, pivot, text, lp, opts.max_context, dispatch);
            splits.push(from);
        }
        splits.push(sa.len());

        for j in 0..p {
            let (lo, hi) = (splits[j], splits[j + 1]);
            if lo >= hi {
                continue;
            }
            let mut bucket = partition_buckets[j].lock().unwrap();
            bucket.add_soa_reset_first_lcp(&sa[lo..hi], &lcp_arr[lo..hi])?;
            bucket.mark_boundary();
        }
        Ok(())
    })?;

    Ok(partition_buckets
        .into_iter()
        .map(|m| m.into_inner().expect("partition mutex poisoned"))
        .collect())
}

/// Upper bound of `pivot` in `records`, searched forward from the previous
/// pivot's split. Pivots are sorted, so upper bounds never move backward.
fn upper_bound_positions_from<S, I, L>(
    positions: &[I],
    from: usize,
    pivot: I,
    text: &[S],
    lp: &L,
    max_ctx: usize,
    dispatch: LcpDispatch,
) -> usize
where
    S: Symbol,
    I: Index,
    L: LimitProvider,
{
    let n = positions.len();
    let greater = |i: usize| -> bool {
        dispatch.suffix_cmp_with(text, lp, positions[i].to_usize(), pivot.to_usize(), max_ctx)
            == Ordering::Greater
    };

    if from >= n {
        return n;
    }
    if greater(from) {
        return from;
    }

    let mut lo = from;
    let mut step = 1usize;
    loop {
        let probe = from.saturating_add(step);
        if probe >= n {
            break;
        }
        if greater(probe) {
            let mut hi = probe;
            while lo + 1 < hi {
                let mid = lo + (hi - lo) / 2;
                if greater(mid) {
                    hi = mid;
                } else {
                    lo = mid;
                }
            }
            return hi;
        }
        lo = probe;
        step = step.saturating_mul(2);
    }

    let mut hi = n;
    while lo + 1 < hi {
        let mid = lo + (hi - lo) / 2;
        if greater(mid) {
            hi = mid;
        } else {
            lo = mid;
        }
    }
    hi
}

/// Upper-bound binary search: returns the first index `i` such that the
/// suffix at `records[i].pos` is **strictly greater than** the suffix at
/// `pivot`.
fn upper_bound_by_pivot<S, I, L>(
    records: &[SaLcp<I>],
    pivot: I,
    text: &[S],
    lp: &L,
    max_ctx: usize,
    dispatch: LcpDispatch,
) -> usize
where
    S: Symbol,
    I: Index,
    L: LimitProvider,
{
    let mut lo = 0;
    let mut hi = records.len();
    while lo < hi {
        let mid = lo + (hi - lo) / 2;
        match dispatch.suffix_cmp_with(
            text,
            lp,
            records[mid].pos.to_usize(),
            pivot.to_usize(),
            max_ctx,
        ) {
            Ordering::Greater => hi = mid,
            Ordering::Equal | Ordering::Less => lo = mid + 1,
        }
    }
    lo
}

/// Phase 4 + 5: parallel-merge partitions in chunks of `num_threads`,
/// emitting each chunk's results in lex order before starting the next.
///
/// Up to `4 × T` partitions are dispatched per chunk so workers can steal
/// around partition-size skew. Each active merge holds two position/LCP sides;
/// completed position results are drained sequentially via `emit`. Between
/// chunks all merge storage is dropped, so residency scales with the bounded
/// in-flight chunk rather than the total suffix-array length.
#[allow(clippy::too_many_arguments)]
fn phase4_merge_and_emit<S, I, L, B, E, F>(
    text: &[S],
    lp: &L,
    partition_buckets: &mut [B],
    max_ctx: usize,
    ordered_emit: bool,
    memo_config: Option<MemoConfig>,
    collect_memo_stats: bool,
    emit: &mut F,
    dispatch: LcpDispatch,
) -> Result<(), BuildError<E>>
where
    S: Symbol,
    I: Index,
    L: LimitProvider,
    SaLcp<I>: BucketRecord,
    B: SaLcpBucketStore<I> + Send,
    F: FnMut(u64) -> Result<(), E>,
{
    let n_partitions = partition_buckets.len();
    if n_partitions == 0 {
        return Ok(());
    }
    // `chunk_size = 4 × num_threads` (not `= num_threads`): with one
    // partition per thread per chunk, rayon's `par_iter_mut` assigns
    // 1-to-1 with no opportunity to steal, and the chunk's wall is
    // set by its slowest partition. Sample-sort partition sizes vary
    // ~2× from random sampling, so the slow tail leaves ~half the
    // cores idle waiting (observed: 52% parallel efficiency on
    // GRCh38 / 32 t).
    //
    // Bumping the chunk to `4 × num_threads` gives rayon four
    // partitions per thread to dispatch — fast threads can steal from
    // slow neighbours, smoothing out the size variance. Peak RAM
    // grows linearly: each in-flight merged partition holds its
    // result `Vec<I>` (~3 MB at human-genome scale with `u32`
    // indices), so the chunk's transient cost goes from `32 × 3 MB =
    // 96 MB` to `128 × 3 MB = 384 MB` — well within the budget we
    // already spend on phase 1.
    let chunk_size = rayon::current_num_threads().max(1) * 4;

    // Per-thread CPU-µs accumulators for the two parallel sub-steps. They
    // add across threads, so the printed values are CPU-time (sum), not
    // wall-time; the ratio between them still tells us where the work is.
    use std::sync::atomic::{AtomicU64, Ordering as AtomicOrdering};
    let profile = std::env::var_os("CAPS_SA_PROFILE").is_some();
    let memo_profiled = profile && collect_memo_stats && memo_config.is_some();
    let load_us = AtomicU64::new(0);
    let merge_us = AtomicU64::new(0);
    let memo_stats = Mutex::new(MemoStats::default());
    let mut emit_secs: f64 = 0.0;

    let mut start = 0;
    while start < n_partitions {
        let end = (start + chunk_size).min(n_partitions);
        let chunk = &mut partition_buckets[start..end];
        if ordered_emit {
            phase4_merge_chunk_ordered_emit(
                text,
                lp,
                chunk,
                max_ctx,
                emit,
                dispatch,
                memo_config,
                &memo_stats,
                memo_profiled,
                profile,
                &load_us,
                &merge_us,
                &mut emit_secs,
            )?;
        } else {
            phase4_merge_chunk_collect_emit(
                text,
                lp,
                chunk,
                max_ctx,
                emit,
                dispatch,
                memo_config,
                &memo_stats,
                memo_profiled,
                profile,
                &load_us,
                &merge_us,
                &mut emit_secs,
            )?;
        }
        start = end;
    }
    if profile {
        profile_log(&format!(
            "phase4 breakdown CPU: load {:.3}s merge {:.3}s; wall emit {:.3}s",
            load_us.load(AtomicOrdering::Relaxed) as f64 * 1e-6,
            merge_us.load(AtomicOrdering::Relaxed) as f64 * 1e-6,
            emit_secs,
        ));
        if let Some(config) = memo_config.filter(|_| memo_profiled) {
            let stats = *memo_stats.lock().expect("memo profile mutex poisoned");
            profile_log(&format!(
                "geometric memo probe={} min_lcp={} cap={} activate_entries={} tables={} active_tables={} table_bins=[{},{},{},{},{},{}] calls={} training_direct={} probe_resolved={} lookups={} direct_hits={} gap_hits={} gap_mismatches={} gap_caps={} misses={} inserts={} extensions={} cap_rejects={} final_entries={} max_entries={} unique_diagonals={} singleton_diagonals={} max_entries_per_diagonal={} lookup_steps={} insert_steps={} insert_shifts={} scanned_matches={} skipped_matches={}",
                config.probe,
                config.min_lcp,
                config.capacity,
                config.activate_entries,
                stats.tables,
                stats.active_tables,
                stats.tables_0_15,
                stats.tables_16_31,
                stats.tables_32_63,
                stats.tables_64_127,
                stats.tables_128_255,
                stats.tables_256_plus,
                stats.calls,
                stats.cold_direct,
                stats.probe_resolved,
                stats.lookups,
                stats.direct_hits,
                stats.gap_hits,
                stats.gap_mismatches,
                stats.gap_caps,
                stats.misses,
                stats.inserts,
                stats.extensions,
                stats.capacity_rejects,
                stats.final_entries,
                stats.max_entries,
                stats.unique_diagonals,
                stats.singleton_diagonals,
                stats.max_entries_per_diagonal,
                stats.lookup_steps,
                stats.insert_steps,
                stats.insert_shifts,
                stats.scanned_matches,
                stats.skipped_matches,
            ));
        }
    }
    Ok(())
}

#[allow(clippy::too_many_arguments)]
fn phase4_merge_chunk_collect_emit<S, I, L, B, E, F>(
    text: &[S],
    lp: &L,
    chunk: &mut [B],
    max_ctx: usize,
    emit: &mut F,
    dispatch: LcpDispatch,
    memo_config: Option<MemoConfig>,
    memo_stats: &Mutex<MemoStats>,
    memo_profiled: bool,
    profile: bool,
    load_us: &std::sync::atomic::AtomicU64,
    merge_us: &std::sync::atomic::AtomicU64,
    emit_secs: &mut f64,
) -> Result<(), BuildError<E>>
where
    S: Symbol,
    I: Index,
    L: LimitProvider,
    SaLcp<I>: BucketRecord,
    B: SaLcpBucketStore<I> + Send,
    F: FnMut(u64) -> Result<(), E>,
{
    // Default fast path: let rayon merge the whole chunk with minimal
    // coordination, then emit the collected partition results in order.
    let merged: Vec<Vec<I>> = chunk
        .par_iter_mut()
        .map(|bucket| -> io::Result<Vec<I>> {
            merge_one_partition(
                text,
                lp,
                bucket,
                max_ctx,
                dispatch,
                memo_config,
                memo_stats,
                memo_profiled,
                profile,
                load_us,
                merge_us,
            )
        })
        .collect::<Result<Vec<_>, io::Error>>()?;

    let t = Instant::now();
    for positions in merged {
        for pos in positions {
            emit(pos.to_usize() as u64).map_err(BuildError::Emit)?;
        }
    }
    if profile {
        *emit_secs += t.elapsed().as_secs_f64();
    }
    Ok(())
}

#[allow(clippy::too_many_arguments)]
fn phase4_merge_chunk_ordered_emit<S, I, L, B, E, F>(
    text: &[S],
    lp: &L,
    chunk: &mut [B],
    max_ctx: usize,
    emit: &mut F,
    dispatch: LcpDispatch,
    memo_config: Option<MemoConfig>,
    memo_stats: &Mutex<MemoStats>,
    memo_profiled: bool,
    profile: bool,
    load_us: &std::sync::atomic::AtomicU64,
    merge_us: &std::sync::atomic::AtomicU64,
    emit_secs: &mut f64,
) -> Result<(), BuildError<E>>
where
    S: Symbol,
    I: Index,
    L: LimitProvider,
    SaLcp<I>: BucketRecord,
    B: SaLcpBucketStore<I> + Send,
    F: FnMut(u64) -> Result<(), E>,
{
    let n_jobs = chunk.len();
    let channel_bound = (rayon::current_num_threads().max(1) * 2).min(n_jobs).max(1);
    let (tx, rx) = std::sync::mpsc::sync_channel::<(usize, io::Result<Vec<I>>)>(channel_bound);
    let mut pending = std::collections::BTreeMap::<usize, Vec<I>>::new();
    let mut next_to_emit = 0usize;
    let mut received = 0usize;
    let mut io_err: Option<io::Error> = None;
    let mut emit_err: Option<E> = None;

    std::thread::scope(|thread_scope| {
        let worker = thread_scope.spawn(|| {
            chunk
                .par_iter_mut()
                .enumerate()
                .for_each_with(tx, |tx, (local_idx, bucket)| {
                    let result = merge_one_partition(
                        text,
                        lp,
                        bucket,
                        max_ctx,
                        dispatch,
                        memo_config,
                        memo_stats,
                        memo_profiled,
                        profile,
                        load_us,
                        merge_us,
                    );
                    let _ = tx.send((local_idx, result));
                });
        });

        while received < n_jobs {
            let (local_idx, result) = rx
                .recv()
                .expect("phase4 worker channel closed before all partitions completed");
            received += 1;
            match result {
                Ok(positions) => {
                    pending.insert(local_idx, positions);
                }
                Err(err) => {
                    if io_err.is_none() {
                        io_err = Some(err);
                    }
                }
            }

            while let Some(positions) = pending.remove(&next_to_emit) {
                if io_err.is_none() && emit_err.is_none() {
                    let t = Instant::now();
                    for pos in positions {
                        // Widen back to the public `u64` emit contract.
                        if let Err(err) = emit(pos.to_usize() as u64) {
                            emit_err = Some(err);
                            break;
                        }
                    }
                    if profile {
                        *emit_secs += t.elapsed().as_secs_f64();
                    }
                }
                next_to_emit += 1;
            }
        }
        worker.join().expect("phase4 merge worker panicked");
    });

    if let Some(err) = io_err {
        return Err(BuildError::Io(err));
    }
    if let Some(err) = emit_err {
        return Err(BuildError::Emit(err));
    }
    Ok(())
}

#[allow(clippy::too_many_arguments)]
fn merge_one_partition<S, I, L, B>(
    text: &[S],
    lp: &L,
    bucket: &mut B,
    max_ctx: usize,
    dispatch: LcpDispatch,
    memo_config: Option<MemoConfig>,
    memo_stats: &Mutex<MemoStats>,
    memo_profiled: bool,
    profile: bool,
    load_us: &std::sync::atomic::AtomicU64,
    merge_us: &std::sync::atomic::AtomicU64,
) -> io::Result<Vec<I>>
where
    S: Symbol,
    I: Index,
    L: LimitProvider,
    SaLcp<I>: BucketRecord,
    B: SaLcpBucketStore<I>,
{
    use std::sync::atomic::Ordering as AtomicOrdering;

    if bucket.total_records() == 0 {
        return Ok(Vec::new());
    }
    let t = Instant::now();
    let (positions, lcps) = bucket.load_all_soa()?;
    let boundaries: Vec<usize> = bucket.boundaries().to_vec();
    if profile {
        load_us.fetch_add(t.elapsed().as_micros() as u64, AtomicOrdering::Relaxed);
    }

    let t = Instant::now();
    let workspace = CascadeWorkspace::<I>::from_soa(positions, lcps);
    let result = if let Some(config) = memo_config {
        let mut memo = GeometricMemo::new(config);
        let result = if memo_profiled {
            workspace.cascade_merge_memoized_profiled(
                text,
                lp,
                &boundaries,
                max_ctx,
                dispatch,
                &mut memo,
            )
        } else {
            workspace.cascade_merge_memoized(text, lp, &boundaries, max_ctx, dispatch, &mut memo)
        };
        if memo_profiled {
            memo_stats
                .lock()
                .expect("memo profile mutex poisoned")
                .add_assign(memo.finish());
        }
        result
    } else {
        workspace.cascade_merge(text, lp, &boundaries, max_ctx, dispatch)
    };
    if profile {
        merge_us.fetch_add(t.elapsed().as_micros() as u64, AtomicOrdering::Relaxed);
    }
    Ok(result)
}

/// Reusable ping-pong scratch for the partition cascade merge.
///
/// Holds two `(sa, lcp)` buffers each sized to the largest partition seen.
/// The cascade alternates reads from one side and writes to the other,
/// flipping a `src_is_a` flag after each level. Avoids the
/// per-level allocations that the previous immutable-`Vec` cascade
/// performed for every pair of sub-subarrays.
struct CascadeWorkspace<I> {
    a_sa: Vec<I>,
    a_lcp: Vec<I>,
    b_sa: Vec<I>,
    b_lcp: Vec<I>,
}

impl<I: Index> CascadeWorkspace<I> {
    fn from_soa(a_sa: Vec<I>, a_lcp: Vec<I>) -> Self {
        assert_eq!(a_sa.len(), a_lcp.len());
        let n = a_sa.len();
        Self {
            a_sa,
            a_lcp,
            b_sa: vec![I::zero(); n],
            b_lcp: vec![I::zero(); n],
        }
    }

    /// Cascade 2-way LCP-enhanced merges across the sub-subarrays of one
    /// partition (delimited by `boundaries`) until a single sorted run
    /// remains. **Consumes the workspace** and returns the result side
    /// as a `Vec<I>`; the other three buffers (`a_lcp`, the opposing
    /// `*_sa`, the opposing `*_lcp`) drop immediately. This shape lets
    /// the caller skip the per-partition `to_vec()` round-trip that
    /// would otherwise sit briefly alongside all four workspace buffers
    /// at peak.
    fn cascade_merge<S, L>(
        self,
        text: &[S],
        lp: &L,
        boundaries: &[usize],
        max_ctx: usize,
        dispatch: LcpDispatch,
    ) -> Vec<I>
    where
        S: Symbol,
        L: LimitProvider,
    {
        self.cascade_merge_impl(text, lp, boundaries, max_ctx, dispatch, None, false)
    }

    #[allow(clippy::too_many_arguments)]
    fn cascade_merge_memoized<S, L>(
        self,
        text: &[S],
        lp: &L,
        boundaries: &[usize],
        max_ctx: usize,
        dispatch: LcpDispatch,
        memo: &mut GeometricMemo,
    ) -> Vec<I>
    where
        S: Symbol,
        L: LimitProvider,
    {
        self.cascade_merge_impl(text, lp, boundaries, max_ctx, dispatch, Some(memo), false)
    }

    #[allow(clippy::too_many_arguments)]
    fn cascade_merge_memoized_profiled<S, L>(
        self,
        text: &[S],
        lp: &L,
        boundaries: &[usize],
        max_ctx: usize,
        dispatch: LcpDispatch,
        memo: &mut GeometricMemo,
    ) -> Vec<I>
    where
        S: Symbol,
        L: LimitProvider,
    {
        self.cascade_merge_impl(text, lp, boundaries, max_ctx, dispatch, Some(memo), true)
    }

    #[allow(clippy::too_many_arguments)]
    fn cascade_merge_impl<S, L>(
        mut self,
        text: &[S],
        lp: &L,
        boundaries: &[usize],
        max_ctx: usize,
        dispatch: LcpDispatch,
        mut memo: Option<&mut GeometricMemo>,
        memo_profiled: bool,
    ) -> Vec<I>
    where
        S: Symbol,
        L: LimitProvider,
    {
        let n = self.a_sa.len();
        if n == 0 {
            return Vec::new();
        }

        // Side A was decoded directly from the bucket in SOA form. Collect the
        // lengths of the non-empty sub-subarrays.
        let mut run_lens: Vec<usize> = boundaries
            .windows(2)
            .filter_map(|w| {
                let l = w[1] - w[0];
                if l > 0 { Some(l) } else { None }
            })
            .collect();
        let mut src_is_a = true;
        while run_lens.len() > 1 {
            run_lens = self.merge_one_level(
                src_is_a,
                &run_lens,
                text,
                lp,
                max_ctx,
                dispatch,
                memo.as_deref_mut(),
                memo_profiled,
            );
            src_is_a = !src_is_a;
        }

        // Take ownership of the buffer holding the result, truncate to
        // the actual record count, drop the other three buffers with
        // `self` going out of scope.
        let mut result = if src_is_a { self.a_sa } else { self.b_sa };
        result.truncate(n);
        result
    }

    /// Pair the runs in `run_lens` (last odd one passes through unchanged),
    /// running each pair through the LCP-enhanced 2-way merge from the
    /// `src_is_a`-selected buffer side into the other. Returns the new
    /// run-length list (each entry is the sum of the two it replaced, or
    /// the carry-over for an odd tail).
    #[allow(clippy::too_many_arguments)]
    fn merge_one_level<S, L>(
        &mut self,
        src_is_a: bool,
        run_lens: &[usize],
        text: &[S],
        lp: &L,
        max_ctx: usize,
        dispatch: LcpDispatch,
        mut memo: Option<&mut GeometricMemo>,
        memo_profiled: bool,
    ) -> Vec<usize>
    where
        S: Symbol,
        L: LimitProvider,
    {
        // Destructure self so the borrow checker can see the two sides as
        // disjoint locals — we borrow one immutably and the other mutably.
        let Self {
            a_sa,
            a_lcp,
            b_sa,
            b_lcp,
        } = self;
        let (src_sa, src_lcp, dst_sa, dst_lcp) = if src_is_a {
            (
                a_sa.as_slice(),
                a_lcp.as_slice(),
                b_sa.as_mut_slice(),
                b_lcp.as_mut_slice(),
            )
        } else {
            (
                b_sa.as_slice(),
                b_lcp.as_slice(),
                a_sa.as_mut_slice(),
                a_lcp.as_mut_slice(),
            )
        };

        let mut new_lens = Vec::with_capacity(run_lens.len().div_ceil(2));
        let mut src_off = 0usize;
        let mut dst_off = 0usize;
        let mut i = 0;
        while i < run_lens.len() {
            let l1 = run_lens[i];
            if i + 1 < run_lens.len() {
                let l2 = run_lens[i + 1];
                let x_end = src_off + l1;
                let xy_end = x_end + l2;
                let dst_end = dst_off + l1 + l2;
                if let Some(memo) = memo.as_deref_mut() {
                    if memo.is_active() && memo_profiled {
                        sample_sort::merge_memoized_profiled(
                            text,
                            lp,
                            &src_sa[src_off..x_end],
                            &src_sa[x_end..xy_end],
                            &src_lcp[src_off..x_end],
                            &src_lcp[x_end..xy_end],
                            &mut dst_sa[dst_off..dst_end],
                            &mut dst_lcp[dst_off..dst_end],
                            max_ctx,
                            dispatch,
                            memo,
                        );
                    } else if memo.is_active() {
                        sample_sort::merge_memoized(
                            text,
                            lp,
                            &src_sa[src_off..x_end],
                            &src_sa[x_end..xy_end],
                            &src_lcp[src_off..x_end],
                            &src_lcp[x_end..xy_end],
                            &mut dst_sa[dst_off..dst_end],
                            &mut dst_lcp[dst_off..dst_end],
                            max_ctx,
                            dispatch,
                            memo,
                        );
                    } else if memo_profiled {
                        sample_sort::merge_memoized_training_profiled(
                            text,
                            lp,
                            &src_sa[src_off..x_end],
                            &src_sa[x_end..xy_end],
                            &src_lcp[src_off..x_end],
                            &src_lcp[x_end..xy_end],
                            &mut dst_sa[dst_off..dst_end],
                            &mut dst_lcp[dst_off..dst_end],
                            max_ctx,
                            dispatch,
                            memo,
                        );
                    } else {
                        sample_sort::merge_memoized_training(
                            text,
                            lp,
                            &src_sa[src_off..x_end],
                            &src_sa[x_end..xy_end],
                            &src_lcp[src_off..x_end],
                            &src_lcp[x_end..xy_end],
                            &mut dst_sa[dst_off..dst_end],
                            &mut dst_lcp[dst_off..dst_end],
                            max_ctx,
                            dispatch,
                            memo,
                        );
                    }
                } else {
                    sample_sort::merge(
                        text,
                        lp,
                        &src_sa[src_off..x_end],
                        &src_sa[x_end..xy_end],
                        &src_lcp[src_off..x_end],
                        &src_lcp[x_end..xy_end],
                        &mut dst_sa[dst_off..dst_end],
                        &mut dst_lcp[dst_off..dst_end],
                        max_ctx,
                        dispatch,
                    );
                }
                new_lens.push(l1 + l2);
                src_off = xy_end;
                dst_off = dst_end;
                i += 2;
            } else {
                // Odd run carries over unchanged.
                let end = dst_off + l1;
                dst_sa[dst_off..end].copy_from_slice(&src_sa[src_off..src_off + l1]);
                dst_lcp[dst_off..end].copy_from_slice(&src_lcp[src_off..src_off + l1]);
                new_lens.push(l1);
                src_off += l1;
                dst_off = end;
                i += 1;
            }
        }
        new_lens
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::build_in_memory;
    use std::ffi::OsString;
    use tempfile::tempdir;

    static ENV_LOCK: Mutex<()> = Mutex::new(());

    struct EnvGuard(Vec<(&'static str, Option<OsString>)>);

    impl EnvGuard {
        fn capture(keys: &[&'static str]) -> Self {
            Self(
                keys.iter()
                    .map(|&key| (key, std::env::var_os(key)))
                    .collect(),
            )
        }

        fn set(&self, key: &'static str, value: &str) {
            // The test serializes every mutation through ENV_LOCK and restores
            // all touched variables before releasing it.
            unsafe { std::env::set_var(key, value) };
        }
    }

    impl Drop for EnvGuard {
        fn drop(&mut self) {
            for (key, value) in self.0.drain(..) {
                // See EnvGuard::set: this runs while the serial test lock is
                // still held and restores the process environment exactly.
                unsafe {
                    if let Some(value) = value {
                        std::env::set_var(key, value);
                    } else {
                        std::env::remove_var(key);
                    }
                }
            }
        }
    }

    fn ext_mem_sa(text: &[u8], p: usize) -> Vec<u64> {
        let dir = tempdir().unwrap();
        let opts = ExtMemOpts {
            subproblem_count: p,
            physical_file_count: 1,
            work_dir: dir.path().to_path_buf(),
            ..ExtMemOpts::default()
        };
        let mut out: Vec<u64> = Vec::with_capacity(text.len());
        build_ext_mem(text, &opts, |pos| {
            out.push(pos);
            Ok(())
        })
        .unwrap();
        out
    }

    fn ext_mem_sa_with_policy(
        text: &[u8],
        p: usize,
        lcp_memoization: LcpMemoizationPolicy,
    ) -> Vec<u64> {
        let dir = tempdir().unwrap();
        let opts = ExtMemOpts {
            subproblem_count: p,
            physical_file_count: 1,
            work_dir: dir.path().to_path_buf(),
            lcp_memoization,
            ..ExtMemOpts::default()
        };
        let mut out = Vec::with_capacity(text.len());
        build_ext_mem(text, &opts, |pos| {
            out.push(pos);
            Ok(())
        })
        .unwrap();
        out
    }

    #[test]
    fn memoization_policy_defaults_to_disabled() {
        let opts = ExtMemOpts::default();
        assert_eq!(opts.lcp_memoization, LcpMemoizationPolicy::Disabled);
        assert!(!opts.collect_lcp_memoization_stats);

        let config = GeometricMemoizationConfig::default();
        assert_eq!(config.probe_symbols(), 256);
        assert_eq!(config.min_lcp_symbols(), 1_024);
        assert_eq!(config.activate_after_entries(), 64);
        assert_eq!(config.max_entries_per_partition(), 4_096);
        assert_eq!(
            LcpMemoizationPolicy::geometric(),
            LcpMemoizationPolicy::Geometric(config)
        );
    }

    #[test]
    fn geometric_policy_matches_direct_output() {
        let mut text = Vec::new();
        for i in 0..600 {
            text.extend_from_slice(b"ACGTACGTACGTACGTACGTACGTACGT");
            text.push((i % 5) as u8);
        }
        text.push(200);

        let direct = ext_mem_sa_with_policy(&text, 16, LcpMemoizationPolicy::Disabled);
        let config = GeometricMemoizationConfig::default()
            .with_probe_symbols(NonZeroUsize::new(8).unwrap())
            .with_min_lcp_symbols(NonZeroUsize::new(16).unwrap())
            .with_activate_after_entries(NonZeroUsize::new(1).unwrap())
            .with_max_entries_per_partition(NonZeroUsize::new(128).unwrap());
        let memoized = ext_mem_sa_with_policy(&text, 16, LcpMemoizationPolicy::Geometric(config));
        assert_eq!(memoized, direct);
    }

    #[test]
    fn from_env_parses_memoization_policy_and_rejects_zero_values() {
        let _lock = ENV_LOCK.lock().unwrap();
        let keys = [
            "CAPS_SA_GEOMETRIC_MEMO",
            "CAPS_SA_MEMO_PROBE",
            "CAPS_SA_MEMO_MIN_LCP",
            "CAPS_SA_MEMO_ACTIVATE_ENTRIES",
            "CAPS_SA_MEMO_CAPACITY",
            "CAPS_SA_MEMO_STATS",
        ];
        let env = EnvGuard::capture(&keys);
        env.set("CAPS_SA_GEOMETRIC_MEMO", "true");
        env.set("CAPS_SA_MEMO_PROBE", "32");
        env.set("CAPS_SA_MEMO_MIN_LCP", "256");
        env.set("CAPS_SA_MEMO_ACTIVATE_ENTRIES", "8");
        env.set("CAPS_SA_MEMO_CAPACITY", "512");
        env.set("CAPS_SA_MEMO_STATS", "yes");

        let opts = ExtMemOpts::from_env();
        let LcpMemoizationPolicy::Geometric(config) = opts.lcp_memoization else {
            panic!("environment should enable geometric memoization");
        };
        assert_eq!(config.probe_symbols(), 32);
        assert_eq!(config.min_lcp_symbols(), 256);
        assert_eq!(config.activate_after_entries(), 8);
        assert_eq!(config.max_entries_per_partition(), 512);
        assert!(opts.collect_lcp_memoization_stats);

        env.set("CAPS_SA_MEMO_PROBE", "0");
        env.set("CAPS_SA_MEMO_MIN_LCP", "0");
        env.set("CAPS_SA_MEMO_ACTIVATE_ENTRIES", "0");
        env.set("CAPS_SA_MEMO_CAPACITY", "0");
        let opts = ExtMemOpts::from_env();
        let LcpMemoizationPolicy::Geometric(config) = opts.lcp_memoization else {
            panic!("environment should still enable geometric memoization");
        };
        assert_eq!(config, GeometricMemoizationConfig::default());
    }

    fn assert_matches_in_memory(text: &[u8], p: usize) {
        let want: Vec<u32> = build_in_memory(text);
        let want64: Vec<u64> = want.iter().map(|&x| x as u64).collect();
        let got = ext_mem_sa(text, p);
        assert_eq!(got, want64, "mismatch on text {text:?} with p={p}");
    }

    #[test]
    fn ext_mem_empty() {
        let got = ext_mem_sa(b"", 4);
        assert!(got.is_empty());
    }

    #[test]
    fn ext_mem_single_partition() {
        assert_matches_in_memory(b"banana", 1);
    }

    #[test]
    fn ext_mem_p_greater_than_n() {
        assert_matches_in_memory(b"abc", 10);
    }

    #[test]
    fn ext_mem_banana_p4() {
        assert_matches_in_memory(b"banana", 4);
    }

    #[test]
    fn ext_mem_mississippi_p3() {
        assert_matches_in_memory(b"mississippi", 3);
    }

    #[test]
    fn ext_mem_random_byte_texts() {
        use rand::{RngExt, SeedableRng};
        let mut rng = rand::rngs::StdRng::seed_from_u64(0xCAFE);
        for &n in &[16usize, 100, 1000, 5000] {
            for &p in &[1usize, 2, 4, 16] {
                let text: Vec<u8> = (0..n).map(|_| rng.random_range(0..6u8)).collect();
                assert_matches_in_memory(&text, p);
            }
        }
    }

    #[test]
    fn ext_mem_with_unique_terminator() {
        use rand::{RngExt, SeedableRng};
        let mut rng = rand::rngs::StdRng::seed_from_u64(0xF00D);
        for &n in &[10usize, 200, 2000] {
            for &p in &[1usize, 3, 8] {
                let mut text: Vec<u8> = (0..n).map(|_| rng.random_range(0..5u8)).collect();
                text.push(200);
                assert_matches_in_memory(&text, p);
            }
        }
    }

    fn ext_mem_for_positions(text: &[u8], positions: Vec<u64>, p: usize) -> Vec<u64> {
        let dir = tempdir().unwrap();
        let opts = ExtMemOpts {
            subproblem_count: p,
            physical_file_count: 1,
            work_dir: dir.path().to_path_buf(),
            ..ExtMemOpts::default()
        };
        let mut out: Vec<u64> = Vec::with_capacity(positions.len());
        build_ext_mem_for_positions(text, positions, &opts, |pos| {
            out.push(pos);
            Ok(())
        })
        .unwrap();
        out
    }

    #[test]
    fn ext_mem_for_positions_full_set_matches_ext_mem() {
        let text = b"mississippi";
        let want = ext_mem_sa(text, 3);
        let positions: Vec<u64> = (0..text.len() as u64).collect();
        let got = ext_mem_for_positions(text, positions, 3);
        assert_eq!(got, want);
    }

    #[test]
    fn ext_mem_for_positions_subset_matches_brute_force() {
        let text = b"mississippi";
        let positions: Vec<u64> = (0..text.len() as u64).step_by(2).collect();
        let mut want = positions.clone();
        want.sort_by(|&a, &b| text[a as usize..].cmp(&text[b as usize..]));
        let got = ext_mem_for_positions(text, positions, 4);
        assert_eq!(got, want);
    }

    #[test]
    fn ext_mem_for_positions_random_subsets() {
        use rand::{RngExt, SeedableRng};
        let mut rng = rand::rngs::StdRng::seed_from_u64(0xC0DE);
        for &n in &[50usize, 500, 2000] {
            for &p in &[1usize, 3, 8] {
                let text: Vec<u8> = (0..n).map(|_| rng.random_range(0..5u8)).collect();
                let mut positions: Vec<u64> = (0..n as u64).collect();
                positions.retain(|_| rng.random_range(0..10) < 7);
                let mut want = positions.clone();
                want.sort_by(|&a, &b| text[a as usize..].cmp(&text[b as usize..]));
                let got = ext_mem_for_positions(&text, positions, p);
                assert_eq!(got, want, "subset ext-mem mismatch n={n} p={p}");
            }
        }
    }

    fn in_memory_sample_sort(text: &[u8], p: usize) -> Vec<u64> {
        let dir = tempdir().unwrap();
        let opts = ExtMemOpts {
            subproblem_count: p,
            physical_file_count: 1,
            work_dir: dir.path().to_path_buf(),
            ..ExtMemOpts::default()
        };
        let mut out: Vec<u64> = Vec::with_capacity(text.len());
        build_in_memory_sample_sort(text, &opts, |pos| {
            out.push(pos);
            Ok(())
        })
        .unwrap();
        out
    }

    #[test]
    fn in_memory_sample_sort_matches_in_memory() {
        for text in [b"banana" as &[u8], b"mississippi", b"abracadabra"] {
            let want: Vec<u32> = build_in_memory(text);
            let want64: Vec<u64> = want.iter().map(|&x| x as u64).collect();
            let got = in_memory_sample_sort(text, 0);
            assert_eq!(got, want64, "in-mem sample-sort mismatch on {text:?}");
        }
    }

    #[test]
    fn in_memory_sample_sort_random_byte_texts() {
        use rand::{RngExt, SeedableRng};
        let mut rng = rand::rngs::StdRng::seed_from_u64(0xC0DE_C0DE);
        for &n in &[16usize, 200, 2000] {
            for &p in &[1usize, 4, 16] {
                let text: Vec<u8> = (0..n).map(|_| rng.random_range(0..6u8)).collect();
                let want: Vec<u32> = build_in_memory(&text);
                let want64: Vec<u64> = want.iter().map(|&x| x as u64).collect();
                let got = in_memory_sample_sort(&text, p);
                assert_eq!(got, want64, "in-mem ss mismatch n={n} p={p}");
            }
        }
    }

    /// Helper that drives [`build_ext_mem_for_filter`] and collects the
    /// emitted positions.
    fn ext_mem_for_filter<Pred>(text: &[u8], keep: Pred, p: usize) -> Vec<u64>
    where
        Pred: Fn(u64) -> bool + Send + Sync,
    {
        let dir = tempdir().unwrap();
        let opts = ExtMemOpts {
            subproblem_count: p,
            physical_file_count: 1,
            work_dir: dir.path().to_path_buf(),
            ..ExtMemOpts::default()
        };
        let mut out: Vec<u64> = Vec::new();
        build_ext_mem_for_filter(text, keep, &opts, |pos| {
            out.push(pos);
            Ok(())
        })
        .unwrap();
        out
    }

    #[test]
    fn ext_mem_for_filter_matches_for_positions_on_full_set() {
        // Filter that accepts every position → must equal the
        // identity-positions ext-mem build.
        let text = b"mississippi";
        let want = ext_mem_sa(text, 3);
        let got = ext_mem_for_filter(text, |_p| true, 3);
        assert_eq!(got, want);
    }

    #[test]
    fn ext_mem_for_filter_matches_for_positions_on_dna_subset() {
        // STAR-style "keep ACGT (`< 4`), drop N (`4`)/spacer (`5`)"
        // filter. The filter API must produce exactly the same SA as
        // pre-materialising the kept positions and going through the
        // _for_positions path.
        use rand::{RngExt, SeedableRng};
        let mut rng = rand::rngs::StdRng::seed_from_u64(0xCA_755A);
        for &n in &[50usize, 500, 2000] {
            for &p in &[1usize, 3, 8] {
                let text: Vec<u8> = (0..n).map(|_| rng.random_range(0..6u8)).collect();
                let positions: Vec<u64> = (0..n as u64).filter(|&i| text[i as usize] < 4).collect();
                let want = ext_mem_for_positions(&text, positions, p);
                let got = ext_mem_for_filter(&text, |i| text[i as usize] < 4, p);
                assert_eq!(got, want, "filter vs positions mismatch n={n} p={p}");
            }
        }
    }

    #[test]
    fn ext_mem_for_filter_handles_block_aligned_boundaries() {
        // Exercise the bitmap word/block boundaries by using a text
        // longer than one popcount block (1024 × 64 bits = 64 K
        // positions) — but stay under that to keep the test fast.
        // 200 K positions touches the cumsum's second block too.
        use rand::{RngExt, SeedableRng};
        let mut rng = rand::rngs::StdRng::seed_from_u64(0xB10C_C0DE);
        let n = 200_000usize;
        let text: Vec<u8> = (0..n).map(|_| rng.random_range(0..6u8)).collect();
        let positions: Vec<u64> = (0..n as u64).filter(|&i| text[i as usize] < 4).collect();
        let want = ext_mem_for_positions(&text, positions, 8);
        let got = ext_mem_for_filter(&text, |i| text[i as usize] < 4, 8);
        assert_eq!(got, want, "filter API mismatch across block boundaries");
    }

    #[test]
    fn ext_mem_for_filter_sparse_predicate() {
        // ~5% acceptance — exercises long runs of zero-bits in the
        // bitmap (skip-loop across whole words).
        use rand::{RngExt, SeedableRng};
        let mut rng = rand::rngs::StdRng::seed_from_u64(0x5_AA_55);
        let n = 50_000usize;
        let text: Vec<u8> = (0..n).map(|_| rng.random_range(0..20u8)).collect();
        let positions: Vec<u64> = (0..n as u64).filter(|&i| text[i as usize] < 1).collect();
        let want = ext_mem_for_positions(&text, positions, 4);
        let got = ext_mem_for_filter(&text, |i| text[i as usize] < 1, 4);
        assert_eq!(got, want, "filter API mismatch on sparse predicate");
    }

    #[derive(Debug, PartialEq, Eq)]
    enum EmitTestError {
        Stop,
    }

    #[test]
    fn try_ext_mem_returns_typed_emit_error() {
        let dir = tempdir().unwrap();
        let opts = ExtMemOpts {
            subproblem_count: 2,
            physical_file_count: 1,
            work_dir: dir.path().to_path_buf(),
            ..ExtMemOpts::default()
        };
        let mut seen = 0usize;
        let err = try_build_ext_mem(b"banana", &opts, |_pos| {
            seen += 1;
            if seen == 2 {
                Err(EmitTestError::Stop)
            } else {
                Ok(())
            }
        })
        .unwrap_err();

        assert!(matches!(err, BuildError::Emit(EmitTestError::Stop)));
    }

    #[test]
    fn ext_mem_repetitive_does_not_blow_up() {
        // Many copies of a long repeat — what killed the Phase 2 v1
        // linear-scan merge. The sample-sort + LCP-enhanced cascade
        // should handle it in proportional time.
        use std::time::Instant;
        let unit = b"ACGTACGTACGTACGTACGTACGTACGT"; // 28 bases
        let mut text: Vec<u8> = Vec::new();
        for _ in 0..100 {
            text.extend_from_slice(unit);
        }
        text.push(200);
        let start = Instant::now();
        let got = ext_mem_sa(&text, 8);
        let elapsed = start.elapsed();
        let want: Vec<u32> = build_in_memory(&text);
        let want64: Vec<u64> = want.iter().map(|&x| x as u64).collect();
        assert_eq!(got, want64);
        // Sanity: should finish in well under a second on this input.
        assert!(
            elapsed.as_secs() < 2,
            "ext-mem build on a tiny repetitive text took {elapsed:?}"
        );
    }
}