fix(index): position-bin placed-but-unmapped reads in the BAI

crates/fgumi-bam-io/src/writer.rs

@@ -366,8 +366,14 @@ struct CachedIndexEntry {
 
 /// Extract alignment context from raw BAM bytes.
 ///
-/// Returns `Some((ref_id, start, end, is_mapped))` for mapped reads, or `None`
-/// for unmapped reads (the indexer still counts them via a `None` context).
+/// 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.
 #[inline]
 #[allow(clippy::cast_sign_loss, clippy::cast_possible_wrap)]
 pub(crate) fn extract_alignment_context(bam: &[u8]) -> Option<(usize, Position, Position, bool)> {
@@ -376,22 +382,26 @@ pub(crate) fn extract_alignment_context(bam: &[u8]) -> Option<(usize, Position,
     let pos = v.pos();
     let flags = v.flags();
 
-    let is_unmapped = (flags & fgumi_raw_bam::flags::UNMAPPED) != 0;
-
-    // Unmapped reads: return None (indexer handles them specially).
-    if tid < 0 || is_unmapped {
+    // Only a read with no reference / no position is unplaced. A read carrying a
+    // valid tid+pos is binned by position regardless of its mapped flag — this is
+    // what htslib does, and what makes placed-but-unmapped mates queryable.
+    if tid < 0 || pos < 0 {
         return None;
     }
 
-    // Calculate alignment end from CIGAR using byte-safe access
-    // (doesn't require 4-byte alignment like get_cigar_ops).
-    let ref_len = fgumi_raw_bam::reference_length_from_raw_bam(bam);
+    let is_mapped = (flags & fgumi_raw_bam::flags::UNMAPPED) == 0;
+
+    // Reference span from the CIGAR (byte-safe access; no 4-byte alignment
+    // needed). A placed-but-unmapped read has no CIGAR, so this is 0; htslib
+    // bins such a read over `[pos, pos + 1)`, so floor the span at one base
+    // (matching `bam_endpos`, which returns `pos + max(rlen, 1)`).
+    let ref_len = fgumi_raw_bam::reference_length_from_raw_bam(bam).max(1);
 
     // BAM positions are 0-based, Position is 1-based.
     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, true))
+    Some((tid as usize, start, end, is_mapped))
 }
 
 /// Incrementally builds a BAI index from per-record positions plus per-block
@@ -1413,6 +1423,62 @@ mod tests {
         record
     }
 
+    /// Build a raw BAM record for a *placed-but-unmapped* read: a valid
+    /// reference/position, the UNMAPPED flag set, and no CIGAR — as produced for
+    /// the unmapped mate of a mapped read.
+    #[allow(clippy::cast_possible_truncation)]
+    fn create_placed_unmapped_bam_record(ref_id: i32, pos: i32, read_name: &[u8]) -> Vec<u8> {
+        let name_with_null = read_name.len() + 1;
+        let padding = (4 - (name_with_null % 4)) % 4;
+        let l_read_name = (name_with_null + padding) as u8;
+
+        let mut record = Vec::new();
+        record.extend_from_slice(&ref_id.to_le_bytes());
+        record.extend_from_slice(&pos.to_le_bytes());
+        record.push(l_read_name);
+        record.push(0); // mapq
+        record.extend_from_slice(&0u16.to_le_bytes()); // bin
+        record.extend_from_slice(&0u16.to_le_bytes()); // n_cigar_op = 0 (no CIGAR)
+        record.extend_from_slice(&fgumi_raw_bam::flags::UNMAPPED.to_le_bytes()); // flag
+        record.extend_from_slice(&0u32.to_le_bytes()); // l_seq = 0
+        record.extend_from_slice(&(-1i32).to_le_bytes()); // next_ref_id
+        record.extend_from_slice(&(-1i32).to_le_bytes()); // next_pos
+        record.extend_from_slice(&0i32.to_le_bytes()); // tlen
+        record.extend_from_slice(read_name);
+        record.push(0);
+        record.extend(std::iter::repeat_n(0u8, padding));
+        record
+    }
+
+    #[test]
+    fn extract_alignment_context_bins_placed_unmapped_reads() {
+        // Placed-but-unmapped (valid ref+pos, UNMAPPED flag, no CIGAR) must be
+        // 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");
+
+        // 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);
+
+        // A truly unplaced read (no reference) is not binned.
+        assert!(
+            extract_alignment_context(&create_placed_unmapped_bam_record(-1, -1, b"unplaced"))
+                .is_none(),
+            "unplaced read -> None (counted as unplaced by the indexer)"
+        );
+    }
+
     #[test]
     fn test_create_indexing_bam_writer() -> Result<()> {
         let temp_file = NamedTempFile::new()?;

tests/integration/test_sort_write_index.rs

@@ -192,10 +192,9 @@ fn create_large_test_bam(dir: &Path, mapped_per_contig: usize) -> PathBuf {
         let _ = writeln!(sam, "u_{i}\t4\t*\t0\t0\t*\t*\t0\t0\tACGTACGTAC\tIIIIIIIIII");
     }
     // Placed-but-unmapped reads: FUNMAP (0x4) set, but with a reference + position
-    // (e.g. the unmapped mate of a mapped read). samtools position-bins these so
-    // they appear in region queries; fgumi's indexer classifies them as unmapped
-    // and does not position-bin them (pre-existing behavior, unchanged by the
-    // pooled-writer work). The test therefore compares MAPPED-only retrieval.
+    // (e.g. the unmapped mate of a mapped read). Both samtools and fgumi
+    // position-bin these, so they appear in region queries and idxstats; the test
+    // asserts fgumi matches samtools over ALL reads, these included.
     for i in 0..2000 {
         let p1 = 1 + (i * 101) % 999_000;
         let p2 = 1 + (i * 149 + 7) % 999_000;
@@ -216,39 +215,51 @@ fn create_large_test_bam(dir: &Path, mapped_per_contig: usize) -> PathBuf {
 
 /// `samtools view -c [extra] <bam> <region>` — record count in a region via the
 /// `.bai`. `extra` carries filter flags (e.g. `-F 4` to count mapped reads only).
-fn region_count(bam: &Path, region: &str, extra: &[&str]) -> i64 {
+/// `samtools view <bam> <region>` — the records a region query returns, verbatim.
+///
+/// Returns the records rather than a count. The two BAMs being compared are
+/// byte-identical copies that differ only in which `.bai` retrieves from them, so a
+/// matching count cannot distinguish "returned the right records" from "returned the
+/// right *number* of the wrong records" — which is exactly what a mis-recovered
+/// virtual offset, or a bin pointing at a neighbouring chunk, would produce.
+fn region_records(bam: &Path, region: &str, extra: &[&str]) -> String {
     let out = Command::new("samtools")
         .arg("view")
-        .arg("-c")
         .args(extra)
         .args([bam.to_str().unwrap(), region])
         .output()
-        .expect("run samtools view -c");
+        .expect("run samtools view");
     assert!(
         out.status.success(),
-        "samtools view -c {region} failed on {}: {}",
+        "samtools view {region} failed on {}: {}",
         bam.display(),
         String::from_utf8_lossy(&out.stderr)
     );
-    String::from_utf8_lossy(&out.stdout).trim().parse().expect("parse count")
+    String::from_utf8_lossy(&out.stdout).into_owned()
 }
 
-/// Per-reference MAPPED counts (`ref\tlen\tmapped`) from idxstats. Excludes the
-/// unmapped column, which legitimately differs for placed-but-unmapped reads.
-fn idxstats_mapped(bam: &Path) -> String {
-    idxstats(bam)
-        .lines()
-        .map(|line| {
-            let mut f = line.split('\t');
-            format!(
-                "{}\t{}\t{}",
-                f.next().unwrap_or_default(),
-                f.next().unwrap_or_default(),
-                f.next().unwrap_or_default()
-            )
-        })
-        .collect::<Vec<_>>()
-        .join("\n")
+/// Assert two region-query outputs hold the identical records, reporting the first
+/// divergence compactly — these regions return tens of 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 .bai differ from samtools' .bai \
+         ({} vs {} records); {detail}",
+        fgumi_lines.len(),
+        samtools_lines.len()
+    );
 }
 
 /// `samtools idxstats <bam>` — per-reference mapped/unmapped counts from the `.bai`.
@@ -313,7 +324,8 @@ fn test_sort_write_index_matches_samtools_index() {
         .expect("run samtools index");
     assert!(status.success(), "samtools index failed");
 
-    // Region counts via fgumi's .bai must equal those via samtools' .bai.
+    // The records a region query returns via fgumi's .bai must be identical to those
+    // returned via samtools' .bai.
     let regions = [
         "chr1",
         "chr2",
@@ -324,27 +336,25 @@ fn test_sort_write_index_matches_samtools_index() {
         "chr1:999500-1000000",
         "chr1:1-1000000",
     ];
-    // MAPPED-read retrieval (`-F 4`) via fgumi's .bai must exactly match samtools'
-    // .bai over the same bytes — this validates the recovered virtual offsets.
-    // (Placed-but-unmapped reads are position-binned by samtools but not by fgumi;
-    // see create_large_test_bam. That is a pre-existing indexing choice, not a
-    // property of the pooled writer, so it is deliberately excluded here.)
-    let mut total = 0i64;
+    // Retrieval via fgumi's .bai must exactly match samtools' .bai over the same
+    // bytes — for ALL reads, including the placed-but-unmapped reads
+    // (`create_large_test_bam` seeds many): fgumi now position-bins them like
+    // htslib, so they are returned by region queries too. This validates both the
+    // recovered virtual offsets and the placed-unmapped binning.
+    let mut total = 0usize;
     for r in regions {
-        let via_fgumi = region_count(&sorted, r, &["-F", "4"]);
-        let via_samtools = region_count(&reference, r, &["-F", "4"]);
-        assert_eq!(
-            via_fgumi, via_samtools,
-            "region {r}: fgumi .bai mapped count {via_fgumi} != samtools .bai {via_samtools}"
-        );
-        total += via_fgumi;
+        let via_fgumi = region_records(&sorted, 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 mapped region queries");
+    assert!(total > 0, "expected non-empty region queries");
 
-    // Per-reference mapped counts must match exactly.
+    // Per-reference mapped AND unmapped counts must match exactly (the unmapped
+    // column now agrees because placed-but-unmapped reads are position-binned).
     assert_eq!(
-        idxstats_mapped(&sorted),
-        idxstats_mapped(&reference),
-        "per-reference mapped counts via fgumi .bai must match samtools"
+        idxstats(&sorted),
+        idxstats(&reference),
+        "idxstats via fgumi .bai must match samtools"
     );
 }