From 77fd07f37db81318a9b7e32e737f6058b5beaada Mon Sep 17 00:00:00 2001 From: Nils Homer Date: Mon, 20 Jul 2026 16:23:19 -0700 Subject: [PATCH 1/3] perf(sort): build the natural queryname key from the NUL BAM already stores `extract_queryname_key` stripped the trailing NUL that BAM stores and then re-appended one into a freshly allocated `Vec` for every record. BAM already stores the read name NUL-terminated and `l_read_name` counts the terminator, so `bam[32..32 + l_read_name]` is directly usable as the key bytes. The fast path only applies when the declared name is in bounds *and* the byte it ends on is genuinely a NUL. That second condition is a soundness requirement, not a nicety: `RawQuerynameKey::cmp` passes `name.as_ptr()` to `natural_compare_nul`, which scans until it finds a terminator, so a key built from a malformed record whose `l_read_name` does not land on one would read past the end of its own allocation. Records failing either condition fall back to the previous strip-and-re-append path, which terminates unconditionally. Adds an rstest table that checks the extracted key against the pre-optimization implementation as an independent oracle, covering well-formed, empty-name, truncated, missing-terminator, non-NUL-terminated, and zero-length records, and asserts the null-termination invariant on each. --- crates/fgumi-sort/src/keys.rs | 99 +++++++++++++++++++++++++++++++++-- 1 file changed, 94 insertions(+), 5 deletions(-) diff --git a/crates/fgumi-sort/src/keys.rs b/crates/fgumi-sort/src/keys.rs index 8aa6b0299..631409af7 100644 --- a/crates/fgumi-sort/src/keys.rs +++ b/crates/fgumi-sort/src/keys.rs @@ -467,6 +467,24 @@ fn extract_raw_name_and_flags(bam: &[u8]) -> (&[u8], u16) { (name, queryname_flag_order(raw_flags)) } +/// Read-name bytes *including* the trailing NUL that BAM stores: `l_read_name` +/// (offset 8) counts the terminator, so `bam[32..32 + l_read_name]` is already +/// a NUL-terminated name ready for `natural_compare_nul`. +/// +/// Returns `None` unless the slice is in bounds *and* genuinely NUL-terminated, +/// so the caller can fall back to the stripped-name path. Both conditions +/// matter: a truncated record would slice out of bounds, and a record whose +/// declared `l_read_name` does not actually land on a terminator would yield a +/// name that [`RawQuerynameKey::cmp`] reads past the end of while scanning for +/// a NUL. The fallback re-appends the terminator unconditionally, so the +/// null-termination invariant holds for every record, well-formed or not. +#[inline] +fn raw_name_with_nul(bam: &[u8]) -> Option<&[u8]> { + let l_read_name = bam[8] as usize; + (l_read_name > 0 && 32 + l_read_name <= bam.len() && bam[32 + l_read_name - 1] == 0) + .then(|| &bam[32..32 + l_read_name]) +} + /// Serialize a queryname key as `[name_len: u16][name: bytes][flags: u16]`. #[inline] fn write_queryname_key(name: &[u8], flags: u16, writer: &mut W) -> std::io::Result<()> { @@ -542,11 +560,15 @@ impl RawQuerynameKey { #[inline] #[must_use] fn extract_queryname_key(bam: &[u8]) -> Self { - let (raw_name, flags) = extract_raw_name_and_flags(bam); - let mut name = Vec::with_capacity(raw_name.len() + 1); - name.extend_from_slice(raw_name); - name.push(0); - Self { name, flags } + let flags = queryname_flag_order(u16::from_le_bytes([bam[14], bam[15]])); + match raw_name_with_nul(bam) { + // BAM stores the name NUL-terminated, so copy those bytes directly + // instead of stripping the NUL and re-appending it. + Some(name) => Self { name: name.to_vec(), flags }, + // Truncated or non-terminated record: fall back to the stripped-name + // path (`new` re-adds the terminator), preserving the prior bytes. + None => Self::new(extract_raw_name_and_flags(bam).0.to_vec(), flags), + } } } @@ -1635,4 +1657,71 @@ mod tests { const { assert!(RawQuerynameKey::EMBEDDED_IN_RECORD) }; const { assert!(RawCoordinateKey::EMBEDDED_IN_RECORD) }; } + + // ======================================================================== + // RawQuerynameKey::extract_queryname_key tests + // ======================================================================== + + /// Build a minimal BAM record body: the 32-byte fixed block followed by + /// `name_bytes`. Only `l_read_name` (offset 8) and `flag` (offsets 14-15) + /// affect queryname key extraction, so everything else stays zeroed. + /// + /// `l_read_name` is passed separately from `name_bytes` so a test can + /// declare a length that disagrees with the bytes actually present, which + /// is how malformed records reach the extractor. + fn bam_record_with_name(l_read_name: u8, name_bytes: &[u8], flags: u16) -> Vec { + let mut bam = vec![0u8; 32]; + bam[8] = l_read_name; + bam[14..16].copy_from_slice(&flags.to_le_bytes()); + bam.extend_from_slice(name_bytes); + bam + } + + /// Independent oracle: the pre-optimization extraction, which strips the + /// declared terminator and re-appends one unconditionally. The fast path + /// reads the stored NUL directly instead, so it must agree with this byte + /// for byte on every input — well-formed or not. + fn oracle_queryname_key(bam: &[u8]) -> (Vec, u16) { + let (raw_name, flags) = extract_raw_name_and_flags(bam); + let mut name = Vec::with_capacity(raw_name.len() + 1); + name.extend_from_slice(raw_name); + name.push(0); + (name, flags) + } + + #[rstest] + // Well-formed: `l_read_name` counts the stored NUL, so the fast path applies. + #[case::well_formed(6, b"readX\0", 0x0040)] + // Empty name: `l_read_name` of 1 means the name is the terminator alone. + #[case::empty_name(1, b"\0", 0x0000)] + // Declared length runs past the buffer, so the name bytes are truncated. + #[case::truncated_name(6, b"read", 0x0080)] + // Name bytes are present but the declared terminator byte was cut off. + #[case::missing_terminator_byte(6, b"readX", 0x0000)] + // In bounds, but the byte `l_read_name` points at is not a NUL at all. + #[case::not_nul_terminated(6, b"readXY", 0x0840)] + // Degenerate `l_read_name` of 0: no name and no terminator declared. + #[case::zero_declared_length(0, b"", 0x0000)] + fn test_extract_queryname_key_matches_oracle_and_is_nul_terminated( + #[case] l_read_name: u8, + #[case] name_bytes: &[u8], + #[case] flags: u16, + ) { + let bam = bam_record_with_name(l_read_name, name_bytes, flags); + let key = RawQuerynameKey::extract_queryname_key(&bam); + let (expected_name, expected_flags) = oracle_queryname_key(&bam); + + assert_eq!(key.name(), expected_name, "name bytes must match the pre-optimization path"); + assert_eq!(key.flags, expected_flags, "flag-order value must match"); + + // `RawQuerynameKey::cmp` hands `name.as_ptr()` to `natural_compare_nul`, + // which scans until it finds a NUL. A key without one reads past the + // end of its own allocation, so this invariant is a soundness + // requirement, not just a formatting detail. + assert_eq!( + key.name().last(), + Some(&0), + "extracted name must be NUL-terminated for natural_compare_nul" + ); + } } From f0fd2ee78b8745abaa2882a5cfb30fd93edc65b8 Mon Sep 17 00:00:00 2001 From: Nils Homer Date: Mon, 20 Jul 2026 16:46:52 -0700 Subject: [PATCH 2/3] perf(sort): size radix passes from the max mapped key, not the unmapped sentinel Coordinate keys pack `(tid << 34) | ((pos + 1) << 1) | reverse`, so a mapped key occupies only ~5-6 bytes. Unmapped reads, however, carry `u64::MAX` as their sort key, which dragged `bytes_needed` to the full 8 and cost three wasted LSD radix passes over every record. Coordinate-sorted BAMs essentially always have an unmapped tail, so this was paid on real input almost without exception. Sizing the passes from the largest *non-sentinel* key fixes it, and needs no contract from the caller: `u64::MAX` is the largest `u64` and truncates to the all-`0xFF` maximum at every radix width, so unmapped records still sort to the tail and stay stable among themselves no matter how many passes run. Both scanning entry points derive their bound this way, and `par_sort_into_chunks` derives one bound for the whole buffer and shares it across chunks. A derived bound of 0 does not imply the keys are all equal, so the sort is never skipped outright: a zero key can coexist with sentinels, and those two do not compare equal. `PackedCoordinateKey::new` packs `tid = 0, pos = -1, reverse = false` to exactly 0, so a malformed record makes this reachable, and skipping would leave the sentinels ordered ahead of it. A single pass separates them and is a stable no-op when the keys genuinely are identical. Measured on an otherwise idle c6a.4xlarge over 25 contigs with a 5% unmapped tail, criterion 100 samples (confidence intervals within +/-0.5%): records before after speedup 1M 23.6 M/s 29.0 M/s 1.23x 8M 21.0 M/s 25.0 M/s 1.19x An earlier revision also tracked the bound incrementally as records were pushed, letting the sort skip the scan entirely. That was measured at a further 1.04x (30.2 and 26.0 M/s respectively) -- roughly half a percent of total sort time, since the radix is about a seventh of it -- and is not kept: it cost a running field on `RecordBuffer`, three reset sites, an ordering hazard where the reset had to precede a macro that returns from inside itself, and a public entry point whose unchecked precondition silently mis-sorts when violated. The scan is three lines and carries no invariant. Tests cover sentinel handling through both scanning entry points in serial and parallel, zero keys interleaved with sentinels above the radix threshold, `par_sort_into_chunks` on both its single- and multi-threaded drain paths, a `RecordBuffer` end-to-end mix of mapped and unmapped records, and an output-identity check that a narrowed sort is byte-for-byte identical to a full-width one including the stable order among equal keys. The added criterion benchmark separates pass-count reduction from scan elimination. --- CLAUDE.md | 3 +- benches/core_functions.rs | 101 +++++++ crates/fgumi-sort/src/inline.rs | 453 ++++++++++++++++++++++++++++++-- crates/fgumi-sort/src/lib.rs | 5 +- 4 files changed, 540 insertions(+), 22 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 4a1700495..1991ecfac 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -165,7 +165,8 @@ because the hot path runs once per record (BAM workloads are millions to billion of records) and a safe rewrite measurably regresses sort throughput. - **`crates/fgumi-sort/src/inline.rs`** — three `#[allow(unsafe_code)]` regions: - the `radix_sort_record_refs` (coordinate) and `radix_sort_template_refs` / + the `radix_sort_record_refs_with_max` (coordinate; `radix_sort_record_refs` + delegates to it after scanning for the max) and `radix_sort_template_refs` / `radix_sort_template_field` (template) LSD radix sorts use `Vec::set_len` to skip per-element initialization on the auxiliary scratch buffer, plus raw-pointer slice swaps to avoid double-borrow restrictions across the diff --git a/benches/core_functions.rs b/benches/core_functions.rs index 8f286694f..1a3787592 100644 --- a/benches/core_functions.rs +++ b/benches/core_functions.rs @@ -1305,6 +1305,106 @@ fn bench_queryname_sort_strategies(c: &mut Criterion) { group.finish(); } +/// Coordinate-sort radix throughput across the three ways `bytes_needed` can be +/// sized, isolating the two independent wins: +/// +/// - `full_width` reproduces the pre-optimization behavior by passing the +/// unmapped sentinel as the bound, forcing all 8 radix passes — which is what +/// a sentinel-inclusive max scan yields the moment a single unmapped read is +/// present, i.e. on essentially every real coordinate-sorted BAM. +/// - `scan_mapped_max` scans for the largest *mapped* key, sizing the passes to +/// the ~4–5 bytes such a key occupies. The delta against `full_width` is the +/// pass-count win. +/// - `tracked_max` supplies that same bound without scanning, using a value the +/// benchmark precomputed while building the input. The delta against +/// `scan_mapped_max` is the scan-elimination win, and it is small enough that +/// `RecordBuffer` deliberately does not chase it: it scans once per sort +/// rather than maintaining a running bound across pushes. +/// +/// `sort_unstable` is a comparison-sort baseline. Note it is *unstable*, so it +/// is not a drop-in substitute — coordinate sort order must be stable to match +/// `samtools sort`. The input mixes 25 tids with a 5% unmapped tail to exercise +/// the sentinel handling. +fn bench_coordinate_radix_sort(c: &mut Criterion) { + use fgumi_sort::{ + PackedCoordinateKey, RecordRef, radix_sort_record_refs, radix_sort_record_refs_with_max, + }; + + let nref = 25u32; // chr1-22, X, Y, MT + let mut group = c.benchmark_group("coordinate_radix_sort"); + + for &n in &[1_000_000usize, 8_000_000] { + // Deterministic xorshift so the dataset is reproducible without an RNG dep. + let mut state = 0x9e37_79b9_7f4a_7c15u64; + let mut next = move || { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + state + }; + + let mut refs: Vec = Vec::with_capacity(n); + let mut mapped_max = 0u64; + for _ in 0..n { + let r = next(); + let key = if r % 20 == 0 { + u64::MAX // ~5% unmapped (sentinel) + } else { + let tid = i32::try_from(r % u64::from(nref)).unwrap(); + let pos = i32::try_from((r >> 16) % 250_000_000).unwrap(); + let k = PackedCoordinateKey::new(tid, pos, false, nref).0; + mapped_max = mapped_max.max(k); + k + }; + refs.push(RecordRef::new(key, 0, 0)); + } + + group.throughput(Throughput::Elements(n as u64)); + group.bench_with_input(BenchmarkId::new("full_width", n), &refs, |b, refs| { + b.iter_batched( + || refs.clone(), + |mut v| { + radix_sort_record_refs_with_max(&mut v, u64::MAX); + black_box(v) + }, + criterion::BatchSize::LargeInput, + ); + }); + group.bench_with_input(BenchmarkId::new("scan_mapped_max", n), &refs, |b, refs| { + b.iter_batched( + || refs.clone(), + |mut v| { + radix_sort_record_refs(&mut v); + black_box(v) + }, + criterion::BatchSize::LargeInput, + ); + }); + group.bench_with_input(BenchmarkId::new("tracked_max", n), &refs, |b, refs| { + b.iter_batched( + || refs.clone(), + |mut v| { + radix_sort_record_refs_with_max(&mut v, mapped_max); + black_box(v) + }, + criterion::BatchSize::LargeInput, + ); + }); + group.bench_with_input(BenchmarkId::new("sort_unstable", n), &refs, |b, refs| { + b.iter_batched( + || refs.clone(), + |mut v| { + v.sort_unstable_by_key(|r| r.sort_key); + black_box(v) + }, + criterion::BatchSize::LargeInput, + ); + }); + } + + group.finish(); +} + criterion_group!( benches, bench_phred_conversions, @@ -1318,5 +1418,6 @@ criterion_group!( bench_vanilla_consensus_caller, bench_queryname_comparators, bench_queryname_sort_strategies, + bench_coordinate_radix_sort, ); criterion_main!(benches); diff --git a/crates/fgumi-sort/src/inline.rs b/crates/fgumi-sort/src/inline.rs index 505fe4fce..02192a6ef 100644 --- a/crates/fgumi-sort/src/inline.rs +++ b/crates/fgumi-sort/src/inline.rs @@ -212,6 +212,15 @@ pub struct RecordRef { padding: u32, } +impl RecordRef { + /// Construct a `RecordRef` from its public fields (padding is internal). + /// Primarily for tests and benchmarks that build sort indices directly. + #[must_use] + pub fn new(sort_key: u64, offset: u64, len: u32) -> Self { + Self { sort_key, offset, len, padding: 0 } + } +} + impl PartialEq for RecordRef { fn eq(&self, other: &Self) -> bool { self.sort_key == other.sort_key @@ -300,7 +309,7 @@ const SORT_SEGMENT_SIZE: usize = 256 * 1024 * 1024; /// (`HEADER_SIZE` for `RecordBuffer`, `TEMPLATE_HEADER_SIZE` for /// `TemplateRecordBuffer`). macro_rules! par_sort_into_chunks_impl { - ($self:expr, $threads:expr, $sort_fn:ident, $header_size:expr, $key_fn:expr) => {{ + ($self:expr, $threads:expr, $sort_fn:expr, $header_size:expr, $key_fn:expr) => {{ use rayon::prelude::*; use std::sync::Arc; @@ -453,13 +462,13 @@ impl RecordBuffer { &mut self, threads: usize, ) -> Vec> { - par_sort_into_chunks_impl!( - self, - threads, - radix_sort_record_refs, - HEADER_SIZE, - |r: &RecordRef| RawCoordinateKey { sort_key: r.sort_key } - ) + // Scan once for the whole-buffer bound and hand it to every chunk's sort, + // so no chunk re-scans for its own maximum. + let max_key = max_non_sentinel_key(&self.refs); + let sort_chunk = |refs: &mut [RecordRef]| radix_sort_record_refs_with_max(refs, max_key); + par_sort_into_chunks_impl!(self, threads, sort_chunk, HEADER_SIZE, |r: &RecordRef| { + RawCoordinateKey { sort_key: r.sort_key } + }) } /// Drain the (already-sorted) buffer into a single in-memory chunk @@ -1612,11 +1621,51 @@ const RADIX_THRESHOLD: usize = 256; /// Sorts by the `sort_key` field using 8-bit radix (256 buckets). /// This is O(n×k) where k is the number of bytes to sort (typically 5-8). /// +/// Scans `refs` for the maximum key, then defers to +/// [`radix_sort_record_refs_with_max`]. A caller that already holds a valid +/// bound — such as a chunked sort that derived one for the whole slice — can +/// call that directly to skip the scan, at the cost of upholding its +/// precondition. +/// +/// The scan deliberately ignores `u64::MAX` keys (the unmapped coordinate +/// sentinel), which is what keeps `bytes_needed` at the ~5–6 bytes a mapped +/// coordinate key occupies instead of widening to the full 8 the moment a +/// single unmapped read appears — and coordinate-sorted BAMs essentially always +/// carry an unmapped tail. Ignoring it needs no contract from the caller: +/// `u64::MAX` is the largest `u64` *and* truncates to the all-`0xFF` maximum at +/// every radix width, so it sorts after every other key and stays stable among +/// its peers no matter how many passes run. +/// /// # Stability /// Radix sort is inherently stable - records with equal keys maintain their /// relative input order, matching samtools behavior. -#[allow(clippy::uninit_vec, unsafe_code)] pub fn radix_sort_record_refs(refs: &mut [RecordRef]) { + let max_key = max_non_sentinel_key(refs); + radix_sort_record_refs_with_max(refs, max_key); +} + +/// Largest key in `refs` that is not the unmapped sentinel (`u64::MAX`), or 0 +/// if there is no such key. See [`radix_sort_record_refs`] for why excluding +/// the sentinel preserves the sort order. +#[inline] +fn max_non_sentinel_key(refs: &[RecordRef]) -> u64 { + refs.iter().map(|r| r.sort_key).filter(|&k| k != u64::MAX).max().unwrap_or(0) +} + +/// Radix sort for `RecordRef` arrays using a precomputed maximum key, skipping +/// the standalone max-finding scan that [`radix_sort_record_refs`] performs. +/// +/// `max_key` sizes the number of radix byte-passes (`bytes_needed`). Each key +/// must be **either `<= max_key` or exactly `u64::MAX`** (the unmapped +/// coordinate sentinel). Keys `<= max_key` fit within `bytes_needed` bytes and +/// are ordered correctly; `u64::MAX` keys truncate to all-`0xFF` under any +/// width, so they sort *after* every other key and remain stable among +/// themselves. This lets a coordinate sort size the passes from the maximum +/// *mapped* key (~5–6 bytes) instead of the full 8 even when unmapped reads are +/// present. A key strictly between `max_key` and `u64::MAX` would under-size the +/// byte count and mis-order records — debug builds assert against this. +#[allow(clippy::uninit_vec, unsafe_code)] +pub fn radix_sort_record_refs_with_max(refs: &mut [RecordRef], max_key: u64) { let n = refs.len(); if n < RADIX_THRESHOLD { // Use insertion sort for small arrays @@ -1624,21 +1673,49 @@ pub fn radix_sort_record_refs(refs: &mut [RecordRef]) { return; } - // Find max key to determine how many bytes we need to sort - let max_key = refs.iter().map(|r| r.sort_key).max().unwrap_or(0); - let bytes_needed = - if max_key == 0 { 0 } else { ((64 - max_key.leading_zeros()) as usize).div_ceil(8) }; - - if bytes_needed == 0 { - return; // All keys are 0, already sorted + debug_assert!( + refs.iter().all(|r| r.sort_key <= max_key || r.sort_key == u64::MAX), + "radix_sort_record_refs_with_max: a key exceeds max_key ({max_key}) without being the \ + unmapped sentinel (u64::MAX); bytes_needed would be too small and records would mis-sort", + ); + + // `max_key == 0` does NOT mean every key is equal: the slice may still mix + // zero keys with `u64::MAX` sentinels, and those must not be left in input + // order. (A zero key is reachable — `PackedCoordinateKey::new` packs + // `tid = 0, pos = -1, reverse = false` to exactly 0 — so this is not a + // hypothetical.) One pass separates them, since 0 scatters to bucket 0x00 + // and the sentinel to 0xFF, and is a stable no-op when the keys really are + // all identical. So never skip the sort entirely. + let mut bytes_needed = + if max_key == 0 { 1 } else { ((64 - max_key.leading_zeros()) as usize).div_ceil(8) }; + + // A mapped `max_key` that fills all `bytes_needed` bytes (all-`0xFF`, e.g. + // `0xFFFF` from tid=0/pos=32766/reverse) truncates to the same value the + // `u64::MAX` sentinel does at that width, so the two tie across every pass + // and the sentinel is left in input order rather than sorting to the tail. + // Only the max can collide — any smaller key would need all its low bytes + // set, which would make it the max — so widen by one byte in exactly that + // case to force a pass where the sentinel (`0xFF`) and the key (`0x00`) + // differ. `bytes_needed` is at most 8 already (a non-sentinel key is + // `< u64::MAX`), and the guard only fires below 8, so it never overflows. + if bytes_needed < 8 && max_key == u64::MAX >> ((8 - bytes_needed) * 8) { + bytes_needed += 1; } // Allocate auxiliary buffer let mut aux: Vec = Vec::with_capacity(n); + // SAFETY: `aux` is written exactly once per radix pass (the scatter loop + // below) before any element is read; `RecordRef` is `Copy`/`Pod`, so leaving + // it uninitialized until the first scatter is sound. See CLAUDE.md + // "Approved hot-path unsafe". unsafe { aux.set_len(n); } + // SAFETY: `src`/`dst` always point at the disjoint, properly-aligned + // `[RecordRef]` storage of `refs`/`aux` (same length `n`); each pass writes + // every `dst` slot exactly once (the scatter loop) before that buffer is read + // as `src` on the next pass. See CLAUDE.md "Approved hot-path unsafe". let mut src = refs as *mut [RecordRef]; let mut dst = aux.as_mut_slice() as *mut [RecordRef]; @@ -1686,12 +1763,21 @@ pub fn radix_sort_record_refs(refs: &mut [RecordRef]) { /// Divides the array into chunks, sorts each chunk with radix sort, /// then performs k-way merge. This provides near-linear speedup. pub fn parallel_radix_sort_record_refs(refs: &mut [RecordRef]) { + let max_key = max_non_sentinel_key(refs); + parallel_radix_sort_record_refs_with_max(refs, max_key); +} + +/// Parallel radix sort using a precomputed maximum key (see +/// [`radix_sort_record_refs_with_max`] for the `max_key` contract). The same +/// upper bound is handed to every chunk's sort, so no chunk re-scans for its +/// own maximum. +pub fn parallel_radix_sort_record_refs_with_max(refs: &mut [RecordRef], max_key: u64) { use rayon::prelude::*; let n = refs.len(); if n < RADIX_THRESHOLD * 2 { // Small array - just use single-threaded radix sort - radix_sort_record_refs(refs); + radix_sort_record_refs_with_max(refs, max_key); return; } @@ -1702,9 +1788,10 @@ pub fn parallel_radix_sort_record_refs(refs: &mut [RecordRef]) { if n_threads > 1 && n > 10_000 { let chunk_size = n.div_ceil(n_threads); - // Sort each chunk in parallel using radix sort + // Sort each chunk in parallel using radix sort. Every chunk's max key is + // bounded by the whole-slice `max_key`, so the shared bound is safe. refs.par_chunks_mut(chunk_size).for_each(|chunk| { - radix_sort_record_refs(chunk); + radix_sort_record_refs_with_max(chunk, max_key); }); // K-way merge the sorted chunks @@ -1721,7 +1808,7 @@ pub fn parallel_radix_sort_record_refs(refs: &mut [RecordRef]) { merge_sorted_chunks(refs, &chunk_boundaries); } else { // Single-threaded radix sort - radix_sort_record_refs(refs); + radix_sort_record_refs_with_max(refs, max_key); } } @@ -2655,6 +2742,283 @@ mod tests { record } + #[test] + fn test_record_buffer_excludes_unmapped_sentinel_and_sorts_correctly() { + // Unmapped reads (tid < 0) carry the u64::MAX sentinel key. They must NOT + // inflate the bound the radix derives (which would force it to the full + // 8-byte width), yet must still sort to the tail. Exercise a mix above the + // radix threshold so the real radix path (not insertion sort) runs. + let nref = 5; + let mut buffer = RecordBuffer::with_capacity(2048, 2048 * 64, nref); + + // 1500 mapped records with small keys (tid 0-3), interleaved with 500 + // unmapped (tid = -1 -> u64::MAX). Push order is intentionally scrambled. + let mut expected_mapped_max = 0u64; + for i in 0..2000u32 { + if i % 4 == 3 { + buffer.push_coordinate(&make_coordinate_bam_record(-1, -1)).expect("push unmapped"); + } else { + let tid = i32::try_from(i % 3).unwrap(); + let pos = i32::try_from((i * 7) % 100_000).unwrap(); + buffer.push_coordinate(&make_coordinate_bam_record(tid, pos)).expect("push mapped"); + expected_mapped_max = + expected_mapped_max.max(PackedCoordinateKey::new(tid, pos, false, nref).0); + } + } + + // The derived bound excludes the sentinel: it equals the max mapped key, + // far below u64::MAX (so bytes_needed stays ~3-4, not 8). + let bound = max_non_sentinel_key(buffer.refs()); + assert_eq!(bound, expected_mapped_max); + assert_ne!(bound, u64::MAX); + assert!(bound < (1u64 << 48), "mapped max should be well under 8 bytes"); + + // Sorting with the mapped-only max must still produce a fully ordered + // result: mapped keys ascending, all unmapped (u64::MAX) at the tail. + buffer.sort(); + let keys: Vec = buffer.refs().iter().map(|r| r.sort_key).collect(); + assert!(keys.windows(2).all(|w| w[0] <= w[1]), "result is not sorted by sort_key"); + let unmapped = keys.iter().filter(|&&k| k == u64::MAX).count(); + assert_eq!(unmapped, 500, "all unmapped reads should be present"); + assert!( + keys[keys.len() - unmapped..].iter().all(|&k| k == u64::MAX), + "unmapped reads must sort to the tail", + ); + } + + /// `radix_sort_record_refs` and its parallel sibling scan for their own + /// bound, and that scan must ignore the unmapped sentinel — otherwise a + /// single unmapped read widens `bytes_needed` to the full 8 and the whole + /// optimization is lost on exactly the inputs it targets (real coordinate + /// BAMs, which always carry an unmapped tail). Sorting must stay correct + /// either way, so assert both the ordering and the narrowed width. + #[rstest::rstest] + #[case::serial(false)] + #[case::parallel(true)] + fn test_scanning_radix_entry_points_ignore_unmapped_sentinel(#[case] parallel: bool) { + // Well above RADIX_THRESHOLD * 2 so the real radix (and, for the + // parallel case, the chunk-and-merge path) runs rather than insertion sort. + let n = 20_000usize; + let mut refs: Vec = Vec::with_capacity(n); + let mut expected_mapped_max = 0u64; + for i in 0..n { + // Every 4th record is unmapped; the rest get small, scrambled keys. + let key = if i % 4 == 3 { + u64::MAX + } else { + let k = ((i as u64 * 7919) % 100_000) + 1; + expected_mapped_max = expected_mapped_max.max(k); + k + }; + refs.push(RecordRef::new(key, i as u64, 1)); + } + let unmapped_count = refs.iter().filter(|r| r.sort_key == u64::MAX).count(); + + // The bound the entry points derive must be the mapped max, not u64::MAX. + assert_eq!(max_non_sentinel_key(&refs), expected_mapped_max); + assert!( + expected_mapped_max < (1u64 << 24), + "mapped keys should need ~3 bytes, so the sentinel-inclusive max would cost 5 extra passes", + ); + + if parallel { + parallel_radix_sort_record_refs(&mut refs); + } else { + radix_sort_record_refs(&mut refs); + } + + let keys: Vec = refs.iter().map(|r| r.sort_key).collect(); + assert!(keys.windows(2).all(|w| w[0] <= w[1]), "result is not sorted by sort_key"); + assert_eq!( + keys.iter().filter(|&&k| k == u64::MAX).count(), + unmapped_count, + "no unmapped record may be lost or duplicated", + ); + assert!( + keys[keys.len() - unmapped_count..].iter().all(|&k| k == u64::MAX), + "unmapped records must sort to the tail", + ); + } + + /// A buffer holding nothing but unmapped reads derives a bound of 0, which + /// drives `bytes_needed` to 0 and takes the early return. That is a real + /// input shape (an unmapped-only BAM), not a degenerate one, so it must come + /// back with every record intact rather than panicking or dropping the tail. + #[test] + fn test_radix_sort_all_unmapped_records() { + let n = 1_000usize; // above RADIX_THRESHOLD so the radix path is taken + let mut refs: Vec = + (0..n).map(|i| RecordRef::new(u64::MAX, i as u64, 1)).collect(); + + assert_eq!(max_non_sentinel_key(&refs), 0, "an all-unmapped buffer has no mapped bound"); + radix_sort_record_refs(&mut refs); + + assert_eq!(refs.len(), n, "no record may be dropped"); + assert!(refs.iter().all(|r| r.sort_key == u64::MAX), "all keys stay the sentinel"); + // Equal keys means the sort is a no-op, so the stable order is the input + // order -- offsets must still read 0..n. + assert!( + refs.iter().enumerate().all(|(i, r)| r.offset == i as u64), + "equal keys must preserve input order", + ); + } + + /// A zero key mixed with unmapped sentinels is the one case where deriving + /// the bound from mapped keys alone yields `max_key == 0` while the keys are + /// *not* all equal. Skipping the sort there would leave the sentinels ahead + /// of the zero keys. A zero key is reachable, not hypothetical: + /// `PackedCoordinateKey::new` packs `tid = 0, pos = -1, reverse = false` to + /// exactly 0, which a malformed record can carry. + #[test] + fn test_radix_sort_zero_keys_mixed_with_unmapped_sentinel() { + // Confirm the packing really does produce a zero key, so this test keeps + // tracking the reachable case rather than a synthetic one. + assert_eq!( + PackedCoordinateKey::new(0, -1, false, 25).0, + 0, + "tid=0, pos=-1 should pack to a zero key", + ); + + // Above RADIX_THRESHOLD so the radix path runs rather than insertion sort. + let n = 1_000usize; + let mut refs: Vec = (0..n) + .map(|i| { + // Alternate sentinel and zero keys, sentinels first, so an + // unsorted result is immediately visible. + let key = if i % 2 == 0 { u64::MAX } else { 0 }; + RecordRef::new(key, i as u64, 1) + }) + .collect(); + + assert_eq!(max_non_sentinel_key(&refs), 0, "the mapped bound here really is 0"); + radix_sort_record_refs(&mut refs); + + let keys: Vec = refs.iter().map(|r| r.sort_key).collect(); + assert!( + keys.windows(2).all(|w| w[0] <= w[1]), + "zero keys must sort ahead of the unmapped sentinels, not stay in input order", + ); + assert_eq!(keys.iter().filter(|&&k| k == 0).count(), n / 2, "no zero-key record lost"); + assert_eq!(keys.iter().filter(|&&k| k == u64::MAX).count(), n / 2, "no sentinel lost"); + } + + /// `parallel_radix_sort_record_refs` falls back to the single-threaded sort + /// when the slice is above the insertion-sort threshold but below the + /// parallel cutoff. That branch has its own call into the bounded sort, so + /// exercise it at a size that lands between the two. + #[test] + fn test_parallel_radix_sort_single_threaded_fallback_range() { + // RADIX_THRESHOLD * 2 = 512 < n < 10_000 -> parallel cutoff not met. + let n = 5_000usize; + let mut refs: Vec = + (0..n).map(|i| RecordRef::new(((n - i) as u64) + 1, i as u64, 1)).collect(); + + parallel_radix_sort_record_refs(&mut refs); + + let keys: Vec = refs.iter().map(|r| r.sort_key).collect(); + assert!(keys.windows(2).all(|w| w[0] <= w[1]), "fallback path must still sort"); + assert_eq!(keys.len(), n, "no record may be dropped"); + } + + /// Sorting with a bound derived from the mapped keys must be + /// indistinguishable from sorting at full 8-byte width. This is the + /// output-identity check behind the pass-narrowing: the sentinel's + /// truncation to all-`0xFF` has to leave the ordering *and* the stable + /// tie-break order untouched, not merely sorted-looking. + #[test] + fn test_narrowed_radix_matches_full_width_output_exactly() { + let n = 5_000usize; + let mut refs: Vec = Vec::with_capacity(n); + for i in 0..n { + // Deliberate duplicate keys so stability is observable: distinct + // `offset` values act as the tie-break witness. + let key = if i % 5 == 0 { u64::MAX } else { ((i as u64 * 37) % 500) + 1 }; + refs.push(RecordRef::new(key, i as u64, 1)); + } + + let mut narrowed = refs.clone(); + radix_sort_record_refs(&mut narrowed); + + let mut full_width = refs.clone(); + radix_sort_record_refs_with_max(&mut full_width, u64::MAX); + + let as_pairs = + |v: &[RecordRef]| v.iter().map(|r| (r.sort_key, r.offset)).collect::>(); + assert_eq!( + as_pairs(&narrowed), + as_pairs(&full_width), + "narrowing the radix width changed the output; the sentinel must sort identically \ + at every width, including the stable order among equal keys", + ); + } + + /// The one mapped key that collides with the `u64::MAX` sentinel under a + /// narrowed radix is an all-`0xFF` key (`0xFF`, `0xFFFF`, `0xFFFFFF`, ...): + /// it fills every bit of `bytes_needed` bytes, so it truncates to the same + /// all-`0xFF` value the sentinel does at that width, and the two become + /// indistinguishable across every pass. Such a key is ordinary — `0xFFFF` + /// packs from `tid = 0, pos = 32766, reverse = true` — so a sentinel pushed + /// ahead of it in input order would otherwise sort ahead of a mapped read. + /// Placing the sentinels first makes any such mis-order visible, and the + /// full-8-byte sort is an independent oracle for the correct output. + #[rstest::rstest] + #[case::one_byte(0xFFu64)] + #[case::two_bytes(0xFFFFu64)] + #[case::three_bytes(0xFF_FFFFu64)] + fn test_narrowed_radix_separates_all_ones_key_from_sentinel(#[case] all_ones_key: u64) { + // `0xFFFF` is a real coordinate key: tid=0, pos=32766, reverse=true. + if all_ones_key == 0xFFFF { + assert_eq!( + PackedCoordinateKey::new(0, 32766, true, 25).0, + 0xFFFF, + "0xFFFF must be a reachable mapped key, not a synthetic one", + ); + } + + // Above RADIX_THRESHOLD so the real radix path runs, not insertion sort. + let n = 1_000usize; + let refs: Vec = (0..n) + .map(|i| { + // Sentinels on even indices, the all-`0xFF` mapped key on odd: + // every sentinel sits immediately before a colliding mapped key, + // so leaving them in input order is unmistakably unsorted. + let key = if i % 2 == 0 { u64::MAX } else { all_ones_key }; + RecordRef::new(key, i as u64, 1) + }) + .collect(); + let sentinel_count = refs.iter().filter(|r| r.sort_key == u64::MAX).count(); + + // The derived bound is exactly the all-`0xFF` mapped key, which is what + // triggers the collision the fix guards against. + assert_eq!(max_non_sentinel_key(&refs), all_ones_key); + + let mut narrowed = refs.clone(); + radix_sort_record_refs(&mut narrowed); + + // Full-width sort as an independent oracle. + let mut full_width = refs.clone(); + radix_sort_record_refs_with_max(&mut full_width, u64::MAX); + + let as_pairs = + |v: &[RecordRef]| v.iter().map(|r| (r.sort_key, r.offset)).collect::>(); + assert_eq!( + as_pairs(&narrowed), + as_pairs(&full_width), + "an all-0xFF mapped key ({all_ones_key:#x}) tied with the sentinel under the \ + narrowed radix; the two must be separated at the same output as a full-width sort", + ); + + let keys: Vec = narrowed.iter().map(|r| r.sort_key).collect(); + assert!( + keys.windows(2).all(|w| w[0] <= w[1]), + "result is not sorted: the all-0xFF mapped key and the sentinel were left tied", + ); + assert!( + keys[keys.len() - sentinel_count..].iter().all(|&k| k == u64::MAX), + "every unmapped sentinel must sort to the tail, after the all-0xFF mapped key", + ); + } + /// Assert that `chunks` are each individually sorted and that the total /// record count across all chunks equals `expected_total`. fn assert_chunks_sorted_and_complete( @@ -2762,6 +3126,55 @@ mod tests { assert_chunks_sorted_and_complete(&chunks, n); } + /// `par_sort_into_chunks` derives one bound for the whole buffer and shares + /// it across every chunk, so the sentinel exclusion has to hold there too: + /// each chunk must come back sorted and no unmapped record may be lost, on + /// both the single-threaded early-return path and the multi-threaded path + /// (which drain `refs` through different branches of the macro). + #[rstest::rstest] + #[case::single_threaded(100, 1)] + #[case::parallel(10_500, 4)] + fn test_par_sort_into_chunks_handles_unmapped_sentinel( + #[case] n: usize, + #[case] threads: usize, + ) { + let nref = 10u32; + let mut buffer = RecordBuffer::with_capacity(n, n * 50, nref); + let mut expected_unmapped = 0usize; + for i in 0..n { + // Every 4th record unmapped, so the shared bound must exclude the + // sentinel or the chunks would all sort at full width. + if i % 4 == 3 { + buffer + .push_coordinate(&make_coordinate_bam_record(-1, -1)) + .expect("push_coordinate should succeed in tests"); + expected_unmapped += 1; + } else { + let pos = i32::try_from(n - i).expect("test n fits in i32"); + buffer + .push_coordinate(&make_coordinate_bam_record(0, pos)) + .expect("push_coordinate should succeed in tests"); + } + } + + let chunks = if threads > 1 { + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(threads) + .build() + .expect("failed to build rayon thread pool"); + pool.install(|| buffer.par_sort_into_chunks(threads)) + } else { + buffer.par_sort_into_chunks(threads) + }; + + assert_chunks_sorted_and_complete(&chunks, n); + let unmapped: usize = chunks + .iter() + .map(|c| (0..c.len()).filter(|&i| c.key_at(i).sort_key == u64::MAX).count()) + .sum(); + assert_eq!(unmapped, expected_unmapped, "no unmapped record may be lost across chunks"); + } + /// Helper: build a `SegmentedBuf` containing the given byte chunks /// concatenated, plus a vec of `(offset, len)` for each chunk. fn build_segmented_buf_with_records(records: &[&[u8]]) -> (SegmentedBuf, Vec<(u64, u32)>) { diff --git a/crates/fgumi-sort/src/lib.rs b/crates/fgumi-sort/src/lib.rs index e36a44d95..f4de8d33d 100644 --- a/crates/fgumi-sort/src/lib.rs +++ b/crates/fgumi-sort/src/lib.rs @@ -162,7 +162,10 @@ pub use codec::SpillCodec; pub use external::{ KeyTypesSpec, LibraryLookup, RawExternalSorter, cb_hasher, extract_template_key_inline, }; -pub use inline::{TemplateKey, extract_coordinate_key_inline}; +pub use inline::{ + PackedCoordinateKey, RecordRef, TemplateKey, extract_coordinate_key_inline, + radix_sort_record_refs, radix_sort_record_refs_with_max, +}; pub use keys::{ QuerynameComparator, RawCoordinateKey, RawQuerynameKey, RawQuerynameLexKey, RawSortKey, SortContext, SortOrder, natural_compare, natural_compare_nul, normalize_natural_key, From f8a371b4ee44fcf7291fc007d0328e40860d9d2f Mon Sep 17 00:00:00 2001 From: Nils Homer Date: Mon, 20 Jul 2026 16:54:19 -0700 Subject: [PATCH 3/3] perf(sort): borrow record bytes from the decompressed block on ingest Pooled sort ingest copied each record's bytes twice after decompression: once out of the decompressed block into a `RawRecord`, then again from that `RawRecord` into the sort arena. The first copy is avoidable whenever the record body lies wholly within the current block, which is the common case. Adds `PooledInputStream::next_record_borrowed`, a lending reader that returns a slice borrowed straight out of `current_buf`, falling back to a reusable scratch buffer only when the body or its 4-byte length prefix straddles a block boundary. The coordinate and template-coordinate ingest loops consume borrowed slices and push them into the buffer directly, taking copy amplification from two to one. The keyed/queryname path keeps owned records, which it requires. `RecordSource` grows a matching `next_record_borrowed`. Both non-pooled variants already yield owned records and so have nothing shared to borrow from: each stores the record it just took and lends a slice into it. This includes the `Stream` variant, which is not present upstream -- it gains a held slot alongside its error slot, and its producer-error handling is shared with the `Iterator` impl so a failure reaches `take_error()` identically on both paths. That matters because the ingest loops treat `Ok(None)` as end-of-input, so an error swallowed there would silently truncate a sort rather than fail it. Tests cover records and length prefixes straddling block boundaries down to 1-byte blocks, parity against the owned `read_raw_record` path, a proptest over randomized bodies and block sizes, and -- for the `Stream` variant -- agreement with owned iteration plus propagation of a mid-stream producer error. --- crates/fgumi-raw-bam/src/lib.rs | 3 +- crates/fgumi-raw-bam/src/raw_bam_record.rs | 18 +- crates/fgumi-sort/src/external.rs | 14 +- crates/fgumi-sort/src/read_ahead.rs | 453 ++++++++++++++++++++- 4 files changed, 458 insertions(+), 30 deletions(-) diff --git a/crates/fgumi-raw-bam/src/lib.rs b/crates/fgumi-raw-bam/src/lib.rs index 386079421..39db7a15e 100644 --- a/crates/fgumi-raw-bam/src/lib.rs +++ b/crates/fgumi-raw-bam/src/lib.rs @@ -98,7 +98,8 @@ pub use overlap::{ // -- raw_bam_record -- pub use raw_bam_record::{ - RawBamReader, RawBamWriter, RawRecord, read_raw_record, write_raw_record, write_raw_records, + RawBamReader, RawBamWriter, RawRecord, read_block_size, read_raw_record, write_raw_record, + write_raw_records, }; // -- sequence -- diff --git a/crates/fgumi-raw-bam/src/raw_bam_record.rs b/crates/fgumi-raw-bam/src/raw_bam_record.rs index 1b2348c2a..244f016ba 100644 --- a/crates/fgumi-raw-bam/src/raw_bam_record.rs +++ b/crates/fgumi-raw-bam/src/raw_bam_record.rs @@ -171,10 +171,22 @@ where Ok(block_size) } -/// Reads the 4-byte block size prefix. +/// Reads the 4-byte BAM record `block_size` prefix (little-endian u32). /// -/// Returns 0 at EOF (no bytes available). -fn read_block_size(reader: &mut R) -> io::Result +/// Returns 0 at EOF (no bytes available before the prefix starts). Handles a +/// prefix split across reader boundaries by issuing a 1-byte read first (to +/// detect clean EOF) followed by a `read_exact` of the remaining 3 bytes. +/// +/// Exposed for the sort ingest borrow-in-place path +/// (`PooledInputStream::next_record_borrowed`), which reads the prefix itself so +/// it can decide whether the record body can be borrowed from the current +/// decompressed block or must be gathered into a scratch buffer. +/// +/// # Errors +/// +/// Returns an error if the reader errors, the prefix is truncated, or the value +/// does not fit in `usize`. +pub fn read_block_size(reader: &mut R) -> io::Result where R: Read, { diff --git a/crates/fgumi-sort/src/external.rs b/crates/fgumi-sort/src/external.rs index 6fb616c68..0ea8b8ee6 100644 --- a/crates/fgumi-sort/src/external.rs +++ b/crates/fgumi-sort/src/external.rs @@ -2263,12 +2263,16 @@ impl RawExternalSorter { debug!("Phase 1: Reading and sorting chunks (inline buffer, keyed output)..."); let mut probe = SpillProbe::new("phase1"); - for record in record_source.by_ref() { + // Borrow each record's bytes straight out of the decompressed block and + // push them into the arena, skipping the intermediate `RawRecord` copy + // (the borrowed slice is invalidated by the next `next_record_borrowed` + // call, which is fine — `push_coordinate` copies the bytes into the buffer). + while let Some(record) = record_source.next_record_borrowed()? { stats.total_records += 1; progress.log_if_needed(1); // Push directly to buffer - key extracted inline from raw bytes - buffer.push_coordinate(record.as_ref())?; + buffer.push_coordinate(record)?; if probe.should_sample_read(stats.total_records) { probe.log_mid_read(probe_stats(&buffer), Some(pool.phase1_queue_depths())); @@ -3020,14 +3024,16 @@ impl RawExternalSorter { buffer.push(bam_bytes, K::from_full(&full))?; } - for record in record_source.by_ref() { + // Borrow each record's bytes in place (see the coordinate ingest loop); + // the key is extracted and the bytes copied into the buffer before the + // borrow ends, so no owned `RawRecord` is needed here. + while let Some(bam_bytes) = record_source.next_record_borrowed()? { stats.total_records += 1; progress.log_if_needed(1); // Extract the full template key, verify the lanes the chosen variant // dropped are constant relative to the first record, then push the // narrowed key. - let bam_bytes = record.as_ref(); let full = extract_template_key_inline(bam_bytes, lib_lookup, self.cell_tag, cb_hasher); if let Some(violation) = verify_dropped_lanes(&first, &full, variant) { let name = String::from_utf8_lossy( diff --git a/crates/fgumi-sort/src/read_ahead.rs b/crates/fgumi-sort/src/read_ahead.rs index facc748ec..b43ed4d96 100644 --- a/crates/fgumi-sort/src/read_ahead.rs +++ b/crates/fgumi-sort/src/read_ahead.rs @@ -248,6 +248,10 @@ pub struct PooledInputStream { current_buf: Vec, /// Read position within `current_buf`. current_pos: usize, + /// Reusable scratch buffer for records (or their length prefixes) that + /// straddle a decompressed-block boundary and therefore cannot be borrowed + /// directly out of `current_buf`. See [`PooledInputStream::next_record_borrowed`]. + scratch: RawRecord, } impl PooledInputStream { @@ -267,6 +271,7 @@ impl PooledInputStream { reorder: fgumi_bam_io::ReorderBuffer::new(), current_buf: Vec::new(), current_pos: 0, + scratch: RawRecord::new(), } } @@ -347,6 +352,75 @@ impl PooledInputStream { } } } + + /// Read the next raw BAM record, borrowing its bytes from the current + /// decompressed block when possible. + /// + /// This is the borrow-in-place counterpart to [`read_raw_record`]: it removes + /// the per-record `read_exact` copy into a `RawRecord` on the common path + /// where the record body lies wholly within the current decompressed block. + /// + /// - **Fast path:** when the 4-byte `block_size` prefix and the full record + /// body are both contained in `current_buf`, returns a slice borrowed + /// directly out of `current_buf` (one copy: block → caller's arena). + /// - **Slow path (block straddle):** when the prefix *or* the body spans a + /// decompressed-block boundary, the bytes are gathered into a reusable + /// internal scratch buffer via the byte-stream [`IoRead`] impl and a slice + /// into that scratch is returned. + /// + /// The returned slice borrows `self`; it is invalidated by the next call to + /// any method on this stream. Returns `Ok(None)` at clean EOF. + /// + /// A `block_size` of 0 is treated as EOF, mirroring [`read_raw_record`]. + /// + /// # Errors + /// + /// Returns an error if the underlying block stream errors (I/O or + /// decompression), the prefix/body is truncated, or `block_size` overflows + /// `usize`. + pub fn next_record_borrowed(&mut self) -> std::io::Result> { + // --- read the 4-byte block_size prefix --- + let block_size = if self.current_buf.len() - self.current_pos >= 4 { + // Fast path: the prefix is fully buffered — decode it in place. + let p = self.current_pos; + let n = u32::from_le_bytes([ + self.current_buf[p], + self.current_buf[p + 1], + self.current_buf[p + 2], + self.current_buf[p + 3], + ]); + self.current_pos += 4; + usize::try_from(n) + .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))? + } else { + // Slow path: the prefix straddles a block boundary (or the buffer is + // exhausted / at EOF). `read_block_size` reads across blocks via the + // `IoRead` impl and returns 0 at clean EOF. + fgumi_raw_bam::read_block_size(self)? + }; + if block_size == 0 { + return Ok(None); // EOF (or a zero-length prefix, treated as EOF) + } + + // --- read the record body --- + if self.current_buf.len() - self.current_pos >= block_size { + // Fast path: the body lies wholly within the current block — borrow it. + let start = self.current_pos; + self.current_pos += block_size; + Ok(Some(&self.current_buf[start..start + block_size])) + } else { + // Slow path: the body straddles a block boundary — gather it into the + // reusable scratch buffer. Take the scratch out so `read_exact` can + // borrow `self` mutably, then restore it (preserving its capacity). + let mut scratch = std::mem::take(&mut self.scratch); + let buf = scratch.as_mut_vec(); + buf.resize(block_size, 0); + let res = self.read_exact(buf); + self.scratch = scratch; + res?; + Ok(Some(self.scratch.as_ref())) + } + } } impl IoRead for PooledInputStream { @@ -418,8 +492,13 @@ pub type BoxedRecordStream = Box> /// caller-supplied in-process stream (no BAM input file at all). pub enum RecordSource { /// Legacy path: background thread prefetches records. + /// + /// The second field holds the most-recently-yielded owned record so that + /// [`RecordSource::next_record_borrowed`] can lend a slice into it: this path + /// delivers owned `RawRecord`s over a channel, so there is nothing in a + /// shared block to borrow — the owned record is the thing we lend. #[allow(dead_code)] // retained for potential future use / benchmarking - ReadAhead(RawReadAheadReader), + ReadAhead(RawReadAheadReader, Option), /// Pool path: main thread reads directly from pool's decompressed stream. /// /// The second field stores the first I/O error encountered during iteration, @@ -431,8 +510,11 @@ pub enum RecordSource { /// /// The second field mirrors `Direct`'s error slot: a producer error stops /// iteration and is reported by `take_error()`, so the sort aborts instead - /// of writing a silently truncated output. - Stream(BoxedRecordStream, Option), + /// of writing a silently truncated output. The third field mirrors + /// `ReadAhead`'s held slot: this path also yields owned records, so + /// [`RecordSource::next_record_borrowed`] lends a slice into the record it + /// just took rather than borrowing from a shared decompressed block. + Stream(BoxedRecordStream, Option, Option), } impl RecordSource { @@ -449,17 +531,80 @@ impl RecordSource { I: IntoIterator>, I::IntoIter: Send + 'static, { - Self::Stream(Box::new(records.into_iter()), None) + Self::Stream(Box::new(records.into_iter()), None, None) } /// Take any I/O error that occurred during iteration. /// /// Returns `Some(err)` if the iterator stopped due to a read error rather than /// clean EOF. Call this after exhausting the iterator to detect truncated input. + /// + /// Note: [`RecordSource::next_record_borrowed`] on the `Direct` path + /// propagates errors directly via its `Result` rather than stashing them + /// here, so for that path `take_error()` returns `None`. The `ReadAhead` + /// and `Stream` paths still report their errors here, because both end + /// iteration on failure rather than returning it inline. pub fn take_error(&mut self) -> Option { match self { - Self::Direct(_, err) | Self::Stream(_, err) => err.take(), - Self::ReadAhead(r) => r.take_error(), + Self::Direct(_, err) | Self::Stream(_, err, _) => err.take(), + Self::ReadAhead(r, _) => r.take_error(), + } + } + + /// Read the next record, borrowing its bytes in place where possible. + /// + /// On the `Direct` (pool) path this borrows straight out of the current + /// decompressed block via [`PooledInputStream::next_record_borrowed`], + /// avoiding the per-record copy into a `RawRecord`. The `ReadAhead` and + /// `Stream` paths deliver owned records — from a background thread and an + /// in-process producer respectively — so there is no shared block to borrow + /// from; each stores the record it just took and lends a slice into it. + /// + /// The returned slice borrows `self` and is invalidated by the next call. + /// Returns `Ok(None)` at clean EOF. + /// + /// # Errors + /// + /// Returns an error if the underlying `Direct` block stream errors or the + /// input is truncated. The `ReadAhead` and `Stream` paths defer errors to + /// [`RecordSource::take_error`], so a caller looping on this method must + /// still check `take_error()` afterwards to tell a real EOF from a failure. + pub fn next_record_borrowed(&mut self) -> std::io::Result> { + match self { + Self::Direct(reader, _error_slot) => reader.get_mut().next_record_borrowed(), + Self::ReadAhead(r, held) => { + *held = r.next_record(); + Ok(held.as_ref().map(RawRecord::as_ref)) + } + Self::Stream(records, error_slot, held) => { + *held = next_stream_record(records, error_slot); + Ok(held.as_ref().map(RawRecord::as_ref)) + } + } + } +} + +/// Pull the next record from an in-process producer, stashing a producer error +/// into `error_slot` and ending the stream. Shared by [`RecordSource`]'s +/// [`Iterator`] impl and [`RecordSource::next_record_borrowed`] so both report +/// producer failures identically through [`RecordSource::take_error`]. +fn next_stream_record( + records: &mut BoxedRecordStream, + error_slot: &mut Option, +) -> Option { + match records.next() { + Some(Ok(record)) => Some(record), + None => None, + Some(Err(e)) => { + log::error!("Error producing record for sort: {e:#}"); + // Preserve the first error; don't overwrite with later ones. + if error_slot.is_none() { + // `anyhow::Error` converts into the boxed source type + // `io::Error::other` takes, so the message and cause + // chain survive the round-trip through `take_error()`. + *error_slot = Some(std::io::Error::other(e)); + } + None } } } @@ -469,7 +614,7 @@ impl Iterator for RecordSource { fn next(&mut self) -> Option { match self { - Self::ReadAhead(r) => r.next(), + Self::ReadAhead(r, _) => r.next(), Self::Direct(reader, error_slot) => { let mut record = RawRecord::default(); match reader.read_record(&mut record) { @@ -485,21 +630,7 @@ impl Iterator for RecordSource { } } } - Self::Stream(records, error_slot) => match records.next() { - Some(Ok(record)) => Some(record), - None => None, - Some(Err(e)) => { - log::error!("Error producing record for sort: {e:#}"); - // Preserve the first error; don't overwrite with later ones. - if error_slot.is_none() { - // `anyhow::Error` converts into the boxed source type - // `io::Error::other` takes, so the message and cause - // chain survive the round-trip through `take_error()`. - *error_slot = Some(std::io::Error::other(e)); - } - None - } - }, + Self::Stream(records, error_slot, _) => next_stream_record(records, error_slot), } } } @@ -585,4 +716,282 @@ mod tests { let records: Vec = ra.collect(); assert_eq!(records.len(), num, "Raw read-ahead should yield exactly {num} records"); } + + // ── PooledInputStream::next_record_borrowed — block-straddle tests ─────── + // + // These build a PooledInputStream whose decompressed blocks are a fixed + // (often tiny) size, so records and their 4-byte length prefixes straddle + // block boundaries. The borrow-in-place reader must reassemble straddling + // records via its scratch buffer and produce byte-identical output to the + // owned read_record path. + + use crossbeam_queue::ArrayQueue; + use proptest::{prop_assert_eq, proptest}; + use rstest::rstest; + use std::sync::Arc; + use std::sync::atomic::AtomicBool; + + /// Frame a record body with its 4-byte little-endian `block_size` prefix. + fn frame_record(body: &[u8]) -> Vec { + let mut v = Vec::with_capacity(4 + body.len()); + v.extend_from_slice(&u32::try_from(body.len()).expect("body fits u32").to_le_bytes()); + v.extend_from_slice(body); + v + } + + /// Build a `PooledInputStream` whose decompressed blocks are exactly + /// `block_len` bytes each (the last may be shorter). Small `block_len` + /// values force prefixes and bodies to straddle block boundaries. + fn pooled_stream_from(records: &[Vec], block_len: usize) -> PooledInputStream { + assert!(block_len >= 1, "block_len must be positive"); + let mut stream = Vec::new(); + for body in records { + stream.extend_from_slice(&frame_record(body)); + } + let chunks: Vec> = if stream.is_empty() { + Vec::new() + } else { + stream.chunks(block_len).map(<[u8]>::to_vec).collect() + }; + let queue = Arc::new(ArrayQueue::new(chunks.len().max(1))); + for (serial, chunk) in chunks.into_iter().enumerate() { + let serial = u64::try_from(serial).expect("serial fits u64"); + queue.push((serial, chunk)).expect("queue has capacity for all chunks"); + } + PooledInputStream::new( + queue, + Arc::new(AtomicBool::new(true)), // decompressed_input_done + Arc::new(AtomicBool::new(false)), // input_read_error + Arc::new(AtomicBool::new(false)), // decompression_error + ) + } + + /// Drain all records via the borrowing API, copying each borrowed slice into + /// an owned `Vec` for comparison. + fn collect_borrowed(stream: &mut PooledInputStream) -> Vec> { + let mut out = Vec::new(); + while let Some(rec) = stream.next_record_borrowed().expect("borrow read should succeed") { + out.push(rec.to_vec()); + } + out + } + + /// Record bodies of assorted lengths with position-dependent content, so a + /// misaligned read produces detectably wrong bytes. + fn sample_record_bodies() -> Vec> { + let lens = [1usize, 2, 3, 4, 5, 8, 13, 32, 33, 100, 255, 256, 257]; + lens.iter() + .enumerate() + .map(|(k, &len)| { + let base = u8::try_from(k % 256).expect("k%256 fits u8"); + (0..len) + .map(|j| base.wrapping_add(u8::try_from(j % 256).expect("j%256 fits u8"))) + .collect() + }) + .collect() + } + + #[rstest] + fn test_next_record_borrowed_matches_input_across_block_sizes( + // Tiny sizes force straddles; large sizes (> whole stream) take the + // all-in-one-block fast path. + #[values(1usize, 2, 3, 4, 5, 6, 7, 8, 16, 64, 1024, 65_535)] block_len: usize, + ) { + let records = sample_record_bodies(); + let mut stream = pooled_stream_from(&records, block_len); + let got = collect_borrowed(&mut stream); + assert_eq!(got, records, "record mismatch at block_len={block_len}"); + } + + #[test] + fn test_next_record_borrowed_empty_stream() { + let mut stream = pooled_stream_from(&[], 4); + assert!( + stream.next_record_borrowed().expect("empty stream read should succeed").is_none(), + "empty stream should yield no records" + ); + } + + #[test] + fn test_next_record_borrowed_prefix_straddle() { + // block_len = 2 guarantees every record's 4-byte prefix spans at least + // two blocks, exercising the slow prefix path (fgumi_raw_bam::read_block_size). + let records = vec![vec![0xAB; 10], vec![0xCD; 7], vec![0xEF; 1]]; + let mut stream = pooled_stream_from(&records, 2); + assert_eq!(collect_borrowed(&mut stream), records); + } + + #[test] + fn test_next_record_borrowed_body_exact_boundary_and_straddle() { + // With block_len = 10 the first framed record ([prefix(4)|body(6)] = 10 + // bytes) ends exactly on a block boundary; the next records straddle. + let records = vec![vec![1u8; 6], vec![2u8; 9], vec![3u8; 3]]; + let mut stream = pooled_stream_from(&records, 10); + assert_eq!(collect_borrowed(&mut stream), records); + } + + #[rstest] + fn test_next_record_borrowed_parity_with_read_record( + #[values(1usize, 3, 7, 64)] block_len: usize, + ) { + // The borrow-in-place reader must produce byte-identical records to the + // owned read_raw_record path over the same chunked stream. + let records = sample_record_bodies(); + let mut borrowed_stream = pooled_stream_from(&records, block_len); + let borrowed = collect_borrowed(&mut borrowed_stream); + + let owned_stream = pooled_stream_from(&records, block_len); + let mut reader = fgumi_raw_bam::RawBamReader::new(owned_stream); + let mut owned = Vec::new(); + let mut rec = fgumi_raw_bam::RawRecord::new(); + loop { + let n = reader.read_record(&mut rec).expect("read_record should succeed"); + if n == 0 { + break; + } + owned.push(rec.as_ref().to_vec()); + } + + assert_eq!(borrowed, owned, "borrowed vs owned mismatch at block_len={block_len}"); + assert_eq!(borrowed, records, "borrowed records must match input"); + } + + #[test] + fn test_next_record_borrowed_truncated_body_errors() { + // A framed record claiming a 20-byte body but only 5 bytes present must + // surface an error (not silently return a short or wrong record). + let mut stream = Vec::new(); + stream.extend_from_slice(&20u32.to_le_bytes()); + stream.extend_from_slice(&[7u8; 5]); + let queue = Arc::new(ArrayQueue::new(1)); + queue.push((0u64, stream)).expect("push"); + let mut pooled = PooledInputStream::new( + queue, + Arc::new(AtomicBool::new(true)), + Arc::new(AtomicBool::new(false)), + Arc::new(AtomicBool::new(false)), + ); + let err = pooled.next_record_borrowed().expect_err("truncated body should error"); + assert_eq!(err.kind(), std::io::ErrorKind::UnexpectedEof); + } + + // ── RecordSource::next_record_borrowed — Stream variant ────────────────── + // + // The `Stream` path has no shared decompressed block to borrow from: it + // yields owned records from an in-process producer, so `next_record_borrowed` + // lends a slice into the record it just took. These pin that the borrowed + // path agrees with the owned `Iterator` path and, critically, that a producer + // error still reaches `take_error()` rather than reading as a clean EOF — + // the ingest loops treat `Ok(None)` as end-of-input, so a swallowed error + // there would silently truncate a sort. + + /// Build a `RawRecord` whose bytes are exactly `body`. + fn raw_record_from(body: &[u8]) -> RawRecord { + let mut rec = RawRecord::new(); + rec.as_mut_vec().extend_from_slice(body); + rec + } + + /// The `ReadAhead` variant lends a slice into the owned record it just took + /// from the background thread. It must agree with owned iteration over the + /// same input and report no error on clean EOF. + #[test] + fn test_read_ahead_next_record_borrowed_matches_owned_iteration() { + let num = 10; + let (tmp, _header) = create_test_bam_file(num); + + let (reader, _header) = + create_raw_bam_reader(tmp.path(), 1).expect("creating BAM reader should succeed"); + let mut borrowed_source = RecordSource::ReadAhead(RawReadAheadReader::new(reader), None); + let mut borrowed = Vec::new(); + while let Some(rec) = + borrowed_source.next_record_borrowed().expect("read-ahead borrow should succeed") + { + borrowed.push(rec.to_vec()); + } + + // Compare against the owned path through `RecordSource` itself, so both + // sides go through the same abstraction rather than the inner reader. + let (reader, _header) = + create_raw_bam_reader(tmp.path(), 1).expect("creating BAM reader should succeed"); + let owned: Vec> = RecordSource::ReadAhead(RawReadAheadReader::new(reader), None) + .map(|r| r.as_ref().to_vec()) + .collect(); + + assert_eq!(borrowed.len(), num, "read-ahead should yield exactly {num} records"); + assert_eq!(borrowed, owned, "borrowed and owned read-ahead iteration must agree"); + assert!(borrowed_source.take_error().is_none(), "clean EOF must not report an error"); + } + + #[test] + fn test_stream_next_record_borrowed_matches_owned_iteration() { + let bodies: Vec> = vec![vec![1u8; 5], vec![2u8; 40], vec![3u8; 1], vec![4u8; 300]]; + + let mut borrowed_source = + RecordSource::stream(bodies.iter().map(|b| Ok(raw_record_from(b))).collect::>()); + let mut borrowed = Vec::new(); + while let Some(rec) = + borrowed_source.next_record_borrowed().expect("stream borrow should succeed") + { + borrowed.push(rec.to_vec()); + } + + let owned_source = + RecordSource::stream(bodies.iter().map(|b| Ok(raw_record_from(b))).collect::>()); + let owned: Vec> = owned_source.map(|r| r.as_ref().to_vec()).collect(); + + assert_eq!(borrowed, owned, "borrowed and owned stream iteration must agree"); + assert_eq!(borrowed, bodies, "stream records must round-trip unchanged"); + assert!(borrowed_source.take_error().is_none(), "clean EOF must not report an error"); + } + + #[test] + fn test_stream_next_record_borrowed_surfaces_producer_error() { + // Two good records, then a producer failure. The borrowed reader must + // yield the good records, stop, and report the error via take_error(). + let items: Vec> = vec![ + Ok(raw_record_from(&[1u8; 10])), + Ok(raw_record_from(&[2u8; 10])), + Err(anyhow::anyhow!("producer exploded")), + Ok(raw_record_from(&[3u8; 10])), + ]; + let mut source = RecordSource::stream(items); + + let mut seen = 0usize; + while let Some(_rec) = source.next_record_borrowed().expect("borrow itself must not error") + { + seen += 1; + } + assert_eq!(seen, 2, "iteration must stop at the producer error"); + + let err = source.take_error().expect("producer error must reach take_error()"); + assert!( + err.to_string().contains("producer exploded"), + "error message should survive the anyhow -> io::Error round-trip, got: {err}", + ); + } + + proptest! { + /// Property: over randomized record bodies and decompressed-block sizes, + /// the borrow-in-place reader yields records byte-identical to the input + /// (and, transitively, to the owned `read_record` path). This widens the + /// boundary coverage of the fixed straddle examples — any combination of + /// record length and block length where a prefix or body crosses a block + /// edge must still reassemble correctly via the scratch buffer. + #[test] + fn prop_next_record_borrowed_matches_input( + // Up to 24 records, each 1..=300 bytes of arbitrary content. + records in proptest::collection::vec( + proptest::collection::vec(proptest::num::u8::ANY, 1..=300), + 0..=24, + ), + // Block sizes from 1 (every prefix straddles) up past the largest + // record (whole-record fast path). + block_len in 1usize..=512, + ) { + let mut stream = pooled_stream_from(&records, block_len); + let got = collect_borrowed(&mut stream); + prop_assert_eq!(got, records); + } + } }