diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..34f36e3 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,74 @@ +name: CI + +on: + # Every branch, not just `main`. A pull request opened from a fork by a + # first-time contributor does not run workflows until a maintainer approves + # them, so `pull_request` alone leaves reviewers with no check runs to look + # at. Building on push means the contributor's own fork produces evidence + # that can be linked from the PR. + push: + pull_request: + workflow_dispatch: + +env: + CARGO_TERM_COLOR: always + RUSTFLAGS: -D warnings + +jobs: + test: + name: test (${{ matrix.os }}) + runs-on: ${{ matrix.os }} + strategy: + fail-fast: false + matrix: + # macOS is aarch64/NEON, Ubuntu is x86_64/AVX2. The LCP kernel and the + # pooled external-memory bucket path are the parts that differ per + # platform, so both need to run. + os: [ubuntu-latest, macos-latest] + steps: + - uses: actions/checkout@v4 + - uses: dtolnay/rust-toolchain@stable + - uses: Swatinem/rust-cache@v2 + + # Debug catches the `debug_assert`s that guard the unchecked scatter in + # `radix.rs` and the buffer-length invariants in the merge kernel. + - name: Test (debug) + run: cargo test --all-targets + - name: Test (release) + run: cargo test --release --all-targets + - name: Doc tests + run: cargo test --doc + + lint: + name: fmt + clippy + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: dtolnay/rust-toolchain@stable + with: + components: rustfmt, clippy + - uses: Swatinem/rust-cache@v2 + - run: cargo fmt --all --check + - run: cargo clippy --all-targets -- -D warnings + + msrv: + name: MSRV (1.89) + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + # Pinned to the `rust-version` in Cargo.toml, which is set by the + # stabilised AVX-512 intrinsics the LCP fast path uses. + - uses: dtolnay/rust-toolchain@1.89 + - uses: Swatinem/rust-cache@v2 + - run: cargo check --all-targets + + docs: + name: rustdoc + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: dtolnay/rust-toolchain@stable + - uses: Swatinem/rust-cache@v2 + - run: cargo doc --no-deps + env: + RUSTDOCFLAGS: -D warnings diff --git a/Cargo.toml b/Cargo.toml index 12261dc..a6e777a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -15,6 +15,19 @@ categories = ["algorithms", "data-structures"] repository = "https://github.com/COMBINE-lab/caps-sa" readme = "README.md" +# The published benchmark numbers in `bench/` were taken with fat LTO and a +# single codegen unit. That configuration used to come from a parent workspace +# that no longer exists in this repo, so pin it here. A `[profile.release]` in a +# library crate applies only when this crate is the workspace root — i.e. to +# this repo's own tests, examples and benches — and is invisible to downstream +# consumers, who keep their own profile. +[profile.release] +lto = "fat" +codegen-units = 1 + +[profile.bench] +inherits = "release" + [dependencies] rayon = "1" tempfile = "3" @@ -26,3 +39,7 @@ rand = "0.10" [[example]] name = "caps_sa" path = "examples/caps_sa.rs" + +[[example]] +name = "gsj_bench" +path = "examples/gsj_bench.rs" diff --git a/README.md b/README.md index c018b38..be331a7 100644 --- a/README.md +++ b/README.md @@ -15,8 +15,68 @@ streams the SA out as positions are emitted. ## Status Both the in-memory and external-memory paths are implemented, tested, -and benchmarked. 43 unit tests pass and the SA output is differentially -verified against a brute-force reference on small and random inputs. +and benchmarked. 73 unit tests pass and the SA output is differentially +verified against a brute-force reference on small and random 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 + +On the human genome (GRCh38, 32 threads on AMD EPYC 9575F), caps-sa is +**7% faster than upstream CaPS-SA's ext-mem path** and uses **23% less +RAM**, while beating upstream's in-mem wall time by 3% at 1/10 of the +RAM. See [`bench/README.md`](bench/README.md) for the full methodology +and the optimisation ladder that got us there. Those paths still use the +merge kernel, so they retain the scan cost described above on +repeat-heavy input. On the human genome (GRCh38, 32 threads on AMD EPYC 9575F), caps-sa is **7% faster than upstream CaPS-SA's ext-mem path** and uses **23% less @@ -86,6 +146,119 @@ 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 prefix doubling needs a rank for every position in +the text, which would defeat exactly that. Segmented texts and symbols +wider than `u8` also stay on the merge kernel. + +### Skipping long repeats + +Those paths get the same pathology fixed in the comparator instead, which +costs no extra memory. If `text[s..e)` has period `q` and two suffixes +start at `a < b` inside it with `(b - a) % q == 0`, they agree until the +later one reaches `e`, so + +```text +lcp(a, b) >= e - b +``` + +is known in `O(1)` from the run's bounds with nothing scanned. When the +phase does not match, the two suffixes differ within `q` symbols and the +ordinary scan is already short. Scans are additionally bounded so they +stop at a run's start rather than traversing it. + +Detecting only single-symbol runs would miss the case that actually +occurs: in wrapped FASTA an `N` block is 60 `N`s followed by a newline, +which is period 61, not period 1. Periods up to 64 are considered. +Measured on synthetic periodic inputs, periods 1, 2, 61 and 64 sort +5.8-6.9x faster, while periods 65 and 171 are not detected and run at +parity, so alpha-satellite arrays (canonical monomer 171 bases) fall +outside the detector. + +Detection is two-stage so texts without repeats pay almost nothing: a +sampling pass collects the periods that occur at all, and the full scan +runs only for those. On `N`-free DNA the table comes out empty and every +query short-circuits. The table is a few dozen entries, so the +external-memory path keeps its memory bound. + +| ext-mem input | before | after | CPU before | CPU after | peak RSS | +| ------------- | ------ | ----- | ---------- | --------- | -------- | +| chr21 FASTA, 47.5 MB | 24.2 s | **1.47 s** | 268 s | 17.5 s | 147 → 190 MB | +| chr21 `N`-free, 80 MB | 3.49 s | **1.65 s** | 33.9 s | 15.1 s | 214 → 285 MB | + +Seven changes get there, each measured separately: + +| change | chr21.0123 | chr21 FASTA | +| ------ | ---------- | ----------- | +| baseline | 3.49 s | 24.19 s | +| skip periodic runs | 3.55 s | 2.48 s | +| prefetch the next candidates' text | 2.90 s | 1.99 s | +| subarray target 64Ki → 128Ki records | 2.54 s | 1.86 s | +| seed phase-1 subarrays with the packed key | 2.11 s | 1.65 s | +| merge cascade run pairs in parallel | 2.02 s | 1.55 s | +| pipeline the emit against the next merge | 1.85 s | 1.45 s | +| re-sort run-free partitions by key | **1.65 s** | **1.47 s** | + +Phase 1 goes from 22.85 s to 0.42 s on the FASTA input, and from 1.20 s +to 0.32 s on the `N`-free one. Peak RSS rises by under 30 MB, so the +bounded-memory guarantee the path exists for is intact. + ## 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..aee0614 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,198 @@ 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. + +### Where the external-memory path's remaining time is + +Per-phase thread scaling on chr21 forward ++ revcomp (80 MB), 1 to 12 +threads: + +``` + threads phase 1 phase 3 phase 4 total + 1 2.088 s 0.756 s 6.413 s 9.288 s + 6 0.393 s 0.413 s 1.359 s 2.192 s + 12 0.289 s 0.361 s 0.954 s 1.640 s + speedup 7.2x 2.1x 6.7x 5.7x +``` + +Phases 1 and 4 scale acceptably. **Phase 3 does not**: it saturates by +six threads and gains 14% going from six to twelve. It is 22% of the +wall at twelve threads and would be a larger share on any wider machine. + +Two hypotheses for it were tested and both are wrong: + +- *Physical file contention.* The `p` logical buckets share + `n_phys` physical files, defaulting to the thread count, so every + write might be queueing behind twelve descriptors. Raising it does not + help and mildly hurts: phase 3 measures 0.356 / 0.379 / 0.396 / 0.404 / + 0.420 s at `CAPS_SA_N_PHYS` of 12 / 24 / 48 / 96 / 192. +- *Search cost.* Galloping cut the pivot comparisons roughly eightfold + and bought 8%; an earlier attempt to accelerate the same searches with + packed keys measured nothing at all. + +What is left is the write path itself. Phase 3 reads every record back +and writes every record out again — about 1.3 GB of traffic at this +input size — and those are file writes through the page cache, which +does not parallelise with threads. 0.36 s for that traffic is roughly +3.5 GB/s, far below what this machine's DRAM sustains, which is +consistent with the kernel path rather than memory being the limit. + +So the fix is not to make phase 3 faster but to **not run it**: fuse it +into phase 1. That needs the pivots before phase 1 rather than after, +which a cheap pre-sampling pass over the raw positions can supply. +Correctness does not depend on the pivots being good — any splitters +give a correct sample sort, and only the balance changes — so the +restructuring is safer than it looks. It would remove one full +write-and-read round trip of every record. + +### Measuring the external-memory path reliably + +Wall times on the external-memory path drift by 30% or more between +sessions on the same machine and the same commit, and the cause is +mundane: each build streams several GB of bucket data through the page +cache, so anything that has recently filled that cache makes the next +build slower. + +This was diagnosed the hard way, after a phantom "regression" that +bisected to a commit predating it. The actual sequence was that 32 GB of +benchmark *output* files had accumulated in the working directory over a +long session. Deleting them restored the original numbers immediately: + +``` + wall user CPU + with 32 GB resident 20.5-21.1 s — + after deleting them 15.8-17.5 s 154-166 s +``` + +The temp files themselves are not the problem and do not leak: the +pooled buckets are anonymous and `$TMPDIR` measured 454 MB and 801 +entries both before and after a run. + +So, when measuring here: + +- **Write output to `/dev/null`** unless the run is specifically checking + correctness, and delete any output that is kept. +- **Interleave A and B** rather than measuring all of A then all of B. + An interleaved three-pair comparison of the segmented-key change gave + 20.5/21.0/21.1 s against 30.0/31.3/32.2 s — a ratio of 0.67 — on a + session where the absolute numbers were 30% above their clean-state + values. The ratio was right even though neither side's absolute number + was. +- **Report CPU time alongside wall.** It is measured against the same + drift and makes an I/O-bound artefact visible as a wall/CPU divergence. +- Treat a single reading as a hypothesis. Two figures in this file were + published from unreplicated single runs and both turned out wrong: a + claimed +54% RSS that replication showed to be flat, and an 11.4 s + phase-4 time that was 15-16 s on repetition. + +### chr21 external memory — skipping long repeats + +Same machine and inputs as the section above. The external-memory path +stays on the merge kernel by design (prefix doubling needs a rank per +text position, which would defeat the memory bound it exists to +provide), so it was still paying the full scan cost. Profiling put 94% +of the FASTA run in phase 1: + +``` + phase1 phase2 phase3 phase4 total +chr21.0123 before 1.15 s 0.005 s 0.215 s 2.080 s 3.49 s +chr21.fa before 22.85 s 0.058 s 0.317 s 0.948 s 24.19 s +``` + +Fixing it in the comparator rather than the algorithm keeps the memory +bound intact: + +| ext-mem input | before | after | CPU before | CPU after | peak RSS | +| ------------- | ------ | ----- | ---------- | --------- | -------- | +| `chr21.fa`, 47.5 MB | 24.19 s | **2.48 s** | 268 s | 23.3 s | 147 → 151 MB | +| `chr21.0123`, 80 MB | 3.49 s | 3.55 s | 33.9 s | 33.8 s | 214 → 220 MB | + +Phase 1 goes from 22.85 s to 0.95 s. The `N`-free row is flat, which is +the expected result: the sampling stage finds no periodic window, the +run table comes out empty, and every query short-circuits. + +Output was verified identical to the in-memory suffix array on both +inputs (127.7 M entries). + +The detail worth recording is that **a homopolymer detector would not +have worked**. In 60-column wrapped FASTA the longest run of a single +byte is 60, because each line of `N`s is terminated by a newline. The +real structure is a period-61 repeat spanning 6.6 Mb. Periods up to 64 +are considered, which also covers satellite arrays. + +### 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/bench/gsj_fixture.py b/bench/gsj_fixture.py new file mode 100755 index 0000000..9dcadc2 --- /dev/null +++ b/bench/gsj_fixture.py @@ -0,0 +1,90 @@ +#!/usr/bin/env python3 +"""Build an annotation-shaped splice-junction fixture for `gsj_bench`. + + bench/gsj_fixture.py [overhang] + +Writes `.text` (one byte per symbol, A/C/G/T/N as 0..=4) and +`.seg` (packed little-endian u64 segment lengths). + +The layout mirrors what a STAR-style index constructs: the genome sequence, +then one 2*overhang flank per distinct splice junction, then the reverse +complement of the whole thing. Each junction flank is its own segment, which +is what makes the comparator segmented; the genome is one segment per record. + +Note this is a *shape* reproduction. A genome-wide annotation yields far more +junctions than a single-chromosome one, so segment counts differ accordingly. +""" +import collections +import struct +import sys + +CODE = {"A": 0, "C": 1, "G": 2, "T": 3, "N": 4} +COMP = {0: 3, 1: 2, 2: 1, 3: 0, 4: 4} + + +def main() -> None: + if len(sys.argv) not in (4, 5): + sys.exit(__doc__) + fasta, gtf, prefix = sys.argv[1:4] + overhang = int(sys.argv[4]) if len(sys.argv) == 5 else 100 + + seq = bytearray() + with open(fasta) as fh: + for line in fh: + if line.startswith(">"): + continue + for ch in line.strip().upper(): + seq.append(CODE.get(ch, 4)) + + transcripts = collections.defaultdict(list) + with open(gtf) as fh: + for line in fh: + if line.startswith("#"): + continue + f = line.split("\t") + if len(f) < 9 or f[2] != "exon": + continue + i = f[8].find('transcript_id "') + if i < 0: + continue + tid = f[8][i + 15 : f[8].find('"', i + 15)] + transcripts[tid].append((int(f[3]) - 1, int(f[4]))) + + junctions = set() + for exons in transcripts.values(): + exons.sort() + for k in range(len(exons) - 1): + donor, acceptor = exons[k][1], exons[k + 1][0] + if acceptor > donor: + junctions.add((donor, acceptor)) + + flanks = bytearray() + seglens = [len(seq)] + for donor, acceptor in sorted(junctions): + left = seq[max(0, donor - overhang) : donor] + right = seq[acceptor : acceptor + overhang] + flanks += left + bytearray([4] * (overhang - len(left))) + flanks += right + bytearray([4] * (overhang - len(right))) + seglens.append(2 * overhang) + + forward = seq + flanks + text = forward + bytearray(COMP[b] for b in reversed(forward)) + seglens = seglens + list(reversed(seglens)) + assert sum(seglens) == len(text) + + with open(f"{prefix}.text", "wb") as out: + out.write(bytes(text)) + with open(f"{prefix}.seg", "wb") as out: + out.write(struct.pack(f"<{len(seglens)}Q", *seglens)) + + counts = collections.Counter(text) + print( + f"{prefix}: {len(text)} symbols, {len(seglens)} segments, " + f"{len(junctions)} junctions, " + f"{sum(v for k, v in counts.items() if k < 4)} ACGT-start positions", + file=sys.stderr, + ) + + +if __name__ == "__main__": + main() diff --git a/examples/caps_sa.rs b/examples/caps_sa.rs index a26beba..2744945 100644 --- a/examples/caps_sa.rs +++ b/examples/caps_sa.rs @@ -18,7 +18,10 @@ 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_ext_mem_for_filter, build_in_memory, + build_in_memory_sample_sort, verify_sa, +}; struct Args { input: PathBuf, @@ -27,6 +30,8 @@ struct Args { in_mem_ss: bool, subproblem_count: usize, threads: Option, + verify: bool, + filter_acgt: bool, } fn parse_args() -> Args { @@ -36,6 +41,8 @@ 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 filter_acgt = false; let mut i = 1; while i < argv.len() { match argv[i].as_str() { @@ -47,6 +54,17 @@ fn parse_args() -> Args { in_mem_ss = true; i += 1; } + "--verify" => { + verify = true; + i += 1; + } + // Sort only suffixes starting at a symbol below 4, the shape a + // STAR-style genome index uses: A/C/G/T participate, N and + // spacers do not. + "--filter-acgt" => { + filter_acgt = true; + i += 1; + } "--subproblem-count" => { subproblem_count = argv[i + 1] .parse() @@ -64,7 +82,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 +106,24 @@ fn parse_args() -> Args { in_mem_ss, subproblem_count, threads, + verify, + filter_acgt, + } +} + +/// 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 +160,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() @@ -150,7 +187,18 @@ fn main() -> std::io::Result<()> { let mut count = 0usize; let mode_label = if args.ext_mem { "ext-mem" } else { "in-mem-ss" }; let build_start = Instant::now(); - if args.ext_mem { + if args.ext_mem && args.filter_acgt { + build_ext_mem_for_filter( + &text, + |p| text[p as usize] < 4, + &opts, + |pos| { + count += 1; + writer.borrow_mut().write_all(&pos.to_le_bytes())?; + Ok(()) + }, + )?; + } else if args.ext_mem { build_ext_mem(&text, &opts, |pos| { count += 1; writer.borrow_mut().write_all(&pos.to_le_bytes())?; @@ -177,6 +225,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/examples/gsj_bench.rs b/examples/gsj_bench.rs new file mode 100644 index 0000000..d3d78d9 --- /dev/null +++ b/examples/gsj_bench.rs @@ -0,0 +1,235 @@ +//! Benchmark the annotation-shaped, splice-junction index build. +//! +//! This is the shape a STAR-style genome index actually constructs, and it is +//! the one none of the packed-key work reached until segmented keys existed: +//! +//! * the text is **segmented** — one segment per chromosome plus one per +//! splice-junction flank — so LCP comparisons stop at segment boundaries; +//! * the comparator is STAR's **spacer-as-largest** `boundary_order`, in which +//! the suffix that reaches its boundary first is the *larger* one, with an +//! ascending-position tie-break; +//! * only **ACGT-starting** positions participate, so no suffix beginning +//! inside an `N` block enters the sort at all; +//! * construction goes through the **external-memory** path. +//! +//! Usage: +//! +//! ```text +//! gsj_bench [--threads N] [--plain] [--verify] +//! ``` +//! +//! `` is one byte per symbol with A/C/G/T/N coded `0..=4`. +//! `` is a packed little-endian `u64[]` summing to the text +//! length. `bench/gsj_fixture.py` builds both from a FASTA and a GTF. +//! +//! `--plain` swaps the segmented provider for `PlainText`, which is the +//! comparison worth having: it shows what the same positions cost when the +//! segmented comparator is not required. +//! +//! ## Measured +//! +//! Apple M4 Max, 12 threads, chr21 plus a 698,597-junction library +//! (372,858,766 symbols, 1,397,196 segments, 320,856,244 retained ACGT-start +//! positions) — the same junction and segment counts as a GENCODE v50 +//! primary-assembly fixture. Three runs per configuration: +//! +//! ```text +//! segmented keys off segmented keys on +//! phase 1 9.54 s 3.31 s +//! phase 4 11.62 s 11.42 s +//! total 23.5-24.0 s 15.8-16.2 s -33% +//! peak RSS 3.22-3.94 GB 3.21-3.23 GB +//! ``` +//! +//! Phase 1 drops by 65% at unchanged peak memory. An earlier single-run pair +//! suggested a large RSS increase; replicating both sides showed that was one +//! outlier measured against another, not a real cost. +//! +//! Partition-count sweep on the same fixture, which re-checks the 128 Ki +//! target on a segmented workload rather than the plain one it was tuned on: +//! +//! ```text +//! p = 2845 (128 Ki target, current default) 15.26 s 3.22 GB +//! p = 5690 (64 Ki target, previous default) 19.09 s 3.71 GB +//! p = 11380 26.97 s 5.19 GB +//! ``` +//! +//! The target holds up here: the current default is both faster and smaller +//! than the one it replaced, and the trend continues in the same direction. + +use std::cmp::Ordering; +use std::env; +use std::fs; +use std::path::PathBuf; +use std::process; +use std::time::Instant; + +use caps_sa::{ + BoundaryRank, ExtMemOpts, LimitProvider, PlainText, SegmentedText, + build_ext_mem_for_positions_with, +}; + +/// STAR's convention: whichever suffix hits its boundary first is larger. +struct StarConvention { + inner: SegmentedText, +} + +impl LimitProvider for StarConvention { + #[inline] + fn lim_at(&self, p: usize) -> usize { + self.inner.lim_at(p) + } + #[inline] + fn boundary_order(&self, p_a: usize, lim_a: usize, p_b: usize, lim_b: usize) -> Ordering { + lim_b.cmp(&lim_a).then(p_a.cmp(&p_b)) + } + #[inline] + fn boundary_rank(&self) -> Option { + Some(BoundaryRank::LongerFirst) + } +} + +fn main() -> std::io::Result<()> { + let argv: Vec = env::args().collect(); + let mut positional: Vec = Vec::new(); + let mut threads: Option = None; + let mut plain = false; + let mut verify = false; + let mut subproblems: usize = 0; + let mut i = 1; + while i < argv.len() { + match argv[i].as_str() { + "--threads" => { + threads = Some(argv[i + 1].parse().expect("--threads expects an integer")); + i += 2; + } + "--plain" => { + plain = true; + i += 1; + } + "--verify" => { + verify = true; + i += 1; + } + "--subproblem-count" => { + subproblems = argv[i + 1] + .parse() + .expect("--subproblem-count expects an integer"); + i += 2; + } + _ => { + positional.push(argv[i].clone()); + i += 1; + } + } + } + if positional.len() != 2 { + eprintln!("usage: gsj_bench [--threads N] [--plain] [--verify]"); + process::exit(2); + } + + let text = fs::read(PathBuf::from(&positional[0]))?; + let raw = fs::read(PathBuf::from(&positional[1]))?; + let lengths: Vec = raw + .chunks_exact(8) + .map(|c| u64::from_le_bytes(c.try_into().unwrap()) as usize) + .collect(); + assert_eq!( + lengths.iter().sum::(), + text.len(), + "segment lengths must sum to the text length" + ); + + if let Some(t) = threads { + rayon::ThreadPoolBuilder::new() + .num_threads(t) + .build_global() + .expect("failed to configure rayon"); + } + + // Only ACGT starts participate, as in a STAR index. + let positions: Vec = (0..text.len() as u64) + .filter(|&p| text[p as usize] < 4) + .collect(); + eprintln!( + "fixture: {} symbols, {} segments, {} ACGT-start positions, mode={}", + text.len(), + lengths.len(), + positions.len(), + if plain { "plain" } else { "segmented+STAR" }, + ); + + let opts = ExtMemOpts { + subproblem_count: subproblems, + ..ExtMemOpts::default() + }; + let mut count = 0usize; + let mut last = 0u64; + let mut ordered = true; + + let start = Instant::now(); + if plain { + let lp = PlainText::new(text.len()); + build_ext_mem_for_positions_with(&text, positions, &lp, &opts, |pos| { + count += 1; + ordered &= count == 1 || last <= pos; + last = pos; + Ok(()) + })?; + } else { + let lp = StarConvention { + inner: SegmentedText::from_lengths(text.len(), &lengths), + }; + // Checking order here would need the comparator; `--verify` below does + // that properly on the collected output instead. + build_ext_mem_for_positions_with(&text, positions, &lp, &opts, |_pos| { + count += 1; + Ok(()) + })?; + } + let elapsed = start.elapsed(); + eprintln!("build: {count} positions in {:.3}s", elapsed.as_secs_f64()); + let _ = ordered; + + if verify { + // Re-run collecting, then check every adjacent pair against the + // comparator directly. O(n) comparisons, each bounded by a segment. + let lp = StarConvention { + inner: SegmentedText::from_lengths(text.len(), &lengths), + }; + let positions: Vec = (0..text.len() as u64) + .filter(|&p| text[p as usize] < 4) + .collect(); + let mut out: Vec = Vec::with_capacity(positions.len()); + build_ext_mem_for_positions_with(&text, positions, &lp, &opts, |pos| { + out.push(pos); + Ok(()) + })?; + let t = Instant::now(); + let mut bad = 0usize; + for w in out.windows(2) { + let (a, b) = (w[0] as usize, w[1] as usize); + let (la, lb) = (lp.lim_at(a), lp.lim_at(b)); + let mut ord = Ordering::Equal; + for j in 0..la.min(lb) { + if text[a + j] != text[b + j] { + ord = text[a + j].cmp(&text[b + j]); + break; + } + } + if ord == Ordering::Equal { + ord = lp.boundary_order(a, la, b, lb); + } + if ord == Ordering::Greater { + bad += 1; + } + } + if bad == 0 { + eprintln!("verify: OK in {:.3}s", t.elapsed().as_secs_f64()); + } else { + eprintln!("verify: FAILED, {bad} adjacent pairs out of order"); + process::exit(1); + } + } + Ok(()) +} diff --git a/src/ext_mem.rs b/src/ext_mem.rs index 621a3ab..876f6b3 100644 --- a/src/ext_mem.rs +++ b/src/ext_mem.rs @@ -42,6 +42,7 @@ use crate::ext_bucket::{ }; use crate::lcp::{LcpDispatch, Symbol}; use crate::limits::{LimitProvider, PlainText}; +use crate::runs::Cmp; use crate::sample_sort; /// Emit a phase-timing line to stderr if `CAPS_SA_PROFILE` is set in @@ -49,7 +50,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}"); } @@ -388,8 +389,8 @@ where /// bytes — ~770 MB on the human genome, vs the ~50 GB the equivalent /// `Vec` would take). Phase 1's per-subarray fill is then driven /// by popcount-walking the bitmap; the predicate is **never invoked -/// again** after the initial build. See [`FilteredSource`] for the -/// memory accounting and the inner loop. +/// again** after the initial build. See the crate-internal +/// `FilteredSource` for the memory accounting and the inner loop. /// /// Use this entry when the caller already has the text in RAM and /// the kept positions are described by a cheap per-position @@ -502,7 +503,24 @@ where return Ok(()); } let p = effective_subproblem_count(n, opts.subproblem_count); - let dispatch = LcpDispatch::detect(); + let runs = crate::runs::detect_for(text); + let cmp = Cmp::new(LcpDispatch::detect(), &runs); + // Only build the alphabet map once every eligibility condition holds. + // `Packer::new` may materialise a ranked copy of the whole text, and on a + // segmented or context-bounded build that copy could never be used: the + // packed-key paths all decline those. Constructing it first cost real + // resident memory for nothing. + let plain_text = lp.plain_lex_len() == Some(text.len()); + let seed_packer = + if opts.max_context == usize::MAX && (plain_text || lp.boundary_rank().is_some()) { + // A segmented build needs one spare code for the boundary sentinel; + // a plain one must not pay for it, since widening the field halves + // the symbols a key carries. + crate::radix::seed_params(text, !plain_text) + } else { + None + }; + let seed_params = seed_packer.as_ref(); let work_dir = opts.work_dir.clone(); // Pool the `2 × p` bucket files into one anonymous tempfile per @@ -515,7 +533,8 @@ where // local-disk wall time is neutral or marginally improved. See // `bench/README.md` for the empirical sizing. let n_phys = effective_physical_file_count(opts.physical_file_count); - let phase1_pool = BucketPool::new(n_phys, &work_dir)?; + // One pool now, not two: the fused phase 1 writes partition buckets + // directly, so there are no subarray buckets to hold. let phase3_pool = BucketPool::new(n_phys, &work_dir)?; profile_log(&format!( @@ -523,57 +542,39 @@ where std::mem::size_of::() * 8 )); - let sub_factory = |i: usize| phase1_pool.new_bucket::>(i); let part_factory = |j: usize| phase3_pool.new_bucket::>(j); let t = Instant::now(); - let (mut subarray_buckets, samples) = phase1_sort_sample_spill::( - text, - lp, - &source, - p, - opts, - dispatch, - sub_factory, - )?; - profile_log(&format!( - "phase1 (sort+sample+spill) {:.3}s", - t.elapsed().as_secs_f64() - )); - - // Drop the position source as soon as phase 1 returns — phases - // 2/3/4 don't touch it. For `PositionSource::Subset` this frees - // the caller's `Vec` (e.g. ~47 GB on a human-scale - // _for_positions build); for `PositionSource::Filtered` it - // frees the bitmap + cumsum (~770 MB); for `Identity` it's a - // no-op. The text and the spilled `subarray_buckets` are all - // phase 2+ needs. - drop(source); - - let t = Instant::now(); - let pivots = phase2_select_pivots::(text, lp, samples, p, opts.max_context, dispatch); + let pivots = phase0_presample_pivots::(text, lp, &source, p, opts, cmp); profile_log(&format!( - "phase2 (select pivots) {:.3}s", + "phase0 (presample pivots) {:.3}s", t.elapsed().as_secs_f64() )); let t = Instant::now(); - let mut partition_buckets = phase3_distribute::( + let mut partition_buckets = phase1_sort_and_distribute::( text, lp, - &mut subarray_buckets, + &source, &pivots, p, opts, - dispatch, + cmp, + seed_params, part_factory, )?; profile_log(&format!( - "phase3 (distribute) {:.3}s", + "phase1 (sort+distribute) {:.3}s", t.elapsed().as_secs_f64() )); - drop(subarray_buckets); + // Drop the position source as soon as phase 1 returns — phases + // 2/3/4 don't touch it. For `PositionSource::Subset` this frees + // the caller's `Vec` (e.g. ~47 GB on a human-scale + // _for_positions build); for `PositionSource::Filtered` it + // frees the bitmap + cumsum (~770 MB); for `Identity` it's a + // no-op. Phase 4 needs only the text and the partition buckets. + drop(source); let t = Instant::now(); let result = phase4_merge_and_emit::( @@ -583,7 +584,8 @@ where opts.max_context, opts.ordered_phase4_emit, &mut emit, - dispatch, + cmp, + seed_params, ); profile_log(&format!( "phase4 (merge+emit) {:.3}s", @@ -622,16 +624,41 @@ where return Ok(()); } let p = effective_subproblem_count(n, opts.subproblem_count); - let dispatch = LcpDispatch::detect(); + let runs = crate::runs::detect_for(text); + let cmp = Cmp::new(LcpDispatch::detect(), &runs); + // Only build the alphabet map once every eligibility condition holds. + // `Packer::new` may materialise a ranked copy of the whole text, and on a + // segmented or context-bounded build that copy could never be used: the + // packed-key paths all decline those. Constructing it first cost real + // resident memory for nothing. + let plain_text = lp.plain_lex_len() == Some(text.len()); + let seed_packer = + if opts.max_context == usize::MAX && (plain_text || lp.boundary_rank().is_some()) { + // A segmented build needs one spare code for the boundary sentinel; + // a plain one must not pay for it, since widening the field halves + // the symbols a key carries. + crate::radix::seed_params(text, !plain_text) + } else { + None + }; + let seed_params = seed_packer.as_ref(); let factory = |_i: usize| InMemBucket::>::new(); - let (mut subarray_buckets, samples) = - phase1_sort_sample_spill::(text, lp, &source, p, opts, dispatch, factory)?; + let (mut subarray_buckets, samples) = phase1_sort_sample_spill::( + text, + lp, + &source, + p, + opts, + cmp, + seed_params, + factory, + )?; // Same rationale as in `build_ext_mem_inner` — drop the source // as soon as phase 1's `fill_chunk` calls have stopped. drop(source); - let pivots = phase2_select_pivots::(text, lp, samples, p, opts.max_context, dispatch); + let pivots = phase2_select_pivots::(text, lp, samples, p, opts.max_context, cmp); let mut partition_buckets = phase3_distribute::( text, lp, @@ -639,7 +666,7 @@ where &pivots, p, opts, - dispatch, + cmp, factory, )?; drop(subarray_buckets); @@ -650,7 +677,8 @@ where opts.max_context, opts.ordered_phase4_emit, &mut emit, - dispatch, + cmp, + seed_params, ) } @@ -709,6 +737,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, @@ -1074,9 +1122,38 @@ impl<'a> PositionSource<'a> { /// Target subarray size used by [`effective_subproblem_count`] when /// auto-picking `p`. Smaller means more (smaller) subarrays — lower /// per-task phase-1 scratch, at the cost of more phase-3 distribute -/// work (which scales as `O(p² · log(n/p))`, sequentially) and a -/// higher temp-file count. -const PHASE1_TARGET_CHUNK: usize = 65_536; +/// work (which scales as `O(p² · log(n/p))`) and a higher temp-file +/// count. +/// +/// Raised from 65 536 after measuring the trade-off directly. Total work +/// is `n log n` either way, since a smaller `p` moves levels out of +/// phase 4's per-partition cascade and into phase 1's merge sort, but +/// the constants are not equal: phase 3 shrinks quadratically in `p`, +/// and phase 4's cascade does one full pass over its partition per +/// level. Peak RSS is set by phase 4 holding `4 × threads` partitions of +/// `n / p` records at once, so it only starts growing once `p` is small +/// enough for that product to rival the text itself. +/// +/// Measured on chr21 forward ++ revcomp (80 MB), 12 threads: +/// +/// ```text +/// p total peak RSS +/// 48 2.58 s 987 MB +/// 96 2.49 s 538 MB +/// 306 2.85 s 282 MB +/// 612 2.80 s 202 MB <- 131 072 +/// 1224 3.05 s 205 MB <- 65 536 (previous default) +/// ``` +/// +/// 131 072 is the largest step that costs nothing in memory: same peak +/// RSS as before, ~8% less wall time. Going further trades real memory +/// for speed, which is the opposite of what this path is for, so it is +/// left to the caller via `ExtMemOpts::subproblem_count`. +/// +/// At genome scale this changes nothing: `PHASE1_MAX_PARTITIONS` already +/// binds for any `n` above ~1 GB, so GRCh38 still gets `p = 8192`. +const PHASE1_TARGET_CHUNK: usize = 131_072; + /// Hard cap on the number of subarrays. Matches upstream CaPS-SA's /// default of 8192 — phase 3 is now parallelised across rayon /// workers (each subarray distributes independently into per-partition @@ -1143,7 +1220,8 @@ fn phase1_sort_sample_spill( source: &PositionSource<'_>, p: usize, opts: &ExtMemOpts, - dispatch: LcpDispatch, + cmp: Cmp<'_>, + seed_params: Option<&crate::radix::Packer>, mk_bucket: MkB, ) -> io::Result<(Vec, Vec)> where @@ -1176,16 +1254,33 @@ where let mut sa_w = vec![I::zero(); len]; let mut lcp_arr = vec![I::zero(); len]; let mut lcp_w = vec![I::zero(); len]; - sample_sort::merge_sort( + // Seed with the packed key where the comparator allows it: that + // resolves the first `k` symbols with no text access and yields + // the LCP between runs from the key difference, leaving the merge + // kernel only the suffixes that agree through all `k`. + if !crate::radix::seed_subarray( text, lp, + seed_params, &mut sa, - &mut sa_w, &mut lcp_arr, + &mut sa_w, &mut lcp_w, opts.max_context, - dispatch, - ); + cmp, + ) { + sample_sort::merge_sort( + text, + lp, + &mut sa, + &mut sa_w, + &mut lcp_arr, + &mut lcp_w, + 0, + opts.max_context, + cmp, + ); + } // Pull `samples_per_subarray` evenly-spaced positions out of // the now-sorted subarray. Deterministic — no RNG needed for @@ -1247,7 +1342,7 @@ fn phase2_select_pivots( mut samples: Vec, p: usize, max_ctx: usize, - dispatch: LcpDispatch, + cmp: Cmp<'_>, ) -> Vec where S: Symbol, @@ -1268,8 +1363,9 @@ where &mut sa_w, &mut lcp, &mut lcp_w, + 0, max_ctx, - dispatch, + cmp, ); // p-1 pivots at evenly-spaced ranks across the sorted sample pool. @@ -1300,7 +1396,7 @@ fn phase3_distribute( pivots: &[I], p: usize, opts: &ExtMemOpts, - dispatch: LcpDispatch, + cmp: Cmp<'_>, mk_bucket: MkB, ) -> io::Result> where @@ -1326,15 +1422,10 @@ where // *upper bound* in the sorted subarray. let mut splits = Vec::with_capacity(p + 1); splits.push(0usize); + let mut from = 0usize; for &pivot in pivots { - splits.push(upper_bound_by_pivot( - &records, - pivot, - text, - lp, - opts.max_context, - dispatch, - )); + from = upper_bound_from(&records, from, pivot, text, lp, opts.max_context, cmp); + splits.push(from); } splits.push(records.len()); @@ -1362,38 +1453,268 @@ where .collect()) } -/// Upper-bound binary search: returns the first index `i` such that the -/// suffix at `records[i].pos` is **strictly greater than** the suffix at -/// `pivot`. -fn upper_bound_by_pivot( +/// Phase 0: choose the `p - 1` pivots *before* sorting anything. +/// +/// The old flow sampled from the already-sorted subarrays, which forced an +/// ordering: sort everything and spill it, pick pivots, then read all of it +/// back to distribute. That round trip is the single worst-scaling part of the +/// build — phase 3 gained 14% going from six threads to twelve, because it is +/// bound by the page-cache write path rather than by anything threads help +/// with. +/// +/// Pivots do not have to come from sorted data. **Any** splitters produce a +/// correct sample sort; only the balance of the partitions changes. So a cheap +/// pre-pass over the raw positions can supply them, and phase 1 can then sort +/// and distribute in one go, never materialising the subarrays at all. +/// +/// Sampling is by strided blocks rather than strided singletons: a +/// [`PositionSource`] fills a contiguous run cheaply but pays per call, and +/// the `Filtered` variant especially so. Blocks give the same coverage of the +/// position space for a fraction of the calls, and splitter quality is not +/// sensitive to the difference. +fn phase0_presample_pivots( + text: &[S], + lp: &L, + source: &PositionSource<'_>, + p: usize, + opts: &ExtMemOpts, + cmp: Cmp<'_>, +) -> Vec +where + S: Symbol, + I: Index, + L: LimitProvider, +{ + let n = source.len(); + if p <= 1 || n == 0 { + return Vec::new(); + } + const BLOCK: usize = 64; + let target = sample_target_total(n, p).min(n); + let n_blocks = target.div_ceil(BLOCK).max(1); + let stride = (n / n_blocks).max(1); + + let mut sample: Vec = Vec::with_capacity(n_blocks * BLOCK); + let mut start = 0usize; + while start < n && sample.len() < target { + let len = BLOCK.min(n - start); + let base = sample.len(); + sample.resize(base + len, I::zero()); + source.fill_chunk(start, &mut sample[base..]); + start += stride; + } + if sample.is_empty() { + return Vec::new(); + } + + let m = sample.len(); + let mut sa_w = vec![I::zero(); m]; + let mut lcp = vec![I::zero(); m]; + let mut lcp_w = vec![I::zero(); m]; + sample_sort::merge_sort( + text, + lp, + &mut sample, + &mut sa_w, + &mut lcp, + &mut lcp_w, + 0, + opts.max_context, + cmp, + ); + (1..p).map(|i| sample[(i * m / p).min(m - 1)]).collect() +} + +/// Phase 1, fused with the old phase 3: sort each subarray and write its +/// pieces straight into the partition buckets. +/// +/// Because [`phase0_presample_pivots`] has already chosen the splitters, a +/// subarray never has to be spilled and read back. That removes one complete +/// write-and-read round trip of every record from the build. +#[allow(clippy::too_many_arguments)] +fn phase1_sort_and_distribute( + text: &[S], + lp: &L, + source: &PositionSource<'_>, + pivots: &[I], + p: usize, + opts: &ExtMemOpts, + cmp: Cmp<'_>, + seed_params: Option<&crate::radix::Packer>, + mk_bucket: MkB, +) -> io::Result> +where + S: Symbol, + I: Index, + L: LimitProvider, + SaLcp: BucketRecord, + B: SaLcpBucketStore + Send, + MkB: Fn(usize) -> B + Send + Sync, +{ + let n = source.len(); + let chunk_size = n.div_ceil(p); + let partition_buckets: Vec> = (0..p).map(|j| Mutex::new(mk_bucket(j))).collect(); + + (0..p).into_par_iter().try_for_each(|i| -> io::Result<()> { + let start = (i * chunk_size).min(n); + let end = ((i + 1) * chunk_size).min(n); + let len = end - start; + if len == 0 { + return Ok(()); + } + + let mut sa: Vec = vec![I::zero(); len]; + source.fill_chunk(start, &mut sa); + let mut sa_w = vec![I::zero(); len]; + let mut lcp_arr = vec![I::zero(); len]; + let mut lcp_w = vec![I::zero(); len]; + if !crate::radix::seed_subarray( + text, + lp, + seed_params, + &mut sa, + &mut lcp_arr, + &mut sa_w, + &mut lcp_w, + opts.max_context, + cmp, + ) { + sample_sort::merge_sort( + text, + lp, + &mut sa, + &mut sa_w, + &mut lcp_arr, + &mut lcp_w, + 0, + opts.max_context, + cmp, + ); + } + drop(sa_w); + drop(lcp_w); + + // Split the sorted subarray at the pivots and hand each piece to its + // partition. Splits are non-decreasing, so each search gallops from + // the previous one. + let records: Vec> = sa + .iter() + .zip(lcp_arr.iter()) + .map(|(&pos, &lcp)| SaLcp { pos, lcp }) + .collect(); + drop(sa); + drop(lcp_arr); + + let mut splits = Vec::with_capacity(p + 1); + splits.push(0usize); + let mut from = 0usize; + for &pivot in pivots { + from = upper_bound_from(&records, from, pivot, text, lp, opts.max_context, cmp); + splits.push(from); + } + splits.push(records.len()); + + for j in 0..p { + let (lo, hi) = (splits[j], splits[j + 1]); + if lo >= hi { + continue; + } + let mut bucket = partition_buckets[j].lock().unwrap(); + bucket.add_slice_reset_first_lcp(&records[lo..hi])?; + bucket.mark_boundary(); + } + Ok(()) + })?; + + Ok(partition_buckets + .into_iter() + .map(|m| m.into_inner().expect("partition mutex poisoned")) + .collect()) +} + +/// Upper bound of `pivot` in `records`, searched forward from `from`. +/// +/// Returns the first index at or after `from` whose suffix is strictly +/// greater than `pivot`'s. +/// +/// Phase 3 asks this `p` times per subarray, once per pivot, and the pivots +/// are sorted, so the answers are non-decreasing. Searching the whole array +/// each time wasted that: at `p = 612` over 131 K-record subarrays a plain +/// binary search costs about 17 suffix comparisons per pivot regardless of how +/// close the answer is to the previous one. Galloping from the previous split +/// costs one comparison when the bracket is empty, which it usually is once +/// `p` is large enough that consecutive pivots land within a few records of +/// each other. +/// +/// Each probe is a random access into `records` plus a suffix comparison, so +/// the count is what matters here, not the constant. +fn upper_bound_from( records: &[SaLcp], + from: usize, pivot: I, text: &[S], lp: &L, max_ctx: usize, - dispatch: LcpDispatch, + cmp: Cmp<'_>, ) -> usize where S: Symbol, I: Index, L: LimitProvider, { - let mut lo = 0; - let mut hi = records.len(); - while lo < hi { - let mid = lo + (hi - lo) / 2; - match dispatch.suffix_cmp_with( + let n = records.len(); + let greater = |i: usize| -> bool { + cmp.suffix_cmp_with( text, lp, - records[mid].pos.to_usize(), + records[i].pos.to_usize(), pivot.to_usize(), max_ctx, - ) { - Ordering::Greater => hi = mid, - Ordering::Equal | Ordering::Less => lo = mid + 1, + ) == Ordering::Greater + }; + + if from >= n { + return n; + } + // The common case: this pivot's split is the previous one. + if greater(from) { + return from; + } + + // Gallop to bracket the answer, then bisect inside the bracket. + let mut lo = from; + let mut step = 1usize; + loop { + let probe = from + step; + if probe >= n { + break; + } + if greater(probe) { + let mut hi = probe; + while lo + 1 < hi { + let mid = lo + (hi - lo) / 2; + if greater(mid) { + hi = mid; + } else { + lo = mid; + } + } + return hi; + } + lo = probe; + step *= 2; + } + // Never greater within the array: bisect the tail. + let mut hi = n; + while lo + 1 < hi { + let mid = lo + (hi - lo) / 2; + if greater(mid) { + hi = mid; + } else { + lo = mid; } } - lo + hi } /// Phase 4 + 5: parallel-merge partitions in chunks of `num_threads`, @@ -1412,6 +1733,7 @@ where /// subarrays the per-partition size is `≈ n / p`, so this stays /// proportional to `n / 4 = 0.25 n` even at the peak — well below the /// in-memory path's `~4 n` working set. +#[allow(clippy::too_many_arguments)] // buckets + text + lp + ctx + emit + cmp + seed + flag fn phase4_merge_and_emit( text: &[S], lp: &L, @@ -1419,7 +1741,8 @@ fn phase4_merge_and_emit( max_ctx: usize, ordered_emit: bool, emit: &mut F, - dispatch: LcpDispatch, + cmp: Cmp<'_>, + seed_params: Option<&crate::radix::Packer>, ) -> Result<(), BuildError> where S: Symbol, @@ -1442,7 +1765,7 @@ where // GRCh38 / 32 t). // // Bumping the chunk to `4 × num_threads` gives rayon four - // partitions per thread to dispatch — fast threads can steal from + // partitions per thread to cmp — fast threads can steal from // slow neighbours, smoothing out the size variance. Peak RAM // grows linearly: each in-flight merged partition holds its // result `Vec` (~3 MB at human-genome scale with `u32` @@ -1460,38 +1783,86 @@ where let merge_us = AtomicU64::new(0); let mut emit_secs: f64 = 0.0; - let mut start = 0; - while start < n_partitions { - let end = (start + chunk_size).min(n_partitions); - let chunk = &mut partition_buckets[start..end]; - if ordered_emit { + if ordered_emit { + let mut start = 0; + while start < n_partitions { + let end = (start + chunk_size).min(n_partitions); phase4_merge_chunk_ordered_emit( text, lp, - chunk, - max_ctx, - emit, - dispatch, - profile, - &load_us, - &merge_us, - &mut emit_secs, - )?; - } else { - phase4_merge_chunk_collect_emit( - text, - lp, - chunk, + &mut partition_buckets[start..end], max_ctx, emit, - dispatch, + cmp, + seed_params, profile, &load_us, &merge_us, &mut emit_secs, )?; + start = end; } - start = end; + } else { + // Emitting is `Θ(n)` and single-threaded by construction: the caller's + // closure is `FnMut` and its ordering is the whole point. Previously + // it also did not overlap anything, so every worker sat idle once per + // chunk while the main thread drained the merged results. + // + // Pipeline it instead. A scoped producer merges chunk `c + 1` (itself + // rayon-parallel) while the main thread emits chunk `c`. A bound of + // one keeps at most two chunks resident, so the transient cost is one + // extra chunk of merged positions rather than the unbounded queue an + // unsynchronised producer would build. + // + // This is not the `ordered_phase4_emit` path, which coordinates at + // *partition* granularity through a `BTreeMap` and measured slower. + // Here the producer emits whole chunks already in order, so the + // consumer just drains them. + let (tx, rx) = std::sync::mpsc::sync_channel::>>>(1); + let load_ref = &load_us; + let merge_ref = &merge_us; + std::thread::scope(|scope| -> Result<(), BuildError> { + scope.spawn(move || { + let mut start = 0; + while start < n_partitions { + let end = (start + chunk_size).min(n_partitions); + let merged: io::Result>> = partition_buckets[start..end] + .par_iter_mut() + .map(|bucket| { + merge_one_partition( + text, + lp, + bucket, + max_ctx, + cmp, + seed_params, + profile, + load_ref, + merge_ref, + ) + }) + .collect(); + let failed = merged.is_err(); + if tx.send(merged).is_err() || failed { + return; + } + start = end; + } + }); + + for merged in rx { + let t = Instant::now(); + for positions in merged? { + for pos in positions { + emit(pos.to_usize() as u64).map_err(BuildError::Emit)?; + } + } + if profile { + emit_secs += t.elapsed().as_secs_f64(); + } + } + Ok(()) + })?; } if profile { profile_log(&format!( @@ -1504,50 +1875,6 @@ where Ok(()) } -#[allow(clippy::too_many_arguments)] -fn phase4_merge_chunk_collect_emit( - text: &[S], - lp: &L, - chunk: &mut [B], - max_ctx: usize, - emit: &mut F, - dispatch: LcpDispatch, - profile: bool, - load_us: &std::sync::atomic::AtomicU64, - merge_us: &std::sync::atomic::AtomicU64, - emit_secs: &mut f64, -) -> Result<(), BuildError> -where - S: Symbol, - I: Index, - L: LimitProvider, - SaLcp: BucketRecord, - B: BucketStore> + Send, - F: FnMut(u64) -> Result<(), E>, -{ - // Default fast path: let rayon merge the whole chunk with minimal - // coordination, then emit the collected partition results in order. - let merged: Vec> = chunk - .par_iter_mut() - .map(|bucket| -> io::Result> { - merge_one_partition( - text, lp, bucket, max_ctx, dispatch, profile, load_us, merge_us, - ) - }) - .collect::, io::Error>>()?; - - let t = Instant::now(); - for positions in merged { - for pos in positions { - emit(pos.to_usize() as u64).map_err(BuildError::Emit)?; - } - } - if profile { - *emit_secs += t.elapsed().as_secs_f64(); - } - Ok(()) -} - #[allow(clippy::too_many_arguments)] fn phase4_merge_chunk_ordered_emit( text: &[S], @@ -1555,7 +1882,8 @@ fn phase4_merge_chunk_ordered_emit( chunk: &mut [B], max_ctx: usize, emit: &mut F, - dispatch: LcpDispatch, + cmp: Cmp<'_>, + seed_params: Option<&crate::radix::Packer>, profile: bool, load_us: &std::sync::atomic::AtomicU64, merge_us: &std::sync::atomic::AtomicU64, @@ -1585,7 +1913,15 @@ where .enumerate() .for_each_with(tx, |tx, (local_idx, bucket)| { let result = merge_one_partition( - text, lp, bucket, max_ctx, dispatch, profile, load_us, merge_us, + text, + lp, + bucket, + max_ctx, + cmp, + seed_params, + profile, + load_us, + merge_us, ); let _ = tx.send((local_idx, result)); }); @@ -1642,7 +1978,8 @@ fn merge_one_partition( lp: &L, bucket: &mut B, max_ctx: usize, - dispatch: LcpDispatch, + cmp: Cmp<'_>, + seed_params: Option<&crate::radix::Packer>, profile: bool, load_us: &std::sync::atomic::AtomicU64, merge_us: &std::sync::atomic::AtomicU64, @@ -1667,8 +2004,55 @@ where } let t = Instant::now(); - let workspace = CascadeWorkspace::::new(); - let result = workspace.cascade_merge(text, lp, &records, &boundaries, max_ctx, dispatch); + // A partition arrives as `p` sorted sub-subarrays, and the cascade merges + // them pairwise in `log2(p)` levels, each a full pass with LCP-enhanced + // comparisons. That was the single largest cost in the whole ext-mem + // build (15.4 CPU-seconds of a 15.1-second run on 80 MB of DNA). + // + // When the comparator allows a packed key, throwing the existing + // sortedness away and re-sorting the partition outright is dramatically + // cheaper: one key sort resolves the leading `k` symbols with no text + // access, and only suffixes agreeing through all of them need the merge + // kernel. `log2(p)` passes collapse into one. + // + // Only when the text has no long periodic runs, though. A run is exactly + // a stretch where a fixed-depth key resolves nothing, since every suffix + // inside it shares the whole key, so the re-sort would hand the merge + // kernel one enormous tied group and lose the ordering phase 1 had + // already established. Measured on chr21 FASTA (6.6 Mb of `N`) the + // unconditional version was 2.23 s against the cascade's 1.45 s, while on + // run-free DNA it is 1.65 s against 1.95 s. So: key re-sort when the run + // table is empty, cascade otherwise. + let result = match seed_params { + Some(_) + if lp.plain_lex_len() == Some(text.len()) + && max_ctx == usize::MAX + && !cmp.has_long_runs() => + { + let mut sa: Vec = records.iter().map(|r| r.pos).collect(); + let len = sa.len(); + let mut lcp = vec![I::zero(); len]; + let mut sa_w = vec![I::zero(); len]; + let mut lcp_w = vec![I::zero(); len]; + let seeded = crate::radix::seed_subarray( + text, + lp, + seed_params, + &mut sa, + &mut lcp, + &mut sa_w, + &mut lcp_w, + max_ctx, + cmp, + ); + debug_assert!(seeded, "guards agreed but seed_subarray declined"); + sa + } + _ => { + let workspace = CascadeWorkspace::::new(); + workspace.cascade_merge(text, lp, &records, &boundaries, max_ctx, cmp) + } + }; if profile { merge_us.fetch_add(t.elapsed().as_micros() as u64, AtomicOrdering::Relaxed); } @@ -1725,7 +2109,7 @@ impl CascadeWorkspace { records: &[SaLcp], boundaries: &[usize], max_ctx: usize, - dispatch: LcpDispatch, + cmp: Cmp<'_>, ) -> Vec where S: Symbol, @@ -1753,7 +2137,7 @@ impl CascadeWorkspace { let mut src_is_a = true; while run_lens.len() > 1 { - run_lens = self.merge_one_level(src_is_a, &run_lens, text, lp, max_ctx, dispatch); + run_lens = self.merge_one_level(src_is_a, &run_lens, text, lp, max_ctx, cmp); src_is_a = !src_is_a; } @@ -1777,7 +2161,7 @@ impl CascadeWorkspace { text: &[S], lp: &L, max_ctx: usize, - dispatch: LcpDispatch, + cmp: Cmp<'_>, ) -> Vec where S: Symbol, @@ -1807,44 +2191,66 @@ impl CascadeWorkspace { ) }; - let mut new_lens = Vec::with_capacity(run_lens.len().div_ceil(2)); - let mut src_off = 0usize; - let mut dst_off = 0usize; - let mut i = 0; - while i < run_lens.len() { - let l1 = run_lens[i]; - if i + 1 < run_lens.len() { - let l2 = run_lens[i + 1]; - let x_end = src_off + l1; - let xy_end = x_end + l2; - let dst_end = dst_off + l1 + l2; - sample_sort::merge( - text, - lp, - &src_sa[src_off..x_end], - &src_sa[x_end..xy_end], - &src_lcp[src_off..x_end], - &src_lcp[x_end..xy_end], - &mut dst_sa[dst_off..dst_end], - &mut dst_lcp[dst_off..dst_end], - max_ctx, - dispatch, - ); - new_lens.push(l1 + l2); - src_off = xy_end; - dst_off = dst_end; - i += 2; - } else { - // Odd run carries over unchanged. - let end = dst_off + l1; - dst_sa[dst_off..end].copy_from_slice(&src_sa[src_off..src_off + l1]); - dst_lcp[dst_off..end].copy_from_slice(&src_lcp[src_off..src_off + l1]); - new_lens.push(l1); - src_off += l1; - dst_off = end; - i += 1; + // The pairs at one level are independent and write to disjoint + // destination ranges, so the only thing that made this sequential was + // the running `src_off` / `dst_off`. Both are prefix sums, so compute + // them up front and hand each pair its own sub-slices. + // + // This matters because the cascade's last level is a single merge over + // the whole partition. With `p` well above the thread count there is + // enough partition-level parallelism to hide that most of the time, + // but it is what caps phase 4's efficiency: it was running at ~8x on + // 12 threads. + let n_pairs = run_lens.len() / 2; + let mut new_lens: Vec = (0..n_pairs) + .map(|j| run_lens[2 * j] + run_lens[2 * j + 1]) + .collect(); + if run_lens.len() % 2 == 1 { + new_lens.push(run_lens[run_lens.len() - 1]); + } + + // `dst` ranges are exactly `new_lens`; `src` ranges are the pairs. + let mut jobs: Vec<(usize, usize, &mut [I], &mut [I])> = Vec::with_capacity(new_lens.len()); + { + let mut sa_rest: &mut [I] = dst_sa; + let mut lcp_rest: &mut [I] = dst_lcp; + let mut src_off = 0usize; + for (j, &out_len) in new_lens.iter().enumerate() { + let (sa_head, sa_tail) = sa_rest.split_at_mut(out_len); + let (lcp_head, lcp_tail) = lcp_rest.split_at_mut(out_len); + jobs.push((j, src_off, sa_head, lcp_head)); + sa_rest = sa_tail; + lcp_rest = lcp_tail; + src_off += out_len; } } + + jobs.into_par_iter() + .for_each(|(j, src_off, out_sa, out_lcp)| { + if 2 * j + 1 < run_lens.len() { + let l1 = run_lens[2 * j]; + let l2 = run_lens[2 * j + 1]; + let x_end = src_off + l1; + let xy_end = x_end + l2; + sample_sort::merge( + text, + lp, + &src_sa[src_off..x_end], + &src_sa[x_end..xy_end], + &src_lcp[src_off..x_end], + &src_lcp[x_end..xy_end], + out_sa, + out_lcp, + max_ctx, + cmp, + ); + } else { + // Odd run carries over unchanged. + let l1 = run_lens[2 * j]; + out_sa.copy_from_slice(&src_sa[src_off..src_off + l1]); + out_lcp.copy_from_slice(&src_lcp[src_off..src_off + l1]); + } + }); new_lens } } @@ -1853,8 +2259,228 @@ impl CascadeWorkspace { mod tests { use super::*; use crate::build_in_memory; + use crate::limits::SegmentedText; use tempfile::tempdir; + /// Build with the external-memory path over an arbitrary `Symbol`. + fn ext_mem_sa_of(text: &[S], p: usize) -> Vec { + let dir = tempdir().unwrap(); + let opts = ExtMemOpts { + subproblem_count: p, + work_dir: dir.path().to_path_buf(), + ..ExtMemOpts::default() + }; + let mut out: Vec = Vec::with_capacity(text.len()); + build_ext_mem(text, &opts, |pos| { + out.push(pos); + Ok(()) + }) + .unwrap(); + out + } + + /// Reference order: sort suffixes with the slice comparator, which uses + /// `S`'s own `Ord`. + fn direct_sa(text: &[S]) -> Vec { + let mut sa: Vec = (0..text.len() as u64).collect(); + sa.sort_by(|&a, &b| text[a as usize..].cmp(&text[b as usize..])); + sa + } + + /// `Symbol` is implemented for `i8`, and a packed key orders its fields as + /// unsigned, so `-1` (byte `0xFF`) would sort above `1`. The packed-key + /// paths must decline signed symbols; these cover the two entry points + /// that reach them. + #[test] + fn ext_mem_signed_i8_matches_direct_order() { + let fixtures: Vec> = vec![ + vec![-1, 0, -1, 1, -2, 0, 1, -1, 0, -2, 1, 0], + vec![i8::MIN, 0, i8::MAX, -1, 1, i8::MIN, i8::MAX, 0], + (0..200).map(|i: i32| (i % 7 - 3) as i8).collect(), + ]; + for text in fixtures { + for p in [1usize, 3, 8] { + assert_eq!( + ext_mem_sa_of(&text, p), + direct_sa(&text), + "ext-mem p={p} disagrees on {text:?}" + ); + } + } + } + + #[test] + fn in_memory_sample_sort_signed_i8_matches_direct_order() { + let fixtures: Vec> = vec![ + vec![-1, 0, -1, 1, -2, 0, 1, -1, 0, -2, 1, 0], + (0..300).map(|i: i32| (i % 5 - 2) as i8).collect(), + ]; + for text in fixtures { + let mut out: Vec = Vec::new(); + build_in_memory_sample_sort(&text, &ExtMemOpts::default(), |pos| { + out.push(pos); + Ok(()) + }) + .unwrap(); + assert_eq!(out, direct_sa(&text), "sample sort disagrees on {text:?}"); + } + } + + #[test] + fn in_memory_signed_i8_matches_direct_order() { + let text: Vec = vec![-1, 0, -1, 1, -2, 0, 1, -1, 0, -2, 1, 0]; + let got: Vec = crate::build_in_memory(&text); + assert_eq!(got, direct_sa(&text)); + } + + /// A `LimitProvider` with STAR's spacer-as-largest convention: the suffix + /// that reaches its boundary first is *larger*, with an ascending-position + /// tie-break. This is the comparator a splice-junction index uses. + struct StarConvention { + inner: SegmentedText, + } + + impl LimitProvider for StarConvention { + fn lim_at(&self, p: usize) -> usize { + self.inner.lim_at(p) + } + fn boundary_order( + &self, + p_a: usize, + lim_a: usize, + p_b: usize, + lim_b: usize, + ) -> std::cmp::Ordering { + lim_b.cmp(&lim_a).then(p_a.cmp(&p_b)) + } + fn boundary_rank(&self) -> Option { + Some(crate::limits::BoundaryRank::LongerFirst) + } + } + + /// The provider's own comparator, spelled out. + fn direct_cmp(text: &[u8], lp: &L, a: u64, b: u64) -> std::cmp::Ordering { + let (pa, pb) = (a as usize, b as usize); + let (la, lb) = (lp.lim_at(pa), lp.lim_at(pb)); + for i in 0..la.min(lb) { + if text[pa + i] != text[pb + i] { + return text[pa + i].cmp(&text[pb + i]); + } + } + lp.boundary_order(pa, la, pb, lb) + } + + /// Assert `got` is a valid sort of `positions` under `lp`'s comparator. + /// + /// Checked as a property rather than against a canonical permutation: + /// `SegmentedText`'s default `boundary_order` returns `Equal` for suffixes + /// that end together with equal content, so their relative order is free + /// and no single answer is "the" right one. The merge kernel is not a + /// stable sort either. + fn assert_sorted_under( + text: &[u8], + lp: &L, + positions: &[u64], + got: &[u64], + what: &str, + ) { + let mut want = positions.to_vec(); + want.sort_unstable(); + let mut have = got.to_vec(); + have.sort_unstable(); + assert_eq!(have, want, "{what}: not a permutation of the input"); + for w in got.windows(2) { + assert_ne!( + direct_cmp(text, lp, w[0], w[1]), + std::cmp::Ordering::Greater, + "{what}: {} precedes {} but compares greater", + w[0], + w[1], + ); + } + } + + fn ext_mem_sa_with( + text: &[u8], + lp: &L, + p: usize, + positions: Vec, + ) -> Vec { + let dir = tempdir().unwrap(); + let opts = ExtMemOpts { + subproblem_count: p, + work_dir: dir.path().to_path_buf(), + ..ExtMemOpts::default() + }; + let mut out: Vec = Vec::with_capacity(positions.len()); + build_ext_mem_for_positions_with(text, positions, lp, &opts, |pos| { + out.push(pos); + Ok(()) + }) + .unwrap(); + out + } + + /// Segmented texts now get packed keys too: the key stops at the segment + /// boundary and pads with a reserved sentinel on the side the provider's + /// `boundary_order` demands. These check both conventions, and the + /// ACGT-start filter a splice-junction index applies, against the direct + /// comparator. + #[test] + fn ext_mem_segmented_matches_direct_comparator() { + use rand::{RngExt, SeedableRng}; + let mut rng = rand::rngs::StdRng::seed_from_u64(0x5D1CE); + for trial in 0..12 { + let n_seg = rng.random_range(2..8usize); + let lengths: Vec = (0..n_seg).map(|_| rng.random_range(4..60usize)).collect(); + let n: usize = lengths.iter().sum(); + // A/C/G/T/N plus a spacer code, the ruSTAR alphabet. + let text: Vec = (0..n).map(|_| rng.random_range(0..6u8)).collect(); + let seg = SegmentedText::from_lengths(n, &lengths); + let all: Vec = (0..n as u64).collect(); + let acgt: Vec = all + .iter() + .copied() + .filter(|&p| text[p as usize] < 4) + .collect(); + + for p in [1usize, 2, 5] { + assert_sorted_under( + &text, + &seg, + &all, + &ext_mem_sa_with(&text, &seg, p, all.clone()), + &format!("shorter-first, all positions, trial {trial} p={p}"), + ); + assert_sorted_under( + &text, + &seg, + &acgt, + &ext_mem_sa_with(&text, &seg, p, acgt.clone()), + &format!("shorter-first, ACGT filter, trial {trial} p={p}"), + ); + + let star = StarConvention { + inner: SegmentedText::from_lengths(n, &lengths), + }; + assert_sorted_under( + &text, + &star, + &all, + &ext_mem_sa_with(&text, &star, p, all.clone()), + &format!("longer-first, all positions, trial {trial} p={p}"), + ); + assert_sorted_under( + &text, + &star, + &acgt, + &ext_mem_sa_with(&text, &star, p, acgt.clone()), + &format!("longer-first, ACGT filter, trial {trial} p={p}"), + ); + } + } + } + fn ext_mem_sa(text: &[u8], p: usize) -> Vec { let dir = tempdir().unwrap(); let opts = ExtMemOpts { diff --git a/src/lib.rs b/src/lib.rs index 70684f4..3a42376 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -20,6 +20,8 @@ mod ext_bucket; mod ext_mem; mod lcp; mod limits; +mod radix; +mod runs; mod sample_sort; pub use ext_mem::{ @@ -33,12 +35,112 @@ pub use ext_mem::{ try_build_in_memory_sample_sort_for_positions_with, try_build_in_memory_sample_sort_with, }; pub use lcp::{LcpDispatch, Symbol, lcp, lcp_scalar, lcp_u8, suffix_cmp}; -pub use limits::{LimitProvider, PlainText, SegmentedText}; +pub use limits::{BoundaryRank, LimitProvider, PlainText, SegmentedText}; pub use sample_sort::{ Opts, build_in_memory, build_in_memory_for_positions, build_in_memory_for_positions_with, 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(()) +} + +/// The LCP array of a byte text's suffix array, in `O(n)`. +/// +/// `lcp[i]` is the number of symbols `text[sa[i - 1]..]` and `text[sa[i]..]` +/// share; `lcp[0]` is `0`. `sa` must be the suffix array of `text` — check it +/// with [`verify_sa`] first if it came from elsewhere. +/// +/// The merge kernel produces an LCP array as a byproduct, but nothing exposed +/// it, and the fast path does not produce one at all: prefix doubling answers +/// comparisons from ranks and never computes an LCP. This derives one from the +/// suffix array instead, by Kasai's algorithm, in a single linear pass. +/// +/// The bound is worth stating because it is exactly what the scanning merge +/// lacks: `h` falls by at most one per position and rises only while matching, +/// so the total symbol comparisons are at most `2n` no matter how repetitive +/// the text is. +/// +/// ``` +/// let text = b"banana"; +/// let sa: Vec = caps_sa::build_in_memory(text); +/// let lcp = caps_sa::lcp_array(text, &sa); +/// // sa is [5, 3, 1, 0, 4, 2] = a, ana, anana, banana, na, nana +/// assert_eq!(lcp, vec![0u32, 1, 3, 0, 0, 2]); +/// ``` +pub fn lcp_array(text: &[u8], sa: &[I]) -> Vec { + assert_eq!( + sa.len(), + text.len(), + "lcp_array: sa must be the suffix array of the whole text", + ); + radix::kasai_lcp(text, sa) +} + /// 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 9fde6f1..608660e 100644 --- a/src/limits.rs +++ b/src/limits.rs @@ -24,6 +24,24 @@ //! rationale and the comparison against the `[u8; 3]` (24-bit-text) //! alternative. +/// How a provider's [`boundary_order`][LimitProvider::boundary_order] ranks a +/// suffix that ends at its segment boundary against one that keeps going. +/// +/// This is the one fact a fixed-depth packed key needs in order to represent +/// a segmented comparator: a key pads a short suffix out to full width, and +/// the padding symbol has to fall on the correct side of every real symbol. +#[derive(Copy, Clone, Debug, PartialEq, Eq)] +pub enum BoundaryRank { + /// The suffix that ends first is smaller, the standard generalised-SA + /// convention and the default `boundary_order`. Keys pad with a symbol + /// below every real one. + ShorterFirst, + /// The suffix that ends first is *larger*, equivalently the longer one is + /// smaller. STAR's spacer-as-largest ordering. Keys pad with a symbol + /// above every real one. + LongerFirst, +} + /// Per-suffix length provider. The merge and cascade-merge code use /// `lp.lim_at(p)` instead of `text.len() - p`; the LCP function itself /// is unchanged (the merge passes the appropriately-capped @@ -73,6 +91,56 @@ 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 + } + + /// Which side of a real symbol this provider's boundary convention puts + /// the end of a suffix on, when that convention is expressible. + /// + /// Returning `Some` lets the crate build packed fixed-depth keys for a + /// *segmented* text: a key packs `min(k, lim_at(p))` symbols and pads the + /// rest with a sentinel placed according to this answer, so key order + /// agrees with `boundary_order` whenever the key decides at all. Keys that + /// tie still defer to `boundary_order` itself, which is what lets a + /// convention with a position tie-break (STAR's + /// `lim_b.cmp(&lim_a).then(p_a.cmp(&p_b))`) work: the key never has to + /// express the tie-break, only never to contradict it. + /// + /// The contract is exactly this: for suffixes `a` and `b` whose shared + /// prefix ends because one of them reached its limit, + /// `boundary_order(a, .., b, ..)` must be `Less` iff `a` is the one that + /// ended, under `ShorterFirst`, and `Greater` iff `a` is the one that + /// ended, under `LongerFirst`. + /// + /// The default is `None`, which keeps every existing implementation on the + /// comparison path. Answer it only if your `boundary_order` decides purely + /// by which suffix ended first, with at most a tie-break between suffixes + /// that end at the same offset. + #[inline] + fn boundary_rank(&self) -> Option { + None + } } /// Default provider for non-segmented texts: `lim_at(p) = n - p`. @@ -99,6 +167,16 @@ impl LimitProvider for PlainText { fn lim_at(&self, p: usize) -> usize { self.n - p } + + #[inline] + fn plain_lex_len(&self) -> Option { + Some(self.n) + } + + #[inline] + fn boundary_rank(&self) -> Option { + Some(BoundaryRank::ShorterFirst) + } } /// Provider for texts partitioned into segments at known cumulative @@ -204,6 +282,13 @@ impl LimitProvider for SegmentedText { self.n - p } } + + /// `SegmentedText` keeps the default `boundary_order`, which is + /// shorter-is-smaller. + #[inline] + fn boundary_rank(&self) -> Option { + Some(BoundaryRank::ShorterFirst) + } } #[cfg(test)] diff --git a/src/radix.rs b/src/radix.rs new file mode 100644 index 0000000..1381934 --- /dev/null +++ b/src/radix.rs @@ -0,0 +1,1171 @@ +//! 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 crate::lcp::Symbol; +use crate::limits::{BoundaryRank, LimitProvider}; +use crate::runs::Cmp; +use crate::sample_sort; +use rayon::prelude::*; +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, + /// Number of distinct codes in use. Codes are `0..alphabet`; `alphabet` + /// itself is free for use as a boundary sentinel when it fits the field. + alphabet: u32, +} + +impl Packer { + /// Build the map for `text`. + fn new(text: &[u8], need_sentinel: bool) -> 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; + } + } + // A segmented key needs one code above the alphabet for its boundary + // sentinel, so the field must hold `alphabet`, not `alphabet - 1`. + // Plain builds do not pay for that: widening the field would halve the + // symbols per key, which is the whole point of packing. + let widest = if need_sentinel { + next + } else { + next.saturating_sub(1) + }; + let bits: u32 = match widest { + 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, + alphabet: next as u32, + } + } + + #[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, + } + } + + /// Whether a boundary sentinel fits alongside the alphabet in one field. + /// + /// With 8-bit fields and 256 distinct symbols there is no spare code, so + /// segmented keys are unavailable and the caller must fall back. + #[inline] + pub(crate) fn has_sentinel(&self) -> bool { + (self.alphabet as u64) < (1u64 << self.bits) + } + + /// Pack the `min(k, lim)` symbols at `text[p..]`, padding the rest with a + /// boundary sentinel placed according to `rank`. + /// + /// This is the segmented counterpart to [`Self::key_at`]. It never reads + /// past `p + lim`, so a key cannot see into the next segment, and the + /// sentinel falls below every real code under + /// [`BoundaryRank::ShorterFirst`] and above every real code under + /// [`BoundaryRank::LongerFirst`]. That is what makes key order agree with + /// the provider's `boundary_order` whenever the key decides at all. + /// + /// Under `ShorterFirst` the sentinel is code `0` and every real code is + /// shifted up by one, so a padded field is strictly below any real symbol. + /// Under `LongerFirst` the sentinel is `alphabet`, strictly above every + /// real code, and no shift is needed. + #[inline] + pub(crate) fn key_at_bounded( + &self, + text: &[u8], + p: usize, + lim: usize, + rank: BoundaryRank, + ) -> u64 { + debug_assert!(self.has_sentinel()); + let src = self.ranked.as_deref().unwrap_or(text); + let take = self.k.min(lim).min(src.len() - p); + let (bias, pad) = match rank { + BoundaryRank::ShorterFirst => (1u64, 0u64), + BoundaryRank::LongerFirst => (0u64, self.alphabet as u64), + }; + let mut key = 0u64; + for &c in &src[p..p + take] { + key = (key << self.bits) | (c as u64 + bias); + } + for _ in take..self.k { + key = (key << self.bits) | pad; + } + key + } + + /// 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 + } +} + +/// The LCP array of `sa`, computed from the suffix array itself in `O(n)`. +/// +/// Kasai's algorithm. `lcp[i]` is the number of symbols +/// `text[sa[i - 1]..]` and `text[sa[i]..]` share, and `lcp[0]` is `0`. +/// +/// This exists because prefix doubling answers comparisons from ranks and so +/// never produces the LCP array the merge kernel yields as a byproduct, which +/// is the structural reason the external-memory path could not be routed +/// through it. Deriving it afterwards costs one linear pass. +/// +/// The pass is sequential and looks random-access, but it is not quadratic: +/// `h` falls by at most one per position and rises only while matching, so the +/// total number of symbol comparisons is at most `2n`. That bound holds +/// regardless of how repetitive the text is, which is the property the +/// scanning merge lacks. +pub(crate) fn kasai_lcp(text: &[u8], sa: &[I]) -> Vec { + let n = sa.len(); + let mut lcp = vec![I::zero(); n]; + if n == 0 { + return lcp; + } + let mut rank = vec![0usize; n]; + for (i, entry) in sa.iter().enumerate() { + rank[entry.to_usize()] = i; + } + let mut h = 0usize; + for p in 0..n { + let i = rank[p]; + if i == 0 { + h = 0; + continue; + } + let q = sa[i - 1].to_usize(); + while p + h < n && q + h < n && text[p + h] == text[q + h] { + h += 1; + } + lcp[i] = I::from_usize(h); + h = h.saturating_sub(1); + } + lcp +} + +/// The alphabet map for `text`, or `None` when a packed key cannot represent +/// this text's order. +/// +/// Computed once per build and handed to [`seed_subarray`], which would +/// otherwise re-scan the whole text for every subarray. +pub(crate) fn seed_params(text: &[S], need_sentinel: bool) -> Option { + // Exactly `u8`, not merely one byte wide. `Symbol` is implemented for + // `i8` too, and a packed key orders its fields as unsigned: `-1` has byte + // `0xFF` and would sort above `1`, inverting the text's real order. The + // in-memory doubling guard already required exact `u8`; the packed-key + // paths must match it. + 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), so a byte view over the same memory is valid + // for reads of the same length. + let bytes: &[u8] = + unsafe { std::slice::from_raw_parts(text.as_ptr() as *const u8, text.len()) }; + Some(Packer::new(bytes, need_sentinel)) +} + +/// Sort `sa` into suffix order and fill `lcp`, using a packed fixed-depth key +/// so that most of the ordering costs no text access at all. +/// +/// This is the external-memory and sample-sort counterpart to [`build_sa`]. +/// Those paths cannot use prefix doubling, which needs a rank for every +/// position in the text and would break the memory bound they exist to +/// provide. But they can still avoid the part of the merge kernel that hurts +/// most: sorting a subarray from singletons, where every leaf merge starts at +/// `m = 0` and compares two suffixes by scanning the text. +/// +/// Sorting by the packed key resolves the first `k` symbols with no text +/// access (32 symbols for DNA), and hands back the LCP for adjacent entries +/// for free from `(key_a ^ key_b).leading_zeros()`. Only suffixes that agree +/// through all `k` symbols reach the merge kernel, on the small slice they +/// occupy. +/// +/// Returns `false` without touching anything when the comparator is not plain +/// lexicographic, so the caller falls back to a plain `merge_sort`. +/// +/// `sa_w` and `lcp_w` are the caller's existing merge scratch buffers. +#[allow(clippy::too_many_arguments)] +pub(crate) fn seed_subarray( + text: &[S], + lp: &L, + packer: Option<&Packer>, + sa: &mut [I], + lcp: &mut [I], + sa_w: &mut [I], + lcp_w: &mut [I], + max_ctx: usize, + cmp: Cmp<'_>, +) -> bool { + let Some(packer) = packer else { + return false; + }; + let (bits, k) = (packer.bits(), packer.k()); + if max_ctx != usize::MAX { + return false; + } + // Plain text keys pad past end-of-text and need the visible-length + // tie-break, because a real `0` symbol is indistinguishable from padding. + // A segmented text instead stops the key at `lim_at(p)` and pads with a + // reserved sentinel placed on the side the provider's `boundary_order` + // demands, which encodes the boundary directly and needs no tie-break. + let plain = lp.plain_lex_len() == Some(text.len()); + let seg_rank = match (plain, lp.boundary_rank()) { + (true, _) => None, + (false, Some(r)) if packer.has_sentinel() => Some(r), + _ => return false, + }; + let len = sa.len(); + if len < 2 { + if len == 1 { + lcp[0] = I::zero(); + } + return true; + } + // SAFETY: `packer` is `Some` only when `S` is exactly `u8`. + let bytes: &[u8] = + unsafe { std::slice::from_raw_parts(text.as_ptr() as *const u8, text.len()) }; + + // Order by (key, visible length), the same comparator `build_sa` seeds + // with and for the same reason: zero padding makes a short suffix share a + // key with any suffix continuing in zeros, and `0` is a real symbol. + let n = bytes.len(); + let visible = |p: usize| -> usize { (n - p).min(k) }; + let mut keyed: Vec<(u64, u32, I)> = sa + .iter() + .map(|&p| { + let p = p.to_usize(); + match seg_rank { + None => (packer.key_at(bytes, p), visible(p) as u32, p_as(p)), + Some(r) => (packer.key_at_bounded(bytes, p, lp.lim_at(p), r), 0, p_as(p)), + } + }) + .collect(); + keyed.sort_unstable(); + for (slot, e) in sa.iter_mut().zip(keyed.iter()) { + *slot = I::from_usize(e.2.to_usize()); + } + + // Walk equal-key runs. Between runs the LCP falls straight out of the key + // difference; inside one it needs the merge kernel. + let mut i = 0usize; + while i < len { + let mut j = i + 1; + while j < len && (keyed[j].0, keyed[j].1) == (keyed[i].0, keyed[i].1) { + j += 1; + } + if j - i > 1 { + // Everything in this group agreed through the whole key, so the + // merge can start its scans there instead of at zero. For a plain + // text the key covers `k` symbols unless the suffix ran out first, + // which the visible-length component records. + let base = match seg_rank { + None if keyed[i].1 as usize == k => k, + _ => 0, + }; + sample_sort::merge_sort( + text, + lp, + &mut sa[i..j], + &mut sa_w[i..j], + &mut lcp[i..j], + &mut lcp_w[i..j], + base, + max_ctx, + cmp, + ); + } else { + lcp[i] = I::zero(); + } + // Boundary entry: LCP against the last element of the previous run. + if i > 0 { + let a = sa[i - 1].to_usize(); + let b = sa[i].to_usize(); + let xor = keyed[i - 1].0 ^ keyed[i].0; + debug_assert_ne!(xor, 0, "distinct runs must differ in the key"); + // `leading_zeros / bits` counts whole matching fields. Cap by both + // suffixes' lengths: padding can agree with a real `0` symbol past + // the end of the shorter one. + let shared = (xor.leading_zeros() / bits) as usize; + lcp[i] = I::from_usize(shared.min(lp.lim_at(a)).min(lp.lim_at(b))); + } + i = j; + } + lcp[0] = I::zero(); + true +} + +/// Round-trip a position through the index type used in the seed vector. +#[inline] +fn p_as(p: usize) -> I { + I::from_usize(p) +} + +/// 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 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], + 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(); + + let ht = Instant::now(); + // 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(); + + profile_log(&format!( + " seed histogram {:.3}s", + ht.elapsed().as_secs_f64() + )); + let pt = Instant::now(); + // 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); + } + + profile_log(&format!( + " seed prefixsum {:.3}s", + pt.elapsed().as_secs_f64() + )); + let st = Instant::now(); + // 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 = packer.key_at(text, p); + 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; + } + }); + } + + profile_log(&format!( + " seed scatter {:.3}s", + st.elapsed().as_secs_f64() + )); + let bt = Instant::now(); + // 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; + } + }); + + profile_log(&format!( + " seed bucketsort {:.3}s", + bt.elapsed().as_secs_f64() + )); + (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 packer = Packer::new(text, false); + 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 (keys, 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(); + + // 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() { + 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. + let sa_cell = Scatter::new(&mut sa); + let nr_cell = Scatter::new(&mut next_rank); + let rank_ref = &rank; + let sub: Vec<(usize, usize)> = groups + .par_iter() + .flat_map_iter(|&(start, end)| { + let len = end - start; + // SAFETY: `groups` are disjoint, sorted `sa` ranges, so this + // group is the sole owner of `start..end` in both arrays. + let (sa_g, nr_g) = + unsafe { (sa_cell.slice_mut(start, len), nr_cell.slice_mut(start, 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::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((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(); + let n_groups = groups.len(); + groups = sub; + 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. + 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:?}"); + } + + /// `Symbol` is implemented for `i8`, and a one-byte-wide check alone lets a + /// signed text through a packer that orders bytes as unsigned. `-1` has + /// byte `0xFF`, so it would sort above `1`, inverting the true order. + #[test] + fn signed_symbols_are_not_eligible_for_packing() { + let text: Vec = vec![-1, 0, -1, 1, -2, 0, 1, -1, 0, -2, 1, 0]; + assert!( + seed_params(&text, false).is_none(), + "i8 texts must not get a packed key" + ); + // u8 of the same width still qualifies. + let bytes: Vec = vec![1, 0, 1, 2, 3, 0]; + assert!(seed_params(&bytes, false).is_some()); + } + + /// Kasai's output must match a naive per-pair scan, on the inputs that + /// make the naive version expensive: long runs and periodic text. + #[test] + fn kasai_matches_naive() { + use rand::{RngExt, SeedableRng}; + let mut rng = rand::rngs::StdRng::seed_from_u64(0xCA5A1); + let mut fixtures: Vec> = vec![ + b"banana".to_vec(), + b"mississippi".to_vec(), + vec![7u8; 500], + (0..500).map(|i| (i % 3) as u8).collect(), + (0..500).map(|i| (i % 61) as u8).collect(), + Vec::new(), + vec![1], + ]; + for &sigma in &[2u8, 4, 200] { + for &n in &[7usize, 64, 1000] { + fixtures.push((0..n).map(|_| rng.random_range(0..sigma)).collect()); + } + } + for text in fixtures { + let sa: Vec = build_sa(&text); + let lcp = kasai_lcp(&text, &sa); + assert_eq!(lcp.len(), sa.len()); + if sa.is_empty() { + continue; + } + assert_eq!(lcp[0], 0, "lcp[0] must be 0"); + for i in 1..sa.len() { + let (a, b) = (sa[i - 1] as usize, sa[i] as usize); + let want = (0..) + .take_while(|&j| { + a + j < text.len() && b + j < text.len() && text[a + j] == text[b + j] + }) + .count(); + assert_eq!(lcp[i] as usize, want, "lcp[{i}] 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", false); + assert_eq!((two.bits(), two.k()), (1, 64)); + let four = Packer::new(&[0u8, 1, 2, 3, 3, 2, 1, 0], false); + 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, false); + 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, false); + 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, false); + 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, false); + 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/runs.rs b/src/runs.rs new file mode 100644 index 0000000..64263fe --- /dev/null +++ b/src/runs.rs @@ -0,0 +1,500 @@ +//! Periodic-run detection, and a run-aware suffix comparator. +//! +//! The LCP-enhanced merge resolves most steps without touching the text, but +//! when it does have to compare two suffixes it scans their shared prefix one +//! vector at a time. That is fine until the text contains a long *periodic +//! run*, at which point two suffixes inside the run agree for as far as the +//! run continues and a single comparison scans megabytes. +//! +//! Genome assemblies always contain these. An `N` block is the obvious case, +//! and note that in wrapped FASTA it is **not** a run of one symbol: 60 `N`s +//! followed by a newline is a run of period 61. So detecting only +//! single-symbol runs would miss the representation that actually shows up. +//! +//! The detector covers periods up to [`MAX_PERIOD`] and **not** beyond. +//! Measured on synthetic 1 MiB periodic inputs, periods 1, 2, 61 and 64 are +//! detected with full coverage and sort 5.8-6.9x faster; periods 65 and 171 +//! are not detected at all and run at parity. Alpha-satellite arrays, whose +//! canonical monomer is 171 bases, are therefore *outside* this detector. +//! The sampling stage can also miss a run that is localised enough to fall +//! between its windows. +//! +//! The observation that makes this cheap: if `text[s..e)` has period `q`, and +//! two suffixes start at `a < b` inside it with `(b - a) % q == 0`, then they +//! agree until the later one reaches `e`. That is +//! +//! ```text +//! lcp(a, b) >= e - b +//! ``` +//! +//! known in `O(1)` from the run's bounds, with no scanning at all. The scan +//! resumes at `e`, where the run's guarantee stops. When the phase does not +//! match (`(b - a) % q != 0`) the two suffixes must differ within `q` symbols, +//! so the ordinary scan is already short. +//! +//! Detection is two-stage so that texts without runs pay almost nothing. A +//! sampling pass looks for any periodic window at all and collects the set of +//! periods that actually occur; the full scan then runs only for those +//! periods. On N-free DNA the sample finds nothing and the table is empty, so +//! [`RunTable::skip`] returns immediately on a slice-empty check. +//! +//! This is what the external-memory and sample-sort paths use instead of the +//! prefix doubling in [`crate::radix`]: doubling needs a rank for every +//! position in the text, which would defeat the bounded memory those paths +//! exist to provide, while a run table costs a few dozen entries. + +use crate::lcp::{LcpDispatch, Symbol}; +use crate::limits::LimitProvider; +use rayon::prelude::*; +use std::cmp::Ordering; + +/// Shortest run worth recording. Below this the ordinary SIMD scan crosses +/// the run faster than the binary search that would find it. +const MIN_RUN: usize = 1024; + +/// Longest period considered. Covers homopolymers (period 1) and +/// wrapped-FASTA `N` blocks (period 61), which are the cases that occur in +/// practice in assembly FASTA. +/// +/// It does **not** cover alpha-satellite arrays: their canonical monomer is +/// 171 bases, and a synthetic period-171 input measures at parity with no +/// detection at all. Raising this is a constant change, but the detection +/// scan is `O(periods x n)`, so it is not free. +const MAX_PERIOD: usize = 64; + +/// Window used by the sampling pass to decide whether a period occurs at all. +const SAMPLE_WINDOW: usize = 512; + +/// Symbols an ordinary scan must match before the run table is consulted. +/// +/// Two suffixes of real sequence that agree this far are already unusual, so +/// the table is reached only when it might actually help. Small enough that +/// the probe is a handful of vector compares, and the probe is not wasted +/// work: whatever it matches counts towards the answer. +const RUN_PROBE: usize = 256; + +/// A maximal stretch `[start, end)` of the text with period `period`, meaning +/// `text[i] == text[i + period]` for every `i` in `start..end - period`. +#[derive(Copy, Clone, Debug, PartialEq, Eq)] +struct Run { + start: usize, + end: usize, + period: usize, +} + +/// Long periodic runs of a byte text, sorted by start and non-overlapping. +/// +/// Empty for texts without long repeats, which is the common case for +/// randomised or `N`-free input, and empty for symbol types wider than a byte. +#[derive(Clone, Debug, Default)] +pub(crate) struct RunTable { + runs: Vec, +} + +impl RunTable { + /// An empty table. Every query short-circuits. + pub(crate) fn empty() -> Self { + Self { runs: Vec::new() } + } + + pub(crate) fn is_empty(&self) -> bool { + self.runs.is_empty() + } + + /// Detect the long periodic runs of `text`. + pub(crate) fn detect(text: &[u8]) -> Self { + let n = text.len(); + if n < MIN_RUN { + return Self::empty(); + } + + // Stage 1: which periods occur anywhere? Sample windows across the + // text and record every period that makes one of them periodic. A + // text with no long repeat contributes nothing and stops here. + let n_samples = 4096.min(n / SAMPLE_WINDOW).max(1); + let stride = (n / n_samples).max(1); + let mut seen = [false; MAX_PERIOD + 1]; + let found: Vec> = (0..n_samples) + .into_par_iter() + .map(|s| { + let base = s * stride; + let end = (base + SAMPLE_WINDOW).min(n); + let mut periods = Vec::new(); + if end - base < MAX_PERIOD * 2 { + return periods; + } + for q in 1..=MAX_PERIOD { + if (base..end - q).all(|i| text[i] == text[i + q]) { + periods.push(q); + // The smallest period implies all its multiples; one + // per window is enough to trigger the full scan. + break; + } + } + periods + }) + .collect(); + for q in found.into_iter().flatten() { + seen[q] = true; + } + let periods: Vec = (1..=MAX_PERIOD).filter(|&q| seen[q]).collect(); + if periods.is_empty() { + return Self::empty(); + } + + // Stage 2: for each period that occurs, find its maximal runs. + let mut runs: Vec = periods + .par_iter() + .flat_map_iter(|&q| { + let mut out = Vec::new(); + let mut i = 0usize; + while i + q < n { + if text[i] != text[i + q] { + i += 1; + continue; + } + let start = i; + while i + q < n && text[i] == text[i + q] { + i += 1; + } + // Matching through `i` means the periodic stretch covers + // `start..i + q`. + let end = i + q; + if end - start >= MIN_RUN { + out.push(Run { + start, + end, + period: q, + }); + } + } + out + }) + .collect(); + + // Keep a non-overlapping set, preferring the earliest start and then + // the longest reach, so a lookup is a single binary search. + runs.sort_unstable_by_key(|r| (r.start, std::cmp::Reverse(r.end))); + let mut merged: Vec = Vec::with_capacity(runs.len()); + for r in runs { + match merged.last() { + Some(last) if r.start < last.end => { + // Overlaps the previous run. Extending the previous run + // would break its period guarantee, so drop this one + // unless it reaches strictly further, in which case keep + // only the part past the previous end. + if r.end > last.end && r.end - last.end >= MIN_RUN { + merged.push(Run { + start: last.end, + end: r.end, + period: r.period, + }); + } + } + _ => merged.push(r), + } + } + Self { runs: merged } + } + + /// The run containing `pos`, if any. + #[inline] + fn at(&self, pos: usize) -> Option<&Run> { + if self.runs.is_empty() { + return None; + } + let i = self.runs.partition_point(|r| r.start <= pos); + let r = self.runs.get(i.checked_sub(1)?)?; + (pos < r.end).then_some(r) + } + + /// Start of the first run beginning at or after `pos`, or `usize::MAX`. + #[inline] + fn next_start(&self, pos: usize) -> usize { + if self.runs.is_empty() { + return usize::MAX; + } + let i = self.runs.partition_point(|r| r.start < pos); + self.runs.get(i).map_or(usize::MAX, |r| r.start) + } + + /// Symbols that suffixes `a` and `b` are guaranteed to share starting at + /// their current offset, derived from run structure alone. + /// + /// Returns `0` when nothing can be concluded, which is always the answer + /// for an empty table. + #[inline] + fn skip(&self, a: usize, b: usize) -> usize { + let (lo, hi) = if a <= b { (a, b) } else { (b, a) }; + let Some(r) = self.at(lo) else { return 0 }; + if hi >= r.end || (hi - lo) % r.period != 0 { + return 0; + } + // Both offsets sit in the same run, an exact whole number of periods + // apart, so they agree until the later one leaves the run. + r.end - hi + } +} + +/// A suffix comparator: the SIMD LCP kernel plus the run table that lets it +/// skip long periodic repeats instead of scanning them. +/// +/// Threaded through the merge kernel in place of a bare [`LcpDispatch`]. It is +/// `Copy`, so it still travels through the recursion in registers. +#[derive(Copy, Clone)] +pub(crate) struct Cmp<'a> { + pub(crate) dispatch: LcpDispatch, + pub(crate) runs: &'a RunTable, +} + +impl<'a> Cmp<'a> { + pub(crate) fn new(dispatch: LcpDispatch, runs: &'a RunTable) -> Self { + Self { dispatch, runs } + } + + /// Whether the text contains long periodic repeats. + /// + /// Callers use this to decide whether a fixed-depth key can be expected + /// to resolve most suffixes: a long run is precisely a stretch where it + /// cannot, because every suffix inside it shares the whole key. + #[inline] + pub(crate) fn has_long_runs(&self) -> bool { + !self.runs.is_empty() + } + + /// LCP of `text[p..]` and `text[q..]` in symbols, bounded by `max_ctx`, + /// using run structure to jump over long periodic stretches. + /// + /// With an empty run table this is exactly [`LcpDispatch::lcp`] plus one + /// predictable branch. + /// + /// When a table *is* present, the ordinary bounded scan still runs first. + /// Consulting the table costs up to three binary searches before any + /// comparison happens, and the overwhelming majority of LCP calls in real + /// sequence mismatch within a few symbols and never reach a run at all. + /// Paying the lookup up front taxed every one of them: on a filtered, + /// `N`-containing chr21 that alone turned a 1.80 s build into 2.72 s even + /// though the answers were identical. + /// + /// So: probe first, and only once a match has survived [`RUN_PROBE`] + /// symbols — which ordinary genomic difference does not — is it worth + /// asking whether a run explains it. + #[inline] + pub(crate) fn lcp(&self, text: &[S], p: usize, q: usize, max_ctx: usize) -> usize { + if self.runs.is_empty() || size_of::() != 1 { + return self.dispatch.lcp(text, p, q, max_ctx); + } + + let probe = max_ctx.min(RUN_PROBE); + let got = self.dispatch.lcp(text, p, q, probe); + if got < probe || probe == max_ctx { + // Either a real mismatch, or the caller's bound was reached. No + // run can extend this, so the table is never touched. + return got; + } + + let mut i = got; + while i < max_ctx { + let jump = self.runs.skip(p + i, q + i); + if jump > 0 { + i = (i + jump).min(max_ctx); + continue; + } + // Stop the scan where a run begins, so it never traverses one. + // Scanning into a run is exactly the megabyte-long case. + let next = self + .runs + .next_start(p + i) + .saturating_sub(p + i) + .min(self.runs.next_start(q + i).saturating_sub(q + i)) + .max(1); + let window = max_ctx - i; + let bounded = next.min(window); + let got = self.dispatch.lcp(text, p + i, q + i, bounded); + i += got; + if got < bounded { + // A real mismatch, not a window boundary. + break; + } + } + i.min(max_ctx) + } + + /// Total order on two suffixes, mirroring [`LcpDispatch::suffix_cmp_with`] + /// but going through the run-aware [`Self::lcp`]. + #[inline] + pub(crate) fn suffix_cmp_with( + &self, + text: &[S], + lp: &L, + p: usize, + q: usize, + max_ctx: usize, + ) -> Ordering { + let lim_p = lp.lim_at(p); + let lim_q = lp.lim_at(q); + let lim = lim_p.min(lim_q).min(max_ctx); + let common = self.lcp(text, p, q, lim); + if common < lim { + text[p + common].cmp(&text[q + common]) + } else { + lp.boundary_order(p, lim_p, q, lim_q) + } + } +} + +/// Build a run table for `text` when the symbol type is a byte, otherwise an +/// empty one. +/// +/// Detection is a sequential-read pass and only runs at all if the sampling +/// stage finds a periodic window, so texts without long repeats pay a single +/// sampling sweep. +pub(crate) fn detect_for(text: &[S]) -> RunTable { + if size_of::() != 1 { + return RunTable::empty(); + } + // SAFETY: `S` is one byte wide with no padding and no invalid bit + // patterns (the `Symbol` contract), so a byte view over the same memory is + // valid for reads of the same length. + let bytes: &[u8] = + unsafe { std::slice::from_raw_parts(text.as_ptr() as *const u8, text.len()) }; + RunTable::detect(bytes) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn naive_lcp(text: &[u8], a: usize, b: usize, max_ctx: usize) -> usize { + let lim = (text.len() - a).min(text.len() - b).min(max_ctx); + (0..lim).take_while(|&i| text[a + i] == text[b + i]).count() + } + + /// The run-aware LCP must agree with a byte-at-a-time scan for every pair + /// of positions, whether or not a run is involved. + fn assert_lcp_agrees(text: &[u8]) { + let runs = RunTable::detect(text); + let cmp = Cmp::new(LcpDispatch::detect(), &runs); + let n = text.len(); + let step = (n / 64).max(1); + for a in (0..n).step_by(step) { + for b in (0..n).step_by(step) { + let want = naive_lcp(text, a, b, usize::MAX); + let got = cmp.lcp(text, a, b, usize::MAX); + assert_eq!(got, want, "lcp({a}, {b}) on len-{n} text"); + } + } + } + + #[test] + fn empty_table_for_texts_without_runs() { + use rand::{RngExt, SeedableRng}; + let mut rng = rand::rngs::StdRng::seed_from_u64(0x0DD1); + let text: Vec = (0..50_000).map(|_| rng.random_range(0..4u8)).collect(); + assert!(RunTable::detect(&text).is_empty()); + } + + #[test] + fn detects_homopolymer() { + let mut text: Vec = vec![1, 2, 3]; + text.extend(std::iter::repeat_n(0u8, 5000)); + text.extend([1, 2, 3]); + let table = RunTable::detect(&text); + assert!(!table.is_empty()); + assert!(table.runs.iter().any(|r| r.end - r.start >= 5000)); + } + + /// Wrapped FASTA: 60 `N`s then a newline. The single-symbol runs are only + /// 60 long, so a homopolymer-only detector would find nothing; the real + /// structure is period 61. + #[test] + fn detects_wrapped_fasta_n_block() { + let mut text: Vec = b"ACGT".to_vec(); + for _ in 0..200 { + text.extend(std::iter::repeat_n(b'N', 60)); + text.push(b'\n'); + } + text.extend(b"ACGT"); + let table = RunTable::detect(&text); + assert!(!table.is_empty(), "period-61 N block should be detected"); + assert!(table.runs.iter().any(|r| r.period == 61 || r.period == 1)); + } + + #[test] + fn runs_are_sorted_and_disjoint() { + let mut text: Vec = Vec::new(); + text.extend(std::iter::repeat_n(0u8, 3000)); + text.extend(b"ACGTACGT"); + text.extend((0..3000).map(|i| (i % 7) as u8)); + text.extend(b"TTTT"); + let table = RunTable::detect(&text); + for w in table.runs.windows(2) { + assert!(w[0].end <= w[1].start, "runs overlap: {:?}", w); + assert!(w[0].start < w[1].start); + } + for r in &table.runs { + for i in r.start..r.end - r.period { + assert_eq!(text[i], text[i + r.period], "period claim is wrong"); + } + } + } + + #[test] + fn lcp_agrees_on_homopolymer() { + let mut text: Vec = b"ACGT".to_vec(); + text.extend(std::iter::repeat_n(0u8, 4000)); + text.extend(b"ACGT"); + assert_lcp_agrees(&text); + } + + #[test] + fn lcp_agrees_on_wrapped_fasta() { + let mut text: Vec = b"ACGTAC".to_vec(); + for _ in 0..120 { + text.extend(std::iter::repeat_n(b'N', 60)); + text.push(b'\n'); + } + text.extend(b"GTGTGT"); + assert_lcp_agrees(&text); + } + + #[test] + fn lcp_agrees_on_multi_period_text() { + let mut text: Vec = Vec::new(); + text.extend((0..3000).map(|i| (i % 3) as u8)); + text.extend(b"XYZ"); + text.extend(std::iter::repeat_n(9u8, 2500)); + text.extend((0..2000).map(|i| (i % 5) as u8)); + assert_lcp_agrees(&text); + } + + #[test] + fn lcp_respects_max_ctx_inside_a_run() { + let mut text: Vec = b"AC".to_vec(); + text.extend(std::iter::repeat_n(0u8, 4000)); + let runs = RunTable::detect(&text); + let cmp = Cmp::new(LcpDispatch::detect(), &runs); + for &ctx in &[0usize, 1, 7, 100, 3000] { + assert_eq!(cmp.lcp(&text, 2, 3, ctx), naive_lcp(&text, 2, 3, ctx)); + } + } + + #[test] + fn suffix_cmp_matches_slice_order() { + use crate::limits::PlainText; + let mut text: Vec = b"ACGT".to_vec(); + text.extend(std::iter::repeat_n(5u8, 3000)); + text.extend(b"ACGT"); + let runs = RunTable::detect(&text); + let cmp = Cmp::new(LcpDispatch::detect(), &runs); + let lp = PlainText::new(text.len()); + let step = (text.len() / 40).max(1); + for a in (0..text.len()).step_by(step) { + for b in (0..text.len()).step_by(step) { + let want = text[a..].cmp(&text[b..]); + let got = cmp.suffix_cmp_with(&text, &lp, a, b, usize::MAX); + assert_eq!(got, want, "suffix_cmp({a}, {b})"); + } + } + } +} diff --git a/src/sample_sort.rs b/src/sample_sort.rs index daa9264..8a4e524 100644 --- a/src/sample_sort.rs +++ b/src/sample_sort.rs @@ -33,7 +33,39 @@ use crate::Index; use crate::lcp::{LcpDispatch, Symbol}; use crate::limits::{LimitProvider, PlainText}; +use crate::runs::Cmp; 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 +/// line is still resident when the step that needs it arrives. +const PREFETCH_DISTANCE: usize = 8; + +/// Hint the CPU to start pulling `text[at]` into cache. +/// +/// A no-op on targets without a stable prefetch intrinsic, and harmless when +/// `at` is out of bounds: the address is never dereferenced, only used as a +/// prefetch operand, and prefetch instructions on both supported targets +/// ignore faulting addresses. +#[inline(always)] +fn prefetch_symbol(text: &[S], at: usize) { + let _ = (text, at); + #[cfg(target_arch = "x86_64")] + unsafe { + std::arch::x86_64::_mm_prefetch( + text.as_ptr().add(at.min(text.len())) as *const i8, + std::arch::x86_64::_MM_HINT_T0, + ); + } + #[cfg(target_arch = "aarch64")] + unsafe { + // `core::arch::aarch64::_prefetch` is still unstable, so emit the + // instruction directly. `prfm` never faults. + let p = text.as_ptr().add(at.min(text.len())); + std::arch::asm!("prfm pldl1keep, [{p}]", p = in(reg) p, options(nostack, readonly, preserves_flags)); + } +} /// Tunable options for SA construction. #[derive(Clone, Debug)] @@ -43,12 +75,36 @@ pub struct Opts { /// caller's text doesn't guarantee comparisons terminate via sentinels /// within a known window. pub max_context: usize, + + /// Peak extra bytes `*_for_positions` may spend to sort a *subset* by + /// building the whole suffix array and filtering it. + /// + /// Prefix doubling cannot be restricted to a subset: a round compares + /// `rank[p + d]`, and that successor is generally outside the subset, so + /// ranks have to exist for every position in the text. Building the whole + /// array and filtering it in one `O(n)` pass sidesteps that, and is much + /// faster when the subset is a real fraction of the text — but it is a + /// resource decision, not a speed one: the arrays it needs are sized by + /// the *text*, not by the subset. + /// + /// `None`, the default, never makes that trade: subsets always take the + /// merge kernel, whose footprint stays proportional to the subset. + /// `Some(budget)` allows it when the estimated extra footprint fits, which + /// is `n * (3 * size_of::() + 9)` bytes: three index-wide arrays for + /// the suffix array, the ranks and the round scratch, eight bytes of key + /// per position, and a one-byte membership flag. + /// + /// Set it from what the caller can actually spare. It was previously a + /// silent `subset >= text / 8` rule, which made a large allocation on the + /// caller's behalf without telling them. + pub subset_full_sa_budget: Option, } impl Default for Opts { fn default() -> Self { Self { max_context: usize::MAX, + subset_full_sa_budget: None, } } } @@ -90,11 +146,123 @@ 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. +/// +/// Gated on [`Opts::subset_full_sa_budget`], which defaults to `None` and so +/// declines. The arrays this needs are sized by the *text*, not by the subset, +/// so it is a resource decision the caller has to make: a small subset of a +/// large text can cost far more this way than sorting it directly would. That +/// is a budget, not a correctness condition — unlike the guards in +/// [`try_doubling_fast_path`]. +/// +/// 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 { + return None; + } + // Explicitly budgeted: the arrays below are sized by the text, not by the + // subset, so a small subset of a large text can cost far more than + // sorting it directly would. + let estimate = n.checked_mul(3 * size_of::() + 9)?; + if estimate > opts.subset_full_sa_budget? { + 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. @@ -145,6 +313,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(); @@ -158,7 +330,8 @@ where // Choose the LCP implementation once for the whole build; the captured // function pointer travels through the recursion in a register, so the // inner merge loop pays no atomic load or feature-detection branch. - let dispatch = LcpDispatch::detect(); + let runs = crate::runs::detect_for(text); + let cmp = Cmp::new(LcpDispatch::detect(), &runs); merge_sort( text, @@ -167,8 +340,9 @@ where &mut sa_w, &mut lcp_arr, &mut lcp_w, + 0, opts.max_context, - dispatch, + cmp, ); sa @@ -186,7 +360,7 @@ where /// /// Visible to the rest of the crate so the external-memory path can sort /// individual subarrays of positions using the same kernel. -#[allow(clippy::too_many_arguments)] // 4 buffers + text + lp + ctx + dispatch +#[allow(clippy::too_many_arguments)] // 4 buffers + text + lp + ctx + cmp pub(crate) fn merge_sort( text: &[S], lp: &L, @@ -194,8 +368,9 @@ pub(crate) fn merge_sort( sa_w: &mut [I], lcp_arr: &mut [I], lcp_w: &mut [I], + base: usize, max_ctx: usize, - dispatch: LcpDispatch, + cmp: Cmp<'_>, ) where S: Symbol, I: Index, @@ -208,7 +383,7 @@ pub(crate) fn merge_sort( if n <= 1 { if n == 1 { - lcp_arr[0] = I::zero(); + lcp_arr[0] = I::from_usize(base); } return; } @@ -220,15 +395,15 @@ pub(crate) fn merge_sort( let (lcp_w_l, lcp_w_r) = lcp_w.split_at_mut(mid); join( - || merge_sort(text, lp, sa_l, sa_w_l, lcp_l, lcp_w_l, max_ctx, dispatch), - || merge_sort(text, lp, sa_r, sa_w_r, lcp_r, lcp_w_r, max_ctx, dispatch), + || merge_sort(text, lp, sa_l, sa_w_l, lcp_l, lcp_w_l, base, max_ctx, cmp), + || merge_sort(text, lp, sa_r, sa_w_r, lcp_r, lcp_w_r, base, max_ctx, cmp), ); // Merge the two sorted halves (still living in `sa`) into the workspace, // then copy the workspace back into the destination so the caller's // postcondition holds on `sa` / `lcp_arr`. - merge( - text, lp, sa_l, sa_r, lcp_l, lcp_r, sa_w, lcp_w, max_ctx, dispatch, + merge_from( + text, lp, sa_l, sa_r, lcp_l, lcp_r, sa_w, lcp_w, base, max_ctx, cmp, ); sa.copy_from_slice(sa_w); lcp_arr.copy_from_slice(lcp_w); @@ -242,7 +417,7 @@ pub(crate) fn merge_sort( /// /// Visible to the rest of the crate so the external-memory path can cascade /// 2-way merges across each partition's sub-subarrays during Phase 4. -#[allow(clippy::too_many_arguments)] // CaPS-SA's merge takes 5 buffers + text + lp + ctx + dispatch +#[allow(clippy::too_many_arguments)] // CaPS-SA's merge takes 5 buffers + text + lp + ctx + cmp pub(crate) fn merge( text: &[S], lp: &L, @@ -253,7 +428,44 @@ pub(crate) fn merge( z: &mut [I], lcp_z: &mut [I], max_ctx: usize, - dispatch: LcpDispatch, + cmp: Cmp<'_>, +) where + S: Symbol, + I: Index, + L: LimitProvider, +{ + merge_from(text, lp, x, y, lcp_x, lcp_y, z, lcp_z, 0, max_ctx, cmp) +} + +/// [`merge`], but told that every element of both runs already shares `base` +/// symbols. +/// +/// This is the LCP-reuse the key sort makes available and the merge otherwise +/// throws away: a tied group coming out of a packed-key sort agrees on at +/// least `k` symbols by construction, yet every comparison inside it used to +/// rescan them from zero. +/// +/// It costs one substitution in the existing invariant. The three-case rule is +/// stated against the last-output element `z_last`, initialised to the empty +/// string; here it is initialised to the length-`base` prefix `B` that every +/// element shares. `lcp(B, s) = base` for every `s` in either run, so `m` +/// starts at `base`; `B` is a prefix of every element, so it still precedes +/// them all; and `lcp_x[0] = base` is exactly what the first iteration reads. +/// The loop body is untouched, because every case in it is an argument about +/// *relative* offsets and none of them mentions zero. +#[allow(clippy::too_many_arguments)] +pub(crate) fn merge_from( + text: &[S], + lp: &L, + x: &[I], + y: &[I], + lcp_x: &[I], + lcp_y: &[I], + z: &mut [I], + lcp_z: &mut [I], + base: usize, + max_ctx: usize, + cmp: Cmp<'_>, ) where S: Symbol, I: Index, @@ -289,12 +501,32 @@ pub(crate) fn merge( let mut len_b = len_y; let mut i_a: usize = 0; let mut i_b: usize = 0; - let mut m: usize = 0; + let mut m: usize = base; let mut k: usize = 0; let mut lim_a_cache: Option<(usize, usize)> = None; let mut lim_b_cache: Option<(usize, usize)> = None; while i_a < len_a && i_b < len_b { + // The tied branch below dereferences the text at two *random* + // addresses, and the address for step `i + 1` is not known until step + // `i` retires, so there is no memory-level parallelism to exploit and + // the hardware prefetcher cannot see the pattern either. But the + // candidate positions themselves live in `arr_a` / `arr_b`, which are + // sequential and already in cache, so the addresses a few steps ahead + // *are* known. Issue them now. + // + // This is not the prefetch that was tried and reverted in `lcp.rs`: + // that one sat inside the strided scan loop, which the hardware + // prefetcher already covers. Here the access is random, which is the + // case hardware cannot predict. `m` is the current boundary LCP and a + // good estimate of where the next scans will start. + if i_a + PREFETCH_DISTANCE < len_a { + prefetch_symbol(text, arr_a[i_a + PREFETCH_DISTANCE].to_usize() + m); + } + if i_b + PREFETCH_DISTANCE < len_b { + prefetch_symbol(text, arr_b[i_b + PREFETCH_DISTANCE].to_usize() + m); + } + let l_a = lcp_a[i_a].to_usize(); // (output_a, lcp_for_output, new_m) @@ -333,7 +565,7 @@ pub(crate) fn merge( // intersection — no extra work. let cap = lim_a.min(lim_b).min(max_ctx); let remaining_ctx = cap.saturating_sub(m); - let ext = dispatch.lcp(text, p_a + m, p_b + m, remaining_ctx); + let ext = cmp.lcp(text, p_a + m, p_b + m, remaining_ctx); let total = m + ext; let a_smaller = if total < lim_a && total < lim_b { text[p_a + total] < text[p_b + total] @@ -422,6 +654,120 @@ mod tests { assert_eq!(got, want, "mismatch on text {text:?}"); } + /// Run the production kernel and return **both** the suffix array and + /// the LCP array it computes as a byproduct. + /// + /// The public entry points discard the LCP array, but it is not an + /// incidental artefact: the next merge level *consumes* it in the + /// three-case decision, so a single wrong LCP entry silently reorders + /// suffixes at the level above. It therefore needs direct coverage. + fn build_sa_and_lcp(text: &[u8], max_ctx: usize) -> (Vec, Vec) { + let n = text.len(); + let mut sa: Vec = (0..n as u32).collect(); + let mut sa_w = vec![0u32; n]; + let mut lcp_arr = vec![0u32; n]; + let mut lcp_w = vec![0u32; n]; + merge_sort( + text, + &PlainText::new(n), + &mut sa, + &mut sa_w, + &mut lcp_arr, + &mut lcp_w, + 0, + max_ctx, + Cmp::new(LcpDispatch::detect(), &crate::runs::RunTable::empty()), + ); + (sa, lcp_arr) + } + + /// Byte-at-a-time LCP of `text[a..]` and `text[b..]`, capped at `max_ctx`. + fn naive_lcp(text: &[u8], a: usize, b: usize, max_ctx: usize) -> usize { + let lim = (text.len() - a).min(text.len() - b).min(max_ctx); + (0..lim).take_while(|&i| text[a + i] == text[b + i]).count() + } + + /// Assert the LCP-array postcondition stated on [`merge_sort`]: + /// `lcp[0] == 0` and `lcp[i] == lcp(text[sa[i-1]..], text[sa[i]..])`. + fn assert_lcp_valid(text: &[u8], max_ctx: usize) { + let (sa, lcp) = build_sa_and_lcp(text, max_ctx); + if sa.is_empty() { + return; + } + assert_eq!(lcp[0], 0, "lcp[0] must be 0 (text {text:?})"); + for i in 1..sa.len() { + let want = naive_lcp(text, sa[i - 1] as usize, sa[i] as usize, max_ctx); + assert_eq!( + lcp[i] as usize, + want, + "lcp[{i}] wrong for sa[{}]={} vs sa[{i}]={} (text {text:?})", + i - 1, + sa[i - 1], + sa[i], + ); + } + } + + #[test] + fn lcp_array_matches_naive_on_fixtures() { + for text in [ + b"banana".as_slice(), + b"mississippi", + b"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + b"abababababababababababababab", + b"a", + b"", + ] { + assert_lcp_valid(text, usize::MAX); + } + } + + #[test] + fn lcp_array_matches_naive_on_random() { + use rand::{RngExt, SeedableRng}; + let mut rng = rand::rngs::StdRng::seed_from_u64(0x1CB0); + for &sigma in &[2u8, 4, 6, 255] { + for &n in &[2usize, 3, 7, 16, 17, 63, 64, 65, 200, 1000, 5000] { + let text: Vec = (0..n).map(|_| rng.random_range(0..sigma)).collect(); + assert_lcp_valid(&text, usize::MAX); + } + } + } + + /// Long runs of one symbol are the worst case for the LCP invariant: + /// adjacent suffixes share almost everything, so every `lcp[i]` is + /// large and an off-by-one is easy to miss. + #[test] + fn lcp_array_on_long_runs_and_periodic_text() { + assert_lcp_valid(&vec![7u8; 2000], usize::MAX); + let periodic: Vec = (0..2000).map(|i| (i % 3) as u8).collect(); + assert_lcp_valid(&periodic, usize::MAX); + // A run embedded in noise, the shape a poly-N genome block has. + let mut mixed: Vec = (0..500).map(|i| (i % 4) as u8).collect(); + mixed.extend(std::iter::repeat_n(4u8, 1500)); + mixed.extend((0..500).map(|i| (i % 4) as u8)); + assert_lcp_valid(&mixed, usize::MAX); + } + + #[test] + fn lcp_array_respects_max_context() { + use rand::{RngExt, SeedableRng}; + let mut rng = rand::rngs::StdRng::seed_from_u64(0xC7A); + for &max_ctx in &[1usize, 2, 4, 16] { + for &n in &[64usize, 500] { + let text: Vec = (0..n).map(|_| rng.random_range(0..3u8)).collect(); + let (sa, lcp) = build_sa_and_lcp(&text, max_ctx); + for i in 1..sa.len() { + let want = naive_lcp(&text, sa[i - 1] as usize, sa[i] as usize, max_ctx); + assert_eq!( + lcp[i] as usize, want, + "lcp[{i}] wrong with max_ctx={max_ctx}" + ); + } + } + } + } + #[test] fn empty_text() { let sa: Vec = build_in_memory::(&[]); @@ -488,6 +834,90 @@ 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); + } + + /// The subset full-SA path is now opt-in through a byte budget. Check both + /// that it is correct when allowed, and that it is actually declined when + /// the budget is too small to cover its footprint. + #[test] + fn for_positions_budget_gates_the_full_sa_path() { + use rand::{RngExt, SeedableRng}; + let mut rng = rand::rngs::StdRng::seed_from_u64(0xB0D6E7); + for &n in &[64usize, 500, 4000] { + let text: Vec = (0..n).map(|_| rng.random_range(0..4u8)).collect(); + let positions: Vec = (0..n as u32).filter(|p| p % 3 != 0).collect(); + let mut want = positions.clone(); + want.sort_by(|&a, &b| text[a as usize..].cmp(&text[b as usize..])); + + // Estimate the path advertises: n * (3 * 4 + 9) for `I = u32`. + let need = n * (3 * size_of::() + 9); + + let generous = Opts { + subset_full_sa_budget: Some(need), + ..Opts::default() + }; + assert_eq!( + build_in_memory_for_positions_with_opts(&text, positions.clone(), &generous), + want, + "budgeted path wrong at n={n}" + ); + + let tight = Opts { + subset_full_sa_budget: Some(need - 1), + ..Opts::default() + }; + assert_eq!( + build_in_memory_for_positions_with_opts(&text, positions.clone(), &tight), + want, + "declined path wrong at n={n}" + ); + + // Default declines outright. + assert_eq!( + build_in_memory_for_positions(&text, positions.clone()), + want, + "default path wrong at n={n}" + ); + } + } + #[test] fn for_positions_random_subsets() { use rand::{RngExt, SeedableRng};