diff --git a/README.md b/README.md index 5d9a382..a137913 100644 --- a/README.md +++ b/README.md @@ -16,7 +16,53 @@ streams the SA out as positions are emitted. Both the in-memory and external-memory paths are implemented, tested on Linux, macOS, and Windows, and differentially verified against direct suffix -comparison on small, random, segmented, filtered, and finite-context inputs. +comparison on small, random, segmented, filtered, and finite-context inputs, +and against [`verify_sa`](#verifying-a-suffix-array) at genome scale. + +### In-memory fast path + +`build_in_memory` on a byte text routes through a **radix-seeded prefix +doubling** algorithm rather than the merge kernel. The merge kernel is +still the general path and still backs everything else; the fast path is +taken only when the comparator is provably plain lexicographic (see +[Choosing a path](#choosing-a-path)). + +The reason is that a comparison-based suffix sort pays twice on real +genomic input. It performs `n log n` merge steps, and every tied step +scans the shared prefix of two suffixes from the beginning. Genome FASTA +carries megabyte-scale runs of `N` — period-61 once 60-column line +wrapping is included — so a single comparison can scan millions of +bytes. Measured on chr21 that drives the cost per merge step from 13 ns +to 222 ns, a 16x penalty that is entirely scan time. + +The fast path sorts by a packed fixed-depth key, then resolves what +remains by doubling on ranks. The packing picks the narrowest field +width that holds the alphabet, so DNA over `{0,1,2,3}` resolves 32 +symbols per key rather than the 8 a raw byte key gives. After the seed, +no comparison reads the text again, so a megabyte-long run of `N` costs +exactly what random DNA costs. + +Apple M4 Max (12 P-cores), 12 threads, suffix arrays byte-identical to +the merge kernel's and independently verified: + +| input | before | after | CPU before | CPU after | +| ----- | ------ | ----- | ---------- | --------- | +| chr21 fwd ++ revcomp, `N`-free, 80 MB | 6.08 s | **0.84 s** | 28.1 s | 5.0 s | +| chr21 FASTA, 47.5 MB, 6.6 Mb of `N` | 27.8 s | **1.04 s** | 283.5 s | 5.2 s | +| same, via `build_in_memory_sample_sort` | 3.41 s | **1.18 s** | 34.5 s | 7.2 s | + +Peak RSS on the 80 MB input is 2.21 GB, down from 2.83 GB, because the +seed is an MSD counting sort that recomputes keys from the text rather +than materialising a key array. + +Note the two inputs are different problems; benchmarking one +implementation on the first and another on the second is not a +comparison. `bench/chr21.sh` prepares both. + +### External memory + +The external-memory and sample-sort paths still use the merge kernel, so +they retain the scan cost described above on repeat-heavy input. On the complete ruSTAR-shaped GENCODE Human v50 input (6.56 billion text symbols, 6.18 billion retained suffixes, and 1.40 million segments), caps-sa @@ -105,6 +151,67 @@ build_ext_mem_for_positions(&text, positions, &opts, |sa_pos| { })?; ``` +### Verifying a suffix array + +`verify_sa` checks a candidate in `O(n)` without re-running any +construction algorithm and without depending on LCP length, so it stays +usable on the repetitive inputs that are hardest to trust: + +```rust +use caps_sa::{build_in_memory, verify_sa}; + +let text = b"banana"; +let sa: Vec = build_in_memory(text); +assert!(verify_sa(text, &sa).is_ok()); +``` + +It inverts `sa` to get ranks, then checks that +`(text[p], rank[p + 1])` increases strictly along it, with `rank[n]` +treated as smaller than every real rank. A permutation of `0..n` +satisfies that condition exactly when it is the suffix array. The bench +CLI exposes it as `--verify`. + +## Choosing a path + +`build_in_memory` takes the radix-seeded doubling fast path only when +the requested comparator is provably the plain lexicographic one. All +three conditions are soundness requirements, and each defaults to +declining: + +| Condition | Why | +| --------- | --- | +| `Opts::max_context` unbounded | A finite bound makes the merge comparator fall through to `LimitProvider::boundary_order`, which compares *lengths*, so it is not lexicographic. | +| `LimitProvider::plain_lex_len()` reports the full text | Rules out `SegmentedText`, whose scans stop at segment boundaries, and any custom `boundary_order`. | +| symbol type is exactly `u8` | Packing wider symbols into an order-preserving key is endianness-dependent: on a little-endian host `0x0100 > 0x0001` as `u16` values, but their byte views compare the other way. | + +`plain_lex_len` is a new `LimitProvider` method that defaults to `None`. +An implementation that delegates `lim_at` to `PlainText` but overrides +`boundary_order` for a different convention — STAR's spacer-as-largest +ordering is the motivating example — inherits `None` and keeps today's +semantics without changing a line. + +Given those, the fast path also covers two cases beyond a plain whole-text +build: + +- **`*_for_positions` subsets.** Doubling cannot be restricted to a subset + directly, since a round compares `rank[p + d]` and that successor is + generally outside the subset, so ranks must exist for every text + position. The full array is built and filtered in one `O(n)` pass + instead. Below one eighth of the text this declines and the merge kernel + runs, since building and discarding a whole array would cost more than + sorting a small subset. That ratio is a performance heuristic, not a + correctness condition. Duplicate or out-of-range positions also decline, + because the output is a permutation of the input *multiset* and a + membership filter cannot reproduce that. +- **`build_in_memory_sample_sort`.** This path exists to sort in RAM, so + where doubling applies it is strictly better: same output, no bucket + machinery, none of the scan cost. + +`build_ext_mem` deliberately stays on the merge kernel. Its purpose is to +bound peak memory, and routing it through an in-memory algorithm would +defeat exactly that, so it keeps the scan cost on repeat-heavy input. +Segmented texts and symbols wider than `u8` also stay on the merge kernel. + ## Algorithm The in-memory kernel is a parallel merge-sort whose two-way merge uses diff --git a/bench/README.md b/bench/README.md index 45cc7f7..ff9e06f 100644 --- a/bench/README.md +++ b/bench/README.md @@ -32,8 +32,16 @@ It reports wall time and peak RSS via `/usr/bin/time`. Machine: 64-core x86_64 Linux node, 1 socket, AVX2 enabled. Builds: upstream `cmake -DCMAKE_BUILD_TYPE=Release`, this crate -`cargo build --release --example caps_sa` (release profile in the -workspace `Cargo.toml` enables `lto = "fat"` + `codegen-units = 1`). +`cargo build --release --example caps_sa`. The `[profile.release]` in +this repo's `Cargo.toml` sets `lto = "fat"` + `codegen-units = 1`; for +the SIMD paths also pass `RUSTFLAGS="-C target-cpu=native"`. + +> **Correction.** This paragraph used to attribute the LTO settings to a +> parent workspace. No such workspace exists in the repository, so from +> the commit that made the crate standalone until the profile was added +> here, every build made from this repo actually used `lto = false` and +> `codegen-units = 16`. Numbers taken in that window are not comparable +> with numbers taken now. Re-measure before quoting them. All caps-sa runs include the five optimizations applied incrementally to the Phase 2b sample-sort baseline: @@ -399,6 +407,80 @@ against `build_ext_mem_for_positions` (sort only the kept positions) on the same fixture, showing **6–10× speedups** on padding-dominated inputs. The pool change leaves this gap intact. +### chr21 — the two inputs are different problems (in-memory fast path) + +Machine: Apple M4 Max (12 P-cores + 4 E-cores, aarch64/NEON, no +AVX-512), 12 threads, `RUSTFLAGS="-C target-cpu=native"`. Reproduce with +`bench/chr21.sh`. + +Two inputs are built from hg38 chr21, and the distinction matters more +than any tuning parameter in this document: + +| input | size | alphabet | long runs | +| ----- | ---- | -------- | --------- | +| `chr21.0123` — forward ++ revcomp, codes `0..=3`, ambiguous bases dropped | 80.2 MB | 4 | none | +| `chr21.fa` — the raw FASTA, headers and newlines included | 47.5 MB | 6 | 6.6 Mb of `N`, period-61 after 60-column wrapping | + +Benchmarking one implementation on the first and another on the second +compares two different problems. That is worth stating explicitly +because it is an easy mistake to make: the second input is 40% smaller +yet used to take 4.6x longer. + +Merge kernel vs. the radix-seeded doubling fast path, suffix arrays +byte-identical and independently `--verify`-checked in both cases: + +| input | merge kernel | fast path | speedup | CPU before | CPU after | CPU speedup | +| ----- | ------------ | --------- | ------- | ---------- | --------- | ----------- | +| `chr21.0123` | 6.08 s | **0.89 s** | 6.8x | 28.1 s | 5.0 s | 5.6x | +| `chr21.fa` | 27.8 s | **1.04 s** | 26.7x | 283.5 s | 5.2 s | **54x** | + +The per-merge-step cost is what separates the two rows. Dividing CPU +time by `n log2 n` merge steps: + +``` +chr21.0123 28.1 s / 2.11e9 steps = 13 ns/step +chr21.fa 283.5 s / 1.21e9 steps = 222 ns/step +``` + +Same kernel, same machine, 16x apart. The difference is entirely scan +length: every leaf merge starts with `m = 0`, so two suffixes inside an +`N` block scan their whole shared prefix, which is megabytes. + +Phase breakdown of the fast path (`CAPS_SA_PROFILE=1`): + +``` + chr21.0123 chr21.fa +key extraction 0.086 s 0.020 s +seed sort 0.255 s 0.176 s +grouping 0.075 s 0.037 s +doubling rounds 0.475 s 0.807 s +``` + +The doubling rounds are the larger share on the FASTA input, as +expected: `N` blocks stay tied for many rounds. But each round is +rank-only, so they cost a sort over a shrinking residual rather than a +text scan, which is why the pathology disappears rather than merely +shrinking. + +### Reading the 97.54% LCP profile correctly + +The profile in the next section shows `lcp_u8_avx2` taking 97.54% of +samples on a human-genome slice. That was read at the time as "LCP +scanning is expensive", and it drove several rounds of work on making +the scan wider — AVX2, then AVX-512, then the 32-byte/64-byte hybrid. + +The chr21 numbers above suggest a second reading. The LCP kernel is also +where the two *random* text loads happen, so on short-LCP input those +samples are load stalls rather than scan length, and a wider vector does +not help. That is consistent with the AVX-512 ablation further down, +where the 64-byte-only variant was **16% slower** on `rand100m` and only +the long-LCP human slice gained. Widening the scan helps when the scan +is genuinely long; when it is short, the cost is latency and the fix is +to issue fewer random probes. + +The fast path does the latter: it removes the probes entirely rather +than making each one faster. + ### Where AVX-512 helps and where it doesn't — the measurement A `perf record --call-graph dwarf` run on a 200 MB human-genome slice diff --git a/bench/chr21.sh b/bench/chr21.sh new file mode 100755 index 0000000..5a5398f --- /dev/null +++ b/bench/chr21.sh @@ -0,0 +1,83 @@ +#!/usr/bin/env bash +# Reproduce the chr21 numbers quoted in bench/README.md. +# +# bench/chr21.sh [work-dir] [threads] +# +# Builds two inputs from hg38 chr21, because they exercise different costs and +# conflating them is what made the original slowdown report hard to read: +# +# chr21.0123 forward ++ reverse complement, one byte per base, codes 0..=3, +# ambiguous bases dropped. ~80 MB, alphabet size 4, no long runs. +# This is the input libsais is usually benchmarked on. +# chr21.fa the raw FASTA, headers and newlines included. ~45 MB, and it +# still contains its ~6.6 Mb of `N`. Wrapped at 60 columns, so +# the `N` blocks are a period-61 repeat rather than a plain run. +# +# The second is the realistic one and the one that used to be pathological: a +# comparison-based suffix sort scans the whole shared prefix on every tied +# comparison, so an `N` block costs megabytes per comparison. +set -euo pipefail + +work_dir="${1:-bench/work}" +threads="${2:-$( (command -v nproc >/dev/null && nproc) || sysctl -n hw.perflevel0.logicalcpu 2>/dev/null || echo 4)}" +root="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" + +mkdir -p "$work_dir" +gz="$work_dir/chr21.fa.gz" +fa="$work_dir/chr21.fa" +bin="$work_dir/chr21.0123" + +if [ ! -f "$gz" ]; then + echo "== downloading hg38 chr21 ==" >&2 + curl -sSL -o "$gz" \ + https://hgdownload.soe.ucsc.edu/goldenPath/hg38/chromosomes/chr21.fa.gz +fi +[ -f "$fa" ] || gzip -dc "$gz" > "$fa" + +if [ ! -f "$bin" ]; then + echo "== encoding forward ++ revcomp as codes 0..=3 ==" >&2 + python3 - "$gz" "$bin" <<'PY' +import gzip, sys +code = {"A": 0, "C": 1, "G": 2, "T": 3} +comp = {0: 3, 1: 2, 2: 1, 3: 0} +fwd = bytearray() +with gzip.open(sys.argv[1], "rt") as fh: + for line in fh: + if line.startswith(">"): + continue + for ch in line.strip().upper(): + c = code.get(ch) + if c is not None: + fwd.append(c) +rc = bytearray(comp[b] for b in reversed(fwd)) +with open(sys.argv[2], "wb") as out: + out.write(fwd) + out.write(rc) +print(f"{sys.argv[2]}: {len(fwd) + len(rc)} bytes", file=sys.stderr) +PY +fi + +echo "== building (fat LTO, target-cpu=native) ==" >&2 +RUSTFLAGS="-C target-cpu=native" cargo build --release --example caps_sa --manifest-path "$root/Cargo.toml" +caps_sa="$root/target/release/examples/caps_sa" + +run() { + local label="$1" input="$2" + shift 2 + # `--verify` is an O(n) independent check of the result; it is timed and + # reported separately by the binary, so it never inflates the build time. + printf '%-28s ' "$label" + "$caps_sa" "$input" /dev/null --threads "$threads" --verify "$@" 2>&1 | + awk '/^build:/ { for (i = 1; i <= NF; i++) if ($i ~ /^[0-9.]+s$/) b = $i } + /^verify:/ { v = $0 } + END { printf "build %-9s %s\n", b, (v ~ /OK/ ? "verify OK" : "VERIFY FAILED") }' +} + +echo +echo "threads: $threads" +echo +run "0123 (80 MB, no N)" "$bin" +run "FASTA (45 MB, 6.6 Mb N)" "$fa" +echo +echo "For wall+CPU together, wrap a single run:" >&2 +echo " /usr/bin/time -p $caps_sa $bin /dev/null --threads $threads" >&2 diff --git a/examples/caps_sa.rs b/examples/caps_sa.rs index 17945d8..53f5de3 100644 --- a/examples/caps_sa.rs +++ b/examples/caps_sa.rs @@ -18,7 +18,7 @@ use std::path::PathBuf; use std::process; use std::time::Instant; -use caps_sa::{ExtMemOpts, build_ext_mem, build_in_memory, build_in_memory_sample_sort}; +use caps_sa::{ExtMemOpts, build_ext_mem, build_in_memory, build_in_memory_sample_sort, verify_sa}; struct Args { input: PathBuf, @@ -27,6 +27,7 @@ struct Args { in_mem_ss: bool, subproblem_count: usize, threads: Option, + verify: bool, } fn parse_args() -> Args { @@ -36,6 +37,7 @@ fn parse_args() -> Args { let mut in_mem_ss = false; let mut subproblem_count: usize = 0; let mut threads: Option = None; + let mut verify = false; let mut i = 1; while i < argv.len() { match argv[i].as_str() { @@ -47,6 +49,10 @@ fn parse_args() -> Args { in_mem_ss = true; i += 1; } + "--verify" => { + verify = true; + i += 1; + } "--subproblem-count" => { subproblem_count = argv[i + 1] .parse() @@ -64,7 +70,7 @@ fn parse_args() -> Args { "--help" | "-h" => { eprintln!( "usage: caps_sa [--ext-mem | --in-mem-ss] \ - [--subproblem-count N] [--threads N]" + [--subproblem-count N] [--threads N] [--verify]" ); process::exit(0); } @@ -88,6 +94,23 @@ fn parse_args() -> Args { in_mem_ss, subproblem_count, threads, + verify, + } +} + +/// Independently check the built suffix array in O(n). Off by default so it +/// never contaminates a timing run; the check is reported separately. +fn maybe_verify(enabled: bool, text: &[u8], sa: &[I]) { + if !enabled { + return; + } + let t = Instant::now(); + match verify_sa(text, sa) { + Ok(()) => eprintln!("verify: OK in {:.3}s", t.elapsed().as_secs_f64()), + Err(e) => { + eprintln!("verify: FAILED: {e}"); + process::exit(1); + } } } @@ -124,6 +147,7 @@ fn main() -> std::io::Result<()> { let sa: Vec = build_in_memory(&text); build_elapsed = build_start.elapsed(); n_entries = sa.len(); + maybe_verify(args.verify, &text, &sa); eprintln!( "build: mode=in-mem(u32) n={n_entries} entries in {:.3}s", build_elapsed.as_secs_f64() @@ -174,6 +198,7 @@ fn main() -> std::io::Result<()> { let sa: Vec = build_in_memory(&text); build_elapsed = build_start.elapsed(); n_entries = sa.len(); + maybe_verify(args.verify, &text, &sa); eprintln!( "build: mode=in-mem(u64) n={n_entries} entries in {:.3}s", build_elapsed.as_secs_f64() diff --git a/src/ext_mem.rs b/src/ext_mem.rs index 2a4ebbb..0087e5a 100644 --- a/src/ext_mem.rs +++ b/src/ext_mem.rs @@ -43,7 +43,7 @@ use crate::sample_sort; /// 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) { +pub(crate) fn profile_log(message: &str) { if std::env::var_os("CAPS_SA_PROFILE").is_some() { eprintln!("caps-sa profile {message}"); } @@ -739,6 +739,26 @@ where L: LimitProvider, F: FnMut(u64) -> Result<(), E>, { + // This path exists to sort in RAM, so when the doubling path applies it is + // strictly better here: same output, no bucket machinery, and none of the + // scan cost the merge kernel pays on repeat-heavy text. `build_ext_mem` + // deliberately does *not* do this -- its whole purpose is to bound peak + // memory, and routing it through an in-memory algorithm would defeat that. + if opts.max_context == usize::MAX + && lp.plain_lex_len() == Some(text.len()) + && std::any::TypeId::of::() == std::any::TypeId::of::() + { + // SAFETY: `S` is `u8`, so `&[S]` and `&[u8]` have identical layout. + let bytes: &[u8] = + unsafe { std::slice::from_raw_parts(text.as_ptr() as *const u8, text.len()) }; + let sa: Vec = crate::radix::build_sa(bytes); + let mut emit = emit; + for pos in sa { + emit(pos).map_err(BuildError::Emit)?; + } + return Ok(()); + } + if text.len() <= u32::MAX as usize + 1 { build_in_memory_ss_inner::( text, diff --git a/src/lib.rs b/src/lib.rs index 06fc661..407a822 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -22,6 +22,7 @@ mod ext_mem; mod lcp; mod lcp_memo; mod limits; +mod radix; mod sample_sort; pub use ext_mem::{ @@ -42,6 +43,74 @@ pub use sample_sort::{ build_in_memory_for_positions_with_opts, build_in_memory_with, build_in_memory_with_opts, }; +/// Check that `sa` really is the suffix array of `text`, in `O(n)` time and +/// without re-running any construction algorithm. +/// +/// Comparing a candidate against a second implementation only shows the two +/// agree; comparing adjacent suffixes directly is `O(n · lcp)` and becomes +/// unusable on the repetitive inputs that matter most. This instead uses the +/// standard fixpoint characterisation: let `rank` be the inverse of `sa`, and +/// define `f(p) = (text[p], rank[p + 1])`, with `rank[n]` taken as less than +/// every real rank. A permutation is the suffix array of `text` if and only if +/// `f` is strictly increasing along it, because suffix `p` precedes suffix `q` +/// exactly when `f(p) < f(q)`. +/// +/// Returns `Err` with a description of the first violation found. +/// +/// ``` +/// let text = b"banana"; +/// let sa: Vec = caps_sa::build_in_memory(text); +/// assert!(caps_sa::verify_sa(text, &sa).is_ok()); +/// assert!(caps_sa::verify_sa(text, &[0u32, 1, 2, 3, 4, 5]).is_err()); +/// ``` +pub fn verify_sa(text: &[S], sa: &[I]) -> Result<(), String> +where + S: Ord, + I: Index, +{ + let n = text.len(); + if sa.len() != n { + return Err(format!("sa has {} entries, text has {n} symbols", sa.len())); + } + if n == 0 { + return Ok(()); + } + + // Invert `sa`, checking along the way that it is a permutation of `0..n`. + let mut rank = vec![usize::MAX; n]; + for (i, entry) in sa.iter().enumerate() { + let p = entry.to_usize(); + if p >= n { + return Err(format!("sa[{i}] = {p} is out of range for text length {n}")); + } + if rank[p] != usize::MAX { + return Err(format!( + "position {p} appears at sa[{}] and sa[{i}]", + rank[p] + )); + } + rank[p] = i; + } + + // `None` stands for the end of the text, which sorts before every rank: + // the shorter suffix is the smaller one. + let successor = + |p: usize| -> Option { if p + 1 < n { Some(rank[p + 1]) } else { None } }; + for i in 1..n { + let a = sa[i - 1].to_usize(); + let b = sa[i].to_usize(); + let key_a = (&text[a], successor(a)); + let key_b = (&text[b], successor(b)); + if key_a >= key_b { + return Err(format!( + "suffixes out of order at sa[{}] = {a} and sa[{i}] = {b}", + i - 1, + )); + } + } + Ok(()) +} + /// Trait implemented by integer types usable as suffix array indices. /// /// Provided for `u32`, `u64`, and `usize`. Callers pick the narrowest type diff --git a/src/limits.rs b/src/limits.rs index e4e5ead..3a1d9e8 100644 --- a/src/limits.rs +++ b/src/limits.rs @@ -73,6 +73,29 @@ pub trait LimitProvider: Sync { let _ = (p_a, p_b); lim_a.cmp(&lim_b) } + + /// `Some(n)` iff this provider describes an unsegmented text of `n` + /// symbols under the *standard* comparator: `lim_at(p) == n - p` for + /// every `p`, and `boundary_order` left at its shorter-is-smaller + /// default. + /// + /// Returning `Some` lets the crate substitute a specialised suffix-array + /// algorithm that assumes plain lexicographic order. It is therefore a + /// promise about the *comparator*, not merely about the lengths. + /// + /// The default is `None`, which keeps every existing and third-party + /// implementation on the general merge kernel at today's semantics. In + /// particular, an implementation that delegates `lim_at` to [`PlainText`] + /// but overrides [`boundary_order`][LimitProvider::boundary_order] to get + /// a different convention (STAR's spacer-as-largest ordering is the + /// motivating example) inherits `None` and is safe without doing + /// anything. + /// + /// Override this only if you have *not* overridden `boundary_order`. + #[inline] + fn plain_lex_len(&self) -> Option { + None + } } /// Default provider for non-segmented texts: `lim_at(p) = n - p`. @@ -99,6 +122,11 @@ impl LimitProvider for PlainText { fn lim_at(&self, p: usize) -> usize { self.n - p } + + #[inline] + fn plain_lex_len(&self) -> Option { + Some(self.n) + } } /// Provider for texts partitioned into segments at known cumulative diff --git a/src/radix.rs b/src/radix.rs new file mode 100644 index 0000000..7eab7db --- /dev/null +++ b/src/radix.rs @@ -0,0 +1,601 @@ +//! Radix-seeded prefix doubling for the plain in-memory suffix array. +//! +//! The LCP-enhanced merge sort in [`crate::sample_sort`] is the CaPS-SA +//! kernel and stays the general path: it is the only one that honours a +//! [`LimitProvider`][crate::limits::LimitProvider], a finite `max_context`, +//! and symbol types wider than a byte, and it is the only one that produces +//! an LCP array (which the external-memory path needs). +//! +//! But for the single most common request — the standard lexicographic +//! suffix array of a byte text, with no segmentation and no context bound — +//! that kernel is doing far more work than the problem requires, in two +//! distinct ways that the benchmarks separate cleanly: +//! +//! * **Step count.** The merge sort performs `n log n` merge steps, and a +//! large majority of them are resolved by an actual symbol comparison at a +//! random text address. On 80 MB of N-free DNA that is ~2.1e9 steps at +//! ~13 ns each. +//! * **Scan length.** Every leaf merge starts with `m = 0`, so comparing two +//! suffixes that share a long prefix costs a scan proportional to that +//! prefix. Genome assemblies contain megabyte-scale runs of `N` (and the +//! period-61 `N`-then-newline pattern of wrapped FASTA), where a single +//! comparison scans millions of bytes. On a 47.5 MB chr21 FASTA this +//! pushes the cost per merge step from 13 ns to 222 ns — a 16x penalty +//! that is entirely scan time. +//! +//! This module attacks both. It sorts by a packed fixed-depth key first +//! (killing the step count), then resolves the remainder by **prefix +//! doubling** on ranks (killing the scan length: after the seed, no +//! comparison ever reads the text again, so a megabyte-long run of `N` costs +//! exactly as much as random DNA). +//! +//! ## The algorithm +//! +//! 1. **Pack.** Find the maximum symbol and choose the smallest field width +//! in `{1, 2, 4, 8}` bits that can hold it, so `k = 64 / bits` symbols fit +//! in one `u64` key. DNA over `{0,1,2,3}` gets 2-bit fields and therefore +//! resolves **32 symbols per key** rather than the 8 a raw byte key would. +//! 2. **Seed.** Sort `(key, position)`. This is a full sort of the suffixes +//! by their first `k` symbols, and it touches the text only in one +//! sequential pass. +//! 3. **Double.** Suffixes still tied after depth `d` are ordered by the pair +//! `(rank_d(p), rank_d(p + d))`, which resolves them to depth `2d`. Repeat +//! until every group is a singleton. Each round reads only the rank array. +//! +//! ## Ordering convention +//! +//! Keys are big-endian in the field sense (the first symbol occupies the most +//! significant field) and short suffixes are zero-padded. Since `0` is the +//! minimum of `u8`, a padding field can never exceed a real symbol's field, +//! so a padded key compares less-or-equal to any key it shares a prefix with. +//! That is exactly the crate's "shorter suffix is smaller" convention. A real +//! `0` symbol is indistinguishable from padding *in the key*, which can only +//! make two suffixes tie — never invert them — and ties are resolved by the +//! doubling rounds, which use the true remaining length via the end-of-text +//! sentinel. So `A = 0` DNA encodings and STAR's `0..5` codes are both safe. + +use crate::Index; +use crate::ext_mem::profile_log; +use rayon::prelude::*; +use std::time::Instant; + +/// Field width in bits and the number of symbols that fit in a `u64` key. +/// +/// Restricted to divisors of 64 so a key is an exact number of whole fields +/// and no symbol ever straddles the key boundary. +fn pack_params(max_sym: u8) -> (u32, usize) { + let bits: u32 = match max_sym { + 0..=1 => 1, + 2..=3 => 2, + 4..=15 => 4, + _ => 8, + }; + (bits, 64 / bits as usize) +} + +/// Pack the `k` symbols at `text[p..]` into one order-preserving `u64`, +/// zero-padding past the end of the text. +#[inline] +fn key_at(text: &[u8], p: usize, bits: u32, k: usize) -> u64 { + if bits == 8 { + // Whole-byte fields: this is just a big-endian load, and the common + // case (p + 8 <= n) is a single unaligned u64 read plus a bswap. + let mut buf = [0u8; 8]; + let end = (p + 8).min(text.len()); + buf[..end - p].copy_from_slice(&text[p..end]); + return u64::from_be_bytes(buf); + } + let end = (p + k).min(text.len()); + let mut key: u64 = 0; + for &s in &text[p..end] { + key = (key << bits) | s as u64; + } + // Shift the packed prefix up so the missing trailing fields read as zero. + key << (bits as usize * (k - (end - p))) +} + +/// Bits of the key used for the MSD counting-sort pass, and the resulting +/// bucket count. +/// +/// 2048 buckets keeps the write-combining state at 2048 × 2 streams × 128 B +/// ≈ 512 KB, which stays inside a core's private cache. Going to 16 bits +/// would need 16 MB of open write lines and thrashes the TLB instead. +const RADIX_BITS: u32 = 11; +const RADIX_BUCKETS: usize = 1 << RADIX_BITS; + +/// Sort every suffix position by `(packed key, visible length)`, returning the +/// sorted keys alongside the sorted positions. +/// +/// This is an MSD counting sort rather than a comparison sort, for three +/// reasons that all matter at genome scale: +/// +/// * The source is never materialised. Keys are recomputed from `text` in +/// both the histogram and the scatter pass, which is a sequential read of +/// the text instead of a random read of an `n`-element key array. +/// * Peak memory is the two destination buffers only, 12 bytes per position +/// with `I = u32`, against the 16 a `(u64, u32, I)` record costs. +/// * The top-level partition is a counting pass, so it parallelises evenly. +/// A parallel comparison sort's first partitioning steps are close to +/// serial, which is exactly where a `n log n` sort loses on many cores. +fn seed_sort( + text: &[u8], + bits: u32, + k: usize, + visible_len: &(dyn Fn(usize) -> usize + Sync), +) -> (Vec, Vec) { + let n = text.len(); + let bucket_of = |key: u64| -> usize { (key >> (64 - RADIX_BITS)) as usize }; + + // Chunk the position range so each worker builds a private histogram. + let n_chunks = (rayon::current_num_threads() * 4).clamp(1, 1024); + let chunk_len = n.div_ceil(n_chunks); + let bounds: Vec<(usize, usize)> = (0..n) + .step_by(chunk_len) + .map(|s| (s, (s + chunk_len).min(n))) + .collect(); + + // Pass 1: per-chunk histograms over the top `RADIX_BITS` of each key. + let histograms: Vec> = bounds + .par_iter() + .map(|&(start, end)| { + let mut counts = vec![0u32; RADIX_BUCKETS]; + for p in start..end { + counts[bucket_of(key_at(text, p, bits, k))] += 1; + } + counts + }) + .collect(); + + // Exclusive prefix sum, bucket-major then chunk-minor, so every (chunk, + // bucket) pair gets a disjoint destination range and the buckets come out + // in ascending key order. + let mut offsets = vec![0usize; bounds.len() * RADIX_BUCKETS]; + let mut bucket_start = vec![0usize; RADIX_BUCKETS + 1]; + { + let mut running = 0usize; + for b in 0..RADIX_BUCKETS { + bucket_start[b] = running; + for (c, hist) in histograms.iter().enumerate() { + offsets[c * RADIX_BUCKETS + b] = running; + running += hist[b] as usize; + } + } + bucket_start[RADIX_BUCKETS] = running; + debug_assert_eq!(running, n); + } + + // Pass 2: scatter. Each chunk owns a disjoint slice of every bucket, so + // the writes never collide even though they are not contiguous. + let mut keys: Vec = vec![0; n]; + let mut sa: Vec = vec![I::zero(); n]; + { + let key_out = Scatter::new(&mut keys); + let sa_out = Scatter::new(&mut sa); + bounds + .par_iter() + .enumerate() + .for_each(|(c, &(start, end))| { + let mut cursor: Vec = + offsets[c * RADIX_BUCKETS..(c + 1) * RADIX_BUCKETS].to_vec(); + for p in start..end { + let key = key_at(text, p, bits, k); + let slot = &mut cursor[bucket_of(key)]; + // SAFETY: the prefix sum gives this (chunk, bucket) pair a + // range of exactly its own histogram count, and the cursor + // never leaves it, so no other thread writes this index. + unsafe { + key_out.set(*slot, key); + sa_out.set(*slot, I::from_usize(p)); + } + *slot += 1; + } + }); + } + + // Pass 3: order within each bucket. Buckets share their top `RADIX_BITS`, + // so what remains is the low bits of the key and then the visible-length + // tie-break. Buckets are contiguous and independent. + let mut rest: &mut [u64] = &mut keys; + let mut rest_sa: &mut [I] = &mut sa; + let mut slices: Vec<(&mut [u64], &mut [I])> = Vec::with_capacity(RADIX_BUCKETS); + for b in 0..RADIX_BUCKETS { + let len = bucket_start[b + 1] - bucket_start[b]; + let (kb, kt) = rest.split_at_mut(len); + let (sb, st) = rest_sa.split_at_mut(len); + slices.push((kb, sb)); + rest = kt; + rest_sa = st; + } + slices.into_par_iter().for_each(|(kb, sb)| { + if kb.len() < 2 { + return; + } + let mut pairs: Vec<(u64, I)> = kb.iter().copied().zip(sb.iter().copied()).collect(); + pairs.sort_unstable_by(|a, b| { + a.0.cmp(&b.0) + .then_with(|| visible_len(a.1.to_usize()).cmp(&visible_len(b.1.to_usize()))) + }); + for (i, &(key, pos)) in pairs.iter().enumerate() { + kb[i] = key; + sb[i] = pos; + } + }); + + (keys, sa) +} + +/// Build the standard lexicographic suffix array of `text` by radix-seeded +/// prefix doubling. +/// +/// The caller is responsible for the guards: `text` must be the whole, +/// non-segmented text, the comparator must be plain lexicographic with +/// shorter-is-smaller, and there must be no `max_context` bound. See +/// [`crate::sample_sort::build_in_memory_with`] for where those are checked. +pub(crate) fn build_sa(text: &[u8]) -> Vec { + let n = text.len(); + if n == 0 { + return Vec::new(); + } + if n == 1 { + return vec![I::zero()]; + } + + let t0 = Instant::now(); + let max_sym = text.par_iter().copied().max().unwrap_or(0); + let (bits, k) = pack_params(max_sym); + + // ---- Seed: sort by the first `k` symbols, then by visible length. ---- + // + // The second component is `min(n - p, k)`, and it is load-bearing rather + // than cosmetic. Zero-padding makes a suffix shorter than `k` share a key + // with any suffix whose symbols continue with zeros, and `0` is a real + // symbol in every DNA encoding. Ordering those by visible length puts the + // proper prefix first, which is the shorter-is-smaller convention. For + // suffixes at least `k` long the component is `k` for all of them, so it + // never separates suffixes that the doubling rounds still need to see as + // tied. Without it, `[0, 0]` leaves positions 0 and 1 permanently tied + // and the doubling loop cannot terminate. + // Only the last `k - 1` positions can have a visible length below `k`, so + // the tie-break is a function of the position alone and never has to be + // stored alongside the key. + let visible_len = |p: usize| -> usize { (n - p).min(k) }; + + profile_log(&format!( + "radix setup {:.3}s", + t0.elapsed().as_secs_f64() + )); + let t1 = Instant::now(); + let (keys, mut sa) = seed_sort::(text, bits, k, &visible_len); + profile_log(&format!( + "radix seed sort {:.3}s", + t1.elapsed().as_secs_f64() + )); + let t2 = Instant::now(); + + // Two seeded entries tie exactly when key and visible length both match. + let seed_eq = |a: usize, b: usize| -> bool { + keys[a] == keys[b] && visible_len(sa[a].to_usize()) == visible_len(sa[b].to_usize()) + }; + + // `rank[p]` is the index in `sa` of the first element of `p`'s group, so + // two suffixes tie at the current depth exactly when their ranks match, + // and rank order is the current partial order. + let mut rank: Vec = vec![I::zero(); n]; + // Non-singleton `sa` ranges, the only ones any later round touches. + // + // Each index decides for itself whether it starts a group; the index that + // does then owns the whole group, walks it to find the end, and writes its + // members' ranks. Every group has exactly one owner and groups partition + // `0..n`, so the scattered writes never collide. `collect` on an indexed + // parallel iterator preserves order, so `groups` comes out sorted. + let ranks = Scatter::new(&mut rank); + let groups: Vec<(usize, usize)> = (0..n) + .into_par_iter() + .filter_map(|h| { + if h > 0 && seed_eq(h - 1, h) { + return None; + } + let mut e = h + 1; + while e < n && seed_eq(e, h) { + e += 1; + } + let g = I::from_usize(h); + for entry in &sa[h..e] { + // SAFETY: `sa` is a permutation of `0..n`, and this thread + // owns the whole group `h..e`, so `entry` is a distinct index + // no other thread writes. + unsafe { ranks.set(entry.to_usize(), g) }; + } + (e - h > 1).then_some((h, e)) + }) + .collect(); + let mut groups = groups; + drop(keys); + profile_log(&format!( + "radix grouping {:.3}s", + t2.elapsed().as_secs_f64() + )); + let t3 = Instant::now(); + + // ---- Double: (rank_d(p), rank_d(p + d)) resolves to depth 2d. ---- + let mut depth = k; + // Scratch for the new rank of each `sa` slot, so a round's reads of + // `rank` never observe that same round's writes. + let mut next_rank: Vec = vec![I::zero(); n]; + + while !groups.is_empty() { + // Phase A: sort each tied group by the successor rank, and record the + // ranks it should get. Groups are disjoint `sa` ranges, so this is + // data-parallel with no synchronisation. + let sub: Vec> = split_disjoint(&mut sa, &mut next_rank, &groups) + .into_par_iter() + .zip(groups.par_iter()) + .map(|((sa_g, nr_g), &(start, _))| { + let succ = |p: usize| -> u64 { + // End-of-text sorts first: the shorter suffix is smaller. + match p.checked_add(depth) { + Some(q) if q < n => rank[q].to_usize() as u64 + 1, + _ => 0, + } + }; + // Materialise the successor ranks once. `sort_unstable_by_key` + // re-evaluates its key function O(len log len) times, and each + // evaluation is a random probe into `rank`; paying for it once + // per element turns the sort's memory traffic sequential. + let mut keyed: Vec<(u64, I)> = + sa_g.iter().map(|&e| (succ(e.to_usize()), e)).collect(); + keyed.sort_unstable(); + + let mut fresh = Vec::new(); + let mut i = 0; + while i < keyed.len() { + let key = keyed[i].0; + let mut j = i + 1; + while j < keyed.len() && keyed[j].0 == key { + j += 1; + } + let g = I::from_usize(start + i); + for slot in &mut nr_g[i..j] { + *slot = g; + } + if j - i > 1 { + fresh.push((start + i, start + j)); + } + i = j; + } + for (slot, &(_, e)) in sa_g.iter_mut().zip(keyed.iter()) { + *slot = e; + } + fresh + }) + .collect(); + + // Phase B: publish the new ranks, now that every read is done. + // Groups are disjoint and `sa` is a permutation, so each `rank` slot + // is written by exactly one group. + let ranks = Scatter::new(&mut rank); + groups.par_iter().for_each(|&(start, end)| { + for i in start..end { + // SAFETY: `sa[start..end]` are distinct positions owned solely + // by this group, and the groups partition their index range. + unsafe { ranks.set(sa[i].to_usize(), next_rank[i]) }; + } + }); + + let before: usize = groups.iter().map(|&(s, e)| e - s).sum(); + groups = sub.into_iter().flatten().collect(); + let after: usize = groups.iter().map(|&(s, e)| e - s).sum(); + + // A doubling round can only ever refine, so `after <= before`. If a + // round refines nothing at all the text has a run longer than the + // whole remaining depth budget; doubling still terminates because + // `depth` grows geometrically and every suffix eventually runs off + // the end of the text, which the sentinel orders. Guard against + // overflow rather than against non-progress. + debug_assert!(after <= before); + match depth.checked_mul(2) { + Some(d) if d <= n.saturating_mul(2) => depth = d, + _ => { + debug_assert!(groups.is_empty(), "doubling exhausted with ties left"); + break; + } + } + } + + profile_log(&format!( + "radix doubling {:.3}s", + t3.elapsed().as_secs_f64() + )); + sa +} + +/// Write access to disjoint slots of one slice from several rayon threads. +/// +/// Both users here scatter through a permutation: the target index is +/// `sa[i]`, not `i`, so the writes cannot be expressed as disjoint sub-slices +/// and `split_at_mut` does not apply. What makes them safe is that `sa` is a +/// permutation and the ranges being processed partition its index space, so +/// every slot is written exactly once across all threads. +struct Scatter { + ptr: *mut T, + len: usize, +} + +// SAFETY: `Scatter` hands out writes only through `set`, whose contract is +// that no two calls target the same index. Under that contract there is no +// aliasing between threads, so the pointer is safe to share. +unsafe impl Send for Scatter {} +unsafe impl Sync for Scatter {} + +impl Scatter { + fn new(slice: &mut [T]) -> Self { + Self { + ptr: slice.as_mut_ptr(), + len: slice.len(), + } + } + + /// Write `value` at `index`. + /// + /// # Safety + /// + /// No two concurrent calls may pass the same `index`, and the borrow the + /// `Scatter` was built from must still be live. + #[inline] + unsafe fn set(&self, index: usize, value: T) { + debug_assert!(index < self.len); + unsafe { self.ptr.add(index).write(value) }; + } +} + +/// Borrow each `(start, end)` range of `sa` and `next_rank` mutably and +/// simultaneously. The ranges come from a scan of `sa` so they are sorted and +/// non-overlapping, which is what makes the repeated `split_at_mut` sound. +fn split_disjoint<'a, I: Index>( + sa: &'a mut [I], + next_rank: &'a mut [I], + groups: &[(usize, usize)], +) -> Vec<(&'a mut [I], &'a mut [I])> { + let mut out = Vec::with_capacity(groups.len()); + let mut sa_rest = sa; + let mut nr_rest = next_rank; + let mut consumed = 0usize; + for &(start, end) in groups { + debug_assert!( + start >= consumed, + "group ranges must be sorted and disjoint" + ); + let (_, sa_tail) = sa_rest.split_at_mut(start - consumed); + let (_, nr_tail) = nr_rest.split_at_mut(start - consumed); + let (sa_g, sa_tail) = sa_tail.split_at_mut(end - start); + let (nr_g, nr_tail) = nr_tail.split_at_mut(end - start); + out.push((sa_g, nr_g)); + sa_rest = sa_tail; + nr_rest = nr_tail; + consumed = end; + } + out +} + +#[cfg(test)] +mod tests { + use super::*; + + fn brute(text: &[u8]) -> Vec { + let mut sa: Vec = (0..text.len() as u32).collect(); + sa.sort_by(|&a, &b| text[a as usize..].cmp(&text[b as usize..])); + sa + } + + fn check(text: &[u8]) { + let got: Vec = build_sa(text); + assert_eq!(got, brute(text), "mismatch on {text:?}"); + } + + #[test] + fn fixtures() { + check(b""); + check(b"a"); + check(b"banana"); + check(b"mississippi"); + check(b"abracadabra"); + } + + #[test] + fn pack_params_covers_every_width() { + assert_eq!(pack_params(0), (1, 64)); + assert_eq!(pack_params(1), (1, 64)); + assert_eq!(pack_params(3), (2, 32)); + assert_eq!(pack_params(4), (4, 16)); + assert_eq!(pack_params(15), (4, 16)); + assert_eq!(pack_params(16), (8, 8)); + assert_eq!(pack_params(255), (8, 8)); + } + + /// Texts whose symbols include a real `0`, so padding and a genuine + /// minimum symbol are indistinguishable in the packed key. This is the + /// case DNA encodings hit (`A = 0`) and the one the ordering argument in + /// the module docs turns on. + #[test] + fn real_zero_symbol_is_not_confused_with_padding() { + check(&[0]); + check(&[0, 0, 0, 0, 0, 0, 0, 0, 0, 0]); + check(&[1, 2, 0, 0, 0, 0, 0, 0, 0, 0]); + check(&[3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]); + check(&[0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0]); + let mut t: Vec = (0..200).map(|i| (i % 4) as u8).collect(); + t.extend(std::iter::repeat_n(0u8, 100)); + check(&t); + } + + /// Every text of length <= 10 over a binary alphabet, plus every text of + /// length <= 6 over a ternary one. Total coverage of the padding and + /// end-of-text logic at the sizes where exhaustive checking is free. + #[test] + fn exhaustive_small_alphabets() { + for n in 0..=10u32 { + for mask in 0..(1u32 << n) { + let t: Vec = (0..n).map(|i| ((mask >> i) & 1) as u8).collect(); + check(&t); + } + } + for n in 0..=6u32 { + let total = 3u32.pow(n); + for mut code in 0..total { + let mut t = Vec::with_capacity(n as usize); + for _ in 0..n { + t.push((code % 3) as u8); + code /= 3; + } + check(&t); + } + } + } + + #[test] + fn random_across_alphabet_widths() { + use rand::{RngExt, SeedableRng}; + let mut rng = rand::rngs::StdRng::seed_from_u64(0xADD1); + for &sigma in &[2u8, 3, 4, 6, 16, 17, 255] { + for &n in &[ + 2usize, 7, 31, 32, 33, 63, 64, 65, 127, 128, 129, 1000, 20_000, + ] { + let t: Vec = (0..n).map(|_| rng.random_range(0..sigma)).collect(); + check(&t); + } + } + } + + /// Long runs and periodic text are the inputs that make the merge kernel + /// quadratic. Doubling must handle them and must terminate. + #[test] + fn long_runs_and_periodic_text() { + check(&vec![0u8; 5000]); + check(&vec![7u8; 5000]); + check(&(0..5000).map(|i| (i % 2) as u8).collect::>()); + check(&(0..5000).map(|i| (i % 61) as u8).collect::>()); + // A long run flanked by noise: the shape of a poly-N genome block. + let mut t: Vec = (0..500).map(|i| (i % 4) as u8).collect(); + t.extend(std::iter::repeat_n(4u8, 4000)); + t.extend((0..500).map(|i| (i % 4) as u8)); + check(&t); + // Wrapped-FASTA shape: 60 `N`s then a newline, repeated. + let mut fasta: Vec = Vec::new(); + for _ in 0..100 { + fasta.extend(std::iter::repeat_n(b'N', 60)); + fasta.push(b'\n'); + } + check(&fasta); + } + + #[test] + fn u64_index_matches_u32_index() { + use rand::{RngExt, SeedableRng}; + let mut rng = rand::rngs::StdRng::seed_from_u64(0xD0AB); + let t: Vec = (0..5000).map(|_| rng.random_range(0..4u8)).collect(); + let a: Vec = build_sa(&t); + let b: Vec = build_sa(&t); + assert_eq!(a.len(), b.len()); + assert!(a.iter().zip(&b).all(|(&x, &y)| x as u64 == y)); + } +} diff --git a/src/sample_sort.rs b/src/sample_sort.rs index b3343aa..353fa28 100644 --- a/src/sample_sort.rs +++ b/src/sample_sort.rs @@ -35,6 +35,7 @@ use crate::lcp::{LcpDispatch, Symbol}; use crate::lcp_memo::GeometricMemo; use crate::limits::{LimitProvider, PlainText}; use rayon::join; +use rayon::prelude::*; /// How many merge steps ahead the text prefetch runs. Large enough to cover a /// DRAM round trip at the merge's step rate, small enough that the prefetched @@ -121,11 +122,117 @@ where I: Index, L: LimitProvider, { + if let Some(sa) = try_doubling_fast_path::(text, lp, opts) { + return sa; + } let n = text.len(); let positions: Vec = (0..n).map(I::from_usize).collect(); build_in_memory_for_positions_with(text, positions, lp, opts) } +/// Route a whole-text byte build through [`crate::radix`]'s radix-seeded +/// prefix doubling, which is dramatically faster on real genomic input, or +/// return `None` to fall back to the CaPS-SA merge kernel. +/// +/// Every condition below is a soundness requirement, not a heuristic. The +/// doubling path implements exactly one comparator — plain lexicographic over +/// bytes with shorter-is-smaller and no context bound — so anything that can +/// change the comparator has to decline. +/// +/// * `max_context` must be unbounded. With a finite bound the merge's +/// comparator stops being lexicographic: once a scan hits the cap it falls +/// through to [`LimitProvider::boundary_order`], which compares *lengths*. +/// * `lp` must report [`plain_lex_len`][LimitProvider::plain_lex_len]. That +/// rules out `SegmentedText`, whose LCP scans stop at segment boundaries, +/// and any custom `boundary_order` such as STAR's spacer-as-largest. +/// * `S` must be exactly `u8`. Wider symbols are excluded because packing +/// them into an order-preserving key is endianness-dependent: for `u16` on +/// a little-endian host, `0x0100 > 0x0001` as values but their byte views +/// compare the other way. The rest of the crate is immune to this only +/// because [`LcpDispatch`] resolves *equality* over bytes and recovers +/// ordering through `S: Ord`. +fn try_doubling_fast_path(text: &[S], lp: &L, opts: &Opts) -> Option> +where + S: Symbol, + I: Index, + L: LimitProvider, +{ + if opts.max_context != usize::MAX { + return None; + } + if lp.plain_lex_len() != Some(text.len()) { + return None; + } + if std::any::TypeId::of::() != std::any::TypeId::of::() { + return None; + } + // SAFETY: `S` is `u8` (just checked by `TypeId`, and `Symbol: 'static` + // so the comparison is exact), hence `&[S]` and `&[u8]` have identical + // layout, length and validity. + let bytes: &[u8] = + unsafe { std::slice::from_raw_parts(text.as_ptr() as *const u8, text.len()) }; + Some(crate::radix::build_sa(bytes)) +} + +/// Sort a *subset* of positions by building the full suffix array with the +/// doubling path and then keeping only the requested positions. +/// +/// Doubling cannot be restricted to a subset directly: a round compares +/// `rank[p + d]`, and that successor is generally not in the subset, so ranks +/// have to be defined for every position in the text. Building the whole array +/// and filtering it sidesteps that, and the filter is a single `O(n)` pass +/// because the full SA is already in the right order. +/// +/// Worth it when the subset is a decent fraction of the text, which is the +/// case this API exists for (STAR-style indexing keeps every ACGT position and +/// drops only spacers). For a small subset, `O(n)` to build the full array +/// would dwarf the `O(m log m)` the merge kernel needs, so below one eighth of +/// the text this declines and the merge kernel runs. That ratio is a +/// performance heuristic, unlike the guards in [`try_doubling_fast_path`], +/// which are correctness conditions. +/// +/// Also declines on duplicate or out-of-range positions, which a +/// membership filter cannot reproduce faithfully. +fn try_doubling_subset(text: &[S], positions: &[I], lp: &L, opts: &Opts) -> Option> +where + S: Symbol, + I: Index, + L: LimitProvider, +{ + let n = text.len(); + let m = positions.len(); + if m == 0 || m.checked_mul(8)? < n { + return None; + } + if opts.max_context != usize::MAX + || lp.plain_lex_len() != Some(n) + || std::any::TypeId::of::() != std::any::TypeId::of::() + { + return None; + } + + let mut wanted = vec![false; n]; + for p in positions { + let p = p.to_usize(); + // Out of range, or the same position twice: a membership filter emits + // each position at most once, so it cannot reproduce either faithfully. + if p >= n || wanted[p] { + return None; + } + wanted[p] = true; + } + + // SAFETY: `S` is `u8`, so `&[S]` and `&[u8]` have identical layout. + let bytes: &[u8] = unsafe { std::slice::from_raw_parts(text.as_ptr() as *const u8, n) }; + let full: Vec = crate::radix::build_sa(bytes); + let kept: Vec = full + .into_par_iter() + .filter(|p| wanted[p.to_usize()]) + .collect(); + debug_assert_eq!(kept.len(), m); + Some(kept) +} + /// Sort the caller-supplied `positions` by the lexicographic order of /// their suffixes in `text`. Returns the positions reordered so that /// `text[output[i]..]` is the i-th smallest suffix among the input set. @@ -176,6 +283,10 @@ where I: Index, L: LimitProvider, { + if let Some(sa) = try_doubling_subset::(text, &positions, lp, opts) { + return sa; + } + let n = positions.len(); if n == 0 { return Vec::new(); @@ -854,6 +965,45 @@ mod tests { assert_eq!(got, want); } + /// Duplicated positions must survive: the output is a permutation of the + /// *input* multiset, which a membership filter cannot reproduce, so the + /// subset fast path has to decline and let the merge kernel run. + #[test] + fn for_positions_with_duplicates_keeps_multiplicity() { + let text = b"mississippi"; + let positions: Vec = vec![0, 1, 1, 4, 4, 4, 7]; + let mut want = positions.clone(); + want.sort_by(|&a, &b| text[a as usize..].cmp(&text[b as usize..])); + let got = build_in_memory_for_positions(text, positions); + assert_eq!(got, want); + } + + /// A subset far smaller than the text takes the merge kernel, since + /// building the whole suffix array to throw nearly all of it away would + /// cost more than sorting the subset directly. Correctness is identical + /// either way; this pins the behaviour. + #[test] + fn for_positions_tiny_subset_of_large_text() { + use rand::{RngExt, SeedableRng}; + let mut rng = rand::rngs::StdRng::seed_from_u64(0x5AB0); + let text: Vec = (0..20_000).map(|_| rng.random_range(0..4u8)).collect(); + let positions: Vec = (0..20_000u32).step_by(500).collect(); + let mut want = positions.clone(); + want.sort_by(|&a, &b| text[a as usize..].cmp(&text[b as usize..])); + let got = build_in_memory_for_positions(&text, positions); + assert_eq!(got, want); + } + + /// Positions out of range are the caller's error, but the subset fast path + /// must not turn them into a silently wrong answer or an unsafe index. + #[test] + #[should_panic] + fn for_positions_out_of_range_still_panics() { + let text = b"banana"; + let positions: Vec = vec![0, 1, 99]; + let _ = build_in_memory_for_positions(text, positions); + } + #[test] fn for_positions_random_subsets() { use rand::{RngExt, SeedableRng};