diff --git a/README.md b/README.md index 5d9a382..7f57881 100644 --- a/README.md +++ b/README.md @@ -16,7 +16,62 @@ 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, ASCII `ACGT`, 80 MB | 6.08 s | **0.61 s** | 28.1 s | 5.0 s | +| same, pre-coded to `0..3`, 80 MB | 6.08 s | **0.57 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. + +The two DNA rows above land in the same place because the alphabet is +ranked to a dense code range before packing. Without that step the ASCII +row would use 8-bit fields — its largest byte is `'T'` (84) even though +it has four symbols — fitting 8 symbols per key instead of 32, and would +take 1.40 s rather than 0.61 s. Keys are then built by a SWAR gather +over the ranked text, so eight symbols cost three shift-or-mask pairs +instead of eight dependent shift-or-lookup steps. + +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 +160,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..3b35d45 --- /dev/null +++ b/src/radix.rs @@ -0,0 +1,911 @@ +//! 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.** Rank the bytes that actually occur onto a dense code range, +//! then choose the smallest field width in `{1, 2, 4, 8}` bits that holds +//! the alphabet, so `k = 64 / bits` symbols fit in one `u64` key. Ranking +//! matters: raw FASTA uses six symbols but its largest byte is `'T'` (84), +//! so packing raw bytes would force 8-bit fields and 8 symbols per key, +//! against 16 after ranking. DNA over `{0,1,2,3}` gets 2-bit fields and +//! **32 symbols per key**. +//! 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::sync::atomic::{AtomicU64, Ordering}; +use std::time::Instant; + +/// An order-preserving remap of the bytes that actually occur in a text onto +/// a dense code range, plus the resulting key geometry. +/// +/// The field width is driven by how many *distinct* symbols a text uses, not +/// by the largest byte value in it, and the difference is not academic. A raw +/// FASTA uses six symbols, but the largest is `'T'` (84), so packing raw bytes +/// forces 8-bit fields and fits only 8 symbols per key. Ranking those six +/// bytes to `0..6` gives 4-bit fields and 16 symbols per key, which halves the +/// number of doubling rounds needed downstream. The DNA-coded input is already +/// dense, so it is unaffected. +/// +/// The map is monotone by construction, since codes are assigned in ascending +/// byte order. That is what keeps a packed key order-preserving: `key_a < +/// key_b` still implies `suffix_a < suffix_b`, and the zero-padding argument +/// carries over because code `0` remains the minimum. +pub(crate) struct Packer { + /// The text with every byte replaced by its code, when the identity map + /// does not already do that. Materialising it once removes a dependent + /// table load from the packing loop, which is otherwise the chain that + /// sets the cost of building a key. `None` when the text is already dense + /// (a `0..3` DNA encoding, say), so the common pre-coded input pays no + /// extra memory. + ranked: Option>, + /// Bits per packed field. + bits: u32, + /// Symbols per `u64` key. + k: usize, +} + +impl Packer { + /// Build the map for `text`. + fn new(text: &[u8]) -> Self { + // Which bytes occur? One parallel pass, folded into a 256-entry set. + let present = text + .par_chunks(1 << 16) + .map(|c| { + let mut seen = [false; 256]; + for &b in c { + seen[b as usize] = true; + } + seen + }) + .reduce( + || [false; 256], + |mut a, b| { + for i in 0..256 { + a[i] |= b[i]; + } + a + }, + ); + + let mut code = [0u8; 256]; + let mut next = 0u16; + let mut identity = true; + for (b, &seen) in present.iter().enumerate() { + if seen { + code[b] = next as u8; + identity &= next as usize == b; + next += 1; + } + } + let bits: u32 = match next.saturating_sub(1) { + 0..=1 => 1, + 2..=3 => 2, + 4..=15 => 4, + _ => 8, + }; + // An identity map, or 8-bit fields where the code never changes the + // packed value's order, both let the original text be read directly. + let ranked = if identity { + None + } else { + let mut out = vec![0u8; text.len()]; + out.par_chunks_mut(1 << 16) + .zip(text.par_chunks(1 << 16)) + .for_each(|(dst, src)| { + for (d, &s) in dst.iter_mut().zip(src) { + *d = code[s as usize]; + } + }); + Some(out) + }; + Self { + ranked, + bits, + k: 64 / bits as usize, + } + } + + /// Bits per packed field. Used by the key-geometry tests. + #[cfg(test)] + #[inline] + pub(crate) fn bits(&self) -> u32 { + self.bits + } + + #[inline] + pub(crate) fn k(&self) -> usize { + self.k + } + + /// Gather the low `bits` of each of 8 ranked bytes into one contiguous + /// field, most-significant byte first. + /// + /// A binary-tree SWAR shuffle: each step folds neighbouring fields + /// together and halves the stride, so eight symbols cost three + /// shift-or-mask pairs instead of eight dependent shift-or steps. The + /// input is a big-endian load, so the text's first byte lands in the + /// result's most significant field, which is the order the key needs. + #[inline(always)] + fn gather8(v: u64, bits: u32) -> u64 { + match bits { + 1 => { + let mut x = v & 0x0101_0101_0101_0101; + x = (x | (x >> 7)) & 0x0003_0003_0003_0003; + x = (x | (x >> 14)) & 0x0000_000F_0000_000F; + (x | (x >> 28)) & 0xFF + } + 2 => { + let mut x = v & 0x0303_0303_0303_0303; + x = (x | (x >> 6)) & 0x000F_000F_000F_000F; + x = (x | (x >> 12)) & 0x0000_00FF_0000_00FF; + (x | (x >> 24)) & 0xFFFF + } + 4 => { + let mut x = v & 0x0F0F_0F0F_0F0F_0F0F; + x = (x | (x >> 4)) & 0x00FF_00FF_00FF_00FF; + x = (x | (x >> 8)) & 0x0000_FFFF_0000_FFFF; + (x | (x >> 16)) & 0xFFFF_FFFF + } + _ => v, + } + } + + /// Pack the `k` symbols at `text[p..]` into one order-preserving `u64`, + /// zero-padding past the end of the text. + #[inline] + pub(crate) fn key_at(&self, text: &[u8], p: usize) -> u64 { + let src = self.ranked.as_deref().unwrap_or(text); + let n = src.len(); + + // Fast path: a whole key's worth of symbols is available, so it is + // `k / 8` big-endian loads and their gathers, with no bounds fuss. + if p + self.k <= n { + return self + .fold(|i| u64::from_be_bytes(src[p + 8 * i..p + 8 * i + 8].try_into().unwrap())); + } + + // Tail: fewer than `k` symbols remain. Pad with zero codes, which are + // the alphabet's minimum, matching shorter-is-smaller. + let mut buf = [0u8; 64]; + buf[..n - p].copy_from_slice(&src[p..n]); + self.fold(|i| u64::from_be_bytes(buf[8 * i..8 * i + 8].try_into().unwrap())) + } + + /// Concatenate the gathers of the `k / 8` words produced by `word`. + /// + /// The first group is assigned rather than shifted in. With 8-bit fields + /// there is exactly one group and `8 * bits` is 64, which is not a legal + /// shift distance for `u64`; release builds mask it to 0 and happen to + /// give the right answer, debug builds panic. Assigning avoids relying on + /// either behaviour. + #[inline(always)] + fn fold(&self, word: impl Fn(usize) -> u64) -> u64 { + let shift = 8 * self.bits; + let mut key = 0u64; + for i in 0..self.k / 8 { + let g = Self::gather8(word(i), self.bits); + key = if i == 0 { g } else { (key << shift) | g }; + } + key + } +} + +/// Largest tied group whose key vector is built on the stack. Groups average +/// about four elements, so nearly every one avoids the allocator entirely. +const DOUBLING_STACK_GROUP: usize = 32; + +/// 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 positions and a bit per slot marking where a tied group starts. +/// +/// The key array does not come back. Grouping only needs to know where one +/// group ends and the next begins, which is one bit per slot rather than the +/// eight bytes a key costs, and dropping the keys here rather than after the +/// grouping pass is what keeps the doubling path's peak below the merge +/// kernel's. See [`build_sa`] for the accounting. +/// +/// 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], + packer: &Packer, + 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(packer.key_at(text, p))] += 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 positions into their buckets. Each chunk owns a + // disjoint slice of every bucket, so the writes never collide even though + // they are not contiguous. + // + // The key is computed here to pick the bucket and then thrown away. An + // `n`-entry key array would be the largest allocation in the whole build, + // 8 bytes per position against the 4 that `sa` costs, and it would still + // be resident when the doubling rounds need their own arrays. Pass 3 + // recomputes the keys of one bucket at a time instead, which is one extra + // key per position spread across the workers. + let mut sa: Vec = vec![I::zero(); n]; + { + 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 slot = &mut cursor[bucket_of(packer.key_at(text, p))]; + // 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 { sa_out.set(*slot, I::from_usize(p)) }; + *slot += 1; + } + }); + } + + // Pass 3: order within each bucket, and record where tied groups start. + // + // 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, and two positions in different buckets differ in the + // key by construction, so a bucket's first slot always starts a group and + // the group-start bits can be filled in here, per bucket, while that + // bucket's keys exist. + // Atomic words: two buckets can share the word their boundary bits live + // in, so the bits are set with relaxed fetch-or. Reads are relaxed loads + // and cost nothing once the fill is done. + let starts: Vec = (0..n.div_ceil(64)).map(|_| AtomicU64::new(0)).collect(); + let mut rest_sa: &mut [I] = &mut sa; + let mut slices: Vec<(usize, &mut [I])> = Vec::with_capacity(RADIX_BUCKETS); + for b in 0..RADIX_BUCKETS { + let len = bucket_start[b + 1] - bucket_start[b]; + let (sb, st) = rest_sa.split_at_mut(len); + slices.push((bucket_start[b], sb)); + rest_sa = st; + } + { + let start_bits = &starts; + let set_bit = |h: usize| { + start_bits[h / 64].fetch_or(1 << (h % 64), Ordering::Relaxed); + }; + slices.into_par_iter().for_each(|(base, sb)| { + if sb.is_empty() { + return; + } + set_bit(base); + if sb.len() > 1 { + let mut pairs: Vec<(u64, I)> = sb + .iter() + .map(|&p| (packer.key_at(text, p.to_usize()), p)) + .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, &(_, pos)) in pairs.iter().enumerate() { + sb[i] = pos; + } + for i in 1..pairs.len() { + let (ka, pa) = pairs[i - 1]; + let (kb, pb) = pairs[i]; + if ka != kb || visible_len(pa.to_usize()) != visible_len(pb.to_usize()) { + set_bit(base + i); + } + } + } + }); + } + + (starts, 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 packer = Packer::new(text); + let k = packer.k(); + + // ---- 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 (starts, mut sa) = seed_sort::(text, &packer, &visible_len); + profile_log(&format!( + "radix seed sort {:.3}s", + t1.elapsed().as_secs_f64() + )); + let t2 = Instant::now(); + + // Slot `h` begins a tied group exactly when the seed marked it, so two + // adjacent slots tie exactly when the later one is not a start. + let starts_group = + |h: usize| -> bool { starts[h / 64].load(Ordering::Relaxed) >> (h % 64) & 1 == 1 }; + + // `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<(I, I)> = (0..n) + .into_par_iter() + .filter_map(|h| { + if !starts_group(h) { + return None; + } + let mut e = h + 1; + while e < n && !starts_group(e) { + 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((I::from_usize(h), I::from_usize(e))) + }) + .collect(); + let mut groups = groups; + profile_log(&format!( + "radix groups {} groups, {} MB", + groups.len(), + groups.len() * std::mem::size_of::<(I, I)>() / (1 << 20) + )); + drop(starts); + 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 *tied* `sa` slot, so a round's reads + // of `rank` never observe that same round's writes. + // + // Sized to the tied population rather than to `n`. Only slots inside a + // group are rewritten, and even the first round has far fewer of those + // than the text has positions (18% of them on chr21), so a full second + // rank array would be mostly untouched pages. Each group gets a + // contiguous window at its prefix-sum offset, and the buffer is reused + // across rounds, which shrink monotonically. + let mut next_rank: Vec = Vec::new(); + let mut offsets: Vec = Vec::new(); + + while !groups.is_empty() { + let round_t = Instant::now(); + // 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. + // Groups average about four elements, and there are millions of them + // per round, so the two things that dominated here were not the sort + // or the rank probes but the bookkeeping around them: one heap + // allocation per group for the key vector, and a sequential + // `split_at_mut` chain to hand each group its sub-slices. + // + // Both go. `Scatter` already encodes "disjoint ranges, one owner + // each", which is exactly the property the groups have, so each group + // takes its own sub-slices directly with no sequential prepass. And a + // group that fits the stack buffer never touches the allocator. + // Window each group takes in the tied-slot buffer. + offsets.clear(); + offsets.reserve(groups.len() + 1); + let mut tied = 0usize; + for &(start, end) in &groups { + offsets.push(tied); + tied += end.to_usize() - start.to_usize(); + } + offsets.push(tied); + if next_rank.len() < tied { + next_rank.resize(tied, I::zero()); + } + + let sa_cell = Scatter::new(&mut sa); + let nr_cell = Scatter::new(&mut next_rank); + let rank_ref = &rank; + let offsets_ref = &offsets; + let sub: Vec<(I, I)> = groups + .par_iter() + .enumerate() + .flat_map_iter(|(gi, &(start, end))| { + let (start, end) = (start.to_usize(), end.to_usize()); + let len = end - start; + // SAFETY: `groups` are disjoint, sorted `sa` ranges, so this + // group is the sole owner of `start..end` in `sa` and of its + // own prefix-sum window in the tied-slot buffer. + let (sa_g, nr_g) = unsafe { + ( + sa_cell.slice_mut(start, len), + nr_cell.slice_mut(offsets_ref[gi], len), + ) + }; + 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_ref[q].to_usize() as u64 + 1, + _ => 0, + } + }; + + let mut stack = [(0u64, I::zero()); DOUBLING_STACK_GROUP]; + let mut heap: Vec<(u64, I)>; + let keyed: &mut [(u64, I)] = if len <= DOUBLING_STACK_GROUP { + let slot = &mut stack[..len]; + for (dst, &e) in slot.iter_mut().zip(sa_g.iter()) { + *dst = (succ(e.to_usize()), e); + } + slot + } else { + heap = sa_g.iter().map(|&e| (succ(e.to_usize()), e)).collect(); + &mut heap + }; + keyed.sort_unstable(); + + let mut fresh: Vec<(I, I)> = Vec::new(); + let mut i = 0; + while i < len { + let key = keyed[i].0; + let mut j = i + 1; + while j < 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((I::from_usize(start + i), I::from_usize(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() + .enumerate() + .for_each(|(gi, &(start, end))| { + let base = offsets[gi]; + for (i, slot) in (start.to_usize()..end.to_usize()).enumerate() { + // SAFETY: `sa[start..end]` are distinct positions owned + // solely by this group, and the groups partition their + // index range. + unsafe { ranks.set(sa[slot].to_usize(), next_rank[base + i]) }; + } + }); + + let before: usize = groups + .iter() + .map(|&(s, e)| e.to_usize() - s.to_usize()) + .sum(); + let n_groups = groups.len(); + groups = sub; + let after: usize = groups + .iter() + .map(|&(s, e)| e.to_usize() - s.to_usize()) + .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. + profile_log(&format!( + " doubling round depth={depth}: {before} tied in {} groups (avg {:.1}) -> {after} tied, {:.3}s", + n_groups, + before as f64 / n_groups.max(1) as f64, + round_t.elapsed().as_secs_f64() + )); + 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(), + } + } + + /// Borrow `len` elements starting at `index` mutably. + /// + /// # Safety + /// + /// No other live borrow may overlap `index..index + len`, and the borrow + /// the `Scatter` was built from must still be live. + #[inline] + unsafe fn slice_mut<'a>(&self, index: usize, len: usize) -> &'a mut [T] { + debug_assert!(index + len <= self.len); + unsafe { std::slice::from_raw_parts_mut(self.ptr.add(index), 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) }; + } +} + +#[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"); + } + + /// The field width follows the number of *distinct* symbols, not the + /// largest byte value. Raw FASTA is the case that matters: six symbols + /// whose largest is `'T'` (84) would force 8-bit fields without ranking, + /// fitting only 8 symbols per key instead of 16. + #[test] + fn packer_width_follows_alphabet_size_not_byte_value() { + let two = Packer::new(b"abababab"); + assert_eq!((two.bits(), two.k()), (1, 64)); + let four = Packer::new(&[0u8, 1, 2, 3, 3, 2, 1, 0]); + assert_eq!((four.bits(), four.k()), (2, 32)); + + let mut fasta: Vec = b"ACGTN".to_vec(); + fasta.push(b'\n'); + let f = Packer::new(&fasta); + assert_eq!((f.bits(), f.k()), (4, 16), "6 symbols should pack 4 bits"); + + let dense: Vec = (0..=255u8).collect(); + let d = Packer::new(&dense); + assert_eq!((d.bits(), d.k()), (8, 8)); + } + + /// The remap must be monotone, or a packed key would stop being + /// order-preserving and the whole seed would be wrong. Checked through + /// the observable behaviour: over a text whose bytes ascend and are all + /// distinct, successive suffixes must produce strictly increasing keys. + #[test] + fn packer_keys_follow_byte_order() { + for text in [ + b"\nACGNTZq".to_vec(), + (0..40u8) + .map(|i| i.wrapping_mul(6).wrapping_add(3)) + .collect(), + b"ACGT".to_vec(), + ] { + let mut ascending: Vec = text.clone(); + ascending.sort_unstable(); + ascending.dedup(); + let p = Packer::new(&ascending); + let keys: Vec = (0..ascending.len()) + .map(|i| p.key_at(&ascending, i)) + .collect(); + for (i, w) in keys.windows(2).enumerate() { + assert!( + w[0] < w[1], + "key({i}) = {:#x} should precede key({}) = {:#x} for {ascending:?}", + w[0], + i + 1, + w[1], + ); + } + } + } + + /// The SWAR gather must agree with the obvious shift-or loop for every + /// field width and every alignment, including the zero-padded tail. + #[test] + fn swar_gather_matches_scalar_packing() { + use rand::{RngExt, SeedableRng}; + let mut rng = rand::rngs::StdRng::seed_from_u64(0x5AA5); + for &sigma in &[2u8, 4, 16, 200] { + for &n in &[1usize, 7, 8, 9, 31, 32, 33, 63, 64, 65, 200] { + let text: Vec = (0..n).map(|_| rng.random_range(0..sigma)).collect(); + let p = Packer::new(&text); + let (bits, k) = (p.bits(), p.k()); + let ranked = p.ranked.as_deref().unwrap_or(&text); + for pos in 0..n { + let end = (pos + k).min(n); + let mut want: u64 = 0; + for &c in &ranked[pos..end] { + want = (want << bits) | c as u64; + } + want <<= bits as usize * (k - (end - pos)); + assert_eq!( + p.key_at(&text, pos), + want, + "sigma={sigma} n={n} pos={pos} bits={bits}", + ); + } + } + } + } + + /// 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};