diff --git a/crates/fgumi-bam-io/Cargo.toml b/crates/fgumi-bam-io/Cargo.toml index 9d60eb086..311333bde 100644 --- a/crates/fgumi-bam-io/Cargo.toml +++ b/crates/fgumi-bam-io/Cargo.toml @@ -27,12 +27,12 @@ bstr = { workspace = true } bytes = { workspace = true } crossbeam-channel = { workspace = true } log = { workspace = true } +tempfile = { workspace = true } [target.'cfg(target_os = "linux")'.dependencies] nix = { workspace = true } [dev-dependencies] -tempfile = { workspace = true } rstest = { workspace = true } proptest = { workspace = true } criterion = { workspace = true } diff --git a/crates/fgumi-bam-io/src/lib.rs b/crates/fgumi-bam-io/src/lib.rs index 134c47c80..5d929d79c 100644 --- a/crates/fgumi-bam-io/src/lib.rs +++ b/crates/fgumi-bam-io/src/lib.rs @@ -44,8 +44,8 @@ pub use reader::{ }; pub use reorder::{DrainReady, ReorderBuffer}; pub use writer::{ - BaiBuilder, BamWriter, BgzfWriterEnum, IndexingBamWriter, RawBamWriter, bai_sidecar_path, - create_bam_writer, create_indexing_bam_writer, create_optional_bam_writer, - create_raw_bam_writer, open_output_writer, write_bai_index, write_bai_sidecar, - write_bam_header, + AlignmentContext, BaiBuilder, BamWriter, BgzfWriterEnum, IndexingBamWriter, RawBamWriter, + bai_sidecar_path, create_bam_writer, create_indexing_bam_writer, create_optional_bam_writer, + create_raw_bam_writer, extract_alignment_context, open_output_writer, write_bai_index, + write_bai_sidecar, write_bam_header, }; diff --git a/crates/fgumi-bam-io/src/writer.rs b/crates/fgumi-bam-io/src/writer.rs index 5dcac0c55..362200d03 100644 --- a/crates/fgumi-bam-io/src/writer.rs +++ b/crates/fgumi-bam-io/src/writer.rs @@ -43,6 +43,9 @@ use crate::vendored::{BlockInfoRx, MultithreadedWriter, MultithreadedWriterBuild /// `BGZF_BLOCK_SIZE` (`crate::vendored::bgzf_multithreaded`), a deliberately /// independent copy that keeps the vendored module self-contained. It is equal /// to this value today; keep the two in step, since nothing enforces it. +/// - The arena chain sink's inline indexer, which derives physical-block numbers +/// as `uoffset / fgumi_bgzf::BGZF_MAX_BLOCK_SIZE` — statically equal to this +/// value (`fgumi-bgzf/src/writer.rs:17,325`). const MAX_BLOCK_SIZE: usize = bgzf::BGZF_BLOCK_SIZE; /// Fast, non-cryptographic hasher for the dense sequential `u64` block numbers @@ -360,23 +363,41 @@ struct CachedIndexEntry { offset_in_block: usize, /// Length of record (including 4-byte size prefix). record_len: usize, - /// Alignment context: (reference ID, start, end, mapped flag). - alignment_context: Option<(usize, Position, Position, bool)>, + /// Alignment context, if the record is placed. + alignment_context: Option, +} + +/// A placed record's reference, span, and mapped status. +/// +/// Public so pipeline steps in other crates can extract context once and carry +/// it in the BAM index manifest, keeping this the single source of the logic. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct AlignmentContext { + /// Reference sequence ID (0-based). + pub ref_id: usize, + /// 1-based start position (inclusive). + pub start: Position, + /// 1-based end position (exclusive), spanning the reference length + /// consumed by the CIGAR (see [`extract_alignment_context`]). + pub end: Position, + /// Whether the record is mapped (the `UNMAPPED` flag is unset). + pub is_mapped: bool, } /// Extract alignment context from raw BAM bytes. /// -/// Returns `Some((ref_id, start, end, is_mapped))` for any read that has a -/// reference and a position — **including a placed-but-unmapped read** (e.g. the -/// unmapped mate of a mapped read, which carries its mate's `tid`/`pos` so it -/// sorts alongside it). Such reads are position-binned exactly as htslib/samtools -/// do, so region queries and `idxstats` see them; the `is_mapped` flag records -/// that they are unmapped without excluding them from the index. Only a truly -/// unplaced read (no reference or no position) returns `None`, which the indexer -/// counts toward the unplaced total. +/// Returns `Some(AlignmentContext)` for any read that has a reference and a +/// position — **including a placed-but-unmapped read** (e.g. the unmapped mate +/// of a mapped read, which carries its mate's `tid`/`pos` so it sorts alongside +/// it). Such reads are position-binned exactly as htslib/samtools do, so region +/// queries and `idxstats` see them; the `is_mapped` flag records that they are +/// unmapped without excluding them from the index. Only a truly unplaced read +/// (no reference or no position) returns `None`, which the indexer counts +/// toward the unplaced total. #[inline] +#[must_use] #[allow(clippy::cast_sign_loss, clippy::cast_possible_wrap)] -pub(crate) fn extract_alignment_context(bam: &[u8]) -> Option<(usize, Position, Position, bool)> { +pub fn extract_alignment_context(bam: &[u8]) -> Option { let v = fgumi_raw_bam::RawRecordView::new(bam); let tid = v.ref_id(); let pos = v.pos(); @@ -401,7 +422,7 @@ pub(crate) fn extract_alignment_context(bam: &[u8]) -> Option<(usize, Position, let start = Position::try_from((pos + 1) as usize).ok()?; let end = Position::try_from((pos + ref_len) as usize).ok()?; - Some((tid as usize, start, end, is_mapped)) + Some(AlignmentContext { ref_id: tid as usize, start, end, is_mapped }) } /// Incrementally builds a BAI index from per-record positions plus per-block @@ -476,12 +497,33 @@ impl BaiBuilder { record_len: usize, record_bytes: &[u8], ) { - let alignment_context = extract_alignment_context(record_bytes); + self.record_with_context( + block_number, + offset_in_block, + record_len, + extract_alignment_context(record_bytes), + ); + } + + /// Record a written BAM record's position for later index resolution, given + /// an already-extracted alignment context. + /// + /// Identical to [`Self::record`] except the caller supplies `ctx` directly + /// instead of raw record bytes — for callers (e.g. the arena writer) that + /// have already extracted [`AlignmentContext`] for other purposes and would + /// otherwise redo the same parse. + pub fn record_with_context( + &mut self, + block_number: u64, + offset_in_block: usize, + record_len: usize, + ctx: Option, + ) { self.entry_cache.push_back(CachedIndexEntry { block_number, offset_in_block, record_len, - alignment_context, + alignment_context: ctx, }); } @@ -534,7 +576,7 @@ impl BaiBuilder { // starts a fresh run. The record's own context still drives bin // routing and the linear index inside `add_record`; only the chunk's // start position is widened. - let chunk_start = if let Some((reference_sequence_id, start, end, _)) = + let chunk_start = if let Some(AlignmentContext { ref_id, start, end, .. }) = alignment_context { let bin = reg2bin(start, end); @@ -542,13 +584,15 @@ impl BaiBuilder { // `self.run` below doesn't overlap the borrow. let current = self.run.as_ref().map(|r| (r.reference_sequence_id, r.bin, r.start)); match current { - Some((run_ref, run_bin, run_start)) - if run_ref == reference_sequence_id && run_bin == bin => - { + Some((run_ref, run_bin, run_start)) if run_ref == ref_id && run_bin == bin => { run_start } _ => { - self.run = Some(ChunkRun { reference_sequence_id, bin, start: start_vpos }); + self.run = Some(ChunkRun { + reference_sequence_id: ref_id, + bin, + start: start_vpos, + }); start_vpos } } @@ -559,7 +603,10 @@ impl BaiBuilder { start_vpos }; self.indexer - .add_record(alignment_context, Chunk::new(chunk_start, end_vpos)) + .add_record( + alignment_context.map(|c| (c.ref_id, c.start, c.end, c.is_mapped)), + Chunk::new(chunk_start, end_vpos), + ) .map_err(io::Error::other)?; self.entry_cache.pop_front(); } @@ -582,10 +629,31 @@ impl BaiBuilder { /// Number of records still awaiting block-position resolution. #[must_use] - pub(crate) fn pending(&self) -> usize { + pub fn pending(&self) -> usize { self.entry_cache.len() } + /// Drop noted block offsets no future record can reference and advance the + /// watermark. Called by the arena writer after each fully-drained batch + /// (`pending() == 0`); safe because records never cross batch boundaries, so + /// nothing recorded later references a block `< block_number`. Keeps + /// `block_positions` bounded to the in-flight window (see the field doc on + /// [`Self`]) — without it, the arena's per-batch full drain never triggers + /// `resolve`'s own front-gated prune and the map grows for the whole file. + pub fn prune_below(&mut self, block_number: u64) { + while self.pruned_below < block_number { + self.block_positions.remove(&self.pruned_below); + self.pruned_below += 1; + } + } + + /// Number of entries currently retained in `block_positions`, for test + /// assertions that the pruning window stays bounded. + #[cfg(test)] + pub(crate) fn block_positions_len_for_test(&self) -> usize { + self.block_positions.len() + } + /// Build the final BAI index. /// /// The index is compacted (see `compact_bai_index`) so its on-disk size @@ -1062,6 +1130,11 @@ pub fn write_bai_sidecar>(bam_path: P) -> Result { /// /// # Errors /// Returns an error if the file cannot be created or writing the index fails. +/// +/// The write is atomic: the serialized index is written to a temporary file in +/// the destination directory and then renamed onto `path`, so a failure partway +/// through never leaves a truncated or partial `.bai` at the final path (a +/// half-written index parses as valid but silently mis-answers queries). pub fn write_bai_index>(path: P, index: &bai::Index) -> Result<()> { let path_ref = path.as_ref(); // `bai::io::Writer` serializes the index as a great many tiny fields (each @@ -1073,8 +1146,21 @@ pub fn write_bai_index>(path: P, index: &bai::Index) -> Result<() bai::io::Writer::new(&mut buf) .write_index(index) .with_context(|| format!("Failed to serialize BAI index for: {}", path_ref.display()))?; - std::fs::write(path_ref, &buf) + // Write to a temp file in the destination directory, then rename onto the + // final path. A same-directory rename is atomic, so a reader never observes + // a partial index, and a failed write leaves no file at `path` at all. + let dir = match path_ref.parent() { + Some(parent) if !parent.as_os_str().is_empty() => parent, + _ => Path::new("."), + }; + let mut tmp = tempfile::Builder::new() + .prefix(".bai-") + .tempfile_in(dir) + .with_context(|| format!("Failed to create temp file for index in: {}", dir.display()))?; + tmp.write_all(&buf) .with_context(|| format!("Failed to write index to: {}", path_ref.display()))?; + tmp.persist(path_ref) + .with_context(|| format!("Failed to persist index to: {}", path_ref.display()))?; Ok(()) } @@ -1508,20 +1594,18 @@ mod tests { // binned over a 1-base span at its position, flagged unmapped — so region // queries and idxstats see it, matching htslib. let rec = create_placed_unmapped_bam_record(0, 100, b"unmapped_mate"); - let (ref_id, start, end, is_mapped) = - extract_alignment_context(&rec).expect("placed-unmapped read must be binned"); - assert_eq!(ref_id, 0); - assert_eq!(usize::from(start), 101, "0-based pos 100 -> 1-based 101"); - assert_eq!(usize::from(end), 101, "no CIGAR -> 1-base span [pos, pos+1)"); - assert!(!is_mapped, "flag must still record it as unmapped"); + let ctx = extract_alignment_context(&rec).expect("placed-unmapped read must be binned"); + assert_eq!(ctx.ref_id, 0); + assert_eq!(usize::from(ctx.start), 101, "0-based pos 100 -> 1-based 101"); + assert_eq!(usize::from(ctx.end), 101, "no CIGAR -> 1-base span [pos, pos+1)"); + assert!(!ctx.is_mapped, "flag must still record it as unmapped"); // A mapped read is binned over its CIGAR span and flagged mapped. let mapped = create_test_bam_record(0, 100, b"mapped"); - let (_, m_start, m_end, m_mapped) = - extract_alignment_context(&mapped).expect("mapped read must be binned"); - assert_eq!(usize::from(m_start), 101); - assert_eq!(usize::from(m_end), 110, "10M CIGAR -> span of 10"); - assert!(m_mapped); + let m_ctx = extract_alignment_context(&mapped).expect("mapped read must be binned"); + assert_eq!(usize::from(m_ctx.start), 101); + assert_eq!(usize::from(m_ctx.end), 110, "10M CIGAR -> span of 10"); + assert!(m_ctx.is_mapped); // A truly unplaced read (no reference) is not binned. assert!( @@ -1531,6 +1615,116 @@ mod tests { ); } + /// Encode a CIGAR op (`(length, op_char)`) as the BAM-packed `u32` word + /// [`SamBuilder::cigar_ops`] expects: `(length << 4) | op_code`, per the SAM + /// spec's op-code table (`MIDNSHP=X` -> `0..=8`). + fn cigar_op_word(length: u32, op: char) -> u32 { + let op_code = match op { + 'M' => 0, + 'I' => 1, + 'D' => 2, + 'N' => 3, + 'S' => 4, + 'H' => 5, + 'P' => 6, + '=' => 7, + 'X' => 8, + _ => panic!("unsupported CIGAR op: {op}"), + }; + (length << 4) | op_code + } + + /// Build a record **body** (no 4-byte length prefix) via [`SamBuilder`], for + /// use with [`extract_alignment_context`]. + fn build_record_body(ref_id: i32, pos: i32, flags: u16, cigar: &[(u32, char)]) -> Vec { + let ops: Vec = cigar.iter().map(|&(len, op)| cigar_op_word(len, op)).collect(); + fgumi_raw_bam::SamBuilder::new() + .ref_id(ref_id) + .pos(pos) + .flags(flags) + .cigar_ops(&ops) + .build() + .to_vec() + } + + #[test] + fn extract_alignment_context_classifies_placed_unmapped_and_unplaced() { + use noodles::core::Position; + // Placed + mapped: tid=0, pos=99 (0-based) → start=100 (1-based). + let mapped = build_record_body( + /*ref_id*/ 0, + /*pos*/ 99, + /*flags*/ 0, + /*cigar*/ &[(10, 'M')], + ); + let ctx = extract_alignment_context(&mapped).expect("placed"); + assert_eq!(ctx.ref_id, 0); + assert_eq!(ctx.start, Position::try_from(100).unwrap()); + // 0-based pos=99, 10M -> reference span [99, 109) (0-based, exclusive) -> + // 1-based inclusive end = 109 (matches `bam_endpos`-style pos + ref_len). + assert_eq!(ctx.end, Position::try_from(109).unwrap()); + assert!(ctx.is_mapped); + + // Placed-but-unmapped mate: tid/pos valid, UNMAPPED flag set, no CIGAR → span floored to 1. + let placed_unmapped = build_record_body(0, 99, fgumi_raw_bam::flags::UNMAPPED, &[]); + let ctx = + extract_alignment_context(&placed_unmapped).expect("placed-unmapped is still placed"); + assert!(!ctx.is_mapped); + assert_eq!(ctx.end, Position::try_from(100).unwrap()); + + // Truly unplaced: tid < 0 → None. + let unplaced = build_record_body(-1, -1, fgumi_raw_bam::flags::UNMAPPED, &[]); + assert!(extract_alignment_context(&unplaced).is_none()); + } + + #[test] + fn record_with_context_matches_record_from_bytes() { + // Same inputs via the raw-bytes path and the pre-extracted-context path must + // produce an identical built index. + let body = build_record_body(0, 99, 0, &[(10, 'M')]); + let record_len = 4 + body.len(); + + let mut a = BaiBuilder::new(); + a.note_block(0, 1234); + a.record(0, 0, record_len, &body); + a.resolve().unwrap(); + let ia = a.build(1).unwrap(); + + let mut b = BaiBuilder::new(); + b.note_block(0, 1234); + b.record_with_context(0, 0, record_len, extract_alignment_context(&body)); + b.resolve().unwrap(); + let ib = b.build(1).unwrap(); + + // Serialize both and compare bytes — same index by construction. + let mut ba = Vec::new(); + let mut bb = Vec::new(); + noodles::bam::bai::io::Writer::new(&mut ba).write_index(&ia).unwrap(); + noodles::bam::bai::io::Writer::new(&mut bb).write_index(&ib).unwrap(); + assert_eq!(ba, bb); + } + + #[test] + fn prune_below_bounds_block_positions_across_drained_batches() { + // Simulate the arena writer's per-batch pattern: note this batch's blocks, + // record its records, resolve (fully drains → pending()==0), prune_below the + // next batch's base. block_positions must not accumulate across batches. + let body = build_record_body(0, 0, 0, &[(5, 'M')]); + let record_len = 4 + body.len(); + let mut b = BaiBuilder::new(); + let mut next_block_no = 0u64; + for _ in 0..1000 { + b.note_block(next_block_no, next_block_no * 100); // one physical block/batch + b.record_with_context(next_block_no, 0, record_len, extract_alignment_context(&body)); + b.resolve().unwrap(); + assert_eq!(b.pending(), 0, "each batch drains fully"); + next_block_no += 1; + b.prune_below(next_block_no); + } + // Bounded: at most the current in-flight window, not ~1000 entries. + assert!(b.block_positions_len_for_test() <= 1, "block_positions must stay bounded"); + } + #[test] fn test_create_indexing_bam_writer() -> Result<()> { let temp_file = NamedTempFile::new()?; @@ -1901,4 +2095,31 @@ mod tests { // reg2bin(4681)'s parent must be 585, matching the fold target. assert_eq!(bin_parent(4681), 585); } + + /// `write_bai_index` writes atomically via a same-directory temp file and a + /// rename: the result parses as a valid BAI, and no `.bai-` temp artifact is + /// left behind on success (a partial `.bai` would parse yet mis-answer + /// queries, which is why the write must be all-or-nothing). + #[test] + fn write_bai_index_is_atomic_and_leaves_no_temp_artifact() { + let dir = tempfile::TempDir::new().expect("temp dir"); + let path = dir.path().join("out.bam.bai"); + write_bai_index(&path, &bai::Index::default()).expect("write bai"); + + // The final sidecar exists and parses. + assert!(path.exists(), "sidecar must exist at the final path"); + noodles::bam::bai::fs::read(&path).expect("written .bai must parse"); + + // The temp file was renamed onto the target, not left in the directory. + let temp_leftovers: Vec<_> = std::fs::read_dir(dir.path()) + .expect("read temp dir") + .filter_map(std::result::Result::ok) + .filter(|e| e.file_name().to_string_lossy().starts_with(".bai-")) + .map(|e| e.file_name()) + .collect(); + assert!( + temp_leftovers.is_empty(), + "temp file must be renamed away, found: {temp_leftovers:?}" + ); + } } diff --git a/crates/fgumi-pipeline-io/src/sink/write_bgzf.rs b/crates/fgumi-pipeline-io/src/sink/write_bgzf.rs index 9e6cad0e2..bdfc014fb 100644 --- a/crates/fgumi-pipeline-io/src/sink/write_bgzf.rs +++ b/crates/fgumi-pipeline-io/src/sink/write_bgzf.rs @@ -6,9 +6,10 @@ use std::fs::File; use std::io::{self, BufWriter, Write}; -use std::path::Path; +use std::path::{Path, PathBuf}; -use fgumi_bgzf::{BGZF_EOF, InlineBgzfCompressor}; +use fgumi_bam_io::{BaiBuilder, write_bai_index}; +use fgumi_bgzf::{BGZF_EOF, BGZF_MAX_BLOCK_SIZE, InlineBgzfCompressor}; use noodles::sam::Header; use parking_lot::Mutex; @@ -35,6 +36,25 @@ pub struct WriteBgzfFile { struct WriterState { out: BufWriter, pending_header: Option, + /// Cumulative compressed bytes written so far, header included. Used to + /// attribute each joined block's constituent BGZF blocks to their + /// physical (compressed) file offsets for inline BAI indexing. + coffset: u64, + /// Present only when [`WriteBgzfFile::with_bai_index`] was called. `None` + /// for every writer that doesn't build an inline index. + bai: Option, +} + +/// Per-writer inline BAI indexing state: the [`BaiBuilder`] accumulating +/// resolved chunks, the running physical block-number watermark (continuing +/// across batches so block numbers stay globally unique and monotonic), the +/// reference count to build the final index with, and the sidecar path to +/// write it to. +struct BaiSink { + bai: BaiBuilder, + next_block_no: u64, + num_refs: usize, + sidecar_path: PathBuf, } /// Transform applied to the aligner's runtime-resolved header before it is @@ -72,10 +92,17 @@ impl WriteBgzfFile { let mut hc = InlineBgzfCompressor::new(compression_level); hc.write_all(&header_bytes)?; hc.flush()?; - hc.write_blocks_to(&mut out)?; + // Compress to a `Vec` first (rather than `write_blocks_to(&mut out)` + // directly) so the header's compressed length is known up front — the + // starting point for `coffset`, the cumulative compressed-byte counter + // the inline BAI join uses to attribute blocks to physical offsets. + let mut header_blocks = Vec::new(); + hc.write_blocks_to(&mut header_blocks)?; + let coffset = header_blocks.len() as u64; + out.write_all(&header_blocks)?; Ok(Self { - state: Mutex::new(Some(WriterState { out, pending_header: None })), + state: Mutex::new(Some(WriterState { out, pending_header: None, coffset, bai: None })), name: "WriteBgzfFile", detached_group: None, }) @@ -96,6 +123,59 @@ impl WriteBgzfFile { self } + /// Attach an inline BAI indexer: as each [`BgzfBlock`] carrying a manifest + /// (see [`crate::types::BamIndexManifest`]) is written, its records are + /// joined into a fresh [`BaiBuilder`] using this writer's own cumulative + /// compressed offset, and the resulting `.bai` is written to + /// `sidecar_path` once the BGZF stream is drained and finished. + /// + /// Requires the eager-header constructor ([`Self::new`]) — a + /// deferred-header writer ([`Self::new_with_handle`]) has not yet written + /// its header (and so has no fixed `coffset` starting point) when this is + /// called, and joining a manifest against a `coffset` that later shifts + /// underneath it would produce a `.bai` with wrong virtual offsets. + /// Deferred-header inline indexing is unsupported for now. + /// + /// # Precondition: the record stream MUST already be coordinate-sorted + /// + /// The inline indexer position-bins each record's [`AlignmentContext`] + /// (from its manifest entry) into the [`BaiBuilder`] purely by the order + /// records arrive at this sink — it does **not** read or verify `@HD SO`, + /// and it does not re-derive sort order from the coordinates themselves + /// (unlike the retired `IndexBamFinalizeHook`, which built the `.bai` via + /// `noodles::bam::fs::index` reading back a `SO:coordinate` BAM). Feeding + /// a non-coordinate-sorted stream through this sink produces a `.bai` that + /// parses and looks valid but is semantically wrong (queries against it + /// will silently miss or misattribute records). + /// + /// Coordinate-only is enforced upstream of this sink today, not here: the + /// command layer rejects non-coordinate `--write-index` requests, and the + /// chain builder (`fgumi_lib::pipeline::chains::builder`) only constructs + /// a `BamWithIndex` sink spec for coordinate sorts. Any future caller of + /// `with_bai_index` must preserve (or re-verify) that guarantee before + /// wiring a new record stream into this sink. + /// + /// [`AlignmentContext`]: fgumi_bam_io::AlignmentContext + /// + /// # Panics + /// + /// Panics if `state.pending_header` is `Some` (i.e. this was built via + /// [`Self::new_with_handle`]). + #[must_use] + pub fn with_bai_index(self, sidecar_path: PathBuf, num_refs: usize) -> Self { + { + let mut guard = self.state.lock(); + let state = guard.as_mut().expect("state present"); + assert!( + state.pending_header.is_none(), + "with_bai_index requires the eager-header constructor (deferred-header inline indexing is unsupported)" + ); + state.bai = + Some(BaiSink { bai: BaiBuilder::new(), next_block_no: 0, num_refs, sidecar_path }); + } + self + } + /// Open `path` and return the sink with the BAM header write /// deferred until an upstream step resolves `handle`. /// @@ -119,6 +199,8 @@ impl WriteBgzfFile { state: Mutex::new(Some(WriterState { out, pending_header: Some(PendingHeader { handle, compression_level, transform }), + coffset: 0, + bai: None, })), name: "WriteBgzfFile", detached_group: None, @@ -149,7 +231,15 @@ impl WriteBgzfFile { let mut hc = InlineBgzfCompressor::new(level); hc.write_all(&header_bytes)?; hc.flush()?; - hc.write_blocks_to(&mut state.out)?; + // As in `new`: measure the header's compressed length via a `Vec` + // first so `state.coffset` stays correct for a future fused + // deferred-header + inline-index path (today `with_bai_index` asserts + // against a pending header, so this is currently dead weight — but + // cheap and correct to keep in sync). + let mut header_blocks = Vec::new(); + hc.write_blocks_to(&mut header_blocks)?; + state.coffset += header_blocks.len() as u64; + state.out.write_all(&header_blocks)?; state.pending_header = None; Ok(true) @@ -195,7 +285,40 @@ impl Step for WriteBgzfFile { let header_ready = Self::try_write_pending_header(state)?; if header_ready && let Some(block) = ctx.input.pop() { + // Join this block's manifest into the inline BAI builder, using + // this writer's own cumulative compressed offset, BEFORE writing + // the block's bytes — `state.coffset` at this point is exactly the + // physical file offset the block's first physical BGZF block will + // land at. + if let Some(sink) = state.bai.as_mut() + && let Some(manifest) = &block.index + { + let base = sink.next_block_no; + let mut off = state.coffset; // read by value; no borrow overlap with `sink` + for (j, &clen) in manifest.phys_comp_len.iter().enumerate() { + sink.bai.note_block(base + j as u64, off); + off += u64::from(clen); + } + debug_assert_eq!(off, state.coffset + block.bytes.len() as u64); + for e in &manifest.records { + let block_no = base + (e.uoffset as usize / BGZF_MAX_BLOCK_SIZE) as u64; + let within = e.uoffset as usize % BGZF_MAX_BLOCK_SIZE; + sink.bai.record_with_context(block_no, within, e.len as usize, e.ctx); + } + sink.next_block_no += manifest.phys_comp_len.len() as u64; + sink.bai.resolve().map_err(io::Error::other)?; + debug_assert_eq!( + sink.bai.pending(), + 0, + "resolve drains fully per batch on the arena layout" + ); + // With the cache fully drained, `resolve`'s own front-gated + // prune did not run (nothing left in `entry_cache` to gate on) + // — prune explicitly so `block_positions` stays bounded. + sink.bai.prune_below(sink.next_block_no); + } state.out.write_all(&block.bytes)?; + state.coffset += block.bytes.len() as u64; return Ok(StepOutcome::Progress); } @@ -207,6 +330,11 @@ impl Step for WriteBgzfFile { } state.out.write_all(&BGZF_EOF)?; state.out.flush()?; + if let Some(sink) = state.bai.take() { + let index = sink.bai.build(sink.num_refs).map_err(io::Error::other)?; + write_bai_index(&sink.sidecar_path, &index).map_err(io::Error::other)?; + log::info!("Wrote BAM index: {}", sink.sidecar_path.display()); + } let _ = guard.take(); return Ok(StepOutcome::Finished); } @@ -353,8 +481,8 @@ mod tests { // first, so ordinal 0 is emitted before ordinal 1 for the `ByItemOrdinal` // reorder stage. The bytes are arbitrary non-EOF markers. let blocks = vec![ - BgzfBlock { batch_serial: 1, bytes: vec![0xAB; 16], uncompressed_size: 0 }, - BgzfBlock { batch_serial: 0, bytes: vec![0xCD; 16], uncompressed_size: 0 }, + BgzfBlock { batch_serial: 1, bytes: vec![0xAB; 16], uncompressed_size: 0, index: None }, + BgzfBlock { batch_serial: 0, bytes: vec![0xCD; 16], uncompressed_size: 0, index: None }, ]; let source = BlockSource { blocks, held: HeldSlot::new() }; @@ -384,6 +512,243 @@ mod tests { assert!(bytes.windows(16).any(|w| w == [0xCD; 16]), "second block written"); } + /// L5 positive coverage for the inline-BAI join: drives a real 2-step + /// `BlockSource → WriteBgzfFile` pipeline (mirroring + /// `try_run_drained_finish_writes_exactly_one_eof`) with `with_bai_index` + /// attached and two payload blocks carrying manifests. Block B's record + /// straddles a physical block boundary (`uoffset % B + len > B`) to + /// exercise `compute_end_vpos`'s forward walk across noted blocks. The + /// synthetic compressed bytes are not valid BGZF — validity of the BGZF + /// stream is not under test here, only that the manifest offsets join + /// correctly into a parseable `.bai`. + #[test] + #[allow(clippy::too_many_lines)] // inline test Step + exhaustive offset assertions + fn inline_index_writes_valid_bai_with_straddling_record() { + use fgumi_bam_io::AlignmentContext; + use fgumi_pipeline_core::{ + Unpushed, + builder::{Pipeline, PipelineConfig}, + held::HeldSlot, + outputs::OrderedBytesSingle, + queues::QueueSpec, + reorder::BranchOrdering, + }; + use noodles::core::Position; + + use crate::types::{BamIndexManifest, RecordIndexEntry}; + + /// Exclusive source draining a `Vec`, one block per `try_run`. + struct BlockSource { + blocks: Vec, + held: HeldSlot>, + } + impl Step for BlockSource { + type Input = (); + type Outputs = OrderedBytesSingle; + fn profile(&self) -> StepProfile { + StepProfile { + name: "BlockSource", + kind: StepKind::Exclusive, + sticky: true, + output_queues: vec![QueueSpec::ByteBounded { limit_bytes: 1 << 20 }], + branch_ordering: vec![BranchOrdering::ByItemOrdinal], + } + } + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + if let Some(unpushed) = self.held.take() + && let Err(again) = ctx.outputs.retry(unpushed) + { + self.held.put(again); + return Ok(StepOutcome::Progress); + } + let Some(block) = self.blocks.pop() else { + return Ok(StepOutcome::Finished); + }; + if let Err(unpushed) = ctx.outputs.push(block) { + self.held.put(unpushed); + } + Ok(StepOutcome::Progress) + } + } + + const B: usize = BGZF_MAX_BLOCK_SIZE; + + let path = tempfile::NamedTempFile::new().unwrap().into_temp_path(); + let bai_path = fgumi_bam_io::bai_sidecar_path(&path); + let sink = WriteBgzfFile::new(&path, &empty_header(), 1) + .unwrap() + .with_bai_index(bai_path.clone(), 1); + + let ctx0 = Some(AlignmentContext { + ref_id: 0, + start: Position::try_from(1).unwrap(), + end: Position::try_from(11).unwrap(), + is_mapped: true, + }); + let ctx1 = Some(AlignmentContext { + ref_id: 0, + start: Position::try_from(200).unwrap(), + end: Position::try_from(210).unwrap(), + is_mapped: true, + }); + + // Block A (ordinal 0): one physical block, one record fully inside it. + let block_a = BgzfBlock { + batch_serial: 0, + bytes: vec![0xAA; 50], + uncompressed_size: 30, + index: Some(Box::new(BamIndexManifest { + phys_comp_len: vec![50], + records: vec![RecordIndexEntry { uoffset: 0, len: 20, ctx: ctx0 }], + })), + }; + // Block B (ordinal 1): two physical blocks (block numbers 1 and 2, + // continuing from block A's single block 0); one record straddles the + // boundary between them. + let straddle_uoffset = u32::try_from(B - 10).unwrap(); + let block_b = BgzfBlock { + batch_serial: 1, + bytes: vec![0xBB; 90], + uncompressed_size: u32::try_from(B).unwrap() + 20, + index: Some(Box::new(BamIndexManifest { + phys_comp_len: vec![45, 45], + records: vec![RecordIndexEntry { uoffset: straddle_uoffset, len: 20, ctx: ctx1 }], + })), + }; + + // `pop()` drains the tail first, so ordinal 0 (block_a) is emitted + // before ordinal 1 (block_b) — see the comment in + // `try_run_drained_finish_writes_exactly_one_eof`. + let blocks = vec![block_b, block_a]; + let source = BlockSource { blocks, held: HeldSlot::new() }; + + let builder = Pipeline::builder(); + builder.chain(source).chain(sink).into_sink_marker(); + let pipeline = builder.build().expect("pipeline builds"); + pipeline + .run(PipelineConfig { threads: 2, ..Default::default() }) + .expect("pipeline runs to completion"); + + let index_bytes = std::fs::read(&bai_path).expect("bai sidecar written"); + let index = noodles::bam::bai::io::Reader::new(&index_bytes[..]) + .read_index() + .expect("bai sidecar parses"); + let refseqs = index.reference_sequences(); + assert_eq!(refseqs.len(), 1, "num_refs propagated"); + + // The join's virtual offsets are hand-computable from the fixed manifests. + // The only runtime unknown is the compressed header length; recover it + // independently from the file layout — header (H bytes) + block_a (50) + + // block_b (90) + BGZF EOF — so nothing below is read from the code under test. + let bam_bytes = std::fs::read(&path).expect("bam written"); + let header_len = bam_bytes.len() as u64 - 50 - 90 - BGZF_EOF.len() as u64; + + let vpos = |coffset: u64, uoffset: u16| { + noodles::bgzf::VirtualPosition::try_from((coffset, uoffset)).expect("valid vpos") + }; + // Both records fall in the same bin (positions 1 and 200 are both inside + // the first 16 kbp bin), so their per-record chunks coalesce into one + // chunk spanning the earliest start to the latest end. Its endpoints pin + // the two offsets the join must get right: + // start = record 0's start — block A's single physical block at + // coffset `header_len`, uncompressed offset 0; + // end = record 1's end — the straddling record ends 10 bytes into + // block B's *second* physical block (coffset header_len + 95), + // which is the forward walk `compute_end_vpos` performs. A + // dropped straddling record would end at record 0's end + // (vpos(header_len, 20)) instead. + let expected_start = vpos(header_len, 0); + let expected_end = vpos(header_len + 95, 10); + + let chunks: Vec<_> = refseqs[0] + .bins() + .iter() + .flat_map(|(_, bin)| bin.chunks().iter().map(|c| (c.start(), c.end()))) + .collect(); + assert_eq!( + chunks, + vec![(expected_start, expected_end)], + "the joined bin must hold exactly one chunk spanning record 0's start to the \ + straddling record 1's end" + ); + } + + /// No payload blocks reach the sink (only the header). Drained-finish must + /// still build and write a `.bai` whose references are all present but + /// empty — matching a coordinate-sorted BAM with no records. + #[test] + fn inline_index_empty_input_writes_valid_empty_bai() { + use fgumi_pipeline_core::{ + Unpushed, + builder::{Pipeline, PipelineConfig}, + held::HeldSlot, + outputs::OrderedBytesSingle, + queues::QueueSpec, + reorder::BranchOrdering, + }; + /// Exclusive source that immediately finishes with no blocks. + struct BlockSource { + blocks: Vec, + held: HeldSlot>, + } + impl Step for BlockSource { + type Input = (); + type Outputs = OrderedBytesSingle; + fn profile(&self) -> StepProfile { + StepProfile { + name: "BlockSource", + kind: StepKind::Exclusive, + sticky: true, + output_queues: vec![QueueSpec::ByteBounded { limit_bytes: 1 << 20 }], + branch_ordering: vec![BranchOrdering::ByItemOrdinal], + } + } + fn try_run(&mut self, ctx: &mut StepCtx<'_, Self>) -> io::Result { + if let Some(unpushed) = self.held.take() + && let Err(again) = ctx.outputs.retry(unpushed) + { + self.held.put(again); + return Ok(StepOutcome::Progress); + } + let Some(block) = self.blocks.pop() else { + return Ok(StepOutcome::Finished); + }; + if let Err(unpushed) = ctx.outputs.push(block) { + self.held.put(unpushed); + } + Ok(StepOutcome::Progress) + } + } + + let path = tempfile::NamedTempFile::new().unwrap().into_temp_path(); + let bai_path = fgumi_bam_io::bai_sidecar_path(&path); + let sink = WriteBgzfFile::new(&path, &empty_header(), 1) + .unwrap() + .with_bai_index(bai_path.clone(), 3); + + let source = BlockSource { blocks: vec![], held: HeldSlot::new() }; + + let builder = Pipeline::builder(); + builder.chain(source).chain(sink).into_sink_marker(); + let pipeline = builder.build().expect("pipeline builds"); + pipeline + .run(PipelineConfig { threads: 2, ..Default::default() }) + .expect("pipeline runs to completion"); + + let index_bytes = std::fs::read(&bai_path).expect("bai sidecar written for empty input"); + let index = noodles::bam::bai::io::Reader::new(&index_bytes[..]) + .read_index() + .expect("bai sidecar parses"); + assert_eq!( + index.reference_sequences().len(), + 3, + "empty index still carries an entry per reference" + ); + for rs in index.reference_sequences() { + assert!(rs.bins().is_empty(), "no records means no bins"); + } + } + #[test] fn new_with_handle_defers_header_until_resolved() { let path = tempfile::NamedTempFile::new().unwrap().into_temp_path(); diff --git a/crates/fgumi-pipeline-io/src/sink/write_raw.rs b/crates/fgumi-pipeline-io/src/sink/write_raw.rs index 4f01b08c2..9df116788 100644 --- a/crates/fgumi-pipeline-io/src/sink/write_raw.rs +++ b/crates/fgumi-pipeline-io/src/sink/write_raw.rs @@ -141,7 +141,12 @@ mod tests { fn raw_bytes_block_exposes_payload() { let plain = DecompressedBlock { batch_serial: 0, bytes: b"ACGT".to_vec() }; assert_eq!(RawBytesBlock::bytes(&plain), b"ACGT"); - let bgzf = BgzfBlock { batch_serial: 0, bytes: b"\x1f\x8b".to_vec(), uncompressed_size: 4 }; + let bgzf = BgzfBlock { + batch_serial: 0, + bytes: b"\x1f\x8b".to_vec(), + uncompressed_size: 4, + index: None, + }; assert_eq!(RawBytesBlock::bytes(&bgzf), b"\x1f\x8b"); } diff --git a/crates/fgumi-pipeline-io/src/sort/arena_ingest.rs b/crates/fgumi-pipeline-io/src/sort/arena_ingest.rs index 28cc840cb..dd265d799 100644 --- a/crates/fgumi-pipeline-io/src/sort/arena_ingest.rs +++ b/crates/fgumi-pipeline-io/src/sort/arena_ingest.rs @@ -1789,6 +1789,7 @@ mod tests { batch_serial: 0, bytes: blk0.clone(), uncompressed_size: u32::try_from(isz0).unwrap(), + index: None, }) .is_ok() ); @@ -1797,6 +1798,7 @@ mod tests { batch_serial: 1, bytes: blk1.clone(), uncompressed_size: u32::try_from(isz1).unwrap(), + index: None, }) .is_ok() ); @@ -1869,6 +1871,7 @@ mod tests { batch_serial: 0, bytes: blk0.clone(), uncompressed_size: u32::try_from(isz0).unwrap(), + index: None, }) .is_ok() ); @@ -1877,6 +1880,7 @@ mod tests { batch_serial: 1, bytes: blk1.clone(), uncompressed_size: u32::try_from(isz1).unwrap(), + index: None, }) .is_ok() ); @@ -1922,6 +1926,7 @@ mod tests { batch_serial: 99, bytes: probe_blk, uncompressed_size: u32::try_from(probe_isz).unwrap(), + index: None, }); assert!(result.is_err(), "pool must be exhausted while run 0 Arc is still in flight"); // A failed admit (ensure_arena returned false) must not leak state: the @@ -1940,6 +1945,7 @@ mod tests { batch_serial: 2, bytes: blk2.clone(), uncompressed_size: u32::try_from(isz2).unwrap(), + index: None, }) .is_ok(), "run 1 admit must succeed after pool release" @@ -1949,6 +1955,7 @@ mod tests { batch_serial: 3, bytes: blk3.clone(), uncompressed_size: u32::try_from(isz3).unwrap(), + index: None, }) .is_ok(), "run 1 second admit must succeed" diff --git a/crates/fgumi-pipeline-io/src/sort/tests.rs b/crates/fgumi-pipeline-io/src/sort/tests.rs index 382859b6d..b2f4e0afa 100644 --- a/crates/fgumi-pipeline-io/src/sort/tests.rs +++ b/crates/fgumi-pipeline-io/src/sort/tests.rs @@ -1826,6 +1826,7 @@ fn bgzf_blocks_for( batch_serial: i as u64, bytes: blocks.remove(0).data, uncompressed_size: u32::try_from(payload.len()).expect("payload fits u32"), + index: None, } }) .collect() diff --git a/crates/fgumi-pipeline-io/src/source/read_bam.rs b/crates/fgumi-pipeline-io/src/source/read_bam.rs index 9516cb6e3..4e5a32827 100644 --- a/crates/fgumi-pipeline-io/src/source/read_bam.rs +++ b/crates/fgumi-pipeline-io/src/source/read_bam.rs @@ -141,6 +141,7 @@ impl Step for ReadBgzfBlocks { ) })?, bytes: raw.data, + index: None, }); } diff --git a/crates/fgumi-pipeline-io/src/types.rs b/crates/fgumi-pipeline-io/src/types.rs index bb8ab245a..38cc33893 100644 --- a/crates/fgumi-pipeline-io/src/types.rs +++ b/crates/fgumi-pipeline-io/src/types.rs @@ -8,6 +8,57 @@ use fgumi_pipeline_core::{HeapSize, Ordered}; use fgumi_raw_bam::RawRecord; +// ───────────────────────────────────────────────────────────────────────────── +// BamIndexManifest / RecordIndexEntry — inline BAI sidecar metadata. +// ───────────────────────────────────────────────────────────────────────────── + +/// One record's virtual-offset bookkeeping within a compressed [`BgzfBlock`], +/// as needed to build a BAI index inline with compression. +#[derive(Debug, Clone, Copy)] +pub struct RecordIndexEntry { + /// Offset of the record's 4-byte `block_size` length prefix, measured + /// from the start of the *uncompressed* bytes of the batch that produced + /// this block (i.e. the BGZF virtual offset's uncompressed-offset half, + /// before the physical block offset is added in). + pub uoffset: u32, + /// Total record length in bytes: `4` (the length prefix) plus the body + /// length. + pub len: u32, + /// The record's alignment context (reference id, start/end, mapped + /// flag), or `None` for a record with no reference/position (fully + /// unplaced unmapped reads are not indexable by coordinate). + pub ctx: Option, +} + +/// Sidecar metadata riding alongside a compressed [`BgzfBlock`], carrying +/// everything an inline BAI indexer needs to attribute records to virtual +/// offsets without re-parsing the compressed bytes. +/// +/// # Invariants +/// +/// - `phys_comp_len.iter().map(|&n| n as u64).sum::() == bytes.len() as +/// u64`, where `bytes` is the sibling [`BgzfBlock::bytes`] this manifest +/// describes: the physical lengths of the constituent BGZF blocks sum to +/// the total compressed byte count. +/// - `phys_comp_len.len() == (uncompressed_size as +/// usize).div_ceil(fgumi_bgzf::BGZF_MAX_BLOCK_SIZE)`, where +/// `uncompressed_size` is the sibling [`BgzfBlock::uncompressed_size`]: +/// one physical block per `BGZF_MAX_BLOCK_SIZE`-sized (or smaller, for the +/// final one) chunk of uncompressed input. +/// - `records` is in batch order, which is final coordinate order. +/// - Each entry's `uoffset` is the offset of the record's 4-byte length +/// prefix within the batch's uncompressed bytes; `len` is `4 +` the body +/// length. +#[derive(Debug, Clone)] +pub struct BamIndexManifest { + /// Physical (compressed) length of each constituent BGZF block, in the + /// order those blocks appear in the sibling [`BgzfBlock::bytes`]. + pub phys_comp_len: Vec, + /// Per-record virtual-offset bookkeeping, in batch (= final coordinate) + /// order. + pub records: Vec, +} + // ───────────────────────────────────────────────────────────────────────────── // BgzfBlock — raw compressed BGZF block + read-order serial. // ───────────────────────────────────────────────────────────────────────────── @@ -23,13 +74,22 @@ pub struct BgzfBlock { pub bytes: Vec, /// Decompressed size, parsed from the BGZF block header. pub uncompressed_size: u32, + /// Sidecar index metadata for this block's records, populated only on the + /// inline-BAI-indexed compress path (`None` for every other producer, + /// including sentinel/EOF blocks). + pub index: Option>, } impl HeapSize for BgzfBlock { fn heap_size(&self) -> usize { // Byte-bounded queues budget on resident heap, so account for the full // allocation (`capacity`), not just the populated prefix (`len`). - self.bytes.capacity() + let manifest = self.index.as_ref().map_or(0, |m| { + std::mem::size_of::() + + m.phys_comp_len.capacity() * std::mem::size_of::() + + m.records.capacity() * std::mem::size_of::() + }); + self.bytes.capacity() + manifest } } @@ -228,11 +288,44 @@ mod tests { #[test] fn bgzf_block_heap_size_matches_bytes_capacity() { - let b = BgzfBlock { batch_serial: 0, bytes: vec![0u8; 1024], uncompressed_size: 4096 }; + let b = BgzfBlock { + batch_serial: 0, + bytes: vec![0u8; 1024], + uncompressed_size: 4096, + index: None, + }; assert_eq!(b.heap_size(), 1024); assert_eq!(b.ordinal(), 0); } + #[test] + fn bgzf_block_index_defaults_none_and_heap_counts_manifest() { + let plain = + BgzfBlock { batch_serial: 0, bytes: vec![0u8; 100], uncompressed_size: 0, index: None }; + assert!(plain.index.is_none()); + assert_eq!(plain.heap_size(), plain.bytes.capacity()); + + let manifest = Box::new(BamIndexManifest { + phys_comp_len: vec![10, 20], + records: vec![RecordIndexEntry { uoffset: 0, len: 50, ctx: None }], + }); + let indexed = BgzfBlock { + batch_serial: 1, + bytes: vec![0u8; 30], + uncompressed_size: 0, + index: Some(manifest), + }; + let m = indexed.index.as_ref().expect("manifest present"); + assert_eq!( + indexed.heap_size(), + indexed.bytes.capacity() + + std::mem::size_of::() + + m.phys_comp_len.capacity() * std::mem::size_of::() + + m.records.capacity() * std::mem::size_of::(), + "heap_size must account for the manifest allocation exactly" + ); + } + #[test] fn decompressed_block_heap_size_matches_bytes_capacity() { let b = DecompressedBlock { batch_serial: 7, bytes: vec![0u8; 4096] }; @@ -250,7 +343,7 @@ mod tests { assert!(bytes.capacity() >= 8192 && bytes.len() == 100); let cap = bytes.capacity(); - let block = BgzfBlock { batch_serial: 0, bytes, uncompressed_size: 0 }; + let block = BgzfBlock { batch_serial: 0, bytes, uncompressed_size: 0, index: None }; assert_eq!(block.heap_size(), cap); let mut backing = Vec::with_capacity(4096); diff --git a/src/lib/pipeline/chains/builder.rs b/src/lib/pipeline/chains/builder.rs index 3c64d0baf..40aba4a04 100644 --- a/src/lib/pipeline/chains/builder.rs +++ b/src/lib/pipeline/chains/builder.rs @@ -35,7 +35,7 @@ use std::sync::Arc; use std::sync::atomic::AtomicU64; -use anyhow::{Result, anyhow, bail}; +use anyhow::{Result, anyhow, bail, ensure}; use noodles::sam::Header; use crate::pipeline::chains::{ @@ -401,10 +401,13 @@ pub struct ChainBuilder<'a> { /// Post-pipeline hooks that run **only on a fully successful run**. /// Populated by `add_` methods for actions that publish a derived /// artifact from the run's output (currently only `add_sort`'s - /// `IndexBamFinalizeHook`, which writes the `.bai` sidecar by re-reading - /// the finished BAM). Gating these mirrors the standalone `fgumi sort` - /// flow, where the index write is behind `run_result?`: a failed run - /// leaves a partial BAM, so publishing its index would be stale. + /// `SortSummaryFinalizeHook`, which logs the `=== Summary ===` block from + /// the standalone-sort stats slot). Gating these mirrors the standalone + /// `fgumi sort` flow: a failed run leaves a partial BAM, so publishing a + /// derived artifact or summary from it would be stale. The `.bai` sidecar + /// for `SinkSpec::BamWithIndex` is NOT one of these — it is written inline + /// by the `WriteBgzfFile` sink itself (see `add_sink`), not by a + /// finalize hook. finalize_on_success: Vec>, /// Shared progress counter threaded through source + stage steps for @@ -1145,7 +1148,11 @@ impl<'a> ChainBuilder<'a> { let rejects_branch = (consensus_pt.0, BranchIdx(1)); let rejects_compress_tail = self.pipeline.append_step( - BgzfCompress::new(self.tuning.compression_level, self.tuning.per_step_byte_limit), + BgzfCompress::new( + self.tuning.compression_level, + self.tuning.per_step_byte_limit, + false, + ), rejects_branch, ); self.pipeline.append_step(rejects_write, rejects_compress_tail); @@ -1337,11 +1344,19 @@ impl<'a> ChainBuilder<'a> { /// Add the output sink step(s) to the pipeline. /// /// Appends `BgzfCompress` + `WriteBgzfFile`. Reads `spec.sink` for the - /// output path and uses the output header from `self.header`. + /// output path and uses the output header from `self.header`. For + /// `SinkSpec::BamWithIndex`, also builds `BgzfCompress` with inline + /// indexing enabled and attaches + /// [`WriteBgzfFile::with_bai_index`](crate::pipeline::steps::sink::write_bgzf::WriteBgzfFile::with_bai_index) + /// to the sink, so the `.bai` sidecar is produced as the BAM is written + /// rather than by a post-pipeline re-read. /// /// # Errors /// - /// Returns errors from sink step construction (e.g., cannot create output file). + /// Returns errors from sink step construction (e.g., cannot create output + /// file), or if `SinkSpec::BamWithIndex` targets stdout or a chain with a + /// deferred (not-yet-resolved) header — inline indexing requires an eager + /// header, since it needs a fixed starting compressed-byte offset. /// /// # Panics /// @@ -1367,8 +1382,52 @@ impl<'a> ChainBuilder<'a> { let tail = self.current_tail.expect("add_sink called before add_source"); let output_path = self.spec.sink.path(); - let compress_step = - BgzfCompress::new(self.tuning.compression_level, self.tuning.per_step_byte_limit); + // `SinkSpec::BamWithIndex` selects the inline BAI indexer: `BgzfCompress` + // emits a `BamIndexManifest` alongside each block it produces, and + // `WriteBgzfFile::with_bai_index` joins those manifests against its own + // cumulative compressed offset as it writes, building the `.bai` + // in-band rather than re-reading the finished BAM afterward. + let want_index = matches!(self.spec.sink, SinkSpec::BamWithIndex(_)); + if want_index { + ensure!( + !fgumi_bam_io::is_stdout_path(output_path), + "--write-index cannot target stdout" + ); + ensure!( + self.pending_header_handle.is_none(), + "inline BAI indexing requires a resolved (non-deferred) header" + ); + // The inline BAI indexer (`WriteBgzfFile::with_bai_index`) position-bins + // records purely by arrival order and does not itself verify sort order + // — see that method's doc. Coordinate-only is already enforced above us: + // `validate_cross_stage_constraints` (Rule 3) rejects a non-coordinate + // terminal sort under `BamWithIndex` unconditionally (release included), + // and the command layer rejects `--write-index` for a non-coordinate + // `--order`. This `debug_assert!` is a complementary developer-only + // guardrail on the *produced* @HD SO header tag, catching a future + // producer that advertises the wrong order despite passing spec + // validation; it is not the release-mode enforcement point. + debug_assert!( + self.header + .header() + .and_then(|hd| { + hd.other_fields().get( + &noodles::sam::header::record::value::map::header::tag::SORT_ORDER, + ) + }) + .is_some_and(|so| { + AsRef::<[u8]>::as_ref(so) + == noodles::sam::header::record::value::map::header::sort_order::COORDINATE + }), + "inline BAI indexing requires an @HD SO:coordinate header" + ); + } + + let compress_step = BgzfCompress::new( + self.tuning.compression_level, + self.tuning.per_step_byte_limit, + want_index, + ); let tail = self.pipeline.append_step(compress_step, tail); // When Stage::Align is in the chain, the merged output header is @@ -1388,6 +1447,19 @@ impl<'a> ChainBuilder<'a> { WriteBgzfFile::new(output_path, &self.header, self.tuning.compression_level) .map_err(|e| anyhow!("WriteBgzfFile::new: {e}"))? }; + // `want_index` is only ever true on the eager-header branch above (the + // `ensure!` on `pending_header_handle` guarantees it), so + // `with_bai_index`'s eager-header requirement always holds here. + let write_step = if want_index { + let num_refs = self.header.reference_sequences().len(); + // The derived `.bai` sidecar's collision with the BAM or a rejects + // branch is rejected up front by `validate_cross_stage_constraints` + // (before any sink opens and truncates a file), not here. + let bai_path = fgumi_bam_io::bai_sidecar_path(output_path); + write_step.with_bai_index(bai_path, num_refs) + } else { + write_step + }; // Lever 2: on the standalone-sort terminal only, run the writer on the // sort's I/O driver thread (legacy "N + 2"), sharing it with the phase-1 // `SpillWrite`. Every other chain keeps the pool-scheduled Serial + @@ -1428,8 +1500,11 @@ impl<'a> ChainBuilder<'a> { use crate::pipeline::steps::types::{BgzfBlock, DecompressedBlock}; if path_is_gzip(path) { - let compress = - BgzfCompress::new(self.tuning.compression_level, self.tuning.per_step_byte_limit); + let compress = BgzfCompress::new( + self.tuning.compression_level, + self.tuning.per_step_byte_limit, + false, + ); let compress_tail = self.pipeline.append_step(compress, tail); let writer = WriteRawFile::::new(path, &fgumi_bgzf::BGZF_EOF) .map_err(|e| anyhow!("WriteRawFile: {e}"))?; @@ -1838,7 +1913,11 @@ impl<'a> ChainBuilder<'a> { // Branch 1 = pre-framed rejects bytes → BgzfCompress → WriteBgzfFile. let rejects_branch = (pt.0, BranchIdx(1)); let rejects_compress_tail = self.pipeline.append_step( - BgzfCompress::new(self.tuning.compression_level, self.tuning.per_step_byte_limit), + BgzfCompress::new( + self.tuning.compression_level, + self.tuning.per_step_byte_limit, + false, + ), rejects_branch, ); let _rejects_write_tail = @@ -2281,10 +2360,11 @@ impl<'a> ChainBuilder<'a> { /// ↓ [Intermediate] SortMerge (Serial) → RecordBatch → DecodeFromRecords /// ``` /// - /// For a `SinkSpec::BamWithIndex` (terminal sort only), an - /// [`IndexBamFinalizeHook`] is registered on the success-gated finalize - /// list. `--sort::max-memory=auto` is only honoured for a sole-`[Sort]` - /// chain (it owns the whole budget); a fused chain must pass a fixed value. + /// For a `SinkSpec::BamWithIndex` (terminal sort only), `add_sink` attaches + /// the inline BAI indexer directly to the `WriteBgzfFile` sink — no + /// finalize hook is involved. `--sort::max-memory=auto` is only honoured + /// for a sole-`[Sort]` chain (it owns the whole budget); a fused chain must + /// pass a fixed value. /// /// The chain-level [`StageTimingFinalizeHook`] (registered by [`Self::build`]) /// covers the wall time; sort does not register a per-stage summary hook. @@ -2293,12 +2373,9 @@ impl<'a> ChainBuilder<'a> { /// /// Returns errors if sort options are missing from the spec bag, or if /// `--sort::max-memory=auto` is requested in a fused multi-stage pipeline. - /// - /// [`IndexBamFinalizeHook`]: crate::pipeline::chains::commands::sort::IndexBamFinalizeHook #[allow(clippy::too_many_lines)] fn add_sort(&mut self, position: StagePosition) -> Result<()> { use crate::commands::common::{MemoryLimit, resolve_memory_budget}; - use crate::pipeline::chains::commands::sort::IndexBamFinalizeHook; let sort = self .spec @@ -2676,9 +2753,13 @@ impl<'a> ChainBuilder<'a> { // error path, and the stats slot is unset (default zeros) whenever // SortMerge never reached `Done`, so an always-run summary would // report a bogus "Records processed: 0 ... completed" line on a - // failed run. Registered before the index hook so on success the - // summary still lands before any "Indexing BAM:" line, mirroring - // the legacy SortFinalizeHook ordering. + // failed run. + // + // A `SinkSpec::BamWithIndex` request needs no finalize hook here: + // `add_sink` wires the inline BAI indexer directly onto the + // `WriteBgzfFile` sink (`with_bai_index`), so the `.bai` is + // produced as part of the sink's own drained-finish, not by a + // post-pipeline re-read step queued from this branch. if let Some(slot) = sort_stats_slot { let output_path = self.spec.sink.path().clone(); self.finalize_on_success.push(Box::new( @@ -2689,26 +2770,6 @@ impl<'a> ChainBuilder<'a> { }, )); } - - // When the spec asks for a sidecar BAI, queue the - // `IndexBamFinalizeHook` on the *success-gated* finalize list - // so the indexer runs after the chain has flushed the final - // BAM on a successful run only (a failed run leaves a partial - // BAM whose index would be stale). Every Sort-terminal chain — - // standalone `[Stage::Sort]` or fused (e.g. - // `[Stage::Correct, Stage::Sort]`) — lands here in the unified - // `add_sort` flow, so without this wire a `BamWithIndex` request - // would pass Rule 3 and silently skip indexing. Today only - // standalone `fgumi sort` builds such specs (runall never - // constructs `SinkSpec::BamWithIndex` because its sort output is - // template-coordinate, where BAI is undefined), but the - // architectural invariant — "BamWithIndex triggers the hook in - // every terminal-sort path" — is enforced here for future - // producers. - if let SinkSpec::BamWithIndex(p) = &self.spec.sink { - self.finalize_on_success - .push(Box::new(IndexBamFinalizeHook { output_path: p.clone() })); - } } else { // Intermediate: `SortMerge` (the default) // emits `RecordBatch`, decoded to `DecodedRecordBatch` for the @@ -4101,8 +4162,11 @@ impl<'a> ChainBuilder<'a> { .rejects .as_ref() .ok_or_else(|| anyhow!("rejects path missing (build_for dispatch bug)"))?; - let rejects_compress = - BgzfCompress::new(self.tuning.compression_level, self.tuning.per_step_byte_limit); + let rejects_compress = BgzfCompress::new( + self.tuning.compression_level, + self.tuning.per_step_byte_limit, + false, + ); let rejects_writer = WriteBgzfFile::new(rejects_path, &self.header, self.tuning.compression_level) .map_err(|e| anyhow!("WriteBgzfFile (rejects)::new: {e}"))?; diff --git a/src/lib/pipeline/chains/commands/sort.rs b/src/lib/pipeline/chains/commands/sort.rs index 02e54d0a3..9d268ed00 100644 --- a/src/lib/pipeline/chains/commands/sort.rs +++ b/src/lib/pipeline/chains/commands/sort.rs @@ -3,9 +3,13 @@ //! The stage-by-stage construction lives in //! [`crate::pipeline::chains::builder::ChainBuilder`]'s `add_sort` method. //! This module provides the standalone-sort summary finalize hook -//! (`SortSummaryFinalizeHook`) and the `IndexBamFinalizeHook` BAI indexer -//! (defined locally — see its doc), which `add_sort` registers for a -//! `SinkSpec::BamWithIndex` request. +//! (`SortSummaryFinalizeHook`), registered by `add_sort` for a sole- +//! `[Stage::Sort]` chain. A `SinkSpec::BamWithIndex` request needs no +//! finalize hook here: `ChainBuilder::add_sink` attaches an inline BAI +//! indexer directly to the `WriteBgzfFile` sink +//! (`WriteBgzfFile::with_bai_index`), which builds the `.bai` from the +//! `BamIndexManifest`s each `BgzfCompress`-produced block carries, as the sink +//! drains — no post-pipeline re-read of the finished BAM. //! //! ## Sort pipeline topology //! @@ -55,56 +59,3 @@ impl FinalizeHook for SortSummaryFinalizeHook { Ok(()) } } - -// ───────────────────────────────────────────────────────────────────────────── -// IndexBamFinalizeHook -// ───────────────────────────────────────────────────────────────────────────── - -/// Post-pipeline action that builds a BAI index for a finished -/// coordinate-sorted BAM, writing the `.bai` sidecar next to the BAM. -/// -/// Reads the BAM via `noodles::bam::fs::index` (walking BGZF block offsets) and -/// writes the sidecar via [`fgumi_bam_io::write_bai_sidecar`], which appends -/// `.bai` to the full output path (samtools convention), so `foo.bam` → -/// `foo.bam.bai` and `foo.sorted` → `foo.sorted.bai`. -/// -/// **Invariant:** the BAM at `output_path` must be writer-closed before -/// `finalize` runs — guaranteed by `Pipeline::run` returning (every step's -/// writer has dropped). The hook does not `fsync`: the same process re-reading -/// via the OS page cache sees all flushed bytes on every supported local -/// filesystem. -/// -/// `pub(crate)` so [`crate::pipeline::chains::builder::ChainBuilder::add_sort`] -/// can construct and register it. Defined here (rather than re-exported from a -/// sort-cli crate as upstream did) because `main-runall` keeps the sort command -/// local to `crate::commands::sort`. -pub(crate) struct IndexBamFinalizeHook { - /// Path of the finished coordinate-sorted BAM to index. - pub(crate) output_path: PathBuf, -} - -impl FinalizeHook for IndexBamFinalizeHook { - /// Build and write the BAI index alongside the BAM. - /// - /// # Errors - /// - /// Returns an error if indexing or writing the BAI sidecar fails. - fn finalize(self: Box) -> Result<()> { - use std::time::Instant; - - use crate::logging::format_duration; - - let IndexBamFinalizeHook { output_path } = *self; - - info!("Indexing BAM: {}", output_path.display()); - let start = Instant::now(); - - // Index + write the `.bai` sidecar in one call; the path is - // derived by appending `.bai` to the full BAM path (samtools - // convention), correct for any path — not just `*.bam`. - let index_path = fgumi_bam_io::write_bai_sidecar(&output_path)?; - info!("Wrote BAM index: {} ({})", index_path.display(), format_duration(start.elapsed())); - - Ok(()) - } -} diff --git a/src/lib/pipeline/chains/finalize.rs b/src/lib/pipeline/chains/finalize.rs index 3c01fbb10..192a4d796 100644 --- a/src/lib/pipeline/chains/finalize.rs +++ b/src/lib/pipeline/chains/finalize.rs @@ -82,11 +82,13 @@ pub struct BuiltPipeline { /// completed and every always-drain `finalize` hook succeeded. /// /// Use this for actions that publish a derived artifact from the output - /// the run produced (e.g. writing a `.bai` sidecar by re-reading the - /// finished BAM). On the error path the output is incomplete, so running - /// these would publish a stale/partial artifact — exactly the - /// `IndexBamFinalizeHook` footgun the standalone `fgumi sort` flow guards - /// against by gating the index write behind `run_result?`. + /// the run produced (e.g. a post-run summary computed from stats the run + /// collected). On the error path the output is incomplete, so running + /// these would publish a stale/partial artifact — the footgun the + /// standalone `fgumi sort` flow guards against by gating such publication + /// behind `run_result?`. A `.bai` sidecar for `SinkSpec::BamWithIndex` is + /// NOT an example of this: it is written inline by the `WriteBgzfFile` + /// sink as part of the run itself, not by a finalize hook. pub finalize_on_success: Vec>, } diff --git a/src/lib/pipeline/chains/sink_spec.rs b/src/lib/pipeline/chains/sink_spec.rs index 548b000b2..19c158401 100644 --- a/src/lib/pipeline/chains/sink_spec.rs +++ b/src/lib/pipeline/chains/sink_spec.rs @@ -7,8 +7,9 @@ use std::path::PathBuf; /// **Note for programmatic callers constructing a `ChainSpec` directly:** /// `Bam` is the default — picking it on a sort-terminal chain produces /// a coordinate-sorted BAM with **no** companion `.bai`. If you want a -/// sidecar BAI, use `BamWithIndex`; the chain-builder will register an -/// `IndexBamFinalizeHook` that reads the finished BAM and emits +/// sidecar BAI, use `BamWithIndex`; the chain builder's `add_sink` attaches an +/// inline indexer directly to the output sink, which builds the `.bai` +/// alongside the BAM as it writes (no post-pipeline re-read) and emits /// `.bam.bai`. There is no "auto-index" inference from sort /// order; the caller must opt in via this variant. /// diff --git a/src/lib/pipeline/chains/validate.rs b/src/lib/pipeline/chains/validate.rs index 46a326a66..0019dc4bf 100644 --- a/src/lib/pipeline/chains/validate.rs +++ b/src/lib/pipeline/chains/validate.rs @@ -234,6 +234,7 @@ pub fn validate_stage_opts_present(spec: &ChainSpec) -> Result<()> { /// Returns an error on the first violated constraint. pub fn validate_cross_stage_constraints(spec: &ChainSpec) -> Result<()> { use crate::assigner::Strategy; + use crate::commands::sort::SortOrderArg; use crate::pipeline::chains::{SinkSpec, SourceSpec}; // Rule 1: Zipper requires a PairedBams source. @@ -275,14 +276,22 @@ pub fn validate_cross_stage_constraints(spec: &ChainSpec) -> Result<()> { } // Rule 3: `SinkSpec::BamWithIndex` requires the terminal stage to be - // `Stage::Sort`. The BAI file format indexes BGZF virtual offsets of a - // coordinate-sorted BAM; emitting one for a non-sort terminal stage (or - // for a chain where sort is intermediate, leaving some other stage as - // terminal) would either reference a file that isn't coordinate-sorted - // or attempt to index an output type the hook isn't equipped to handle. - // The per-command CLI check (`Sort::execute` rejects - // `--write-index --order queryname`) handles the "coordinate - // specifically, not template-coordinate or queryname" constraint. + // `Stage::Sort`, sorting by *coordinate*. The BAI file format indexes BGZF + // virtual offsets of a coordinate-sorted BAM; emitting one for a non-sort + // terminal stage (or for a chain where sort is intermediate, leaving some + // other stage as terminal) would either reference a file that isn't + // coordinate-sorted or attempt to index an output type the hook isn't + // equipped to handle. A queryname- or template-coordinate-sorted terminal + // is just as wrong: the BAI would bin records that are not in coordinate + // order, producing a structurally-invalid index. + // + // The per-command CLI check (`Sort::execute` rejects `--write-index` for a + // non-coordinate `--order`) covers the sort command, and `add_sink` carries + // a `debug_assert!` on the produced @HD SO tag — but that assert is skipped + // in release. Enforcing the order here, unconditionally and before any sink + // is constructed, is what protects a *direct* chain (built programmatically + // rather than through the sort command) from writing an invalid BAI in a + // release build. if matches!(spec.sink, SinkSpec::BamWithIndex(_)) { let terminal_is_sort = spec.stages.last().is_some_and(|s| *s == Stage::Sort); if !terminal_is_sort { @@ -302,6 +311,69 @@ pub fn validate_cross_stage_constraints(spec: &ChainSpec) -> Result<()> { got terminal stage {terminal:?}", ); } + + // Terminal is `Stage::Sort`; require it to sort by coordinate. Only + // check when the sort options are present — `validate_stage_opts_present` + // (which `build_for` runs before this validator) guarantees they are on + // the build path, so the `None` arm is reachable only when this + // validator is called in isolation, and skipping it there mirrors how + // Rule 2 guards on `if let Some(group_opts)`. + if let Some(sort_opts) = spec.stage_opts.sort.as_ref() + && !matches!(sort_opts.order, SortOrderArg::Coordinate) + { + bail!( + "SinkSpec::BamWithIndex requires the terminal Stage::Sort to sort by coordinate \ + (the BAI format indexes a coordinate-sorted BAM); got sort order {:?}", + sort_opts.order, + ); + } + } + + // Rule 3b: structural preconditions on `SinkSpec::BamWithIndex` that + // `add_sink` also enforces (as `ensure!` guards) but only *after* earlier + // stages have run. `build_for` adds the sink last, so a rejects-bearing + // stage such as `add_correct` has already created and truncated its rejects + // `WriteBgzfFile` by the time `add_sink` would return one of these errors — + // leaving a clobbered rejects file behind on an input that was invalid from + // the start. Validating here, before any stage or sink is constructed (so + // before any `File::create` truncates a file), preserves the + // check-then-truncate discipline; the `add_sink` guards remain as defense in + // depth. This affects `Correct → Sort` and `Correct → Align → Sort`. + if let SinkSpec::BamWithIndex(output) = &spec.sink { + // An inline BAI is a sidecar file joined against the primary BAM's + // compressed offsets; it cannot be produced for a streamed stdout + // output (there is no seekable file to index). Mirrors `add_sink`'s + // `--write-index cannot target stdout` guard. + if fgumi_bam_io::is_stdout_path(output) { + bail!( + "SinkSpec::BamWithIndex cannot target stdout; an inline BAI index requires a \ + seekable on-disk output" + ); + } + + // `Stage::Align` defers the output header (it is rewritten once the + // aligner has emitted its @SQ lines), and the inline BAI indexer + // requires a resolved (non-deferred) header. Mirrors `add_sink`'s + // `inline BAI indexing requires a resolved (non-deferred) header` + // guard, which fires on `self.pending_header_handle` — set only by + // `add_align`. + if spec.stages.contains(&Stage::Align) { + bail!( + "SinkSpec::BamWithIndex is incompatible with Stage::Align: alignment defers the \ + output header, but inline BAI indexing requires a resolved (non-deferred) header" + ); + } + + // The derived `.bai` sidecar must not collide with the primary BAM or + // any rejects branch. The sidecar path is *derived* (`bai_sidecar_path`), + // so the command-layer output-collision check never saw it: two + // `WriteBgzfFile` sinks aimed at one path would otherwise clobber each + // other byte for byte. + reject_index_sidecar_collisions( + output, + &fgumi_bam_io::bai_sidecar_path(output), + &spec_rejects_paths(spec), + )?; } // Rule 4: `Stage::Extract` requires a `SourceSpec::Fastqs` source. @@ -333,12 +405,98 @@ pub fn validate_cross_stage_constraints(spec: &ChainSpec) -> Result<()> { Ok(()) } +/// The rejects/secondary output paths a chain's stage options declare. +/// +/// These are the writable files a chain produces besides its primary sink; +/// [`validate_cross_stage_constraints`] checks the derived `.bai` sidecar +/// against them (see [`reject_index_sidecar_collisions`]). Only stages that +/// actually emit a rejects branch carry one — the consensus stages are +/// `consensus`-gated, matching their `StageOptionsBag` slots. +fn spec_rejects_paths(spec: &ChainSpec) -> Vec<&std::path::Path> { + let mut paths: Vec<&std::path::Path> = Vec::new(); + let opts = &spec.stage_opts; + if let Some(correct) = &opts.correct + && let Some(rejects) = &correct.rejects_path + { + paths.push(rejects.as_path()); + } + if let Some(filter) = &opts.filter + && let Some(rejects) = &filter.rejects + { + paths.push(rejects.as_path()); + } + #[cfg(feature = "consensus")] + { + if let Some(duplex) = &opts.duplex + && let Some(rejects) = &duplex.rejects_opts.rejects + { + paths.push(rejects.as_path()); + } + if let Some(codec) = &opts.codec + && let Some(rejects) = &codec.rejects_opts.rejects + { + paths.push(rejects.as_path()); + } + if let Some(simplex) = &opts.simplex + && let Some(rejects) = &simplex.rejects_opts.rejects + { + paths.push(rejects.as_path()); + } + } + paths +} + +/// Reject a collision between the inline-index `.bai` sidecar and any other file +/// the chain writes (the primary BAM or a rejects branch). +/// +/// The sidecar path is *derived* ([`fgumi_bam_io::bai_sidecar_path`]), so the +/// command-layer +/// [`reject_output_collisions`](crate::commands::common::reject_output_collisions) +/// check — which only sees the paths the CLI passed — never inspects it. Two +/// `WriteBgzfFile` sinks aimed at one path would truncate and clobber each other +/// byte for byte, exactly the corruption `reject_output_collisions` prevents. +fn reject_index_sidecar_collisions( + output_path: &std::path::Path, + bai_path: &std::path::Path, + rejects_paths: &[&std::path::Path], +) -> Result<()> { + let mut targets: Vec<(&std::path::Path, &str)> = + vec![(output_path, "--output"), (bai_path, "--write-index (.bai sidecar)")]; + targets.extend(rejects_paths.iter().map(|p| (*p, "--rejects"))); + crate::commands::common::reject_output_collisions(&targets) +} + #[cfg(test)] mod tests { use super::*; + use crate::commands::sort::{SortOptions, SortOrderArg}; use crate::pipeline::chains::{SinkSpec, SourceSpec}; + use rstest::rstest; use std::path::PathBuf; + /// A `SortOptions` for a given order, covering every field the type + /// declares (it has no `Default`). Only `order` varies across callers — + /// the rest are the standalone-sort defaults, matching the integration + /// helper in `tests/integration/test_chain_bam_with_index.rs`. + fn sort_opts_with_order(order: SortOrderArg) -> SortOptions { + use crate::commands::common::{MaxTempFiles, MemoryLimit, MemoryReserve}; + SortOptions { + order, + key_types: None, + max_memory: MemoryLimit::Auto, + memory_reserve: MemoryReserve::Auto, + memory_per_thread: true, + tmp_dirs: Vec::new(), + sort_threads: None, + merge_threads: None, + temp_compression: 1, + temp_codec: fgumi_sort::SpillCodec::default(), + max_temp_files: MaxTempFiles::Auto, + block_batch: 4, + file_granularity: false, + } + } + fn empty_spec(stages: Vec) -> ChainSpec { use crate::commands::common::{ CompressionOptions, QueueMemoryOptions, SchedulerOptions, ThreadingOptions, @@ -532,6 +690,67 @@ mod tests { assert!(msg.contains("Group"), "expected stage name in error, got: {msg}"); } + #[rstest] + #[case::queryname_lexicographic(SortOrderArg::Queryname)] + #[case::queryname_natural(SortOrderArg::QuerynameNatural)] + #[case::template_coordinate(SortOrderArg::TemplateCoordinate)] + fn cross_stage_bam_with_index_rejects_non_coordinate_sort_order(#[case] order: SortOrderArg) { + // A terminal `Stage::Sort` under `SinkSpec::BamWithIndex` must sort by + // coordinate — the BAI format indexes a coordinate-sorted BAM. This is + // a plain `#[test]` (not `#[cfg(debug_assertions)]`), so it runs in + // release too, covering the case the `add_sink` `debug_assert!` skips: + // a direct chain built programmatically with a non-coordinate order. + let mut spec = empty_spec(vec![Stage::Sort]); + spec.sink = SinkSpec::BamWithIndex(PathBuf::from("out.bam")); + spec.stage_opts.sort = Some(sort_opts_with_order(order)); + let msg = validate_cross_stage_constraints(&spec).unwrap_err().to_string(); + assert!(msg.contains("BamWithIndex"), "expected BamWithIndex in error, got: {msg}"); + assert!(msg.contains("coordinate"), "expected 'coordinate' in error, got: {msg}"); + assert!( + msg.contains(&format!("{order:?}")), + "expected the rejected order {order:?} in error, got: {msg}" + ); + } + + #[test] + fn cross_stage_bam_with_index_accepts_coordinate_sort_order() { + // The one accepted order, with the sort options present (the `None` + // arm is exercised by `cross_stage_bam_with_index_accepted_on_terminal_sort`). + let mut spec = empty_spec(vec![Stage::Sort]); + spec.sink = SinkSpec::BamWithIndex(PathBuf::from("out.bam")); + spec.stage_opts.sort = Some(sort_opts_with_order(SortOrderArg::Coordinate)); + validate_cross_stage_constraints(&spec) + .expect("BamWithIndex with a coordinate terminal sort is valid"); + } + + #[rstest] + #[case::dash("-")] + #[case::dev_stdout("/dev/stdout")] + fn cross_stage_bam_with_index_rejects_stdout_output(#[case] output: &str) { + // `Correct → Sort` with an indexed sink aimed at stdout. `add_correct` + // would create and truncate its rejects file during stage construction, + // before `add_sink` returns the stdout error — so validation must reject + // this up front, before `build_for` constructs anything. + let mut spec = empty_spec(vec![Stage::Correct, Stage::Sort]); + spec.sink = SinkSpec::BamWithIndex(PathBuf::from(output)); + let msg = validate_cross_stage_constraints(&spec).unwrap_err().to_string(); + assert!(msg.contains("BamWithIndex"), "expected BamWithIndex in error, got: {msg}"); + assert!(msg.contains("stdout"), "expected 'stdout' in error, got: {msg}"); + } + + #[test] + fn cross_stage_bam_with_index_rejects_align_stage() { + // `Correct → Align → Sort` with an indexed sink. `Stage::Align` defers + // the output header, which the inline BAI indexer cannot use; + // `add_sink` rejects the deferred header only after `add_correct` has + // truncated its rejects file, so validation must reject up front. + let mut spec = empty_spec(vec![Stage::Correct, Stage::Align, Stage::Sort]); + spec.sink = SinkSpec::BamWithIndex(PathBuf::from("out.bam")); + let msg = validate_cross_stage_constraints(&spec).unwrap_err().to_string(); + assert!(msg.contains("BamWithIndex"), "expected BamWithIndex in error, got: {msg}"); + assert!(msg.contains("Align"), "expected 'Align' in error, got: {msg}"); + } + #[test] fn opts_present_downsample_not_wired() { // Stage::Downsample has no bag slot yet; should fail at options-presence time. @@ -729,4 +948,101 @@ mod tests { let err = validate_cross_stage_constraints(&spec).unwrap_err(); assert!(err.to_string().contains("Stage::Fastq"), "got: {err}"); } + + // --- Rule 3b: inline-index `.bai` sidecar collision --- + + /// The `reject_index_sidecar_collisions` helper: a sidecar colliding with a + /// rejects branch is rejected, and the error names both flags. + #[test] + fn sidecar_helper_rejects_collision_with_a_rejects_branch() { + use std::path::Path; + let output = Path::new("out.bam"); + let bai = Path::new("out.bam.bai"); + let rejects = [Path::new("out.bam.bai")]; + let err = super::reject_index_sidecar_collisions(output, bai, &rejects) + .expect_err("sidecar/rejects collision must be rejected"); + let msg = err.to_string(); + assert!( + msg.contains("--write-index") && msg.contains("--rejects"), + "error must name both colliding flags, got: {msg}" + ); + } + + /// The helper also rejects a sidecar colliding with the primary BAM output. + #[test] + fn sidecar_helper_rejects_collision_with_the_output() { + use std::path::Path; + let output = Path::new("out.bam"); + let err = super::reject_index_sidecar_collisions(output, output, &[]) + .expect_err("sidecar/output collision must be rejected"); + assert!(err.to_string().contains("--output"), "error must name --output"); + } + + /// The normal case — sidecar derived from the output, distinct or absent + /// rejects — has no collision. + #[test] + fn sidecar_helper_accepts_distinct_paths() { + use std::path::Path; + let output = Path::new("out.bam"); + let bai = Path::new("out.bam.bai"); + assert!(super::reject_index_sidecar_collisions(output, bai, &[]).is_ok()); + assert!( + super::reject_index_sidecar_collisions(output, bai, &[Path::new("rejects.bam")]) + .is_ok(), + "a distinct rejects path must not collide with the sidecar" + ); + } + + /// A `[Correct, Sort]` chain whose `--rejects` names the derived `.bai` + /// sidecar is rejected at validation — before `build_for` opens any sink, so + /// no file is truncated. Pins that Rule 3b is wired into the validator, not + /// just tested as a free helper. + #[test] + fn cross_stage_bam_with_index_rejects_sidecar_rejects_collision() { + let output = PathBuf::from("out.bam"); + let sidecar = fgumi_bam_io::bai_sidecar_path(&output); + let mut spec = empty_spec(vec![Stage::Correct, Stage::Sort]); + spec.sink = SinkSpec::BamWithIndex(output); + spec.stage_opts.correct = Some(correct_opts_with_rejects(Some(sidecar))); + let err = validate_cross_stage_constraints(&spec) + .expect_err("a rejects path colliding with the derived .bai must be rejected"); + let msg = err.to_string(); + assert!( + msg.contains("--write-index") && msg.contains("--rejects"), + "error must name both colliding flags, got: {msg}" + ); + } + + /// The same chain with a distinct `--rejects` path validates cleanly — the + /// rule must not false-positive on a normal rejects-bearing indexed sort. + #[test] + fn cross_stage_bam_with_index_distinct_rejects_is_valid() { + let mut spec = empty_spec(vec![Stage::Correct, Stage::Sort]); + spec.sink = SinkSpec::BamWithIndex(PathBuf::from("out.bam")); + spec.stage_opts.correct = + Some(correct_opts_with_rejects(Some(PathBuf::from("rejects.bam")))); + validate_cross_stage_constraints(&spec) + .expect("a distinct rejects path must not collide with the derived .bai"); + } + + /// `CorrectOptions` carrying only the `rejects_path` under test; all other + /// fields are inert defaults (mirrors `correct::tests::make_default_opts`). + fn correct_opts_with_rejects( + rejects_path: Option, + ) -> crate::commands::correct::CorrectOptions { + use crate::commands::correct::{CorrectOptions, Target}; + CorrectOptions { + metrics: None, + target: Target::Umi, + max_mismatches: 2, + min_distance_diff: 2, + umis: Vec::new(), + umi_files: Vec::new(), + dont_store_original_umis: false, + cache_size: 100_000, + min_corrected: None, + revcomp: false, + rejects_path, + } + } } diff --git a/src/lib/pipeline/steps/bgzf/compress.rs b/src/lib/pipeline/steps/bgzf/compress.rs index ebfdba2b9..ef06ef543 100644 --- a/src/lib/pipeline/steps/bgzf/compress.rs +++ b/src/lib/pipeline/steps/bgzf/compress.rs @@ -26,7 +26,9 @@ use crate::pipeline::core::outputs::OrderedBytesSingle; use crate::pipeline::core::queues::QueueSpec; use crate::pipeline::core::reorder::BranchOrdering; use crate::pipeline::core::step::{Step, StepCtx, StepKind, StepOutcome, StepProfile}; -use crate::pipeline::steps::types::{BgzfBlock, DecompressedBlock}; +use crate::pipeline::steps::types::{ + BamIndexManifest, BgzfBlock, DecompressedBlock, RecordIndexEntry, +}; /// Per-worker BGZF compressed-output scratch capacity. The compressed /// output of `BgzfCompress` is the concatenation of physical 64 KiB BGZF @@ -46,17 +48,23 @@ pub struct BgzfCompress { output_scratch: Vec, held: HeldSlot>, output_byte_limit: u64, + /// When set, `try_run` walks the input batch's framed BAM records to + /// build a `BamIndexManifest` and emits it on the output `BgzfBlock`. + /// When unset, the frame walk and per-physical-block length capture are + /// skipped entirely (zero overhead on the non-indexed path). + index_bam: bool, } impl BgzfCompress { #[must_use] - pub fn new(compression_level: u32, output_byte_limit: u64) -> Self { + pub fn new(compression_level: u32, output_byte_limit: u64, index_bam: bool) -> Self { Self { compressor: InlineBgzfCompressor::new(compression_level), compression_level, output_scratch: Vec::with_capacity(COMPRESS_SCRATCH_CAPACITY), held: HeldSlot::new(), output_byte_limit, + index_bam, } } } @@ -69,10 +77,36 @@ impl Clone for BgzfCompress { output_scratch: Vec::with_capacity(COMPRESS_SCRATCH_CAPACITY), held: HeldSlot::new(), output_byte_limit: self.output_byte_limit, + index_bam: self.index_bam, } } } +/// Walk a batch of `[u32 le len][body]*` frames into per-record index +/// entries. `uoffset` is the frame start (offset of the length prefix, +/// measured from the start of the batch's uncompressed bytes); `len` is `4 +` +/// the body length. +/// +/// Pure and unit-testable without a full pipeline: this is the shape a +/// `SerializeBamRecords` batch always takes, so no `StepCtx` is needed to +/// exercise it. +fn build_record_entries(batch: &[u8]) -> Vec { + let mut out = Vec::new(); + let mut u = 0usize; + while u + 4 <= batch.len() { + let block_size = u32::from_le_bytes(batch[u..u + 4].try_into().expect("4 bytes")) as usize; + let body = &batch[u + 4..u + 4 + block_size]; + out.push(RecordIndexEntry { + uoffset: u32::try_from(u).expect("batch offset fits u32"), + len: u32::try_from(4 + block_size).expect("record len fits u32"), + ctx: fgumi_bam_io::extract_alignment_context(body), + }); + u += 4 + block_size; + } + debug_assert_eq!(u, batch.len(), "frames must exactly cover the batch"); + out +} + impl BgzfCompress { /// Drain the compressor's physical BGZF blocks and assemble the single /// output `BgzfBlock` byte buffer. @@ -116,6 +150,35 @@ impl BgzfCompress { } } } + + /// Like [`Self::assemble_output_bytes`], but also returns each physical + /// block's compressed length (`data.len()`, before it moves or is + /// concatenated away), in emission order — exactly the `phys_comp_len` + /// a `BamIndexManifest` needs. Only called on the `index_bam` path; the + /// non-indexed path uses `assemble_output_bytes` and never allocates a + /// `lens` vec. + fn assemble_output_bytes_with_lens(&mut self) -> (Vec, Vec) { + let mut children = self.compressor.take_blocks(); + let lens: Vec = children + .iter() + .map(|c| u32::try_from(c.data.len()).expect("block len fits u32")) + .collect(); + let bytes = match children.len() { + 0 => Vec::new(), + 1 => children.pop().expect("len == 1").data, + _ => { + for child in children { + self.output_scratch.extend_from_slice(&child.data); + self.compressor.recycle_buffer(child.data); + } + std::mem::replace( + &mut self.output_scratch, + Vec::with_capacity(COMPRESS_SCRATCH_CAPACITY), + ) + } + }; + (bytes, lens) + } } impl Step for BgzfCompress { @@ -154,15 +217,37 @@ impl Step for BgzfCompress { } return Ok(StepOutcome::NoProgress); }; - let DecompressedBlock { batch_serial, bytes } = parent; - let uncompressed_size = u32::try_from(bytes.len()).unwrap_or(u32::MAX); + let DecompressedBlock { batch_serial, bytes: input_bytes } = parent; + let uncompressed_size = u32::try_from(input_bytes.len()).unwrap_or(u32::MAX); // Compress and harvest physical BGZF blocks. - self.compressor.write_all(&bytes)?; + self.compressor.write_all(&input_bytes)?; self.compressor.flush()?; - let bytes = self.assemble_output_bytes(); - let out = BgzfBlock { batch_serial, bytes, uncompressed_size }; + // Build the index manifest from the ORIGINAL uncompressed input + // bytes (the framed BAM records), before they are consumed below — + // never from the compressed output. + let (bytes, index) = if self.index_bam { + let records = build_record_entries(&input_bytes); + let (bytes, phys_comp_len) = self.assemble_output_bytes_with_lens(); + debug_assert_eq!( + phys_comp_len.iter().map(|&n| u64::from(n)).sum::(), + bytes.len() as u64, + "physical block lengths must sum to the total compressed byte count" + ); + debug_assert_eq!( + phys_comp_len.len(), + (uncompressed_size as usize).div_ceil(fgumi_bgzf::BGZF_MAX_BLOCK_SIZE), + "one physical block per BGZF_MAX_BLOCK_SIZE-sized chunk of uncompressed input" + ); + let index = + (!bytes.is_empty()).then(|| Box::new(BamIndexManifest { phys_comp_len, records })); + (bytes, index) + } else { + (self.assemble_output_bytes(), None) + }; + + let out = BgzfBlock { batch_serial, bytes, uncompressed_size, index }; match ctx.outputs.push(out) { Ok(()) => Ok(StepOutcome::Progress), Err(unpushed) => { @@ -183,7 +268,7 @@ mod tests { #[test] fn profile_advertises_parallel_byordinal() { - let s = BgzfCompress::new(1, 1024); + let s = BgzfCompress::new(1, 1024, false); let p = s.profile(); assert_eq!(p.name, "BgzfCompress"); assert_eq!(p.kind, StepKind::Parallel); @@ -192,13 +277,13 @@ mod tests { #[test] fn clone_constructs_fresh_compressor() { - let s = BgzfCompress::new(1, 1024); + let s = BgzfCompress::new(1, 1024, false); let _cloned = s.clone(); } #[test] fn small_input_produces_single_concatenated_output() { - let mut s = BgzfCompress::new(1, 1024); + let mut s = BgzfCompress::new(1, 1024, false); s.compressor.write_all(b"hello bgzf").unwrap(); s.compressor.flush().unwrap(); let blocks = s.compressor.take_blocks(); @@ -224,7 +309,7 @@ mod tests { /// Drive `assemble_output_bytes` on a fresh step over `input`. fn assemble(level: u32, input: &[u8]) -> Vec { - let mut s = BgzfCompress::new(level, 1024); + let mut s = BgzfCompress::new(level, 1024, false); s.compressor.write_all(input).unwrap(); s.compressor.flush().unwrap(); s.assemble_output_bytes() @@ -234,7 +319,7 @@ mod tests { fn single_block_fast_path_is_byte_identical() { // Small input → one physical block → move fast path. let input = b"the quick brown fox jumps over the lazy dog"; - let mut s = BgzfCompress::new(1, 1024); + let mut s = BgzfCompress::new(1, 1024, false); s.compressor.write_all(input).unwrap(); s.compressor.flush().unwrap(); assert_eq!(s.compressor.take_blocks().len(), 1, "input should fit one block"); @@ -252,7 +337,7 @@ mod tests { for i in 0..(200 * 1024u32) { input.push((i % 251) as u8); } - let mut s = BgzfCompress::new(1, 1024); + let mut s = BgzfCompress::new(1, 1024, false); s.compressor.write_all(&input).unwrap(); s.compressor.flush().unwrap(); assert!(s.compressor.take_blocks().len() > 1, "input should span >1 block"); @@ -266,7 +351,7 @@ mod tests { // After the multi-block path recycles buffers, a subsequent compression // reuses them, and output remains byte-identical to a fresh compressor. let big: Vec = (0..(200 * 1024u32)).map(|i| (i % 251) as u8).collect(); - let mut s = BgzfCompress::new(1, 1024); + let mut s = BgzfCompress::new(1, 1024, false); s.compressor.write_all(&big).unwrap(); s.compressor.flush().unwrap(); @@ -286,7 +371,7 @@ mod tests { // assembled output must be a truly empty `Vec`, not the per-worker // scratch — otherwise an "empty" output carries a retained // `COMPRESS_SCRATCH_CAPACITY` allocation and inflates queue memory. - let mut s = BgzfCompress::new(1, 1024); + let mut s = BgzfCompress::new(1, 1024, false); // Fresh compressor, nothing written → no blocks. let fresh = s.assemble_output_bytes(); @@ -310,4 +395,109 @@ mod tests { "empty batch after a multi-block batch must not retain scratch" ); } + + // ──────────────────────────────────────────────────────────────────── + // Frame-walk + manifest emission (`index_bam`). + // ──────────────────────────────────────────────────────────────────── + + /// Encode a CIGAR op (`(length, op_char)`) as the BAM-packed `u32` word + /// [`fgumi_raw_bam::SamBuilder::cigar_ops`] expects: `(length << 4) | + /// op_code`, per the SAM spec's op-code table (`MIDNSHP=X` -> `0..=8`). + fn cigar_op_word(length: u32, op: char) -> u32 { + let op_code = match op { + 'M' => 0, + 'I' => 1, + 'D' => 2, + 'N' => 3, + 'S' => 4, + 'H' => 5, + 'P' => 6, + '=' => 7, + 'X' => 8, + _ => panic!("unsupported CIGAR op: {op}"), + }; + (length << 4) | op_code + } + + /// Build a record **body** (no 4-byte length prefix) via `SamBuilder`, + /// mirroring `fgumi_bam_io::writer`'s Task-1 test helper of the same + /// name — duplicated locally since that helper is private to its own + /// `#[cfg(test)]` module in a different crate. + fn build_record_body(ref_id: i32, pos: i32, flags: u16, cigar: &[(u32, char)]) -> Vec { + let ops: Vec = cigar.iter().map(|&(len, op)| cigar_op_word(len, op)).collect(); + fgumi_raw_bam::SamBuilder::new() + .ref_id(ref_id) + .pos(pos) + .flags(flags) + .cigar_ops(&ops) + .build() + .to_vec() + } + + #[test] + fn build_record_entries_matches_frames() { + // Concatenate framed BAM records: [u32 le len][body]. + let bodies: Vec> = vec![ + build_record_body(0, 0, 0, &[(5, 'M')]), + build_record_body(0, 100, 0, &[(5, 'M')]), + ]; + let mut batch = Vec::new(); + let mut expected_uoffsets = Vec::new(); + for body in &bodies { + expected_uoffsets.push(u32::try_from(batch.len()).unwrap()); + batch.extend_from_slice(&u32::try_from(body.len()).unwrap().to_le_bytes()); + batch.extend_from_slice(body); + } + let entries = build_record_entries(&batch); + assert_eq!(entries.len(), 2); + for (i, e) in entries.iter().enumerate() { + assert_eq!(e.uoffset, expected_uoffsets[i]); + assert_eq!(e.len as usize, 4 + bodies[i].len()); + assert!(e.ctx.is_some()); + } + } + + #[test] + fn assemble_with_lens_sums_to_bytes_and_ceildiv_blocks() { + use fgumi_bgzf::BGZF_MAX_BLOCK_SIZE; + // Small input → one physical block. + let mut s = BgzfCompress::new(1, 1 << 20, true); + let small = b"the quick brown fox"; + s.compressor.write_all(small).unwrap(); + s.compressor.flush().unwrap(); + let (bytes, lens) = s.assemble_output_bytes_with_lens(); + assert_eq!( + lens.iter().map(|&n| u64::from(n)).sum::(), + u64::try_from(bytes.len()).unwrap() + ); + assert_eq!(lens.len(), small.len().div_ceil(BGZF_MAX_BLOCK_SIZE)); // == 1 + + // Input > one block → phys count == ceil(len / 65280). + let big: Vec = (0..(200 * 1024u32)).map(|i| (i % 251) as u8).collect(); + let mut s = BgzfCompress::new(1, 1 << 20, true); + s.compressor.write_all(&big).unwrap(); + s.compressor.flush().unwrap(); + let (bytes, lens) = s.assemble_output_bytes_with_lens(); + assert_eq!( + lens.iter().map(|&n| u64::from(n)).sum::(), + u64::try_from(bytes.len()).unwrap() + ); + assert_eq!(lens.len(), big.len().div_ceil(BGZF_MAX_BLOCK_SIZE)); // > 1 + } + + // `try_run`'s `index_bam=false`/`true` emit contract (whether the + // resulting `BgzfBlock.index` is `None`/`Some` for a real, driven + // `try_run` call) is pinned by + // `compress_step_emits_no_manifest_when_index_bam_false` and + // `compress_step_emits_manifest_matching_frames_when_index_bam_true` in + // `steps/chain_tests.rs`. This module's own tests only drive + // `assemble_output_bytes`/`assemble_output_bytes_with_lens`/ + // `build_record_entries` directly (see the `assemble` helper above) — + // it has no `StepCtx`/pipeline harness, and building one here would + // duplicate the `ReplaySource`/`CollectSink` test-step machinery + // `chain_tests.rs` already has for exactly this purpose. An earlier + // version of this test hand-built a `BgzfBlock { index: None }` and + // asserted `None` on a value it wrote itself, which pinned nothing + // about `try_run`'s actual behavior — replaced by the two chain-level + // tests above. } diff --git a/src/lib/pipeline/steps/chain_tests.rs b/src/lib/pipeline/steps/chain_tests.rs index 09f0ee88a..d2ad64525 100644 --- a/src/lib/pipeline/steps/chain_tests.rs +++ b/src/lib/pipeline/steps/chain_tests.rs @@ -54,7 +54,9 @@ use crate::pipeline::steps::boundaries::bam::FindBamBoundaries; use crate::pipeline::steps::group::position::BatchedProcessedPositionGroups; use crate::pipeline::steps::parse::bam::ParseBamRecords; use crate::pipeline::steps::parse::decode::{DecodeFromRecords, DecodeRecords}; -use crate::pipeline::steps::types::{BgzfBlock, DecodedRecordBatch, RecordBatch}; +use crate::pipeline::steps::types::{ + BgzfBlock, DecodedRecordBatch, DecompressedBlock, RecordBatch, +}; use crate::template::Template; /// Per-edge byte budget. Deliberately small relative to the multi-record @@ -129,6 +131,7 @@ fn bgzf_blocks_for(records: &[RawRecord], payload_bytes: usize) -> Vec Vec { + let mut batch = Vec::new(); + for record in records { + let bytes = record.as_ref(); + let block_size = u32::try_from(bytes.len()).expect("record fits u32"); + batch.extend_from_slice(&block_size.to_le_bytes()); + batch.extend_from_slice(bytes); + } + batch +} + +/// `ReplaySource` → `BgzfCompress` (`index_bam = false`) → +/// `CollectSink`, injecting `DecompressedBlock` directly (skipping +/// decompress/boundary-finding, since it is exactly `BgzfCompress`'s input +/// type). This drives a real `try_run` and asserts on what it actually +/// emits — replacing a prior unit test that hand-built a `BgzfBlock { index: +/// None }` and asserted `None` on a value it wrote itself, which pinned +/// nothing about `try_run`'s actual behavior. +#[rstest] +fn compress_step_emits_no_manifest_when_index_bam_false(#[values(1, 4)] threads: usize) { + let records = test_records(4); + let batch = framed_batch(&records); + let collected: Arc>> = Arc::new(Mutex::new(Vec::new())); + + let sink_handle = Arc::clone(&collected); + let builder = Pipeline::builder(); + builder + .chain(ReplaySource::new(vec![DecompressedBlock { batch_serial: 0, bytes: batch }])) + .chain(BgzfCompress::new(1, EDGE_LIMIT_BYTES, false)) + .chain(CollectSink { collected: sink_handle }) + .into_sink_marker(); + let pipeline = builder.build().expect("chain builds"); + pipeline.run(PipelineConfig { threads, ..Default::default() }).expect("chain runs"); + + let blocks = collected.lock().expect("mutex not poisoned"); + assert_eq!(blocks.len(), 1, "one input DecompressedBlock maps to exactly one output"); + assert!(blocks[0].index.is_none(), "index_bam=false must never emit a manifest"); +} + +/// Symmetric to [`compress_step_emits_no_manifest_when_index_bam_false`]: the +/// same framed batch through a `true` step must emit a manifest whose +/// `phys_comp_len` sums to the compressed byte count, has the expected +/// ceil-div physical-block count, and whose `records` match the input frames +/// — pinning the emit seam `BamIndexManifest` consumers (Task 5's join) rely +/// on. +#[rstest] +fn compress_step_emits_manifest_matching_frames_when_index_bam_true( + #[values(1, 4)] threads: usize, +) { + let records = test_records(4); + let batch = framed_batch(&records); + let batch_len = batch.len(); + let collected: Arc>> = Arc::new(Mutex::new(Vec::new())); + + let sink_handle = Arc::clone(&collected); + let builder = Pipeline::builder(); + builder + .chain(ReplaySource::new(vec![DecompressedBlock { batch_serial: 0, bytes: batch }])) + .chain(BgzfCompress::new(1, EDGE_LIMIT_BYTES, true)) + .chain(CollectSink { collected: sink_handle }) + .into_sink_marker(); + let pipeline = builder.build().expect("chain builds"); + pipeline.run(PipelineConfig { threads, ..Default::default() }).expect("chain runs"); + + let blocks = collected.lock().expect("mutex not poisoned"); + assert_eq!(blocks.len(), 1, "one input DecompressedBlock maps to exactly one output"); + let block = &blocks[0]; + let index = block.index.as_ref().expect("index_bam=true must emit a manifest"); + + assert_eq!( + index.phys_comp_len.iter().map(|&n| u64::from(n)).sum::(), + u64::try_from(block.bytes.len()).expect("compressed len fits u64"), + "phys_comp_len must sum to the total compressed byte count" + ); + assert_eq!( + index.phys_comp_len.len(), + batch_len.div_ceil(fgumi_bgzf::BGZF_MAX_BLOCK_SIZE), + "one physical block per BGZF_MAX_BLOCK_SIZE-sized chunk of uncompressed input" + ); + assert_eq!(index.records.len(), records.len(), "one index entry per framed record"); + let mut expected_uoffset = 0u32; + for (i, (entry, record)) in index.records.iter().zip(records.iter()).enumerate() { + let body_len = u32::try_from(record.as_ref().len()).expect("record fits u32"); + assert_eq!(entry.uoffset, expected_uoffset, "uoffset must be the frame start"); + assert_eq!(entry.len, 4 + body_len, "len must be 4 + body length"); + // `test_records(i)` places record i at ref 0, 0-based pos 100 + i*10 with + // a 4M CIGAR, mapped. `extract_alignment_context` yields the 1-based, + // inclusive span [pos + 1, pos + ref_span]. A join that emitted a + // constant or shifted context would still pass an `is_some()` check. + let ctx = entry.ctx.as_ref().expect("mapped test records carry a context"); + let pos = 100 + i * 10; + assert_eq!(ctx.ref_id, 0, "ref_id must be 0"); + assert_eq!(ctx.start.get(), pos + 1, "1-based start = pos + 1"); + assert_eq!(ctx.end.get(), pos + 4, "end = pos + 4M reference span"); + assert!(ctx.is_mapped, "test records are mapped"); + expected_uoffset += 4 + body_len; + } +} + /// An input with a header but no records must still complete cleanly and /// produce nothing — the empty-BAM case, where every step sees drain before it /// ever sees an item. @@ -561,6 +667,7 @@ fn the_no_header_boundary_variant_consumes_a_header_stripped_stream( batch_serial: i as u64, bytes: blocks.remove(0).data, uncompressed_size: u32::try_from(payload.len()).expect("fits u32"), + index: None, } }) .collect(); @@ -713,6 +820,7 @@ fn a_truncated_record_stream_fails_the_run(#[values(1, 4)] threads: usize) { batch_serial: i as u64, bytes: blocks.remove(0).data, uncompressed_size: u32::try_from(payload.len()).expect("fits u32"), + index: None, } }) .collect(); diff --git a/src/lib/pipeline/steps/tests.rs b/src/lib/pipeline/steps/tests.rs index 3656152bb..81c51a6c7 100644 --- a/src/lib/pipeline/steps/tests.rs +++ b/src/lib/pipeline/steps/tests.rs @@ -128,7 +128,7 @@ fn build_full_chain(input: &std::path::Path, output: &std::path::Path) -> Pipeli .chain(DecodeRecords::new(GroupKeyConfig::default(), bytes_4mb)) .chain(GroupBam::new(64, bytes_4mb)) .chain(SerializeBamRecords::new(bytes_4mb)) - .chain(BgzfCompress::new(1, bytes_4mb)) + .chain(BgzfCompress::new(1, bytes_4mb, false)) .chain(WriteBgzfFile::new(output, &header, 1).unwrap()) .into_sink_marker(); builder.build().unwrap() @@ -297,7 +297,7 @@ fn full_chain_with_process_mutates_mq() { .chain(bump_mq) .chain(GroupBam::new(64, bytes_4mb)) .chain(SerializeBamRecords::new(bytes_4mb)) - .chain(BgzfCompress::new(1, bytes_4mb)) + .chain(BgzfCompress::new(1, bytes_4mb, false)) .chain(WriteBgzfFile::new(&output, &header, 1).unwrap()) .into_sink_marker(); let pipeline = builder.build().unwrap(); diff --git a/src/lib/pipeline/steps/types.rs b/src/lib/pipeline/steps/types.rs index e0224cad2..0278be2ad 100644 --- a/src/lib/pipeline/steps/types.rs +++ b/src/lib/pipeline/steps/types.rs @@ -31,7 +31,10 @@ use fgumi_bam_io::{DecodedRecord, MemoryEstimate}; // `crate::template::Template`. // ───────────────────────────────────────────────────────────────────────────── -pub use fgumi_pipeline_io::types::{BgzfBlock, DecompressedBlock, RecordBatch, RecordBatchBuilder}; +pub use fgumi_pipeline_io::types::{ + BamIndexManifest, BgzfBlock, DecompressedBlock, RecordBatch, RecordBatchBuilder, + RecordIndexEntry, +}; /// Bytes held by a `Vec`'s backing store, counted by **allocated capacity**. /// diff --git a/tests/integration/main.rs b/tests/integration/main.rs index e418f118d..86145e101 100644 --- a/tests/integration/main.rs +++ b/tests/integration/main.rs @@ -9,6 +9,7 @@ mod helpers; mod test_async_reader; mod test_bam_pipeline; mod test_bgzf_eof; +mod test_chain_bam_with_index; mod test_clip_command; #[cfg(feature = "codec")] mod test_codec_command; diff --git a/tests/integration/test_chain_bam_with_index.rs b/tests/integration/test_chain_bam_with_index.rs new file mode 100644 index 000000000..0705f3b83 --- /dev/null +++ b/tests/integration/test_chain_bam_with_index.rs @@ -0,0 +1,693 @@ +//! Chain-direct integration test for the inline BAI indexer wired onto +//! `SinkSpec::BamWithIndex` (Task 6). +//! +//! Builds a `ChainSpec` directly — no command-level driver (`Sort::execute`) +//! involved — with a `[Stage::Sort]` terminal chain and `SinkSpec::BamWithIndex`, +//! runs it via `build_for(spec)?.run()` over a programmatically generated +//! unsorted BAM, and asserts the `.bai` sidecar the *inline* indexer writes +//! (as part of the arena sink's drained-finish, not a post-pipeline re-read +//! hook) exists and parses. `IndexBamFinalizeHook` no longer exists, so a +//! regression that resurrected the re-read hook would show up as a compile +//! failure elsewhere, not here — this test's job is to prove the sink itself +//! produces a working index without it. +//! +//! ## Task 7 additions: the arena-chain correctness gate +//! +//! `fgumi sort --write-index` does **not** route through this chain builder +//! today — only `SinkSpec::BamWithIndex`, built chain-direct as above, reaches +//! the inline indexer (see `task-7-report.md` for the full investigation). +//! The three tests below are therefore the correctness gate for the arena +//! sink's inline `.bai`, run the same chain-direct way as +//! `bam_with_index_produces_inline_sidecar_not_reread` above, but with the +//! full samtools-equivalence and byte-identity assertions Task 7 specifies: +//! +//! - [`arena_bam_with_index_idxstats_equivalent_to_samtools`][]: region-query +//! and `idxstats` parity against samtools' own index, over placed, placed- +//! but-unmapped, and truly-unplaced reads across multiple references. +//! - [`arena_bam_with_index_multi_block_and_straddle`][]: the same parity +//! check forced through many physical BGZF blocks and one record that +//! straddles a physical-block boundary by itself. +//! - [`arena_bam_with_index_detached_writer_is_deterministic_and_correct`][]: +//! Task 7's third requirement (Detached-vs-Serial byte-identity) is only +//! *partially* achievable here — see that test's doc comment for why, and +//! what is asserted instead. + +use std::num::NonZeroUsize; +use std::path::{Path, PathBuf}; +use std::process::Command; + +use fgumi_lib::commands::common::{ + CompressionOptions, MaxTempFiles, MemoryLimit, MemoryReserve, QueueMemoryOptions, + SchedulerOptions, ThreadingOptions, +}; +use fgumi_lib::commands::sort::{SortOptions, SortOrderArg}; +use fgumi_lib::pipeline::chains::{ + ChainSpec, SinkSpec, SourceSpec, Stage, StageOptionsBag, build_for, +}; +use fgumi_raw_bam::{RawRecord, SamBuilder, flags}; +use noodles::sam::Header; +use noodles::sam::header::record::value::Map; +use noodles::sam::header::record::value::map::ReferenceSequence; +use tempfile::TempDir; + +use crate::helpers::bam_generator::{create_minimal_header, write_bam}; + +/// `n` mapped, single-end records on one reference, placed at *descending* +/// positions so the input is deliberately unsorted — the coordinate sort +/// stage has real work to do before the sink writes (and indexes) its output. +fn unsorted_records(n: usize) -> Vec { + (0..n) + .map(|i| { + let pos = i32::try_from((n - i) * 100).expect("pos fits i32"); + let mut b = SamBuilder::new(); + b.read_name(format!("read{i}").as_bytes()) + .ref_id(0) + .pos(pos) + .mapq(60) + .flags(0) + .cigar_ops(&[4u32 << 4]) // 4M + .sequence(b"ACGT") + .qualities(&[30u8; 4]); + b.build() + }) + .collect() +} + +fn samtools_available() -> bool { + Command::new("samtools").arg("--version").output().map(|o| o.status.success()).unwrap_or(false) +} + +#[test] +fn bam_with_index_produces_inline_sidecar_not_reread() { + let dir = TempDir::new().expect("create temp dir"); + let input_bam = dir.path().join("in.bam"); + let output_bam = dir.path().join("out.bam"); + + let header = create_minimal_header("chr1", 1_000_000); + let records = unsorted_records(500); + write_bam(&input_bam, &header, &records); + + let sort_options = SortOptions { + order: SortOrderArg::Coordinate, + key_types: None, + max_memory: MemoryLimit::Fixed(64 * 1024 * 1024), + memory_reserve: MemoryReserve::Auto, + memory_per_thread: true, + tmp_dirs: Vec::new(), + sort_threads: None, + merge_threads: None, + temp_compression: 1, + temp_codec: fgumi_sort::SpillCodec::default(), + max_temp_files: MaxTempFiles::Auto, + block_batch: 4, + file_granularity: false, + }; + + let spec = ChainSpec { + stages: vec![Stage::Sort], + source: SourceSpec::Bam(input_bam.clone()), + sink: SinkSpec::BamWithIndex(output_bam.clone()), + stage_opts: StageOptionsBag { sort: Some(sort_options), ..Default::default() }, + threading: ThreadingOptions { threads: None }, + compression: CompressionOptions::default(), + scheduler: SchedulerOptions::default(), + queue_memory: QueueMemoryOptions::default(), + async_reader: false, + verify_crc: true, + command_line: "fgumi sort --write-index".to_string(), + }; + + build_for(spec) + .expect("build_for should accept a Sort-terminal BamWithIndex chain") + .run() + .expect( + "running the chain should write the coordinate-sorted BAM and its inline .bai sidecar", + ); + + assert!(output_bam.exists(), "sorted output BAM was not written"); + + let bai_path = fgumi_bam_io::bai_sidecar_path(&output_bam); + assert!(bai_path.exists(), "inline indexer did not write a .bai sidecar at {bai_path:?}"); + + // The sidecar must parse as a well-formed BAI: one reference (matching the + // header), covering the records the sort just wrote. + let index = noodles::bam::bai::fs::read(&bai_path) + .expect("inline .bai must parse as a valid BAI index"); + assert_eq!( + index.reference_sequences().len(), + 1, + "expected one reference sequence's worth of bins in the inline .bai" + ); + + // Full samtools-equivalence (idxstats / region-query parity) is Task 7's + // job; here a structural parse plus a basic region query is enough to + // confirm the inline path produces a *usable* index, not just bytes that + // happen to exist. + if samtools_available() { + let status = Command::new("samtools") + .args(["quickcheck", output_bam.to_str().unwrap()]) + .status() + .expect("run samtools quickcheck"); + assert!(status.success(), "samtools quickcheck failed on inline-indexed output BAM"); + + let out = Command::new("samtools") + .args(["view", "-c", output_bam.to_str().unwrap(), "chr1:1-50000"]) + .output() + .expect("run samtools view region query"); + assert!( + out.status.success(), + "samtools view region query against the inline .bai failed: {}", + String::from_utf8_lossy(&out.stderr) + ); + let count: usize = String::from_utf8_lossy(&out.stdout).trim().parse().unwrap_or(0); + assert!(count > 0, "expected at least one record in chr1:1-50000 via the inline .bai"); + } else { + eprintln!("skipping samtools region-query check: samtools not available in PATH"); + } +} + +// ============================================================================= +// Task 7: arena-chain correctness gate — shared helpers +// ============================================================================= + +/// A header with `refs.len()` references, each `(name, length)`. +fn multi_ref_header(refs: &[(&str, usize)]) -> Header { + let mut builder = Header::builder(); + for &(name, len) in refs { + let seq = Map::::new( + NonZeroUsize::new(len).expect("reference length must be non-zero"), + ); + builder = builder.add_reference_sequence(bstr::BString::from(name), seq); + } + builder.build() +} + +/// A deterministic scattered (i.e. deliberately unsorted) position in +/// `[1, modulus]`, salted per caller so different record categories don't +/// collide on the same positions. +fn scattered_pos(i: usize, salt: usize, modulus: usize) -> i32 { + i32::try_from(1 + (i * 97 + salt) % modulus).expect("scattered position fits i32") +} + +/// An unsorted input spanning `refs.len()` references, covering every record +/// shape the inline `.bai` must position-bin the same way samtools does: +/// normally placed mapped reads, placed-but-unmapped reads (`FUNMAP` set but +/// carrying a real ref+pos — e.g. the unmapped mate of a mapped read; both +/// fgumi and samtools bin these by position), and a trailing block of truly +/// unplaced reads (no reference at all — samtools' `n_no_coor`). +fn diverse_multiref_records( + mapped_per_ref: usize, + num_refs: usize, + ref_len: usize, +) -> Vec { + let mut records = Vec::new(); + let mut b = SamBuilder::new(); + let modulus = ref_len.saturating_sub(100).max(1); + + for ref_idx in 0..num_refs { + for i in 0..mapped_per_ref { + let pos = scattered_pos(i, ref_idx * 31 + 13, modulus); + b.clear(); + b.read_name(format!("m{ref_idx}_{i}").as_bytes()) + .ref_id(i32::try_from(ref_idx).expect("ref_id fits i32")) + .pos(pos) + .mapq(60) + .flags(0) + .cigar_ops(&[4u32 << 4]) // 4M + .sequence(b"ACGT") + .qualities(&[30u8; 4]); + records.push(b.build()); + } + } + + // Placed-but-unmapped: FUNMAP set, real ref+pos. + for ref_idx in 0..num_refs { + for i in 0..(mapped_per_ref / 4).max(1) { + let pos = scattered_pos(i, ref_idx * 7 + 5, modulus); + b.clear(); + b.read_name(format!("pu{ref_idx}_{i}").as_bytes()) + .ref_id(i32::try_from(ref_idx).expect("ref_id fits i32")) + .pos(pos) + .mapq(0) + .flags(flags::UNMAPPED) + .sequence(b"ACGT") + .qualities(&[30u8; 4]); + records.push(b.build()); + } + } + + // Trailing truly-unplaced region: no reference, no position. + for i in 0..(mapped_per_ref / 8).max(50) { + b.clear(); + b.read_name(format!("u{i}").as_bytes()) + .ref_id(-1) + .pos(-1) + .mapq(0) + .flags(flags::UNMAPPED) + .sequence(b"ACGT") + .qualities(&[30u8; 4]); + records.push(b.build()); + } + + records +} + +/// `mapped_per_ref` scattered mapped reads on each of `num_refs` references — +/// enough, at default BGZF compression, to span many 65280-byte physical +/// blocks (Task 7 requirement 2). No unmapped reads: this generator's only +/// job is bulk plus one oversized record (below); category coverage is +/// `diverse_multiref_records`'s job. +fn many_mapped_records(mapped_per_ref: usize, num_refs: usize, ref_len: usize) -> Vec { + let mut records = Vec::with_capacity(mapped_per_ref * num_refs); + let mut b = SamBuilder::new(); + let modulus = ref_len.saturating_sub(100).max(1); + for ref_idx in 0..num_refs { + for i in 0..mapped_per_ref { + let pos = scattered_pos(i, ref_idx * 31 + 13, modulus); + b.clear(); + b.read_name(format!("m{ref_idx}_{i}").as_bytes()) + .ref_id(i32::try_from(ref_idx).expect("ref_id fits i32")) + .pos(pos) + .mapq(60) + .flags(0) + .cigar_ops(&[4u32 << 4]) // 4M + .sequence(b"ACGT") + .qualities(&[30u8; 4]); + records.push(b.build()); + } + } + records +} + +/// One record whose SEQ/QUAL/CIGAR are `len` bases long — well over one BGZF +/// physical block (`fgumi_bgzf::BGZF_MAX_BLOCK_SIZE` is 65280 bytes) — so the +/// record's own bytes straddle a physical-block boundary in the compressed +/// output (Task 7 requirement 2's "record spans physical blocks" case). +fn oversized_record(ref_id: i32, pos: i32, len: usize) -> RawRecord { + let bases = vec![b'A'; len]; + let quals = vec![30u8; len]; + let mut b = SamBuilder::new(); + b.read_name(b"oversized") + .ref_id(ref_id) + .pos(pos) + .mapq(60) + .flags(0) + .cigar_ops(&[u32::try_from(len).expect("record length fits u32 CIGAR op") << 4]) + .sequence(&bases) + .qualities(&quals); + b.build() +} + +/// A `SortOptions` covering every field `add_sort` reads, matching +/// `Sort::to_sort_options()`'s field set (mirrors Task 6's test). Only +/// `max_memory` and `tmp_dirs` vary per test below. +fn sort_options(max_memory: MemoryLimit, tmp_dirs: Vec) -> SortOptions { + SortOptions { + order: SortOrderArg::Coordinate, + key_types: None, + max_memory, + memory_reserve: MemoryReserve::Auto, + memory_per_thread: true, + tmp_dirs, + sort_threads: None, + merge_threads: None, + temp_compression: 1, + temp_codec: fgumi_sort::SpillCodec::default(), + max_temp_files: MaxTempFiles::Auto, + block_batch: 4, + file_granularity: false, + } +} + +/// A `[Stage::Sort]`-terminal `SinkSpec::BamWithIndex` chain spec — the exact +/// shape `SinkSpec::BamWithIndex` requires (`validate.rs` Rule 3) and the only +/// shape from which `ChainBuilder::add_sort` sets `detached_writer = true` +/// (`is_standalone_sort`, i.e. `stages == [Stage::Sort]`). See +/// [`arena_bam_with_index_detached_writer_is_deterministic_and_correct`]'s doc +/// comment for why that ties every chain-direct `BamWithIndex` test to the +/// Detached writer. +fn bam_with_index_spec( + input: PathBuf, + output: PathBuf, + sort_options: SortOptions, + threads: Option, +) -> ChainSpec { + ChainSpec { + stages: vec![Stage::Sort], + source: SourceSpec::Bam(input), + sink: SinkSpec::BamWithIndex(output), + stage_opts: StageOptionsBag { sort: Some(sort_options), ..Default::default() }, + threading: ThreadingOptions { threads }, + compression: CompressionOptions::default(), + scheduler: SchedulerOptions::default(), + queue_memory: QueueMemoryOptions::default(), + async_reader: false, + verify_crc: true, + command_line: "fgumi sort --write-index".to_string(), + } +} + +/// `samtools view -c [extra] ` region query via the `.bai`. +/// `extra` carries filter flags (e.g. `-F 4`). +/// +/// Returns the records rather than a count: a matching count alone cannot +/// distinguish "returned the right records" from "returned the right +/// *number* of the wrong records," which a mis-recovered virtual offset or a +/// bin pointing at a neighbouring chunk would produce identically. Mirrors +/// `tests/integration/test_sort_write_index.rs`'s helper of the same name. +fn region_records(bam: &Path, region: &str, extra: &[&str]) -> String { + let out = Command::new("samtools") + .arg("view") + .args(extra) + .args([bam.to_str().unwrap(), region]) + .output() + .expect("run samtools view"); + assert!( + out.status.success(), + "samtools view {region} failed on {}: {}", + bam.display(), + String::from_utf8_lossy(&out.stderr) + ); + String::from_utf8_lossy(&out.stdout).into_owned() +} + +/// Assert two region-query outputs hold the identical records, reporting the +/// first divergence compactly (these regions can return thousands of records, +/// so a whole-output `assert_eq!` diff would be unreadable). +fn assert_same_records(region: &str, via_fgumi: &str, via_samtools: &str) { + if via_fgumi == via_samtools { + return; + } + let fgumi_lines: Vec<&str> = via_fgumi.lines().collect(); + let samtools_lines: Vec<&str> = via_samtools.lines().collect(); + let detail = match fgumi_lines.iter().zip(samtools_lines.iter()).position(|(a, b)| a != b) { + Some(i) => format!( + "first differing record at index {i}:\n via fgumi: {}\n via samtools: {}", + fgumi_lines[i], samtools_lines[i] + ), + None => "one output is a strict prefix of the other".to_string(), + }; + panic!( + "region {region}: records retrieved via fgumi's inline .bai differ from samtools' .bai \ + ({} vs {} records); {detail}", + fgumi_lines.len(), + samtools_lines.len() + ); +} + +/// `samtools idxstats ` — per-reference mapped/unmapped counts (plus the +/// trailing no-coordinate count) from the `.bai`. +fn idxstats(bam: &Path) -> String { + let out = Command::new("samtools") + .args(["idxstats", bam.to_str().unwrap()]) + .output() + .expect("run samtools idxstats"); + assert!( + out.status.success(), + "samtools idxstats failed: {}", + String::from_utf8_lossy(&out.stderr) + ); + String::from_utf8_lossy(&out.stdout).into_owned() +} + +/// Parse the inline `.bai` sidecar as a well-formed BAI and assert it carries +/// one reference sequence's worth of bins per header reference. +/// +/// Called unconditionally — before any `samtools_available()` gate — so the +/// chain, the sidecar write, and the BAI parser are exercised on every run, +/// even the samtools-less CI job that only runs `#[ignore]`d tests. The +/// samtools parity assertions each test adds on top are what the gate skips. +fn assert_inline_bai_valid(bai_path: &Path, expected_refs: usize) { + let index = + noodles::bam::bai::fs::read(bai_path).expect("inline .bai must parse as a valid BAI index"); + assert_eq!( + index.reference_sequences().len(), + expected_refs, + "inline .bai reference-sequence count must match the header", + ); +} + +/// Task 7 requirement 1: the inline `.bai` `SinkSpec::BamWithIndex` writes +/// must be idxstats/region-query equivalent to one samtools builds for the +/// identical output bytes, across the record shapes samtools and fgumi must +/// bin the same way — placed mapped reads, placed-but-unmapped reads, and a +/// trailing truly-unplaced region — spread over three references. +#[test] +fn arena_bam_with_index_idxstats_equivalent_to_samtools() { + let dir = TempDir::new().expect("create temp dir"); + let input_bam = dir.path().join("in.bam"); + let output_bam = dir.path().join("out.bam"); + + let refs = [("chr1", 200_000usize), ("chr2", 200_000usize), ("chr3", 200_000usize)]; + let header = multi_ref_header(&refs); + let records = diverse_multiref_records(800, refs.len(), 200_000); + write_bam(&input_bam, &header, &records); + + let spec = bam_with_index_spec( + input_bam, + output_bam.clone(), + sort_options(MemoryLimit::Fixed(64 * 1024 * 1024), Vec::new()), + None, + ); + + build_for(spec) + .expect("build_for should accept a Sort-terminal BamWithIndex chain") + .run() + .expect("running the chain should write the sorted BAM and its inline .bai"); + + assert!(output_bam.exists(), "sorted output BAM was not written"); + let fgumi_bai = fgumi_bam_io::bai_sidecar_path(&output_bam); + assert!(fgumi_bai.exists(), "inline indexer did not write a .bai sidecar at {fgumi_bai:?}"); + // Exercise the chain, sidecar, and parser regardless of samtools. + assert_inline_bai_valid(&fgumi_bai, refs.len()); + + if !samtools_available() { + eprintln!( + "arena_bam_with_index_idxstats_equivalent_to_samtools: samtools not on PATH; \ + validated the inline .bai only, skipped the samtools parity assertions" + ); + return; + } + + // A samtools-built index over an independent copy of the identical output + // bytes: this isolates the two indexes (each colocated with its own copy + // of the BAM) rather than overwriting fgumi's inline sidecar in place. + let reference = dir.path().join("reference.bam"); + std::fs::copy(&output_bam, &reference).expect("copy output bam for reference indexing"); + let status = Command::new("samtools") + .args(["index", reference.to_str().unwrap()]) + .status() + .expect("run samtools index"); + assert!(status.success(), "samtools index failed"); + + let regions = [ + "chr1", + "chr2", + "chr3", + "chr1:1-100000", + "chr2:50000-150000", + "chr3:150000-200000", + "chr1:199000-200000", + ]; + let mut total = 0usize; + for r in regions { + let via_fgumi = region_records(&output_bam, r, &[]); + let via_samtools = region_records(&reference, r, &[]); + assert_same_records(r, &via_fgumi, &via_samtools); + total += via_fgumi.lines().count(); + } + assert!(total > 0, "expected non-empty region queries"); + + assert_eq!( + idxstats(&output_bam), + idxstats(&reference), + "idxstats via the inline .bai must match samtools' idxstats over the identical bytes" + ); +} + +/// Task 7 requirement 2: force many 65280-byte physical BGZF blocks (a small +/// in-memory budget forces many spill runs and a multi-block k-way merge) +/// plus one record whose own bytes are larger than one physical block, and +/// assert the inline `.bai` still answers region queries identically to +/// samtools — proving the join between the merge's recovered virtual offsets +/// and `BgzfCompress`'s per-record manifest survives both a batch boundary +/// that doesn't line up with a physical block AND a record straddling a +/// physical block by itself. +#[test] +fn arena_bam_with_index_multi_block_and_straddle() { + let dir = TempDir::new().expect("create temp dir"); + let input_bam = dir.path().join("in.bam"); + let output_bam = dir.path().join("out.bam"); + + let refs = [("chr1", 2_000_000usize), ("chr2", 2_000_000usize)]; + let header = multi_ref_header(&refs); + let mut records = many_mapped_records(12_000, refs.len(), 2_000_000); + records.push(oversized_record(0, 100_000, 70_000)); // > BGZF_MAX_BLOCK_SIZE (65280) + write_bam(&input_bam, &header, &records); + + // A 64 KiB in-memory budget over ~24k records forces many real spill + // files, so both the spill k-way merge and the sink's compression cross + // physical-block boundaries repeatedly. Spill temp files go to the test's + // own disk-backed TempDir (`--tmp-dir`), not the default tmp location. + let spec = bam_with_index_spec( + input_bam, + output_bam.clone(), + sort_options(MemoryLimit::Fixed(64 * 1024), vec![dir.path().to_path_buf()]), + Some(4), + ); + + build_for(spec) + .expect("build_for should accept a Sort-terminal BamWithIndex chain") + .run() + .expect("running the chain should write the sorted BAM and its inline .bai"); + + assert!(output_bam.exists(), "sorted output BAM was not written"); + let fgumi_bai = fgumi_bam_io::bai_sidecar_path(&output_bam); + assert!(fgumi_bai.exists(), "inline indexer did not write a .bai sidecar at {fgumi_bai:?}"); + // Exercise the chain, sidecar, and parser regardless of samtools. + assert_inline_bai_valid(&fgumi_bai, refs.len()); + + if !samtools_available() { + eprintln!( + "arena_bam_with_index_multi_block_and_straddle: samtools not on PATH; validated the \ + inline .bai only, skipped the samtools parity assertions" + ); + return; + } + + let reference = dir.path().join("reference.bam"); + std::fs::copy(&output_bam, &reference).expect("copy output bam for reference indexing"); + let status = Command::new("samtools") + .args(["index", reference.to_str().unwrap()]) + .status() + .expect("run samtools index"); + assert!(status.success(), "samtools index failed"); + + let regions = [ + "chr1", + "chr2", + "chr1:1-500000", + "chr1:90000-180000", // spans the oversized record + "chr1:1500000-2000000", + "chr2:1-2000000", + ]; + let mut total = 0usize; + for r in regions { + let via_fgumi = region_records(&output_bam, r, &[]); + let via_samtools = region_records(&reference, r, &[]); + assert_same_records(r, &via_fgumi, &via_samtools); + total += via_fgumi.lines().count(); + } + assert!(total > 0, "expected non-empty region queries"); + + assert_eq!( + idxstats(&output_bam), + idxstats(&reference), + "idxstats via the inline .bai must match samtools' idxstats over the identical bytes" + ); +} + +/// Task 7 requirement 3: Detached-vs-Serial writer byte-identity. +/// +/// **No seam exists to force the Serial (non-detached) writer from a +/// chain-direct test, so this asserts the achievable subset instead of +/// faking the comparison — per the ruling's explicit fallback.** +/// +/// `ChainBuilder::detached_writer` (private, +/// `src/lib/pipeline/chains/builder.rs`) is set to `true` only when +/// `self.spec.is_sort_terminal()` — i.e. `spec.stages == [Stage::Sort]` +/// exactly — and the Sort stage sits at the chain's terminal position +/// (`add_sort`, guarded by `is_standalone_sort`). `SinkSpec::BamWithIndex` +/// only requires `stages.last() == Some(Stage::Sort)` (`validate.rs` Rule 3), +/// which is *weaker* than `is_sort_terminal()` — so in principle a +/// `BamWithIndex` chain with a leading non-Sort stage (e.g. +/// `[Stage::Correct, Stage::Sort]`) would leave `detached_writer` at its +/// default `false` (Serial). But inserting a real upstream stage changes what +/// is being compared: `Correct` (or any other stage) decodes and rewrites +/// every record before Sort ever sees it, so a byte-identity check between +/// that chain's output and a plain `[Stage::Sort]` chain's output would be +/// confounded by the upstream stage's own re-encoding, not isolate the +/// writer-threading variable Task 7 asks about. No `Stage` is a verified +/// byte-preserving no-op, and inventing a test-only production seam to flip +/// `detached_writer` directly from outside `builder.rs` was out of scope for +/// a correctness-gate task (see `task-7-report.md` for the full reasoning). +/// +/// What IS asserted here: every `SinkSpec::BamWithIndex` chain built +/// chain-direct takes the Detached writer, so this runs that (only reachable) +/// configuration twice over the same input and asserts (a) the two runs are +/// byte-identical to each other — the Detached writer is deterministic, not +/// just "correct once" — and (b) the produced `.bai` is idxstats-equivalent +/// to samtools, re-confirming correctness for this configuration +/// independently of the two tests above. This is confirmatory, not +/// load-bearing, for the Detached/Serial claim itself: Task 5's review +/// already reasoned that `with_detached()` only relocates which thread drives +/// the sink without touching the manifest/join math that produces the index +/// bytes, so Serial's BAM/`.bai` bytes are expected — not proven here — to be +/// identical to Detached's. +#[test] +fn arena_bam_with_index_detached_writer_is_deterministic_and_correct() { + let dir = TempDir::new().expect("create temp dir"); + let refs = [("chr1", 100_000usize), ("chr2", 100_000usize)]; + let header = multi_ref_header(&refs); + let records = diverse_multiref_records(300, refs.len(), 100_000); + let input_bam = dir.path().join("in.bam"); + write_bam(&input_bam, &header, &records); + + let run = |name: &str| -> PathBuf { + let output_bam = dir.path().join(name); + let spec = bam_with_index_spec( + input_bam.clone(), + output_bam.clone(), + sort_options(MemoryLimit::Fixed(64 * 1024 * 1024), Vec::new()), + None, + ); + build_for(spec) + .expect("build_for should accept a Sort-terminal BamWithIndex chain") + .run() + .expect("running the chain should write the sorted BAM and its inline .bai"); + output_bam + }; + + let bam_run_a = run("run_a.bam"); + let bam_run_b = run("run_b.bam"); + let index_run_a = fgumi_bam_io::bai_sidecar_path(&bam_run_a); + let index_run_b = fgumi_bam_io::bai_sidecar_path(&bam_run_b); + + assert_eq!( + std::fs::read(&bam_run_a).expect("read run_a.bam"), + std::fs::read(&bam_run_b).expect("read run_b.bam"), + "two runs of the arena chain's Detached-writer sink over the same input must produce \ + byte-identical BAM output" + ); + assert_eq!( + std::fs::read(&index_run_a).expect("read run_a.bam.bai"), + std::fs::read(&index_run_b).expect("read run_b.bam.bai"), + "two runs of the arena chain's Detached-writer sink over the same input must produce a \ + byte-identical inline .bai" + ); + // The deterministic byte-identity checks above need no samtools; the sidecar + // parse also runs regardless of it. + assert_inline_bai_valid(&index_run_a, refs.len()); + + if !samtools_available() { + eprintln!( + "arena_bam_with_index_detached_writer_is_deterministic_and_correct: samtools not on \ + PATH; validated determinism and the inline .bai only, skipped the samtools parity \ + assertions" + ); + return; + } + + let reference = dir.path().join("reference.bam"); + std::fs::copy(&bam_run_a, &reference).expect("copy output bam for reference indexing"); + let status = Command::new("samtools") + .args(["index", reference.to_str().unwrap()]) + .status() + .expect("run samtools index"); + assert!(status.success(), "samtools index failed"); + assert_eq!( + idxstats(&bam_run_a), + idxstats(&reference), + "idxstats via the inline .bai must match samtools' idxstats over the identical bytes" + ); +}