diff --git a/Cargo.lock b/Cargo.lock index 30a6be15..e069b385 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -88,12 +88,6 @@ version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" -[[package]] -name = "block-aligner" -version = "0.5.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8986feb1596f749fdcf1046e1b8068dfe802fc45708e6e293b736e4c0bb3b3a5" - [[package]] name = "bstr" version = "1.12.1" @@ -582,7 +576,6 @@ name = "strobealign" version = "0.18.0-alpha" dependencies = [ "assert_cmd", - "block-aligner", "cc", "clap", "fastrand", diff --git a/Cargo.toml b/Cargo.toml index d5542664..f3be4627 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -21,12 +21,6 @@ flate2 = { version = "1.0", features = ["zlib-rs"]} thiserror = "2.0.11" rayon = "1.11.0" mimalloc = "0.1.48" -[target.'cfg(target_arch = "wasm32-wasi")'.dependencies] -block-aligner = { version = "0.5", features = ["simd_wasm"] } -[target.'cfg(target_arch = "x86_64")'.dependencies] -block-aligner = { version = "0.5", features = ["simd_avx2"] } -[target.'cfg(target_arch = "aarch64")'.dependencies] -block-aligner = { version = "0.5", features = ["simd_neon"] } [dev-dependencies] assert_cmd = "2.0.16" diff --git a/src/aligner.rs b/src/aligner.rs index bdf391c1..e90777dc 100644 --- a/src/aligner.rs +++ b/src/aligner.rs @@ -55,14 +55,14 @@ pub struct Aligner { } impl Aligner { - pub fn new(scores: Scores, k: usize, xdrop: i32) -> Self { + pub fn new(scores: Scores, k: usize, bandwidth: usize) -> Self { let ssw_aligner = SswAligner::new( scores.match_, scores.mismatch, scores.gap_open, scores.gap_extend, ); - let piecewise_aligner = PiecewiseAligner::new(scores, k, xdrop); + let piecewise_aligner = PiecewiseAligner::new(scores, k, bandwidth); Aligner { scores, ssw_aligner, diff --git a/src/cigar.rs b/src/cigar.rs index a45414d7..8dbc2703 100644 --- a/src/cigar.rs +++ b/src/cigar.rs @@ -84,6 +84,12 @@ impl Cigar { Cigar { ops: vec![] } } + pub fn with_capacity(hint: usize) -> Self { + Cigar { + ops: Vec::with_capacity(hint), + } + } + pub fn is_empty(&self) -> bool { self.ops.is_empty() } @@ -124,6 +130,10 @@ impl Cigar { } } + pub fn reverse(&mut self) { + self.ops.reverse(); + } + pub fn extend(&mut self, other: &Cigar) { // TODO use push only for first, then extend for oplen in &other.ops { diff --git a/src/lib.rs b/src/lib.rs index e385f03b..77cc85f7 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -20,4 +20,5 @@ pub mod refseq; pub mod revcomp; pub mod seeding; pub mod shuffle; +pub mod simdaligner; pub mod ssw; diff --git a/src/main.rs b/src/main.rs index 51fccbaf..2c8dbad1 100644 --- a/src/main.rs +++ b/src/main.rs @@ -267,9 +267,10 @@ struct Args { #[arg(long = "ssw", help_heading = "Alignment")] use_ssw: bool, - /// X-drop threshold for piecewise extension - #[arg(long = "xdrop", default_value_t = 500, value_name = "N", help_heading = "Alignment")] - xdrop: i32, + /// Bandwidth for the piecewise start/end extensions; extensions against a longer reference are + /// confined to this many diagonals off the anchor diagonal, shorter ones are aligned exactly + #[arg(long = "bw", default_value_t = 1024, value_name = "N", help_heading = "Alignment")] + bandwidth: usize, /// Path to input reference (in FASTA format) @@ -537,7 +538,7 @@ fn run() -> Result<(), CliError> { debug!("{:?}", &scores); let chainer = Chainer::new(index.k(), chaining_parameters); - let aligner = Aligner::new(scores, index.k(), args.xdrop); + let aligner = Aligner::new(scores, index.k(), args.bandwidth); let cmd_line = env::args().skip(1).collect::>().join(" "); let rg_id = match args.rg_id { diff --git a/src/piecewisealigner.rs b/src/piecewisealigner.rs index a7b4aa59..2b6a1d1f 100644 --- a/src/piecewisealigner.rs +++ b/src/piecewisealigner.rs @@ -1,230 +1,31 @@ use crate::{ aligner::{AlignmentInfo, Scores, hamming_align_global}, chainer::Anchor, - cigar::{Cigar, CigarOperation}, + cigar::CigarOperation, + simdaligner::{AlignmentResult, SimdAligner}, }; -use block_aligner::{ - cigar::{Cigar as BlockCigar, Operation}, - scan_block::{Block, PaddedBytes}, - scores::{AAMatrix, AAProfile, Gaps, NucMatrix, Profile}, -}; - -// Maximum value for blockaligner's block sizes, higher values might cause a crash -const MAXIMUM_BLOCK_SIZE: usize = 8192; - -#[derive(Clone, Copy)] -enum XdropMode { - GlobalStartLocalEnd, - LocalStartGlobalEnd, -} +use std::cell::RefCell; -#[derive(Clone)] pub struct PiecewiseAligner { scores: Scores, k: usize, - gaps: Gaps, - matrix: NucMatrix, - xdrop: i32, + bandwidth: usize, + simd_aligner: RefCell, } -struct AlignmentResult { - pub score: i32, - pub query_start: usize, - pub query_end: usize, - pub ref_start: usize, - pub ref_end: usize, - pub cigar: Cigar, -} - -impl Default for AlignmentResult { - fn default() -> Self { - AlignmentResult { - score: 0, - query_start: 0, - query_end: 0, - ref_start: 0, - ref_end: 0, - cigar: Cigar::new(), - } +impl Clone for PiecewiseAligner { + fn clone(&self) -> Self { + PiecewiseAligner::new(self.scores, self.k, self.bandwidth) } } impl PiecewiseAligner { - pub fn new(scores: Scores, k: usize, xdrop: i32) -> Self { - let gaps = Gaps { - open: -(scores.gap_open as i8), - extend: -(scores.gap_extend as i8), - }; - let matrix = NucMatrix::new_simple(scores.match_ as i8, -(scores.mismatch as i8)); + pub fn new(scores: Scores, k: usize, bandwidth: usize) -> Self { PiecewiseAligner { scores, k, - gaps, - matrix, - xdrop, - } - } - - /// Performs a global alignment between two sequences. - /// - /// The alignment uses the blockaligner crate with block sizes - /// determined based on the largest length of the 2 sequences. - /// - /// # Parameters - /// - /// * `query` - The query sequence as a byte slice - /// * `reference` - The reference sequence as a byte slice - /// - /// # Returns - /// - /// An [`AlignmentResult`] containing: - /// - The alignment score - /// - Start and end positions in both sequences - /// - A CIGAR string representing the alignment operations - fn global_alignment(&self, query: &[u8], reference: &[u8]) -> AlignmentResult { - if query.is_empty() || reference.is_empty() { - return AlignmentResult::default(); - } - - // Blockaligner block ranges, set to maximum values for full global alignment - // Must be powers of 2 - let max_len = query.len().max(reference.len()); - let block_size = max_len.next_power_of_two().clamp(32, MAXIMUM_BLOCK_SIZE); - - // Create padded sequences - let mut q_padded = PaddedBytes::new::(query.len(), block_size); - q_padded.set_bytes::(query, block_size); - let mut r_padded = PaddedBytes::new::(reference.len(), block_size); - r_padded.set_bytes::(reference, block_size); - - // Make a global alignment call with traceback - let mut block = Block::::new(query.len(), reference.len(), block_size); - block.align( - &q_padded, - &r_padded, - &self.matrix, - self.gaps, - block_size..=block_size, - 0, // 0 x-drop threshold for global alignment - ); - let res = block.res(); - - // Retrieve CIGAR - let mut cigar = BlockCigar::new(res.query_idx, res.reference_idx); - block.trace().cigar_eq( - &q_padded, - &r_padded, - res.query_idx, - res.reference_idx, - &mut cigar, - ); - - AlignmentResult { - score: res.score, - query_start: 0, - query_end: res.query_idx, - ref_start: 0, - ref_end: res.reference_idx, - cigar: build_cigar(&cigar), - } - } - - /// Performs an alignment with one end fixed and the other allowed to terminate - /// early based on the x-drop mode. - /// - /// This method computes a semi-global alignment where one end of the alignment is - /// anchored while the other end can terminate early if the alignment score drops - /// too far below the maximum observed score (determined by the x-drop value). - /// - /// The alignment uses the blockaligner crate with block sizes chosen dynamically - /// as 20% of the maximum length of the 2 sequences and 50% of the maximum length - /// of the 2 sequences. - /// - /// # Parameters - /// - /// * `query` - The query sequence as a byte slice - /// * `reference` - The reference sequence as a byte slice - /// * `mode` - Controls which end is fixed: - /// - [`XdropMode::GlobalStartLocalEnd`]: Alignment must start at position 0 of - /// both sequences but may end before the last position - /// - [`XdropMode::LocalStartGlobalEnd`]: Alignment must end at the last position - /// of both sequences but may start after position 0. This is implemented by - /// reversing both sequences, aligning them, and transforming the result back - /// to forward coordinates. - /// - /// # Returns - /// - /// An [`AlignmentResult`] containing: - /// - The alignment score - /// - Start and end positions in both sequences - /// - A CIGAR string representing the alignment operations - fn xdrop_alignment(&self, query: &[u8], reference: &[u8], mode: XdropMode) -> AlignmentResult { - if query.is_empty() || reference.is_empty() { - return AlignmentResult::default(); - } - - // Blockaligner block ranges, set to 20% and 50% of longest sequence - // Must be powers of 2 - let max_len = query.len().max(reference.len()); - let min_size = (max_len / 5) - .next_power_of_two() - .clamp(32, MAXIMUM_BLOCK_SIZE); - let max_size = (max_len / 2) - .next_power_of_two() - .clamp(128, MAXIMUM_BLOCK_SIZE); - - // Create padded sequences - // blockaligner only has a AA profile, no nucleotide profile, so we have to use AA matrix - let mut q_padded = PaddedBytes::new::(query.len(), max_size); - let mut r_padded = PaddedBytes::new::(reference.len(), max_size); - - match mode { - XdropMode::GlobalStartLocalEnd => { - q_padded.set_bytes::(query, max_size); - r_padded.set_bytes::(reference, max_size); - } - XdropMode::LocalStartGlobalEnd => { - q_padded.set_bytes_rev::(query, max_size); - r_padded.set_bytes_rev::(reference, max_size); - } - } - - // Create profile to align the reference on the query with end bonus - // We need to build a position-specific scoring matrix - // blockaligner only has a AA profile, no nucleotide profile - let profile = make_aa_profile(query, &self.scores, max_size, mode); - - // Make an x-drop alignment call with traceback - let mut block = Block::::new(reference.len(), query.len(), max_size); - block.align_profile(&r_padded, &profile, min_size..=max_size, self.xdrop); - let res = block.res(); - - // Retrieve CIGAR - let mut cigar = BlockCigar::new(res.query_idx, res.reference_idx); - block.trace().cigar_eq( - &r_padded, - &q_padded, - res.query_idx, - res.reference_idx, - &mut cigar, - ); - match mode { - XdropMode::GlobalStartLocalEnd => AlignmentResult { - score: res.score, - query_start: 0, - query_end: res.reference_idx, - ref_start: 0, - ref_end: res.query_idx, - cigar: build_cigar_swap_indel(&cigar), - }, - XdropMode::LocalStartGlobalEnd => AlignmentResult { - score: res.score, - query_start: query.len() - res.reference_idx, - query_end: query.len(), - ref_start: reference.len() - res.query_idx, - ref_end: reference.len(), - cigar: build_cigar_reverse_swap_indel(&cigar), - }, + bandwidth, + simd_aligner: RefCell::new(SimdAligner::new(scores)), } } @@ -267,8 +68,11 @@ impl PiecewiseAligner { .saturating_sub(query_part.len() + padding); let ref_part = &refseq[ref_start..first_anchor.ref_start]; - let mut pre_align = - self.xdrop_alignment(query_part, ref_part, XdropMode::LocalStartGlobalEnd); + let mut pre_align = self.simd_aligner.borrow_mut().local_start_alignment( + query_part, + ref_part, + Some(self.bandwidth), + ); if pre_align.score == 0 { AlignmentResult { @@ -335,8 +139,11 @@ impl PiecewiseAligner { .min(last_anchor_ref_end + query_part.len() + padding); let ref_part = &refseq[last_anchor_ref_end..ref_end]; - let mut post_align = - self.xdrop_alignment(query_part, ref_part, XdropMode::GlobalStartLocalEnd); + let mut post_align = self.simd_aligner.borrow_mut().local_end_alignment( + query_part, + ref_part, + Some(self.bandwidth), + ); if post_align.score == 0 { AlignmentResult { @@ -446,7 +253,10 @@ impl PiecewiseAligner { } } - let aligned = self.global_alignment(query_part, ref_part); + let aligned = self + .simd_aligner + .borrow_mut() + .global_alignment(query_part, ref_part, None); score += aligned.score; cigar.extend(&aligned.cigar); @@ -510,119 +320,6 @@ impl PiecewiseAligner { } } -/// Converts a blockaligner CIGAR string to our internal CIGAR format. -fn build_cigar(block_cigar: &BlockCigar) -> Cigar { - let mut result = Cigar::new(); - - for i in 0..block_cigar.len() { - let oplen = block_cigar.get(i); - let op_code = match oplen.op { - Operation::M => CigarOperation::Match, - Operation::Eq => CigarOperation::Eq, - Operation::X => CigarOperation::X, - Operation::I => CigarOperation::Insertion, - Operation::D => CigarOperation::Deletion, - _ => continue, - }; - result.push(op_code, oplen.len); - } - - result -} - -/// Converts a blockaligner CIGAR string to our internal CIGAR format with -/// insertions and deletions swapped. -fn build_cigar_swap_indel(block_cigar: &BlockCigar) -> Cigar { - let mut result = Cigar::new(); - - for i in 0..block_cigar.len() { - let oplen = block_cigar.get(i); - let op_code = match oplen.op { - Operation::M => CigarOperation::Match, - Operation::Eq => CigarOperation::Eq, - Operation::X => CigarOperation::X, - Operation::I => CigarOperation::Deletion, // Swapped - Operation::D => CigarOperation::Insertion, // Swapped - _ => continue, - }; - result.push(op_code, oplen.len); - } - - result -} - -/// Converts a blockaligner CIGAR string to our internal CIGAR format in reverse -/// order with insertions and deletions swapped. -fn build_cigar_reverse_swap_indel(block_cigar: &BlockCigar) -> Cigar { - let mut result = Cigar::new(); - - for i in (0..block_cigar.len()).rev() { - let oplen = block_cigar.get(i); - let op_code = match oplen.op { - Operation::M => CigarOperation::Match, - Operation::Eq => CigarOperation::Eq, - Operation::X => CigarOperation::X, - Operation::I => CigarOperation::Deletion, // Swapped - Operation::D => CigarOperation::Insertion, // Swapped - _ => continue, - }; - result.push(op_code, oplen.len); - } - - result -} - -/// Creates a position-specific scoring matrix (profile) for aligning the reference -/// against the query with end bonuses. -/// -/// # Parameters -/// -/// * `query` - The query sequence to build the profile from -/// * `scores` - Scoring parameters -/// * `max_size` - Maximum block size for alignment -/// * `mode` - Alignment mode determining sequence traversal direction -/// -/// # Returns -/// -/// An [`AAProfile`] configured with position-specific scores for each nucleotide -/// pairing and appropriate end bonuses based on the alignment mode. -fn make_aa_profile(query: &[u8], scores: &Scores, max_size: usize, mode: XdropMode) -> AAProfile { - let mut profile = AAProfile::new(query.len(), max_size, -(scores.gap_extend as i8)); - - // Set scores for each nucleotide combination - for i in 1..=query.len() { - // Select the query character at this position depending on the alignment mode. - // `GlobalStartLocalEnd`: traverse the query from start to end (forward). - // `LocalStartGlobalEnd`: traverse the query from end to start (reversed). - let query_char = match mode { - XdropMode::GlobalStartLocalEnd => query[i - 1], - XdropMode::LocalStartGlobalEnd => query[query.len() - i], - }; - - for &c in b"ACGTN" { - if c == query_char { - profile.set(i, c, scores.match_ as i8); - } else { - profile.set(i, c, -(scores.mismatch as i8)); - }; - } - } - - // Set gap costs - profile.set_all_gap_open_C(-(scores.gap_open as i8) - -(scores.gap_extend as i8)); - profile.set_all_gap_close_C(0); - profile.set_all_gap_open_R(-(scores.gap_open as i8) - -(scores.gap_extend as i8)); - - // Give a bonus score to the end - let bonus_pos = query.len(); - for &c in b"ACGTN" { - let current_score = profile.get(bonus_pos, c); - profile.set(bonus_pos, c, current_score + scores.end_bonus as i8); - } - - profile -} - /// Removes spurious anchors from a chain of anchor points. /// /// This method applies two pruning strategies to filter out unreliable anchors that @@ -686,7 +383,7 @@ pub fn remove_spurious_anchors(anchors: &mut Vec) { // Second pruning: // We remove anchors of the ends of the chain if they create any indel. - // If they were part of the best scoring path, they should be retireved by the x-drop alignment. + // If they were part of the best scoring path, they should be retireved by the local end/start alignment. // the idea is that spurious anchors are much more difficult to detect on the ends of the chain, // so we look at a small ratio of anchors and remove any anchors creating deviations. @@ -722,532 +419,6 @@ pub fn remove_spurious_anchors(anchors: &mut Vec) { mod tests { use super::*; - #[test] - fn global_perfect_match() { - let aligner = PiecewiseAligner::new( - Scores { - match_: 2, - mismatch: 8, - gap_open: 12, - gap_extend: 1, - end_bonus: 10, - }, - 0, - 0, - ); - let result = aligner.global_alignment(b"AAATTT", b"AAATTT"); - assert_eq!(result.score, 6 * 2); - assert_eq!(result.query_start, 0); - assert_eq!(result.query_end, 6); - assert_eq!(result.ref_start, 0); - assert_eq!(result.ref_end, 6); - assert_eq!(result.cigar.to_string(), "6="); - } - - #[test] - fn global_complete_mismatch() { - let aligner = PiecewiseAligner::new( - Scores { - match_: 2, - mismatch: 8, - gap_open: 12, - gap_extend: 1, - end_bonus: 10, - }, - 0, - 0, - ); - let result = aligner.global_alignment(b"AAA", b"TTT"); - assert_eq!(result.score, 3 * -8); - assert_eq!(result.query_start, 0); - assert_eq!(result.query_end, 3); - assert_eq!(result.ref_start, 0); - assert_eq!(result.ref_end, 3); - assert_eq!(result.cigar.to_string(), "3X"); - } - - #[test] - fn global_single_mismatch() { - let aligner = PiecewiseAligner::new( - Scores { - match_: 2, - mismatch: 8, - gap_open: 12, - gap_extend: 1, - end_bonus: 10, - }, - 0, - 0, - ); - let result = aligner.global_alignment(b"AAATAA", b"AAAAAA"); - assert_eq!(result.score, 5 * 2 - 8); - assert_eq!(result.query_start, 0); - assert_eq!(result.query_end, 6); - assert_eq!(result.ref_start, 0); - assert_eq!(result.ref_end, 6); - assert_eq!(result.cigar.to_string(), "3=1X2="); - } - - #[test] - fn global_gap_in_query() { - let aligner = PiecewiseAligner::new( - Scores { - match_: 2, - mismatch: 8, - gap_open: 12, - gap_extend: 1, - end_bonus: 10, - }, - 0, - 0, - ); - let result = aligner.global_alignment(b"AAATTT", b"AAAATTTT"); - assert_eq!(result.score, 6 * 2 - 12 - 1); - assert_eq!(result.query_start, 0); - assert_eq!(result.query_end, 6); - assert_eq!(result.ref_start, 0); - assert_eq!(result.ref_end, 8); - assert_eq!(result.cigar.to_string(), "3=2D3="); - } - - #[test] - fn global_gap_in_reference() { - let aligner = PiecewiseAligner::new( - Scores { - match_: 2, - mismatch: 8, - gap_open: 12, - gap_extend: 1, - end_bonus: 10, - }, - 0, - 0, - ); - let result = aligner.global_alignment(b"AAAATTTT", b"AAATTT"); - assert_eq!(result.score, 6 * 2 - 12 - 1,); - assert_eq!(result.query_start, 0); - assert_eq!(result.query_end, 8); - assert_eq!(result.ref_start, 0); - assert_eq!(result.ref_end, 6); - assert_eq!(result.cigar.to_string(), "3=2I3="); - } - - #[test] - fn global_gap_at_query_start() { - let aligner = PiecewiseAligner::new( - Scores { - match_: 2, - mismatch: 8, - gap_open: 12, - gap_extend: 1, - end_bonus: 10, - }, - 0, - 0, - ); - let result = aligner.global_alignment(b"AAATTT", b"TTAAATTT"); - assert_eq!(result.score, 6 * 2 - 12 - 1); - assert_eq!(result.query_start, 0); - assert_eq!(result.query_end, 6); - assert_eq!(result.ref_start, 0); - assert_eq!(result.ref_end, 8); - assert_eq!(result.cigar.to_string(), "2D6="); - } - - #[test] - fn global_gap_at_query_end() { - let aligner = PiecewiseAligner::new( - Scores { - match_: 2, - mismatch: 8, - gap_open: 12, - gap_extend: 1, - end_bonus: 10, - }, - 0, - 0, - ); - let result = aligner.global_alignment(b"AAATTT", b"AAATTTAA"); - - assert_eq!(result.score, 6 * 2 - 12 - 1); - assert_eq!(result.query_start, 0); - assert_eq!(result.query_end, 6); - assert_eq!(result.ref_start, 0); - assert_eq!(result.ref_end, 8); - assert_eq!(result.cigar.to_string(), "6=2D"); - } - - #[test] - fn global_gap_at_reference_start() { - let aligner = PiecewiseAligner::new( - Scores { - match_: 2, - mismatch: 8, - gap_open: 12, - gap_extend: 1, - end_bonus: 10, - }, - 0, - 0, - ); - let result = aligner.global_alignment(b"TTAAATTT", b"AAATTT"); - - assert_eq!(result.score, 6 * 2 - 12 - 1); - assert_eq!(result.query_start, 0); - assert_eq!(result.query_end, 8); - assert_eq!(result.ref_start, 0); - assert_eq!(result.ref_end, 6); - assert_eq!(result.cigar.to_string(), "2I6="); - } - - #[test] - fn global_gap_at_reference_end() { - let aligner = PiecewiseAligner::new( - Scores { - match_: 2, - mismatch: 8, - gap_open: 12, - gap_extend: 1, - end_bonus: 10, - }, - 0, - 0, - ); - let result = aligner.global_alignment(b"AAATTTAA", b"AAATTT"); - - assert_eq!(result.score, 6 * 2 - 12 - 1); - assert_eq!(result.query_start, 0); - assert_eq!(result.query_end, 8); - assert_eq!(result.ref_start, 0); - assert_eq!(result.ref_end, 6); - assert_eq!(result.cigar.to_string(), "6=2I"); - } - - #[test] - fn global_multiple_mismatches() { - let aligner = PiecewiseAligner::new( - Scores { - match_: 2, - mismatch: 8, - gap_open: 12, - gap_extend: 1, - end_bonus: 10, - }, - 0, - 0, - ); - let result = aligner.global_alignment(b"ATATATATAT", b"AAAAAAAAAA"); - - assert_eq!(result.score, 5 * 2 - 5 * 8); - assert_eq!(result.query_start, 0); - assert_eq!(result.query_end, 10); - assert_eq!(result.ref_start, 0); - assert_eq!(result.ref_end, 10); - assert_eq!(result.cigar.to_string(), "1=1X1=1X1=1X1=1X1=1X"); - } - - #[test] - fn global_complex_alignment() { - let aligner = PiecewiseAligner::new( - Scores { - match_: 2, - mismatch: 8, - gap_open: 12, - gap_extend: 1, - end_bonus: 10, - }, - 0, - 0, - ); - let result = aligner.global_alignment(b"AAACTTAAACCTT", b"AAAATTGAAATT"); - assert_eq!(result.score, 10 * 2 - 8 - 2 * 12 - 1); - assert_eq!(result.query_start, 0); - assert_eq!(result.query_end, 13); - assert_eq!(result.ref_start, 0); - assert_eq!(result.ref_end, 12); - assert_eq!(result.cigar.to_string(), "3=1X2=1D3=2I2="); - } - - #[test] - fn xdrop_forward_perfect_match() { - let aligner = PiecewiseAligner::new( - Scores { - match_: 2, - mismatch: 8, - gap_open: 12, - gap_extend: 1, - end_bonus: 10, - }, - 0, - 100, - ); - let result = aligner.xdrop_alignment(b"AAATTT", b"AAATTT", XdropMode::GlobalStartLocalEnd); - assert_eq!(result.score, 6 * 2 + 10); - assert_eq!(result.query_start, 0); - assert_eq!(result.query_end, 6); - assert_eq!(result.ref_start, 0); - assert_eq!(result.ref_end, 6); - assert_eq!(result.cigar.to_string(), "6="); - } - - #[test] - fn xdrop_forward_with_mismatch() { - let aligner = PiecewiseAligner::new( - Scores { - match_: 2, - mismatch: 8, - gap_open: 12, - gap_extend: 1, - end_bonus: 10, - }, - 0, - 100, - ); - let result = aligner.xdrop_alignment(b"AAATAA", b"AAAAAA", XdropMode::GlobalStartLocalEnd); - assert_eq!(result.score, 5 * 2 - 8 + 10); - assert_eq!(result.query_start, 0); - assert_eq!(result.query_end, 6); - assert_eq!(result.ref_start, 0); - assert_eq!(result.ref_end, 6); - assert_eq!(result.cigar.to_string(), "3=1X2="); - } - - #[test] - fn xdrop_forward_early_termination() { - let aligner = PiecewiseAligner::new( - Scores { - match_: 2, - mismatch: 8, - gap_open: 12, - gap_extend: 1, - end_bonus: 10, - }, - 0, - 100, - ); - let result = aligner.xdrop_alignment( - b"AAAAAAATTTTTTT", - b"AAAAAAACCCCCCC", - XdropMode::GlobalStartLocalEnd, - ); - assert_eq!(result.score, 7 * 2); - assert_eq!(result.query_start, 0); - assert_eq!(result.ref_start, 0); - assert_eq!(result.query_end, 7); - assert_eq!(result.ref_end, 7); - assert_eq!(result.cigar.to_string(), "7="); - } - - #[test] - #[allow(clippy::identity_op)] - fn xdrop_forward_end_bonus_extends_alignment() { - let aligner = PiecewiseAligner::new( - Scores { - match_: 2, - mismatch: 8, - gap_open: 12, - gap_extend: 1, - end_bonus: 50, - }, - 0, - 100, - ); - let result = - aligner.xdrop_alignment(b"AAAAAAAATT", b"AAAAAAAAAA", XdropMode::GlobalStartLocalEnd); - assert_eq!(result.score, 8 * 2 - 8 * 2 + 50); - assert_eq!(result.query_start, 0); - assert_eq!(result.query_end, 10); - assert_eq!(result.ref_start, 0); - assert_eq!(result.ref_end, 10); - assert_eq!(result.cigar.to_string(), "8=2X"); - } - - #[test] - fn xdrop_forward_with_gap() { - let aligner = PiecewiseAligner::new( - Scores { - match_: 2, - mismatch: 8, - gap_open: 12, - gap_extend: 1, - end_bonus: 10, - }, - 0, - 100, - ); - let result = - aligner.xdrop_alignment(b"AAAAACCTTT", b"AAAAATTT", XdropMode::GlobalStartLocalEnd); - assert_eq!(result.score, 8 * 2 - 12 - 1 + 10); - assert_eq!(result.query_start, 0); - assert_eq!(result.query_end, 10); - assert_eq!(result.ref_start, 0); - assert_eq!(result.ref_end, 8); - assert_eq!(result.cigar.to_string(), "5=2I3="); - } - - #[test] - fn xdrop_forward_gap_in_reference() { - let aligner = PiecewiseAligner::new( - Scores { - match_: 2, - mismatch: 8, - gap_open: 12, - gap_extend: 1, - end_bonus: 10, - }, - 0, - 100, - ); - let result = - aligner.xdrop_alignment(b"AAAAATTT", b"AAAAACCTTT", XdropMode::GlobalStartLocalEnd); - assert_eq!(result.score, 8 * 2 - 12 - 1 + 10); - assert_eq!(result.query_start, 0); - assert_eq!(result.query_end, 8); - assert_eq!(result.ref_start, 0); - assert_eq!(result.ref_end, 10); - assert_eq!(result.cigar.to_string(), "5=2D3="); - } - - #[test] - fn xdrop_reverse_perfect_match() { - let aligner = PiecewiseAligner::new( - Scores { - match_: 2, - mismatch: 8, - gap_open: 12, - gap_extend: 1, - end_bonus: 10, - }, - 0, - 100, - ); - let result = aligner.xdrop_alignment(b"AAATTT", b"AAATTT", XdropMode::LocalStartGlobalEnd); - assert_eq!(result.score, 6 * 2 + 10); - assert_eq!(result.query_start, 0); - assert_eq!(result.query_end, 6); - assert_eq!(result.ref_start, 0); - assert_eq!(result.ref_end, 6); - assert_eq!(result.cigar.to_string(), "6="); - } - - #[test] - fn xdrop_reverse_with_mismatch() { - let aligner = PiecewiseAligner::new( - Scores { - match_: 2, - mismatch: 8, - gap_open: 12, - gap_extend: 1, - end_bonus: 10, - }, - 0, - 100, - ); - let result = aligner.xdrop_alignment(b"AATAAA", b"AAAAAA", XdropMode::LocalStartGlobalEnd); - assert_eq!(result.score, 5 * 2 - 8 + 10); - assert_eq!(result.query_start, 0); - assert_eq!(result.query_end, 6); - assert_eq!(result.ref_start, 0); - assert_eq!(result.ref_end, 6); - assert_eq!(result.cigar.to_string(), "2=1X3="); - } - - #[test] - fn xdrop_reverse_early_termination() { - let aligner = PiecewiseAligner::new( - Scores { - match_: 2, - mismatch: 8, - gap_open: 12, - gap_extend: 1, - end_bonus: 10, - }, - 0, - 100, - ); - let result = aligner.xdrop_alignment( - b"TTTTTTTAAAAAAA", - b"CCCCCCCAAAAAAA", - XdropMode::LocalStartGlobalEnd, - ); - assert_eq!(result.score, 7 * 2); - assert_eq!(result.query_start, 7); - assert_eq!(result.query_end, 14); - assert_eq!(result.ref_start, 7); - assert_eq!(result.ref_end, 14); - assert_eq!(result.cigar.to_string(), "7="); - } - - #[test] - #[allow(clippy::identity_op)] - fn xdrop_reverse_end_bonus_extends_alignment() { - let aligner = PiecewiseAligner::new( - Scores { - match_: 2, - mismatch: 8, - gap_open: 12, - gap_extend: 1, - end_bonus: 50, - }, - 0, - 100, - ); - let result = - aligner.xdrop_alignment(b"TTAAAAAAAA", b"AAAAAAAAAA", XdropMode::LocalStartGlobalEnd); - assert_eq!(result.score, 8 * 2 - 8 * 2 + 50); - assert_eq!(result.query_start, 0); - assert_eq!(result.query_end, 10); - assert_eq!(result.ref_start, 0); - assert_eq!(result.ref_end, 10); - assert_eq!(result.cigar.to_string(), "2X8="); - } - - #[test] - fn xdrop_reverse_with_gap() { - let aligner = PiecewiseAligner::new( - Scores { - match_: 2, - mismatch: 8, - gap_open: 12, - gap_extend: 1, - end_bonus: 10, - }, - 0, - 100, - ); - let result = - aligner.xdrop_alignment(b"TTTCCAAAAA", b"TTTAAAAA", XdropMode::LocalStartGlobalEnd); - assert_eq!(result.score, 8 * 2 - 12 - 1 + 10); - assert_eq!(result.query_start, 0); - assert_eq!(result.query_end, 10); - assert_eq!(result.ref_start, 0); - assert_eq!(result.ref_end, 8); - assert_eq!(result.cigar.to_string(), "3=2I5="); - } - - #[test] - fn xdrop_reverse_gap_in_reference() { - let aligner = PiecewiseAligner::new( - Scores { - match_: 2, - mismatch: 8, - gap_open: 12, - gap_extend: 1, - end_bonus: 10, - }, - 0, - 100, - ); - let result = - aligner.xdrop_alignment(b"TTTAAAAA", b"TTTCCAAAAA", XdropMode::LocalStartGlobalEnd); - assert_eq!(result.score, 8 * 2 - 12 - 1 + 10); - assert_eq!(result.query_start, 0); - assert_eq!(result.query_end, 8); - assert_eq!(result.ref_start, 0); - assert_eq!(result.ref_end, 10); - assert_eq!(result.cigar.to_string(), "3=2D5="); - } - #[test] fn remove_spurious_anchors_control() { let expected = vec![ @@ -1523,7 +694,7 @@ mod tests { end_bonus: 10, }, 5, - 100, + 1024, ); let query = b"AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA"; let refseq = b"AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA"; @@ -1572,7 +743,7 @@ mod tests { end_bonus: 10, }, 5, - 100, + 1024, ); let query = b"AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA"; let refseq = b"TTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTT"; @@ -1609,7 +780,7 @@ mod tests { end_bonus: 10, }, 5, - 100, + 1024, ); let query = b"AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA"; let refseq = @@ -1674,7 +845,7 @@ mod tests { end_bonus: 10, }, 5, - 100, + 1024, ); let query = b"CTTTTAAAAATTTTAAAAATGGTTTCAAAAATTCCTAAAAATTTTTCCCCC"; let refseq = b"TTTTTAAAAATTTTTAAAAATTTTTAAAAATTTTTAAAAATTTTTAAAAA"; @@ -1707,6 +878,6 @@ mod tests { assert_eq!(result.query_end, 46); assert_eq!(result.ref_start, 0); assert_eq!(result.ref_end, 45); - assert_eq!(result.cigar.to_string(), "1X13=1D6=2I3=1X7=2X11="); + assert_eq!(result.cigar.to_string(), "1X9=1D10=2I3=1X7=2X11="); } } diff --git a/src/simdaligner/aligner.rs b/src/simdaligner/aligner.rs new file mode 100644 index 00000000..ba5f112b --- /dev/null +++ b/src/simdaligner/aligner.rs @@ -0,0 +1,382 @@ +use super::Scores; + +use super::avx2_u8::{U8Probe, fits_u8_lanes}; +use super::{AlignmentResult, SplitReferenceAlignment}; + +/// A reusable aligner configured with a fixed [`Scores`] scheme. +/// +/// Every alignment method takes `&mut self` because the scratch buffers are held internally and +/// reused; only the query and reference slices change between calls. Construct one per thread +/// and keep it - constructing one per alignment reallocates every buffer. +/// +/// See the [module docs](super) for the guarantees, the alphabet, and what a band means. +pub struct SimdAligner { + workspace: U8Probe, + scores: Scores, +} + +impl Default for SimdAligner { + fn default() -> Self { + SimdAligner::new(Scores::default()) + } +} + +impl SimdAligner { + /// Builds an aligner for the given scoring scheme. + /// + /// Every precondition is checked here, once per scheme, so that the alignment path carries + /// no feature detection and no validation at all. + /// + /// # Panics + /// + /// - If the CPU does not support AVX2. See [`SimdAligner::is_supported`] to check first. + /// - If `scores.gap_open < scores.gap_extend`. The kernel collapses the gap-open and + /// gap-extend cases into a single recurrence on that assumption. + /// - If `match_ + 3 * gap_open - gap_extend > 255`, i.e. roughly `gap_open > 84` at + /// `match_ = 2`. Beyond that bound an intermediate value exceeds the kernel's 8-bit + /// lanes. + pub fn new(scores: Scores) -> Self { + assert!( + Self::is_supported(), + "this CPU doesn't support AVX2, which this aligner requires - check SimdAligner::is_supported() first" + ); + // Hard asserts: once per scheme, never in the inner loop. The messages say why. + assert!( + scores.gap_open >= scores.gap_extend, + "gap_open ({}) must be >= gap_extend ({}): the kernel collapses the open-vs-extend \ + recurrence on that assumption", + scores.gap_open, + scores.gap_extend + ); + assert!( + fits_u8_lanes(scores.match_, scores.gap_open, scores.gap_extend), + "scores do not fit the u8 kernel: match ({}) + 3*gap_open ({}) - gap_extend ({}) \ + must be <= 255, i.e. gap_open <= ~84 at match=2. Past this an intermediate wraps \ + and the score comes back better than optimal, silently.", + scores.match_, + scores.gap_open, + scores.gap_extend + ); + SimdAligner { + workspace: U8Probe::new(), + scores, + } + } + + /// Whether this CPU supports the instruction set the aligner requires (AVX2). + /// + /// [`SimdAligner::new`] panics if this is false. + pub fn is_supported() -> bool { + is_x86_feature_detected!("avx2") + } + + /// The scoring scheme this aligner was built with. + pub fn scores(&self) -> Scores { + self.scores + } + + /// Global (end-to-end) alignment: both sequences are aligned in full, and terminal gaps are + /// penalised. + /// + /// `query_start`/`query_end` and `ref_start`/`ref_end` always span the entire inputs. + /// + /// Empty inputs follow affine-gap semantics: aligning a length-`n` sequence against an empty + /// one yields a single length-`n` gap scoring `-(gap_open + (n - 1) * gap_extend)`, and two + /// empty inputs yield an empty alignment scoring 0. + /// + /// # `bandwidth` + /// + /// Global is the one mode pinned at **both** corners, and `(qlen, rlen)` sits `|qlen - rlen|` + /// off the diagonal a band is centred on, so a narrower band contains no path at all and `w` + /// is widened to reach it. Free when the two lengths are close - the shape global is normally + /// handed - and on a lopsided pair it degrades to correct-but-slow rather than to wrong. + pub fn global_alignment( + &mut self, + query: &[u8], + reference: &[u8], + bandwidth: Option, + ) -> AlignmentResult { + self.workspace.global_alignment( + query, + reference, + self.scores.match_, + self.scores.mismatch, + self.scores.gap_open, + self.scores.gap_extend, + bandwidth, + ) + } + + /// Alignment with a **local reference end**: the query is spanned in full, but the reference + /// may end anywhere, and the unaligned reference tail costs nothing. + /// + /// Formally, the optimum over every end cell `(query.len(), j)` for `j` in + /// `0..=reference.len()`. `query_start`/`query_end` therefore span the whole query, + /// `ref_start` is 0, and `ref_end` is where the alignment chose to stop: the CIGAR covers + /// exactly `reference[..ref_end]` and says nothing about the rest. + /// + /// `j = 0` is a legal end cell, and sometimes the winning one - spanning the query as one + /// gap and aligning no reference at all can beat a run of mismatches, in which case the + /// whole reference is the free end gap. + /// + /// # Tie-break + /// + /// Among equally-scoring end cells the *smallest* `ref_end` wins: don't consume reference + /// you don't have to. + /// + /// # `bandwidth` + /// + /// Anchored at `(0, 0)`, like [`local_end_alignment`](SimdAligner::local_end_alignment), so + /// `Some(w)` bounds `ref_end` to `[qlen - w, qlen + w]`. Size `w` from how much indel + /// imbalance you expect, not from the reference's length. + /// + /// The query must still be spanned, which needs `w >= qlen - rlen`; `w` is widened to that + /// where necessary, which only bites when the query is the longer side. + pub fn local_reference_end_alignment( + &mut self, + query: &[u8], + reference: &[u8], + bandwidth: Option, + ) -> AlignmentResult { + self.workspace.local_reference_end_alignment( + query, + reference, + self.scores.match_, + self.scores.mismatch, + self.scores.gap_open, + self.scores.gap_extend, + bandwidth, + ) + } + + /// Alignment with a **local reference start**: the mirror of + /// [`local_reference_end_alignment`](SimdAligner::local_reference_end_alignment). The query + /// is spanned in full, but the reference may *begin* anywhere, and the unaligned reference + /// prefix costs nothing. + /// + /// Formally, the optimum over every start cell `(0, s)` for `s` in `0..=reference.len()`, + /// ending at `(query.len(), reference.len())`. `query_start`/`query_end` span the whole + /// query, `ref_end` is always `reference.len()`, and `ref_start` is where the alignment + /// chose to begin: the CIGAR covers exactly `reference[ref_start..]`. + /// + /// This is exact rather than an approximation: affine gap costs are symmetric under + /// reversal, and so is literal byte identity, so the reversed problem has the same optimum. + /// + /// # Tie-break + /// + /// Among equally-scoring start cells the *largest* `ref_start` wins - the shortest reference + /// span, mirroring the `_end` variant's rule. + /// + /// # `bandwidth` + /// + /// As [`local_reference_end_alignment`](SimdAligner::local_reference_end_alignment), + /// **mirrored with the mode**: the band is anchored at the *ends* of the two inputs, where + /// this alignment is pinned, and bounds `ref_start` to within `w` of `rlen - qlen`. `None` + /// is exact. + pub fn local_reference_start_alignment( + &mut self, + query: &[u8], + reference: &[u8], + bandwidth: Option, + ) -> AlignmentResult { + self.workspace.local_reference_start_alignment( + query, + reference, + self.scores.match_, + self.scores.mismatch, + self.scores.gap_open, + self.scores.gap_extend, + bandwidth, + ) + } + + /// **Local-end alignment**: both sequences start at 0, and the alignment ends wherever it + /// scores best. Both trailing tails - query and reference - are free. + /// + /// Formally, the optimum over every end cell `(i, j)` of + /// `H[i][j] + (if i == query.len() { scores.end_bonus } else { 0 })`. `query_end` and + /// `ref_end` are wherever it stopped, and the CIGAR covers exactly `query[..query_end]` + /// against `reference[..ref_end]`. + /// + /// # The score is never negative + /// + /// Cell `(0, 0)` is a legal end cell scoring 0, so a hopeless extension returns the empty + /// alignment: spans `0..0` on both sequences, score 0, empty CIGAR. That is a normal result. + /// + /// # `end_bonus` + /// + /// Added to the score only if the alignment covers the whole query. It nudges the alignment + /// towards spanning the query without forcing it, which is the difference between this and + /// [`local_reference_end_alignment`](SimdAligner::local_reference_end_alignment). It never + /// enters the recurrence: it is applied once, at the end, to choose between the best cell + /// anywhere and the best cell that finishes the query. + /// + /// # Tie-break + /// + /// Longest extension: among equally-scoring end cells the largest `query_end` wins, then the + /// largest `ref_end`. + /// + /// # How the CIGAR can end + /// + /// Both of these are properties of the objective, and worth knowing if you consume the + /// coordinates: + /// + /// - A trailing insertion happens only if `end_bonus >= gap_open`. Ending with a `k`-base + /// insertion costs `gap_open + (k - 1) * gap_extend` to reach `(qlen, j)` from + /// `(qlen - k, j)`, and `(qlen - k, j)` is itself a legal end cell, so the bonus must + /// cover that or stopping short would score higher. + /// - With `gap_open > 0` the CIGAR never ends with a deletion, by the same argument: a local + /// maximum cannot end in a gap. + /// + /// When a trailing insertion does occur, `ref_end` does not advance with it. + /// + /// # `bandwidth` + /// + /// This mode and [`local_start_alignment`](SimdAligner::local_start_alignment) are what the + /// band was built for: they *finish* an alignment something else already anchored, so the + /// start is known and the drift is bounded by biology rather than by the matrix. They are + /// also the only two **never widened** - `(0, 0)` is in every band, so `Some(w)` is exactly + /// `w`. Under a band the tie-break is still longest-extension and the score is still never + /// negative. + /// + /// Below roughly `qlen ~ 1.2 * w` the band's bookkeeping costs more than it saves, so pass + /// `None` for short extensions. + pub fn local_end_alignment( + &mut self, + query: &[u8], + reference: &[u8], + bandwidth: Option, + ) -> AlignmentResult { + self.workspace.local_end_alignment( + query, + reference, + self.scores.match_, + self.scores.mismatch, + self.scores.gap_open, + self.scores.gap_extend, + self.scores.end_bonus, + bandwidth, + ) + } + + /// **Local-start alignment**: the mirror of + /// [`local_end_alignment`](SimdAligner::local_end_alignment). Both sequences end at their + /// ends, and the alignment *begins* wherever it scores best; both leading prefixes are free. + /// + /// Formally, the optimum over every start cell `(a, b)` of the best alignment of + /// `query[a..]` against `reference[b..]`, plus `end_bonus` if `a == 0`. `query_start` and + /// `ref_start` are wherever it began; `query_end` and `ref_end` are always the ends of the + /// inputs. + /// + /// # The score is never negative + /// + /// The start cell `(query.len(), reference.len())` is legal and scores 0, so a hopeless + /// extension returns the empty alignment. Note it sits at the *end* of both sequences, the + /// mirror of local-end's, which sits at the start. + /// + /// # Tie-break + /// + /// Longest extension, mirrored: among equally-scoring start cells the *smallest* + /// `query_start` wins, then the smallest `ref_start`. Likewise a *leading* insertion + /// requires `end_bonus >= gap_open`, and with `gap_open > 0` the CIGAR never begins with a + /// deletion. + /// + /// # `bandwidth` + /// + /// As [`local_end_alignment`](SimdAligner::local_end_alignment), **mirrored with the mode**: + /// the band is anchored at the *ends* of the two inputs, where this alignment is pinned, so + /// it constrains start cells to `|(qlen - a) - (rlen - b)| <= w`. In both modes `w` bounds + /// the drift from the end that something else already established. + pub fn local_start_alignment( + &mut self, + query: &[u8], + reference: &[u8], + bandwidth: Option, + ) -> AlignmentResult { + self.workspace.local_start_alignment( + query, + reference, + self.scores.match_, + self.scores.mismatch, + self.scores.gap_open, + self.scores.gap_extend, + self.scores.end_bonus, + bandwidth, + ) + } + + /// **Split-reference alignment**: one query aligned across two references, with a single + /// jump between them. Built for detecting a tandem duplication that a linear aligner reports + /// as an insertion. + /// + /// The alignment starts at `query[0]` against `right_reference[0]` and ends at the end of + /// `left_reference`, jumping once from the right reference to the left. It returns two + /// [`AlignmentResult`]s that **partition the query exactly**: + /// + /// - `right` covers `query[..k]` against a prefix of `right_reference`; + /// - `left` covers `query[k..]` against a suffix of `left_reference`; + /// + /// where `k`, the jump point, is chosen to maximise the total score: + /// + /// ```text + /// score = max over k in 0..=query.len() of F[k] + G[k] + /// F[k] = the best alignment of query[..k] against any prefix of right_reference + /// G[k] = the best alignment of query[k..] against any suffix of left_reference + /// ``` + /// + /// Exact, like everything else here: both full matrices are computed and every jump point is + /// considered. `end_bonus` is ignored - the query is spanned in full by construction, so a + /// bonus for spanning it would be a constant added to every candidate. + /// + /// # The score can be negative, and an arm can be empty + /// + /// Unlike the local modes there is no "give up" option: the two arms must cover the query, + /// so a query that neither reference explains still gets covered, at a cost. Two + /// consequences worth knowing before consuming the coordinates: + /// + /// - `k = 0` and `k = query.len()` are legal. One arm then spans the whole query and the + /// other is the empty alignment. + /// - An arm may consume no reference at all, reporting its query segment as one pure + /// insertion, with `ref_start == ref_end`. That is the right answer, not a failure: it is + /// how an unalignable segment gets covered. + /// + /// # Tie-break + /// + /// Total score decides. Among equal totals the **most balanced** split wins - the smallest + /// `|left.score - right.score|` - because `5 + 5` describes a real duplication and + /// `-20 + 30` does not, though both total 10. Among equally balanced splits, the smallest + /// `k`. Each arm's own reference endpoint breaks ties the way its one-reference counterpart + /// does. + /// + /// # `bandwidth` + /// + /// `Some(w)` is **one band, applied to both arms**, each measured from that arm's own anchor: + /// the right arm from `(query[0], right_reference[0])`, the left arm from the ends of `query` + /// and `left_reference`. + /// + /// **It does not constrain the jump point.** `k` is still chosen over every value the two + /// bands admit, which is what this mode exists to find; the band says only that neither arm + /// wanders far from its own diagonal on the way there. + /// + /// Both arms must together cover the query, and an arm reaches at most `reflen + w` query + /// bases in band, so `w` is widened where necessary to + /// `(qlen - left_reference.len() - right_reference.len()) / 2`. + pub fn split_reference_alignment( + &mut self, + query: &[u8], + left_reference: &[u8], + right_reference: &[u8], + bandwidth: Option, + ) -> SplitReferenceAlignment { + self.workspace.split_reference_alignment( + query, + left_reference, + right_reference, + self.scores.match_, + self.scores.mismatch, + self.scores.gap_open, + self.scores.gap_extend, + bandwidth, + ) + } +} diff --git a/src/simdaligner/avx2_u8.rs b/src/simdaligner/avx2_u8.rs new file mode 100644 index 00000000..aa0a1608 --- /dev/null +++ b/src/simdaligner/avx2_u8.rs @@ -0,0 +1,2441 @@ +//! AVX2 `u8` kernel: the difference recurrence on anti-diagonals, 32 lanes. +//! +//! This is the only kernel; every alignment type routes at it, differing only in which cell the +//! traceback starts from and which maxima the fill tracks. +//! +//! # Where things are +//! +//! A public mode on `U8Probe` picks a band strategy, then goes through one of three wrappers to +//! the fill and the traceback: +//! +//! ```text +//! pub fn -> align | local | row_pass | split_reference +//! -> prepare encode both sequences, size the delta buffers +//! -> size_flags size the tb/nm buffers +//! -> run the fill: deltas forward, flags out +//! -> arg_for_row[_banded] which column won a row (from the nm bits) +//! -> replay walk the flags back, emit the CIGAR +//! ``` +//! +//! See `run`'s "Which mode compiles which flags" table for what each mode switches on. +//! +//! # The difference recurrence +//! +//! Rather than storing the DP cells, this stores the **differences between adjacent cells** +//! (Suzuki & Kasahara 2018, BMC Bioinformatics 19(Suppl 1):45). The differences are bounded by +//! the **scoring scheme alone**, with no length term, so they fit `u8` lanes at any sequence +//! length - which is what makes an exact, unbanded aligner affordable, and what removes any need +//! for an overflow envelope or a wider fallback kernel. +//! +//! The recurrence is the paper's Eq 3, the *offset* form, in this crate's gap convention +//! (`go = gap_open - gap_extend`, `ge = gap_extend`, `C = 2*gap_open`, `s_G(x,y) = s(x,y) + C`). +//! The offset variables fold the constants in so that `ΔE'_G` absorbs `ΔV` and `ΔF'_G` absorbs +//! `ΔH`, which leaves nothing to add before the maxes and gives a critical path of 4: +//! +//! ```text +//! A_G [i,j] = max( s_G(a,b), ΔE'_G[i-1,j], ΔF'_G[i,j-1] ) +//! ΔH_G [i,j] = A_G[i,j] - ΔV_G[i-1,j] +//! ΔV_G [i,j] = A_G[i,j] - ΔH_G[i,j-1] +//! ΔE'_G[i,j] = max( A_G[i,j], ΔE'_G[i-1,j] + go ) - ΔH_G[i,j-1] +//! ΔF'_G[i,j] = max( A_G[i,j], ΔF'_G[i,j-1] + go ) - ΔV_G[i-1,j] +//! ``` +//! +//! with the boundary `ΔV_G[0,j] = ΔE'_G[0,j] = ΔH_G[i,0] = ΔF'_G[i,0] = 0` at index 1 and `go` +//! after. Note the paper's Eq 2 is the same idea in a *naive* difference form; its critical path +//! is 8, so Eq 3 is the one to build on. (Two transcription traps, if comparing against the +//! paper: its gap convention is `G(k) = G_o + k*G_e` against this crate's +//! `gap_open + (k-1)*gap_extend`, so `G_o = gap_open - gap_extend`; and its supplement's Eq 19 +//! is sign-flipped against its own Eq 7.) +//! +//! Everything is unsigned: `0 <= ΔH_G, ΔV_G, ΔE'_G, ΔF'_G <= match + 2*gap_open`. The hard +//! **zero** lower bound is the load-bearing half - it is what permits `u8` lanes with +//! `max_epu8`, and it means the kernel needs **no sentinel at all**. Every value is provably in +//! `[0, upper]`, so the subtractions cannot wrap and plain `_mm256_sub_epi8` suffices. +//! +//! # Why anti-diagonals, and why there are no shifts +//! +//! `ΔE'_G[i,j]` depends on `ΔE'_G[i-1,j]` - same column, previous row - so a column-vectorised +//! layout has an intra-vector serial dependency. On an anti-diagonal `p = i + j`, `A_G`'s inputs +//! are: +//! +//! ```text +//! (i,j) needs ΔV_G[i-1, j] and ΔE'_G[i-1, j] <- anti-diagonal p-1, row i-1 +//! ΔH_G[i, j-1] and ΔF'_G[i, j-1] <- anti-diagonal p-1, row i +//! ``` +//! +//! **Everything comes from `p-1`, and there is no `p-2` term**: `A_G` *is* the diagonal step, so +//! the diagonal predecessor is already folded in. Holding each anti-diagonal in a buffer indexed +//! by **row** makes "row `i-1`" and "row `i`" two `loadu`s one byte apart, so the shifts the +//! paper needs for its fixed-width sliding vector are not needed here - the addressing does it. +//! +//! The sequences fall out the same way. Cell `(i, p-i)` reads `query[i-1]`, contiguous in `i`, +//! and `refseq[p-i-1]`, which runs *backwards* in `i` - so the reference is stored **reversed** +//! and becomes `ref_rev[rlen + i - p]`, contiguous too. (Group it `rlen + i - p`, never +//! `rlen - p + i`: the latter underflows a `usize` once the anti-diagonal starts sliding right.) +//! +//! # Two hazards +//! +//! - **Boundary cells are written after the chunk loop, never before.** A 32-lane store is +//! unconditional, so the partial last chunk scribbles over the row just past the interior - +//! which for `p <= qlen` is the `(p, 0)` boundary cell. +//! - **Garbage lanes need no mask.** Anti-diagonal cells have no cross-lane dependency, so +//! garbage stays in its own lane; and anti-diagonal `p` provably reads only rows +//! `[lo(p-1), hi(p-1)]`, which `p-1`'s interior and boundary writes cover exactly. + +use std::arch::x86_64::*; + +use super::{AlignmentResult, Cigar, CigarOperation, SplitReferenceAlignment}; + +/// Lanes per vector: 32 `u8`s in a `__m256i`. +pub const LANES: usize = 32; + +/// Traceback flag words per 32-lane chunk: `from_diag`, `from_e`, `e_open`, `f_open`. +/// +/// **The alignment state is already in the deltas** - this is the observation the whole +/// traceback rests on, and it is why the difference form is *easier* to trace back than the +/// absolute one, not harder. +/// +/// `A_G = max(s_G, ΔE'_G[i-1,j], ΔF'_G[i,j-1])`, so *which operand achieved the max* is the +/// traceback direction, and the fill already computed it. Likewise +/// `ΔE'_G[i,j] = max(A_G, ΔE'_G[i-1,j] + go) - ΔH_G[i,j-1]`: which operand won says whether the +/// gap opened or extended - our open/extend bits, for free. +/// +/// So every flag is a `_mm256_cmpeq_epi8` against values already in registers, extracted with +/// `_mm256_movemask_epi8`: **4 bits per cell**, 8 ops per 32 cells. +/// +/// Every flag is an *equality*, deliberately. AVX2 has no unsigned byte compare +/// (`_mm256_cmpgt_epi8` is signed, and these values can exceed 127 with a large `gap_open`), so +/// "did `ΔE'_G[i-1,j] + go` win the max" is phrased as `max(A_G, ΔE'_G+go) == A_G` - reusing the +/// max the fill already needs - rather than as a `>`. Equality does not care about signedness. +/// +/// # Only three of the four directions are stored +/// +/// `A_G`'s max has three operands and at least one of them achieves it, so one direction can be +/// left unstored and recovered as the fallthrough. **Which one is not free to choose: the +/// inferred direction is the one the replay checks last, and therefore the one it can never +/// prefer on a tie.** Since the replay prefers the diagonal (see `replay`), `from_diag` is +/// stored and `from_f` is inferred. (`e_open`/`f_open` are a two-way max, so neither is +/// inferable from the other; both stay.) +/// +/// `from_diag` cannot be `cmpeq(A_G, s_G)` on its own - see [`TB_FROM_DIAG`]. +const TB_WORDS: usize = 4; + +/// One `u32` of "this cell took its row's max" per 32-lane chunk - a bit per lane, in the same +/// `(p, chunk)` coordinates as the traceback flags, so `tb_base[p]` indexes both. See +/// `U8Probe::nm`. +const NM_WORDS: usize = 1; + +const TB_FROM_E: usize = 0; +const TB_E_OPEN: usize = 1; +const TB_F_OPEN: usize = 2; +/// Did the **diagonal** achieve `A_G`'s max? +/// +/// `cmpeq(A_G, s_G)` is the obvious answer and it is **wrong for a clamped `s_G`**. Any scheme +/// with `mismatch > 2*gap_open` makes `s_G(mismatch) = C - mismatch` negative, and a u8 lane +/// cannot hold that, so the fill clamps it to zero. Clamping is lossless *for the value* - +/// `A_G` only ever maxes `s_G` against `ΔE'_G`/`ΔF'_G`, which are `>= 0`, so a negative +/// diagonal loses every max it enters and zero loses just as thoroughly. It is **not** lossless +/// for the flag: `A_G == s_G == 0` then holds while the true diagonal is strictly negative and +/// the real max came from a gap. Believing that bit takes a path worth `mismatch` points less +/// than the fill reported - a wrong CIGAR under a right score. +/// +/// The fix costs two ops and no extra state: on a clamped scheme the diagonal is masked off at +/// **mismatch cells**, which is exactly where the clamp bites, using the `q == r` mask the +/// blend already computed. It is sound rather than conservative: where `s_G` clamps, the true +/// diagonal is negative while `ΔE'_G`/`ΔF'_G` are `>= 0`, so at a mismatch cell the diagonal +/// **provably never achieves the max** and masking it off discards nothing. +/// +/// Only a test that re-scores the CIGAR can catch a regression here; one that trusts the +/// reported score cannot, since the score stays right while the CIGAR goes wrong. +const TB_FROM_DIAG: usize = 3; + +/// The interior row range of anti-diagonal `p`, band included. +/// +/// # Band geometry on the anti-diagonal +/// +/// A band is the constraint `|i - j| <= w`. On anti-diagonal `p` the cell in row `i` sits at +/// column `j = p - i`, so `i - j = 2i - p` and the constraint reads +/// +/// ```text +/// |2i - p| <= w <=> ceil((p - w) / 2) <= i <= floor((p + w) / 2) +/// ``` +/// +/// That is **two more clamps on a range the fill already clamps twice**. The matrix corners +/// already force `i >= max(1, p - rlen)` and `i <= min(qlen, p - 1)`; the band is `max` and +/// `min` against one more term each, and every loop below inherits it for free. Nothing else +/// about the fill changes shape: same recurrence, same lanes, same stores, fewer of them. +/// +/// The width is `w + 1` rows when `p + w` is even and `w` when it is odd - the two bounds +/// advance on alternating anti-diagonals, since each moves at half of `p`'s rate. So the fill +/// costs `(qlen + rlen) * (w + 1) / LANES` chunks instead of `qlen * rlen / LANES`, which at +/// the shape this is for (`q = 20000`, `r = 24000`, `w = 100`) is 4.4M cells against 480M. +/// +/// Returns `(lo, hi)`, inclusive. `lo > hi` means the anti-diagonal has no interior, which +/// happens at `p = 1` and nowhere else for `w >= 1`. +#[inline(always)] +fn band_rows(p: usize, qlen: usize, rlen: usize, w: usize) -> (usize, usize) { + let lo = 1.max(p.saturating_sub(rlen)); + let hi = qlen.min(p - 1); + if BANDED { + // `saturating_sub` on both: `p - w` goes negative for every anti-diagonal before the + // band's lower edge lifts off row 0, and `p - rlen` for every one before the band + // starts sliding right. Same trap, same fix, and the same one the unbanded bound + // already needed. + (lo.max(p.saturating_sub(w).div_ceil(2)), hi.min((p + w) / 2)) + } else { + (lo, hi) + } +} + +/// Does a band of `w` reach every cell of a `qlen x rlen` matrix? +/// +/// The furthest any cell sits from the main diagonal is `max(qlen, rlen)` - the corners +/// `(qlen, 0)` and `(0, rlen)`. At or past that the band excludes nothing, so the banded optimum +/// *is* the exact optimum: same candidate set, same tie-breaks, same CIGAR. +/// +/// # Why the callers take the exact path instead +/// +/// Not to change the answer - it cannot - but because the exact path is **cheaper**, and most so +/// exactly where this fires. `global` and the `local_reference_*` modes get their exact score +/// from the cheap `anchor` walk, while every banded path runs `TRACK_ROW_MAX`; and the band's own +/// per-anti-diagonal bookkeeping is pure waste when the strip excludes nothing. +/// +/// This is constant folding, not a policy: `Some(w)` still means exactly what it says, and a +/// band wider than the problem simply stops costing anything. +#[inline] +fn band_reaches_everything(qlen: usize, rlen: usize, w: usize) -> bool { + w >= qlen.max(rlen) +} + +/// The query counterpart of [`ref_span`]. Only the `local_*` modes - the ones that let the +/// query end early - ever pass `consumed != qlen`. +#[inline] +fn query_span(qlen: usize, consumed: usize) -> (usize, usize) { + if REVERSED { + (qlen - consumed, qlen) + } else { + (0, consumed) + } +} + +/// Map "how much reference the CIGAR consumed" onto a span of the caller's reference. +/// +/// Under `REVERSED` the kernel ran on back-to-front codes, so a run of `consumed` bases +/// starting at the reversed origin is a run *ending* at the forward end. One subtraction - the +/// only coordinate work `REVERSED` costs. +#[inline] +fn ref_span(rlen: usize, consumed: usize) -> (usize, usize) { + if REVERSED { + (rlen - consumed, rlen) + } else { + (0, consumed) + } +} + +/// Does this scoring scheme survive `u8` lanes? +/// +/// The bound is `M + 3*gap_open - gap_extend`, **not** the `M + 2*gap_open` that bounds the +/// stored values. The looser rule is a trap: `ΔE'_G`'s recurrence forms `ΔE'_G[i-1,j] + go` +/// **before** subtracting, and that intermediate is not a stored value, so it is not covered. +/// It reaches `upper + go`. +/// +/// Saturating instead of rejecting does not help: `_mm256_adds_epu8` clamping to 255 makes the +/// following `max` pick the wrong operand - a wrong answer rather than a caught one. +/// +/// At default scores this is `2 + 36 - 1 = 37` against a ceiling of 255, and it holds out to +/// `gap_open ~ 84`, so nothing real is near it. +/// +/// Note this takes **no lengths**: the bound depends on the scoring scheme alone, so it is +/// decided once per scheme rather than per call, and no sequence is ever too long for it. +pub fn fits_u8_lanes(match_score: u8, gap_open: u8, gap_extend: u8) -> bool { + let upper = match_score as i64 + 2 * gap_open as i64; + let go = gap_open as i64 - gap_extend as i64; + upper + go <= u8::MAX as i64 +} + +/// One anti-diagonal's four difference vectors, indexed by row `i`. +/// +/// Grow-only, and padded by `LANES` so the partial last chunk can store all 32 lanes without a +/// bounds check. +#[derive(Default)] +struct Deltas { + dh: Vec, + dv: Vec, + de: Vec, + df: Vec, +} + +impl Deltas { + fn grow(&mut self, rows: usize) { + for b in [&mut self.dh, &mut self.dv, &mut self.de, &mut self.df] { + if b.len() < rows { + b.resize(rows, 0); + } + } + } +} + +/// Reusable scratch, owned by the aligner and reused across calls. +/// +/// Every buffer here is **grow-only**: sized up when a call needs more and never shrunk, so a +/// steady stream of similar-sized alignments allocates once. Small alignments dominate the +/// intended workload, where a per-call allocation would dominate the call. +#[derive(Default)] +pub struct U8Probe { + q_codes: Vec, + r_rev: Vec, + prev: Deltas, + cur: Deltas, + /// Traceback flags, `TB_WORDS` `u32`s per 32-lane chunk: **4 bits per cell**, laid out in + /// `(p, lane)` coordinates rather than `(i, j)`. + /// + /// Under a band this stores only the cells actually visited, not the full matrix. + tb: Vec, + /// Chunk index at which anti-diagonal `p` starts within `tb`. Anti-diagonals have unequal + /// lengths - the ramp at both matrix corners - so a fixed stride would waste ~50%. + tb_base: Vec, + + /// The per-row running **absolute** score: `abs[i] == H[i][p - i]` for the anti-diagonal + /// being filled. See `TRACK_ROW_MAX` on `run` for why this is exact and needs no bound. + /// + /// `i32`, not `u8`: the DP itself stays in u8 lanes: this is a reconstruction carried + /// alongside it for the modes that argmax over absolute scores, so `fits_u8_lanes` is + /// untouched and still has no length term. + abs: Vec, + /// `row_max[i] == max over j of H[i][j]`. + row_max: Vec, + + /// **Where** each row's maximum was attained, as one bit per cell: set iff that cell took + /// its row's running max. One `u32` per 32-lane chunk, in the same `(p, chunk)` coordinates + /// as `tb`, so `tb_base[p]` indexes both. + /// + /// The argmax is only ever asked about one row, so carrying a winning `j` per row would + /// compute one for every row on every anti-diagonal and throw all but one away. The bit is + /// enough: `arg_for_row` recovers the row that matters afterwards, scalar. + /// + /// Recovery is exact, and `take` is why: it is the update condition itself, so "the last + /// cell whose bit is set" is by construction the cell a blend would have left, under either + /// tie-break (`ARG_LARGEST` makes `take` non-strict so later ties win; otherwise it is + /// strict, so the first winner sticks). + nm: Vec, + /// Row 0's argmax, which the bit stream cannot carry: row 0 is never *interior* to an + /// anti-diagonal (`i_from >= 1` always), so it has no lane and no bit. It is maintained + /// scalar in the boundary, exactly as it always was, and it is load-bearing - with + /// `gap_open == gap_extend == 0` the whole of row 0 ties at 0 and the longest-extension + /// tie-break has to take the furthest cell. + row0_arg: i32, + + /// A second set of flags and row maxima, for `split_reference_alignment` and nothing else. + /// + /// That mode runs the kernel **twice** - once forward against the right reference, once + /// REVERSED against the left - and then replays **both** arms, so both tracebacks have to + /// be alive at the same time. Empty for every other mode, and grow-only like the rest, so + /// a program that never splits never pays for these. + tb_b: Vec, + tb_base_b: Vec, + row_max_b: Vec, + nm_b: Vec, + row0_arg_b: i32, +} + +/// Recover `argmax over j of H[row][j]` from the new-max bit stream. +/// +/// Takes its buffers explicitly rather than reaching through `self`, because `split_reference` +/// holds two passes' worth at once and has to say which. +/// +/// The stream's bit is the fill's `take` mask, so the cells with the bit set are every cell that +/// *became* row `row`'s max, and the last of them is the one still holding it when the fill +/// ended. Walking backwards therefore needs no knowledge of the tie-break: `ARG_LARGEST` makes +/// `take` non-strict, so later ties re-fire the bit and the last wins; otherwise `take` is +/// strict, so ties do not re-fire and the first winner stays. +/// +/// Row `row` is interior to anti-diagonal `p` exactly for `p` in `[row+1, row+rlen]`, so those +/// are the only `p` scanned - which is also why the garbage the partial chunk tail writes above +/// `i_to` is never read: it lands at `p <= row`, below where this starts. +/// +/// Falling off the bottom means no interior cell ever took the max, so the row's max is still +/// its column-0 seed: `j = 0`. +/// +/// The scan is two phases because `lane = row - max(1, p - rlen)` is constant at `row - 1` while +/// `p <= rlen + 1`, which is most of it; hoisting the lane out of that phase leaves a load, a +/// shift, a load and a test per step. +fn arg_for_row(nm: &[u32], tb_base: &[u32], row0_arg: i32, row: usize, rlen: usize) -> usize { + // Row 0 has no lane and no bit - `i_from >= 1` always - so it is carried scalar. + if row == 0 { + return row0_arg as usize; + } + let word_at = |p: usize, lane: usize| -> bool { + // SAFETY: `p <= row + rlen <= qlen + rlen`, which is what `size_flags` sized `tb_base` + // to, and `tb_base[p] + lane/LANES` is a chunk the fill wrote - `row` is interior to + // every `p` this scans (see below), so its lane is one the group loop covered. + unsafe { + let w = *nm.get_unchecked(*tb_base.get_unchecked(p) as usize + lane / LANES); + w >> (lane % LANES) & 1 != 0 + } + }; + + let mut p = row + rlen; + // Phase 2: `p > rlen + 1`, so `i_from = p - rlen` and the lane slides with `p`. At most + // `qlen` iterations - the ramp off the end of the reference. + while p > rlen + 1 && p > row { + if word_at(p, row + rlen - p) { + return p - row; + } + p -= 1; + } + // Phase 1: `i_from == 1`, so the lane is fixed and so is everything derived from it. + let lane = row - 1; + let off = lane / LANES; + let bit = 1u32 << (lane % LANES); + while p > row { + // SAFETY: as above. + if unsafe { *nm.get_unchecked(*tb_base.get_unchecked(p) as usize + off) } & bit != 0 { + return p - row; + } + p -= 1; + } + 0 +} + +/// [`arg_for_row`] under a band: `argmax over j of H[row][j]`, from the same bit stream. +/// +/// A separate scan rather than a clamp on the other one, because bounding the range here is a +/// **correctness** requirement, not an optimisation. Row `row` is in band only for `p` in +/// `[2*row - w, 2*row + w]`. Above that the row sits above the band, `row - lo(p)` goes negative +/// and the unchecked read lands anywhere; below it the row sits inside the garbage the partial +/// chunk's tail wrote, and `lane / LANES` can walk into the next anti-diagonal's chunks and read +/// a bit that means something else. The unbanded scan needs neither bound because the unbanded +/// matrix has neither edge. +/// +/// Falling off the bottom returns the row's seed column, which under a band is `row - w`, not 0: +/// row `row > w` has no column-0 cell and is seeded at its first in-band column by the fill's +/// band-entry seed. `saturating_sub` covers both cases. +/// +/// That also disposes of the one garbage bit in the stream. At `p = 2*row - w` the row's `take` +/// came from a `ΔV` whose left neighbour is outside the band, so the bit means nothing - but the +/// scan only reaches it after finding no set bit above, which is exactly the case where the +/// answer is the seed column anyway. Set or clear, it returns `row - w`, so the fill leaves it +/// alone rather than paying a read-modify-write to correct it. +fn arg_for_row_banded( + nm: &[u32], + tb_base: &[u32], + row0_arg: i32, + row: usize, + qlen: usize, + rlen: usize, + w: usize, +) -> usize { + if row == 0 { + return row0_arg as usize; + } + // Row `row` is interior to `p` for `j = p - row` in `[1, rlen]` and `|2*row - p| <= w`. + let p_hi = (row + rlen).min(2 * row + w); + let p_lo = (row + 1).max((2 * row).saturating_sub(w)); + + let mut p = p_hi; + while p >= p_lo { + let (lo, _) = band_rows::(p, qlen, rlen, w); + let lane = row - lo; + // SAFETY: `p` is in `[p_lo, p_hi]`, so `row` is interior to it and `lane` is a lane the + // group loop covered; `p <= row + rlen <= qlen + rlen` is what `size_flags` sized + // `tb_base` to. + let set = unsafe { + let word = *nm.get_unchecked(*tb_base.get_unchecked(p) as usize + lane / LANES); + word >> (lane % LANES) & 1 != 0 + }; + if set { + return p - row; + } + p -= 1; + } + row.saturating_sub(w) +} + +/// The alphabet: **A, C, G, T, U**, and one code for everything else. +/// +/// Codes 0-3 are A/C/G/T, case-insensitive, with `U` folded onto `T`. Code 4 is *unknown* - +/// every other byte - and it must never match anything, **not even another unknown byte**. +/// +/// There is no substitution matrix here, and that is the point. The scheme is `match` on the +/// diagonal and `-mismatch` everywhere else: a binary *predicate*, not a table. So scoring a +/// cell is one `cmpeq` against the codes already in registers - no gather, no `pshufb`, no +/// 25-entry lookup. +/// +/// The one thing a plain compare cannot express is `unknown != unknown`, because `x == x` is +/// unavoidably true. Hence the **asymmetry**: [`U8Probe::prepare`] re-encodes unknown to code 5 +/// on the *reference* side only, so the compare says "not equal" and the blend picks +/// `-mismatch`. One `cmpeq` and one `blendv`, exactly what a 2-bit ACGT-only encoding would +/// cost - so the fail-safe is free, and a stray `-` or IUPAC `R` scores a mismatch rather than +/// silently becoming an `A` and scoring a false match. See `mod.rs`, "Alphabet". +fn code(b: u8) -> u8 { + match b.to_ascii_uppercase() { + b'A' => 0, + b'C' => 1, + b'G' => 2, + b'T' | b'U' => 3, + _ => 4, + } +} + +impl U8Probe { + pub fn new() -> Self { + Self::default() + } + + /// Check the preconditions, encode both sequences, and size the delta buffers. + /// + /// Assumes both sequences are non-empty - a zero-length side + /// is one pure gap and never reaches the kernel. + fn prepare( + &mut self, + query: &[u8], + refseq: &[u8], + match_score: u8, + gap_open: u8, + gap_extend: u8, + ) { + assert!( + fits_u8_lanes(match_score, gap_open, gap_extend), + "scores do not fit u8 lanes" + ); + assert!(gap_open >= gap_extend, "gap_open must be >= gap_extend"); + assert!(is_x86_feature_detected!("avx2"), "AVX2 is required"); + + let (qlen, rlen) = (query.len(), refseq.len()); + + // Encoding is O(n) and happens every call regardless, which is why the asymmetric + // unknown-base encoding costs nothing. + // + // Grow-then-overwrite, not `clear()` + `extend()`: the latter re-checks capacity per + // element. + // REVERSED is free, and this is the whole reason: the sequences have to be encoded + // every call anyway, so a `_start` mode just encodes them **back-to-front**. Same pass, + // same cost, no buffer, and the kernel never learns the mode exists. + // + // Note the two axes compose: the kernel always wants the reference laid out backwards + // relative to the query (that is what makes `refseq[p-i-1]` a contiguous load), so + // under REVERSED the reference is reversed *twice* and comes out forward. + if self.q_codes.len() < qlen + LANES { + self.q_codes.resize(qlen + LANES, 4); + } + if REVERSED { + for (dst, &b) in self.q_codes[..qlen].iter_mut().zip(query.iter().rev()) { + *dst = code(b); + } + } else { + for (dst, &b) in self.q_codes[..qlen].iter_mut().zip(query) { + *dst = code(b); + } + } + self.q_codes[qlen..qlen + LANES].fill(4); + + // Code 5 is N on the reference side - the asymmetry that makes `N != N` fall out of + // the compare. + if self.r_rev.len() < rlen + LANES { + self.r_rev.resize(rlen + LANES, 5); + } + let enc_r = |b: u8| -> u8 { + let c = code(b); + if c == 4 { 5 } else { c } + }; + if REVERSED { + for (dst, &b) in self.r_rev[..rlen].iter_mut().zip(refseq.iter()) { + *dst = enc_r(b); + } + } else { + for (dst, &b) in self.r_rev[..rlen].iter_mut().zip(refseq.iter().rev()) { + *dst = enc_r(b); + } + } + self.r_rev[rlen..rlen + LANES].fill(5); + + let rows = qlen + 1 + LANES; + self.prev.grow(rows); + self.cur.grow(rows); + // Each buffer gets its OWN check. Gating all three on `abs.len()` looks equivalent and + // is not: `split_reference` swaps `row_max` with its `_b` twin between its two passes, + // so `abs` can already be grown while the row buffers are the empty ones - the check + // would not fire and the kernel would write off the end. + if self.abs.len() < rows { + self.abs.resize(rows, 0); + } + if self.row_max.len() < rows { + self.row_max.resize(rows, 0); + } + } + + /// Global alignment: both sequences spanned end to end. + #[allow(clippy::too_many_arguments)] // a kernel entry point: the scheme plus the bandwidth + pub fn global_alignment( + &mut self, + query: &[u8], + refseq: &[u8], + m: u8, + mm: u8, + go: u8, + ge: u8, + bandwidth: Option, + ) -> AlignmentResult { + match bandwidth { + // A band that reaches every cell is the exact mode, and the exact mode is cheaper - + // see `band_reaches_everything`. + Some(w) if !band_reaches_everything(query.len(), refseq.len(), w) => { + self.global_banded(query, refseq, m, mm, go, ge, w) + } + _ => self.align::(query, refseq, m, mm, go, ge), + } + } + + /// The banded `global` path: `S[qlen][rlen]`, confined to `|i - j| <= w`. + /// + /// # Widening: global is pinned at both corners + /// + /// Every other mode here has exactly one anchor and lets the far end float, so `|i-j| <= w` + /// is a statement about drift away from that anchor and any `w` is meaningful. Global must + /// also *arrive* at `(qlen, rlen)`, which sits `|qlen - rlen|` off the diagonal the band is + /// centred on. A narrower band does not contain a bad alignment, it contains **no path at + /// all**: `band_rows` returns `lo > hi` for the final anti-diagonal and the corner is never + /// computed. Hence the widening to `|qlen - rlen|` - the narrowest diagonal band that still + /// joins the two corners. + /// + /// That widening is free when `qlen ~ rlen`, which is the shape global is normally handed + /// (two segments already believed to correspond end to end), and it is what makes the band + /// worth having here at all. It bites on a lopsided pair, where it degrades to + /// correct-but-slow rather than to wrong. + /// + /// # The score comes from `abs`, not `anchor` + /// + /// See [`U8Probe::local_reference`]: `align`'s `anchor` walk cannot survive a band. But + /// `TRACK_ROW_MAX` already maintains `abs[i]`, the running **absolute** score of row `i`'s + /// cell on the current anti-diagonal, seeded at the row's band entry and advanced by one + /// `ΔV` per cell. Row `qlen` is live for `p` up to `qlen + rlen`, and the widening guarantees + /// it is still in band there - so when the fill returns, `abs[qlen]` *is* `S[qlen][rlen]`, + /// with nothing to add. The end cell is read, not recomputed. + #[allow(clippy::too_many_arguments)] // ditto, plus the bandwidth + fn global_banded( + &mut self, + query: &[u8], + refseq: &[u8], + m: u8, + mm: u8, + go: u8, + ge: u8, + w: usize, + ) -> AlignmentResult { + let (qlen, rlen) = (query.len(), refseq.len()); + // Band-independent given the widening below: with one side empty the only path is the + // single gap `degenerate` returns, and `w >= |qlen - rlen|` is exactly wide enough to + // hold it. + if qlen == 0 || rlen == 0 { + return self.degenerate::(qlen, rlen, go, ge); + } + let w = w.max(1).max(qlen.abs_diff(rlen)); + + self.row_pass::(query, refseq, m, mm, go, ge, w); + + // `abs[qlen]` after the fill is `S[qlen][rlen]` - see this function's docs. The widening + // is what puts the corner in band, and so what makes this read meaningful. + debug_assert!( + qlen.abs_diff(rlen) <= w, + "the widening should have put the corner in band" + ); + let score = self.abs[qlen]; + + let cigar = self.replay::(query, refseq, qlen, rlen, w); + AlignmentResult { + score, + query_start: 0, + query_end: qlen, + ref_start: 0, + ref_end: rlen, + cigar, + } + } + + /// The query is spanned in full; the **reference may stop anywhere**, and its trailing + /// tail is free and does not appear in the CIGAR. + /// + /// `j = 0` is a legal end cell and sometimes the winning one. Ties resolve to the + /// **smallest** `ref_end`. + #[allow(clippy::too_many_arguments)] // a kernel entry point: the scheme plus the bandwidth + pub fn local_reference_end_alignment( + &mut self, + query: &[u8], + refseq: &[u8], + m: u8, + mm: u8, + go: u8, + ge: u8, + bandwidth: Option, + ) -> AlignmentResult { + match bandwidth { + Some(w) if !band_reaches_everything(query.len(), refseq.len(), w) => { + self.local_reference::(query, refseq, m, mm, go, ge, w) + } + _ => self.align::(query, refseq, m, mm, go, ge), + } + } + + /// The mirror: the reference may **begin** anywhere and its leading prefix is free. + /// + /// This is `local_reference_end_alignment` on reversed inputs, and that is exact rather + /// than an approximation: affine gap costs are symmetric under reversal (a gap of length + /// `L` costs `open + (L-1)*extend` whichever way you read it) and `=`/`X` is literal byte + /// identity, also symmetric. Ties resolve to the **largest** `ref_start`. + #[allow(clippy::too_many_arguments)] // a kernel entry point: the scheme plus the bandwidth + pub fn local_reference_start_alignment( + &mut self, + query: &[u8], + refseq: &[u8], + m: u8, + mm: u8, + go: u8, + ge: u8, + bandwidth: Option, + ) -> AlignmentResult { + match bandwidth { + Some(w) if !band_reaches_everything(query.len(), refseq.len(), w) => { + self.local_reference::(query, refseq, m, mm, go, ge, w) + } + _ => self.align::(query, refseq, m, mm, go, ge), + } + } + + /// The banded `local_reference_*` path: `max over j of H[qlen][j]`, confined to `|i-j| <= w`. + /// + /// # Why not `align` with `BANDED` switched on + /// + /// `align` does not read its score out of the fill. It walks `anchor`, a closed-form + /// recurrence along row `qlen` that reads `ΔV[qlen]` on **every** anti-diagonal past + /// `qlen`, and row `qlen` is a horizontal line while the band is a diagonal strip, so the + /// two meet on only `O(w)` of those. Under a band the rest of that walk reads lanes the + /// fill never wrote. `BANDED` and `TRACK_LAST_ROW` were never compiled together for + /// exactly this reason. + /// + /// `TRACK_ROW_MAX` is the way in: it carries **absolute** per-cell scores, it is already + /// banded (it is what `local` runs), and `row_max[qlen]` *is* this mode's objective. So the + /// banded path is a `row_pass` and a reduce that is one array read, and the anchor never + /// enters it. + /// + /// It costs more per cell than `align` - a `row_max` store and a bit of the new-max stream + /// per chunk - which is why `None` still takes `align`: the two are separate + /// monomorphizations, so the exact path never pays for this one. The band's own win (the + /// fill drops from `qlen*rlen` to `(qlen+rlen)*(w+1)` cells) buys that back many times. + #[allow(clippy::too_many_arguments)] // ditto, plus the bandwidth + fn local_reference( + &mut self, + query: &[u8], + refseq: &[u8], + m: u8, + mm: u8, + go: u8, + ge: u8, + w: usize, + ) -> AlignmentResult { + let (qlen, rlen) = (query.len(), refseq.len()); + // Band-independent: see the widening below, which always admits the whole of the one + // live row or column these cases consist of. + if qlen == 0 || rlen == 0 { + return self.degenerate::(qlen, rlen, go, ge); + } + + // The widening. This mode must span the query, so it must reach a cell `(qlen, j)` with + // `j <= rlen`; in band that needs `qlen - w <= rlen`. Below that the band contains no + // alignment at all - not a bad one, *none* - so `Some(w)` is honoured as "at least w", + // and widened to the narrowest band that has an answer. It only bites when the query is + // longer than the reference, which this mode is not normally handed. + let w = w.max(1).max(qlen.saturating_sub(rlen)); + + self.row_pass::(query, refseq, m, mm, go, ge, w); + + // `last_row == min(qlen, rlen + w) == qlen` after the widening, so row `qlen` is filled + // and this read is in bounds - that is what the widening is for. + debug_assert!( + qlen <= rlen + w, + "the widening should have made row qlen reachable" + ); + let score = self.row_max[qlen]; + // `ARG_LARGEST = false` in `row_pass`, so ties give the smallest `j` - this mode's rule. + let j = arg_for_row_banded(&self.nm, &self.tb_base, self.row0_arg, qlen, qlen, rlen, w); + debug_assert!( + qlen.abs_diff(j) <= w, + "the reduce picked ({qlen}, {j}), which is outside the band of {w}" + ); + + let cigar = self.replay::(query, refseq, qlen, j, w); + let (ref_start, ref_end) = ref_span::(rlen, j); + AlignmentResult { + score, + query_start: 0, + query_end: qlen, + ref_start, + ref_end, + cigar, + } + } + + /// Ends wherever it scores best: `max over (i, j) of H[i][j] + (end_bonus if i == qlen)`. + /// + /// Cell `(0,0)` is a legal end cell, so **the score is never negative** - a hopeless + /// extension returns the empty alignment (spans `0..0`, score 0, empty CIGAR). + /// + /// `end_bonus` never enters the DP. It is applied once, at the end, to choose between "the + /// best cell anywhere" and "the best cell that finishes the query". Ties resolve to the + /// **longest extension**: largest `query_end`, then largest `ref_end`. + #[allow(clippy::too_many_arguments)] // a kernel entry point: the scheme is five scalars + pub fn local_end_alignment( + &mut self, + query: &[u8], + refseq: &[u8], + m: u8, + mm: u8, + go: u8, + ge: u8, + end_bonus: u32, + bandwidth: Option, + ) -> AlignmentResult { + match bandwidth { + Some(w) if !band_reaches_everything(query.len(), refseq.len(), w.max(1)) => { + self.local::(query, refseq, m, mm, go, ge, end_bonus, w.max(1)) + } + _ => self.local::(query, refseq, m, mm, go, ge, end_bonus, 0), + } + } + + /// The mirror: begins wherever it scores best, both leading prefixes free. + /// + /// Its empty alignment sits at the *end* of both sequences (`qlen..qlen`, `rlen..rlen`) - + /// the mirror of `local_end`'s, which sits at the start. Tie-break mirrored too: the + /// **smallest** `query_start`, then the smallest `ref_start`. + #[allow(clippy::too_many_arguments)] // a kernel entry point: the scheme is five scalars + pub fn local_start_alignment( + &mut self, + query: &[u8], + refseq: &[u8], + m: u8, + mm: u8, + go: u8, + ge: u8, + end_bonus: u32, + bandwidth: Option, + ) -> AlignmentResult { + match bandwidth { + Some(w) if !band_reaches_everything(query.len(), refseq.len(), w.max(1)) => { + self.local::(query, refseq, m, mm, go, ge, end_bonus, w.max(1)) + } + _ => self.local::(query, refseq, m, mm, go, ge, end_bonus, 0), + } + } + + /// `w` is read only when `BANDED`, and is the caller's bandwidth already clamped to + /// `>= 1`: a zero-width band is not meaningful, so callers clamp rather than pass it on. + #[allow(clippy::too_many_arguments)] // ditto, plus the bandwidth + fn local( + &mut self, + query: &[u8], + refseq: &[u8], + match_score: u8, + mismatch: u8, + gap_open: u8, + gap_extend: u8, + end_bonus: u32, + w: usize, + ) -> AlignmentResult { + let (qlen, rlen) = (query.len(), refseq.len()); + // A zero-length side skips the kernel - there is no anti-diagonal to walk - but it is + // NOT simply "the empty alignment", and assuming so was a real bug: + // + // * with `qlen == 0`, EVERY cell has `i == qlen`, so `end_bonus` applies to all of + // them. The empty alignment scores `end_bonus`, not 0. + // * with free gaps (`gap_open == gap_extend == 0`) the whole boundary row/column ties + // at 0, and the longest-extension tie-break must then take the FURTHEST cell, not + // the origin. + // + // So scan the one live row or column in closed form. O(qlen + rlen), and only here. + if qlen == 0 || rlen == 0 { + let gap = |k: usize| -> i32 { + if k == 0 { + 0 + } else { + -(gap_open as i32) - (k as i32 - 1) * gap_extend as i32 + } + }; + // The band reaches the live row/column too, and it bites here for the same reason + // it bites anywhere: `H[0][j]` sits `j` off the diagonal, so a band of `w` can only + // see the first `w` of it. Without this the free-gap case (`gap_open == gap_extend + // == 0`, where the whole row ties at 0) would report the furthest cell in the row + // rather than the furthest cell in the *band*. + let reach = |k: usize| if BANDED { k.min(w) } else { k }; + let (mut best, mut bi, mut bj) = (i32::MIN, 0usize, 0usize); + if qlen == 0 { + // Row 0: H[0][j] is j deletions, and i == qlen == 0 throughout, so the bonus + // is on every candidate. `>=` takes the largest j on ties. + for j in 0..=reach(rlen) { + let v = gap(j) + end_bonus as i32; + if v >= best { + best = v; + bj = j; + } + } + } else { + // Column 0: H[i][0] is i insertions; only i == qlen earns the bonus. `>=` + // takes the largest i on ties. + // + // The bonus needs the band's permission: it lands on `i == qlen`, and under a + // band that cell only exists if the band reaches it. Where it does not, the + // loop simply never visits `i == qlen` and no bonus is ever added, which is + // right - the query cannot be spanned inside the band at all. + for i in 0..=reach(qlen) { + let v = gap(i) + if i == qlen { end_bonus as i32 } else { 0 }; + if v >= best { + best = v; + bi = i; + } + } + } + let mut cigar = Cigar::default(); + if bi > 0 { + cigar.push(CigarOperation::Insertion, bi); + } + if bj > 0 { + cigar.push(CigarOperation::Deletion, bj); + } + let (query_start, query_end) = query_span::(qlen, bi); + let (ref_start, ref_end) = ref_span::(rlen, bj); + return AlignmentResult { + score: best, + query_start, + query_end, + ref_start, + ref_end, + cigar, + }; + } + + self.prepare::(query, refseq, match_score, gap_open, gap_extend); + self.size_flags::(qlen, rlen, w); + + unsafe { + self.run::( + qlen, + rlen, + match_score, + mismatch, + gap_open, + gap_extend, + w, + ) + }; + + // How far down the query the band still has cells to offer. + // + // Row `i`'s in-band columns are `[i - w, i + w]` intersected with `[0, rlen]`, which is + // empty once `i > rlen + w`. Those rows are not merely bad candidates, they were never + // filled: `row_max[i]` there is whatever the previous call left, or the garbage the + // partial chunk's tail wrote. So the reduce must not look at them - a wrong-but-huge + // value would win, and would win *silently*. + // + // Only reachable when the query is longer than the reference by more than the band, so + // never in the intended use (an end extension is handed more reference than query) - + // which is exactly why the bound is enforced rather than assumed. + let last_row = if BANDED { qlen.min(rlen + w) } else { qlen }; + + // Reduce the per-row maxima to a single cell. Scalar, and cheap: qlen+1 values against + // the fill's qlen*rlen cells. + // + // Ties resolve to the longest extension - largest i, then largest j. `>=` on i gives the + // largest i directly; the largest j comes from `arg_for_row`, whose bit stream was + // written under `ARG_LARGEST`'s non-strict `take` and so records the last tie. + let (mut best, mut bi) = (0i32, 0usize); + for i in 0..=last_row { + let v = self.row_max[i]; + if v >= best { + best = v; + bi = i; + } + } + // The end_bonus candidate: the best cell that finishes the query. Applied once, here, + // never in the DP. + // + // Under a band it needs `row_max[qlen]` to exist at all. Where the band cannot reach + // row `qlen` there is no cell that finishes the query, so there is nothing for the + // bonus to apply to - and `row_max[qlen]` is unfilled, so asking anyway would read + // garbage and hand it the answer. + if last_row == qlen { + let finish = self.row_max[qlen] + end_bonus as i32; + if finish > best || (finish == best && qlen >= bi) { + best = finish; + bi = qlen; + } + } + // The argmax `j`, recovered for the ONE row that won - which is the whole reason the + // fill stores a bit per cell instead of a `j` per row. `O(rlen)` unbanded, `O(w)` under + // a band. + let bj = if BANDED { + arg_for_row_banded(&self.nm, &self.tb_base, self.row0_arg, bi, qlen, rlen, w) + } else { + arg_for_row(&self.nm, &self.tb_base, self.row0_arg, bi, rlen) + }; + // No early return for `best == 0`, and that is not an oversight. (0,0) is the reduce's + // seed candidate, so the score can never go negative - the "never negative" guarantee + // falls out rather than being clamped. But bailing out at `best <= 0` would ALSO + // discard a *longer* extension that happens to score exactly 0, which the + // longest-extension tie-break says must win over (0,0) - a right score with a wrong + // endpoint, which is what the endpoint assertion below exists to catch. + debug_assert!( + best >= 0, + "(0,0) seeds the reduce, so the score cannot go negative" + ); + debug_assert!( + !BANDED || bi.abs_diff(bj) <= w, + "the reduce picked ({bi}, {bj}), which is outside the band of {w}" + ); + + let cigar = self.replay::(query, refseq, bi, bj, w); + let (query_start, query_end) = query_span::(qlen, bi); + let (ref_start, ref_end) = ref_span::(rlen, bj); + AlignmentResult { + score: best, + query_start, + query_end, + ref_start, + ref_end, + cigar, + } + } + + /// An upper bound on the chunks a **banded** fill writes, and therefore on the `tb`/`nm` it + /// needs. + /// + /// # Derived from the fill's own accounting, not from a cell count + /// + /// `run` advances its chunk counter exactly once per live anti-diagonal, by + /// `(hi - lo) / LANES + 1` (see the `if TRACE && lo <= hi` in the fill). So the total is + /// `sum over live p of ((hi_p - lo_p) / LANES + 1)`, and bounding it needs only a bound on + /// `hi - lo`: + /// + /// `band_rows::` returns `ceil((p - w)/2) ..= floor((p + w)/2)` before the matrix + /// clamps, which spans `w + 1` rows when `p + w` is even and `w` when it is odd - so + /// `hi - lo <= w` always, and the matrix clamps only shrink it. There are at most + /// `qlen + rlen` live anti-diagonals (`p = 0` has no interior). Hence + /// + /// ```text + /// chunks <= (qlen + rlen) * (w / LANES + 1) + /// ``` + /// + /// **`w / LANES + 1`, not `(w + 1) / LANES`.** A strip of `w + 1` rows can straddle a chunk + /// boundary, and the fill starts each anti-diagonal's chunks at the *band's* top row rather + /// than the matrix's, so the `+ 1` is the partial chunk and is not slack to be tuned away. + /// The `+ 2` on top is the same margin the unbanded bound carries. + /// + /// **This is a bound the fill must not exceed, and exceeding it is heap corruption rather + /// than a wrong answer** - the fill's stores are unchecked. `run` debug-asserts the final + /// chunk count against the buffer it was actually handed, which is what makes this derivation + /// checked rather than argued - so run the banded paths in a debug build after touching it. + /// + /// At `q=20000 r=24000 w=64` this is 132,002 chunks against the full matrix's 15,044,002: + /// **2.6 MB instead of 300 MB**, for a fill that only ever visits 2.9M cells. + #[inline] + fn banded_chunks_bound(qlen: usize, rlen: usize, w: usize) -> usize { + (qlen + rlen) * (w / LANES + 1) + 2 + } + + /// Size the traceback flag buffer and its per-anti-diagonal index. Shared by every mode. + /// + /// `NEED_NM` additionally sizes the new-max bit stream, which only the `TRACK_ROW_MAX` modes + /// write. A const generic, not a runtime flag, for the usual reason: `global` must not carry + /// so much as a branch for a buffer it never touches, and grow-only means a program that + /// only ever aligns globally never allocates it at all. + /// + /// `BANDED`/`w` are the same deal, and they are what stops a band allocating for a matrix it + /// will never visit - see [`U8Probe::banded_chunks_bound`]. + fn size_flags( + &mut self, + qlen: usize, + rlen: usize, + w: usize, + ) { + // Size the flag buffer. Anti-diagonals have unequal lengths - the ramp at both matrix + // corners - so a fixed stride would waste ~50%; `tb_base` records where each starts, + // and the fill writes it as it goes. + // + // An exact count would need its own O(qlen + rlen) pass, which at small n is + // comparable to the entire fill. So bound it instead: every cell lives on exactly one + // anti-diagonal, giving `qlen*rlen/LANES` chunks, plus at most one partial chunk per + // anti-diagonal. Over-allocates by ~(qlen+rlen) chunks, and the buffers are grow-only + // anyway. + let full_bound = qlen * rlen / LANES + qlen + rlen + 2; + // Under a band, take whichever bound is smaller. **Both are valid** - a banded fill + // visits a subset of the cells an unbanded one does, so `full_bound` never stops being + // an upper bound - and the `min` matters in both directions: `banded_chunks_bound` + // *exceeds* `full_bound` once `w` approaches the matrix, because it charges a partial + // chunk per anti-diagonal against a strip that is no longer narrow. + let chunks_bound = if BANDED { + full_bound.min(Self::banded_chunks_bound(qlen, rlen, w)) + } else { + full_bound + }; + if self.tb_base.len() < qlen + rlen + 1 { + self.tb_base.resize(qlen + rlen + 1, 0); + } + if self.tb.len() < chunks_bound * TB_WORDS { + self.tb.resize(chunks_bound * TB_WORDS, 0); + } + if NEED_NM && self.nm.len() < chunks_bound * NM_WORDS { + self.nm.resize(chunks_bound * NM_WORDS, 0); + } + } + + /// One query aligned across **two** references, with a single jump point between them. + /// + /// Returns two [`AlignmentResult`]s that **partition the query exactly**: `right` covers + /// `query[..k]` from `right_reference[0]`, `left` covers `query[k..]` to the end of + /// `left_reference`, and `k` is chosen to maximize the total. + /// + /// # Reduction to two ordinary passes + /// + /// Over the same `H` matrix every other mode computes: + /// + /// ```text + /// F[k] = max over j of H(query, right)[k][j] - the best right arm handing over at k + /// G[k] = max over l of (best alignment of query[k..] vs left[l..]) - the best left arm taking over + /// score = max over k in 0..=qlen of F[k] + G[k] + /// ``` + /// + /// `F` is exactly `TRACK_ROW_MAX`'s output. And `G` is `F` on **reversed** inputs, so it is + /// the same pass with `REVERSED` - free, as always. The arms interact *only* through `k`, + /// so the whole "zig-zag" decision is a one-dimensional scan over `qlen + 1` values. + /// Everything before it is the ordinary kernel, run twice. + /// + /// # Ties + /// + /// Resolved to the smallest `|left.score - right.score|`, then the smallest `k`. So if the + /// other arm can pick up a base for free, it will. + /// + /// # Panics + /// If `fits_u8_lanes` rejects the scheme, if `gap_open < gap_extend`, or without AVX2. + #[allow(clippy::too_many_arguments)] // a kernel entry point: two references and the scheme + pub fn split_reference_alignment( + &mut self, + query: &[u8], + left_reference: &[u8], + right_reference: &[u8], + m: u8, + mm: u8, + go: u8, + ge: u8, + bandwidth: Option, + ) -> SplitReferenceAlignment { + match bandwidth { + Some(w) => { + // One band, both arms - each measured from its own arm's anchor. See + // `SimdAligner::split_reference_alignment`. + // + // The widening: the two arms must between them cover the whole query, and an arm + // can reach at most `reflen + w` query bases inside the band, so a `w` below + // `(qlen - llen - rrlen) / 2` admits no split at all. Widening to it is the + // narrowest band that still has an answer. + let deficit = query + .len() + .saturating_sub(left_reference.len() + right_reference.len()); + let w = w.max(1).max(deficit.div_ceil(2)); + // A band that reaches every cell is the exact mode, and the exact mode is + // cheaper - see `band_reaches_everything`. Both arms have to be reached, and + // they align against different references, so the wider of the two decides. + let widest = left_reference.len().max(right_reference.len()); + if band_reaches_everything(query.len(), widest, w) { + return self.split_reference::( + query, + left_reference, + right_reference, + m, + mm, + go, + ge, + 0, + ); + } + self.split_reference::( + query, + left_reference, + right_reference, + m, + mm, + go, + ge, + w, + ) + } + None => self.split_reference::( + query, + left_reference, + right_reference, + m, + mm, + go, + ge, + 0, + ), + } + } + + #[allow(clippy::too_many_arguments)] // ditto, plus the bandwidth + fn split_reference( + &mut self, + query: &[u8], + left_reference: &[u8], + right_reference: &[u8], + m: u8, + mm: u8, + go: u8, + ge: u8, + w: usize, + ) -> SplitReferenceAlignment { + let qlen = query.len(); + + // --- pass A: forward against the right reference. row_max[k] == F[k]. --- + self.row_pass::(query, right_reference, m, mm, go, ge, w); + // Park A's results in the `_b` slots and let pass B use the primary ones. A swap of Vec + // headers, once per call - not per anti-diagonal, where it would matter. + std::mem::swap(&mut self.tb, &mut self.tb_b); + std::mem::swap(&mut self.tb_base, &mut self.tb_base_b); + std::mem::swap(&mut self.row_max, &mut self.row_max_b); + // `nm` and `row0_arg` travel with `tb_base`, and they have to: `arg_for_row` reads all + // three, and reading one pass's bit stream against the other pass's chunk index would + // silently return a `j` from the wrong matrix. + std::mem::swap(&mut self.nm, &mut self.nm_b); + std::mem::swap(&mut self.row0_arg, &mut self.row0_arg_b); + // From here: `*_b` holds pass A (right arm), `*` holds pass B (left arm). + + // --- pass B: REVERSED against the left reference. row_max[qlen - k] == G[k]. --- + self.row_pass::(query, left_reference, m, mm, go, ge, w); + + // --- the jump point: a 1-D scan over qlen+1 values, against the fill's qlen*rlen --- + // + // Under a band the scan must not look at every `k`. Pass A filled rows `0..=min(qlen, + // rrlen + w)` and pass B rows `0..=min(qlen, llen + w)`, and `k` reads `row_max_b[k]` + // from A and `row_max[qlen - k]` from B, so both bounds bite - from opposite ends. + // Outside them the value is not a bad candidate, it is the previous call's leftovers, + // and being unbounded it would win. `split_reference_alignment` widened `w` so this + // range is never empty. + let (k_lo, k_hi) = if BANDED { + ( + qlen.saturating_sub(left_reference.len() + w), + qlen.min(right_reference.len() + w), + ) + } else { + (0, qlen) + }; + debug_assert!( + k_lo <= k_hi, + "no jump point is in band: k_lo={k_lo} > k_hi={k_hi} at w={w}, qlen={qlen}" + ); + let (mut best, mut best_diff, mut best_k) = (i64::MIN, i64::MAX, k_lo); + for k in k_lo..=k_hi { + let f = self.row_max_b[k] as i64; + let g = self.row_max[qlen - k] as i64; + let total = f + g; + let diff = (f - g).abs(); + if total > best || (total == best && diff < best_diff) { + best = total; + best_diff = diff; + best_k = k; + } + } + let k = best_k; + + // --- the right arm: query[..k] from right_reference[0], stopping wherever it scored --- + // + // Pass A's triple, explicitly - this reads before the `tb`/`tb_base` swap below, so the + // `_b` slots are the ones holding A. + let rj = if BANDED { + arg_for_row_banded( + &self.nm_b, + &self.tb_base_b, + self.row0_arg_b, + k, + qlen, + right_reference.len(), + w, + ) + } else { + arg_for_row( + &self.nm_b, + &self.tb_base_b, + self.row0_arg_b, + k, + right_reference.len(), + ) + }; + std::mem::swap(&mut self.tb, &mut self.tb_b); + std::mem::swap(&mut self.tb_base, &mut self.tb_base_b); + let right = AlignmentResult { + score: self.row_max_b[k], + query_start: 0, + query_end: k, + ref_start: 0, + ref_end: rj, + cigar: self.replay::(query, right_reference, k, rj, w), + }; + std::mem::swap(&mut self.tb, &mut self.tb_b); + std::mem::swap(&mut self.tb_base, &mut self.tb_base_b); + + // --- the left arm: query[k..] ending at the end of left_reference --- + // + // The replay runs in the REVERSED frame, walking from (qlen-k, lj) back to the origin. + // Row `i` there is `query[qlen-i]`, so starting at `qlen-k` covers exactly `query[k..]`, + // and the CIGAR comes out in forward order without a flip. + let llen = left_reference.len(); + // Pass B's triple: the swaps above have been undone, so the primary slots hold B. + let lj = if BANDED { + arg_for_row_banded( + &self.nm, + &self.tb_base, + self.row0_arg, + qlen - k, + qlen, + llen, + w, + ) + } else { + arg_for_row(&self.nm, &self.tb_base, self.row0_arg, qlen - k, llen) + }; + let left = AlignmentResult { + score: self.row_max[qlen - k], + query_start: k, + query_end: qlen, + ref_start: llen - lj, + ref_end: llen, + cigar: self.replay::(query, left_reference, qlen - k, lj, w), + }; + + SplitReferenceAlignment { + score: left.score + right.score, + left, + right, + } + } + + /// One `TRACK_ROW_MAX` pass, leaving `row_max`/`nm`/`tb` populated. + /// + /// `ARG_LARGEST = false`: the row maxima pin the **smallest** `j`. + /// + /// `w` is read only when `BANDED`. Under a band **only rows `0..=min(qlen, rlen + w)` are + /// filled** - row `i`'s in-band columns are `[i - w, i + w]` intersected with `[0, rlen]`, + /// which is empty past that - so every caller must bound its scan over `row_max` by that + /// same `last_row`. Reading past it reads whatever the previous call left. This is the same + /// rule `local` states at its own reduce, and it is the caller's to keep because `row_pass` + /// does not know which rows its caller intends to ask about. + #[allow(clippy::too_many_arguments)] // a kernel entry point: the scheme plus the bandwidth + fn row_pass( + &mut self, + query: &[u8], + refseq: &[u8], + match_score: u8, + mismatch: u8, + gap_open: u8, + gap_extend: u8, + w: usize, + ) { + let (qlen, rlen) = (query.len(), refseq.len()); + if rlen == 0 { + // No reference to align against: every row's best is its own column-0 cell, i.e. + // spanning that much query as one pure insertion. Column 0 is load-bearing here - + // for a query segment neither reference explains, it is the ONLY candidate. + if self.row_max.len() < qlen + 1 + LANES { + self.row_max.resize(qlen + 1 + LANES, 0); + } + for i in 0..=qlen { + self.row_max[i] = if i == 0 { + 0 + } else { + -(gap_open as i32) - (i as i32 - 1) * gap_extend as i32 + }; + } + // Every row's best is its column-0 cell, so every argmax is 0 - including row 0's, + // which is carried scalar and would otherwise be whatever the previous call left. + self.row0_arg = 0; + self.tb_base.clear(); + return; + } + self.prepare::(query, refseq, match_score, gap_open, gap_extend); + self.size_flags::(qlen, rlen, w); + unsafe { + self.run::( + qlen, + rlen, + match_score, + mismatch, + gap_open, + gap_extend, + w, + ) + }; + } + + /// The kernel behind every mode above. + /// + /// `TRACK_LAST_ROW` argmaxes `H[qlen][j]` instead of pinning the end cell at + /// `(qlen, rlen)`; `REVERSED` runs the whole thing back-to-front. **Both are `const` + /// generics, never runtime flags**, so that adding a mode costs the others nothing and each + /// mode monomorphizes to exactly the code it needs. + fn align( + &mut self, + query: &[u8], + refseq: &[u8], + match_score: u8, + mismatch: u8, + gap_open: u8, + gap_extend: u8, + ) -> AlignmentResult { + let (qlen, rlen) = (query.len(), refseq.len()); + + // A zero-length side never reaches the kernel: there is no anti-diagonal to walk. + if qlen == 0 || rlen == 0 { + return self.degenerate::(qlen, rlen, gap_open, gap_extend); + } + + self.prepare::(query, refseq, match_score, gap_open, gap_extend); + self.size_flags::(qlen, rlen, 0); + + let (score, end_j) = unsafe { + self.run::( + qlen, + rlen, + match_score, + mismatch, + gap_open, + gap_extend, + 0, + ) + }; + let cigar = self.replay::(query, refseq, qlen, end_j, 0); + let (ref_start, ref_end) = ref_span::(rlen, end_j); + + AlignmentResult { + score, + query_start: 0, + query_end: qlen, + ref_start, + ref_end, + cigar, + } + } + + /// The `qlen == 0` / `rlen == 0` cases, which skip the kernel entirely. + /// + /// Not merely a fast path - the modes genuinely differ here. With an empty query, `global` + /// must still consume the reference (`rlen` deletions), while `local_reference_end` stops + /// at `j = 0` and returns the empty alignment, because the reference tail is free. + fn degenerate( + &mut self, + qlen: usize, + rlen: usize, + gap_open: u8, + gap_extend: u8, + ) -> AlignmentResult { + let gap = |k: usize| -> i32 { + if k == 0 { + 0 + } else { + -(gap_open as i32) - (k as i32 - 1) * gap_extend as i32 + } + }; + let mut cigar = Cigar::default(); + if qlen > 0 { + // The query must always be spanned in full, in every mode here. + cigar.push(CigarOperation::Insertion, qlen); + let (ref_start, ref_end) = ref_span::(rlen, 0); + return AlignmentResult { + score: gap(qlen), + query_start: 0, + query_end: qlen, + ref_start, + ref_end, + cigar, + }; + } + // qlen == 0. Global must still consume the reference; a free reference end need not. + let consumed = if TRACK_LAST_ROW { 0 } else { rlen }; + if consumed > 0 { + cigar.push(CigarOperation::Deletion, consumed); + } + let (ref_start, ref_end) = ref_span::(rlen, consumed); + AlignmentResult { + score: gap(consumed), + query_start: 0, + query_end: 0, + ref_start, + ref_end, + cigar, + } + } + + /// Where cell `(i, j)`'s flags live, as `(word base, lane bit)`, with `p = i + j`. + /// + /// The address, not the flags: the replay's `H` state reads one word and sometimes two, + /// and the `E`/`F` states read exactly one, so returning all four packed would build bits + /// nobody looks at. The lane arithmetic is paid once per cell, the loads once per question, + /// and the four words stay adjacent so they still share a cache line. Most of the replay's + /// cost is the lane arithmetic and the branches rather than the loads, so this is a small + /// win on a serial walk - see `replay`. + #[inline(always)] + fn flag_at( + &self, + p: usize, + i: usize, + qlen: usize, + rlen: usize, + w: usize, + ) -> (*const u32, u32) { + // `saturating_sub`, and it is load-bearing: `p - rlen` underflows for every + // anti-diagonal before the band starts sliding right, which sends `lane` negative and + // the unchecked read below off the end. Caught by a segfault, not by a wrong answer. + // + // The lane is `i - lo(p)`, so it moves with the band: the fill starts each + // anti-diagonal's chunks at the band's top row, not the matrix's. Reading a banded + // fill's flags with the unbanded `lo` would silently address the wrong cell. + let (i_from, _) = band_rows::(p, qlen, rlen, w); + let lane = i - i_from; + debug_assert!( + i >= i_from && (self.tb_base[p] as usize + lane / LANES) * TB_WORDS < self.tb.len(), + "flag_at({p}, {i}) is outside the region the fill wrote" + ); + // SAFETY: `p` is a live anti-diagonal and `i` one of its interior rows, so `base` is + // within the region the fill wrote. The replay only ever asks about cells it walked to, + // and the same argument covers `tb_base[p]`: `p <= qlen + rlen`, which is what + // `size_flags` sized it to. + unsafe { + let base = (*self.tb_base.get_unchecked(p) as usize + lane / LANES) * TB_WORDS; + (self.tb.as_ptr().add(base), 1u32 << (lane % LANES)) + } + } + + /// Walk the flags back from `(qlen, rlen)` to `(0, 0)` and build the CIGAR. + /// + /// Scalar, and unavoidably so - a traceback is a serial walk down one path. Only the flag + /// *generation* vectorises. + /// + /// # Diagonal-first priority minimises the edit distance + /// + /// Many alignments tie for optimal, and the score does not choose between them - this does. + /// Among equally-optimal paths a caller wants the **fewest edits**: `10X` rather than + /// `10I10D` at the same score, because ten substituted bases are what happened and a + /// twenty-base structural event is not. + /// + /// Checking the diagonal first is what delivers that. It is **exact for strongly affine + /// schemes** such as the default, where `gap_open` is large enough against `gap_extend` that + /// the gappy alternative rarely ties the clean one at all. It is a heuristic outside them: a + /// greedy backward walk commits to a tied step knowing nothing about what it costs + /// downstream, so for schemes whose `gap_open` approaches `gap_extend` - where ties are + /// dense - it can occasionally return an optimal alignment that is not the cleanest one. + /// + /// Exactness in general would need the lexicographic `(score, -edits)` DP, and the edit + /// count is the one quantity here with **no scheme-only bound**: its differences grow with + /// sequence length, where every value in the difference recurrence is bounded by the scheme + /// alone. It would need an i32 lane and its own per-row state including anti-diagonal `p-2`, + /// which this kernel never keeps. + /// + /// # The trap, and why `from_diag` is a stored bit + /// + /// Diagonal-first **cannot** be spelled "if neither gap achieved `A_G`'s max, go diagonal": + /// that fallthrough is unreachable on precisely the ties this exists to resolve. It needs + /// its own bit, and that bit needs the clamp correction - see [`TB_FROM_DIAG`]. + /// + /// # Under a band the replay stays inside it, and does not check that it does + /// + /// Every step is decided by a flag, and the fill is what makes the flags safe to follow: + /// an out-of-band predecessor is sealed to `0`, the smallest value a lane can hold, so it + /// loses `A_G`'s max and its direction is never the one recorded. Two of the four + /// directions need an argument beyond "it lost the max", because a sealed `0` can *tie* an + /// `A_G` that is itself `0`: + /// + /// - **The band's lower edge is safe for free.** There `A_G = max(s_G, ΔE'_G, 0)`. If + /// `A_G == 0` then `ΔE'_G == 0` too, so `from_e` is set and the replay leaves via E + /// before it can fall through to F. If `A_G > 0` then either E achieved it (`from_e` + /// set) or `s_G` did, and an `s_G` that is strictly positive is one that did not clamp, + /// so `from_diag` is set and trustworthy. One of the two always fires, and both lead + /// inwards. + /// - **The band's upper edge needs one bit cleared.** There the fallthrough F leads + /// *inwards* and E leads out, so the fill clears `from_e` on that lane - see the mask in + /// `run`. Discarding it costs nothing: it can only ever have been the seal tying `A_G`. + /// + /// So `(i, j)` is always in band, and `|i - j| <= w` with `1 <= i, j` is exactly the + /// condition under which `flag_at` addresses a cell the fill wrote. The two closed-form + /// exits are covered too: the replay can only reach `(i, 0)` or `(0, j)` in band, which + /// bounds `i` and `j` by `w`, and the straight run to the origin it emits there is in band + /// the whole way. + fn replay( + &self, + query: &[u8], + refseq: &[u8], + start_i: usize, + start_j: usize, + w: usize, + ) -> Cigar { + let (qlen, rlen) = (query.len(), refseq.len()); + let mut cigar = Cigar::with_capacity(16); + + enum St { + H, + E, + F, + } + let (mut i, mut j) = (start_i, start_j); + let mut st = St::H; + + // The current run, held in registers and flushed only when the operation changes. + // + // `Cigar::push` already merges a repeated op into the run in progress - but it is the + // **host's** type and reaches it through the `Vec`: two loads, a compare and a store, + // per cell, to add one to a counter. Doing that here makes it a register compare and an + // add, and `push` is then called once per *run* rather than once per *cell*. The + // emitted CIGAR is identical (`push` merges either way; consecutive runs here differ by + // construction, so it only ever takes its append path). + // + // `run_len == 0` means "nothing pending", which is why `run_op`'s initial value is + // never read - guarded by that check, not by an `Option`, so the hot path has no + // discriminant to test. + let mut run_op = CigarOperation::Eq; + let mut run_len = 0usize; + macro_rules! emit { + ($op:expr, $n:expr) => {{ + let op = $op; + if run_len != 0 && op == run_op { + run_len += $n; + } else { + if run_len != 0 { + cigar.push(run_op, run_len); + } + run_op = op; + run_len = $n; + } + }}; + } + + while i > 0 || j > 0 { + // The boundaries carry no flags and need none: down column 0 the only path is all + // insertions, along row 0 all deletions, whatever state we arrived in. And it is + // the *whole* remaining path, so emit it in closed form and stop - the loop used to + // walk it one base at a time to reach the same answer. + if j == 0 { + emit!(CigarOperation::Insertion, i); + break; + } + if i == 0 { + emit!(CigarOperation::Deletion, j); + break; + } + let p = i + j; + // SAFETY: `(w, bit)` addresses cell (i, j)'s flag words, which the fill wrote; see + // `flag_at`. Each `read` below is one of the four adjacent words at that address. + debug_assert!( + !BANDED || i.abs_diff(j) <= w, + "the replay stepped out of the band of {w}, to ({i}, {j})" + ); + let (word, bit) = self.flag_at::(p, i, qlen, rlen, w); + let flag = |which: usize| unsafe { (*word.add(which) & bit) != 0 }; + match st { + St::H => { + // Diagonal FIRST: on a tie it spends one edit where a gap pair would spend + // two or more. See the tie-break note above. + // + // `from_f` is the fallthrough and is never stored: `A_G` is a max of three + // operands, so if neither the diagonal nor E achieved it, `ΔF'_G` did. + if flag(TB_FROM_DIAG) { + // `=` vs `X` is **literal byte identity**, re-derived from the original + // bytes rather than from how the kernel scored the cell. `U` scores as + // `T` but reports `X` against one; an unknown byte never matches, yet + // reports `=` against a byte-identical copy of itself. See `mod.rs`, + // "Alphabet" - this decoupling is deliberate. + // `=` vs `X` in the ORIGINAL frame. Under REVERSED the kernel ran on + // back-to-front codes, so cell (i, j) is the pair + // (query[qlen-i], refseq[rlen-j]) in the caller's bytes. + let (qb, rb) = if REVERSED { + (query[qlen - i], refseq[rlen - j]) + } else { + (query[i - 1], refseq[j - 1]) + }; + let op = if qb == rb { + CigarOperation::Eq + } else { + CigarOperation::X + }; + emit!(op, 1); + i -= 1; + j -= 1; + } else if flag(TB_FROM_E) { + // S[i,j] <- E[i-1,j]: the step consumes query base i. + emit!(CigarOperation::Insertion, 1); + i -= 1; + st = St::E; + } else { + emit!(CigarOperation::Deletion, 1); + j -= 1; + st = St::F; + } + } + St::E => { + if flag(TB_E_OPEN) { + // The run opened here: drop onto the score matrix without moving. + st = St::H; + } else { + emit!(CigarOperation::Insertion, 1); + i -= 1; + } + } + St::F => { + if flag(TB_F_OPEN) { + st = St::H; + } else { + emit!(CigarOperation::Deletion, 1); + j -= 1; + } + } + } + } + // The last run has no following op to flush it. An empty alignment (start_i == start_j + // == 0, which `local_end` returns for a hopeless extension) never enters the loop and + // leaves `run_len == 0`, so this correctly emits nothing. + if run_len != 0 { + cigar.push(run_op, run_len); + } + + // The replay walks backwards and so builds the CIGAR back-to-front, needing one flip. + // In the REVERSED frame, back-to-front IS forward order - so the flip is skipped and + // the two reversals annihilate. This is what makes the `_start` modes free. + if !REVERSED { + cigar.reverse(); + } + cigar + } + + #[target_feature(enable = "avx2")] + /// # `TRACK_ROW_MAX`: per-cell absolute scores, exactly, with nothing to bound + /// + /// `local_end`, `local_start` and `split_reference` argmax over **absolute** `H[i][j]`, + /// which a difference kernel does not have. The paper's answer is a three-level + /// decomposition whose 8-bit small delta **its own authors cannot bound** ("too complex for + /// rigorous analysis"). Under a contract that reads "exact, always", that is unusable. + /// + /// We do not need it. `S[i,j] - S[i,j-1] = ΔV[i,j]` *is* the definition of ΔV, and on + /// anti-diagonal `p` row `i` sits at column `j = p - i` while on `p-1` it sits at `j-1`. So + /// `abs[i] += ΔV[i, p-i]` - one add per cell, from a value the fill already stores. Exact, + /// per-lane (no scan, no cross-lane dependency), seeded from `S[i,0]` in closed form. + /// + /// The accumulator is `i32` and **that is not a fallback to a wider kernel**: the DP stays + /// u8 at 32 lanes and `fits_u8_lanes` still has no length term. It is a reconstruction + /// riding alongside. Its own bound is `(qlen + rlen) * max(match, gap_open) < i32::MAX`, + /// i.e. `qlen + rlen < ~8.4M` at the worst possible u8 scores - and **memory binds first by + /// an enormous margin**: an 8.4M-base alignment needs ~35 TB of traceback flags. There is + /// no reachable input where this is the limit. + /// + /// The seed also disposes of a hazard: the partial chunk tail writes garbage to rows + /// *above* the live range, and those rows **are** read later. But garbage only lands on row + /// `i` at `p <= i`, and the boundary seed at `p = i` scrubs it - so still no mask, only the + /// write ordering already required. + /// + /// `ARG_LARGEST` picks the tie-break on `j`: `local_end` wants the **largest** (longest + /// extension), `split_reference`'s row maxima want the **smallest**. + /// + /// # `BANDED`: the strip, and the two rows either side of it that pay for it + /// + /// `BANDED` narrows each anti-diagonal's interior to `|i - j| <= w` (see [`band_rows`]). + /// The chunk loop does not change at all - it is handed a shorter range - and that is the + /// whole runtime win. What the band costs is at its **edges**, where the recurrence reads + /// cells the fill never visited, and it costs three things: + /// + /// 1. **Two sealed bytes per anti-diagonal.** `A_G = max(s_G, ΔE'_G[i-1,j], ΔF'_G[i,j-1])` + /// reaches one row above the strip's top and one row below its bottom. Both are sealed + /// to `0`, the smallest value a `u8` lane can hold, so they lose every max they enter - + /// which is exactly what "unreachable" has to mean here. **This is not a sentinel**: it + /// is the bottom of the range every delta already lives in, so nothing about the + /// `[0, upper]` invariant, the plain `sub_epi8`s, or `fits_u8_lanes` moves an inch. + /// + /// 2. **Two garbage deltas per anti-diagonal, which are provably dead.** The *other* two + /// reads at each edge - `ΔV_G[i-1,j]` at the top, `ΔH_G[i,j-1]` at the bottom - are + /// subtrahends, and a seal cannot make a difference correct. So they are left as + /// garbage, and the band's geometry retires them: each bound moves at half of `p`'s + /// rate, so an edge that was flat at `p` **must** advance at `p+1`, and advancing is + /// precisely what makes the cell that consumed the garbage fall outside the next + /// anti-diagonal's reads. The values are computed and never read. + /// + /// 3. **One seed per row, replacing the column-0 one.** `abs[i]` accumulates `ΔV` along + /// row `i`, so it needs a starting absolute. Unbanded that is `H[i][0]`, seeded at the + /// matrix corner. Under a band, row `i > w` has no column-0 cell - it is `i` off the + /// diagonal - and enters at column `i - w` instead, where `ΔV`'s left neighbour is + /// outside the strip and its own `ΔV` is therefore the garbage of (2). But `ΔH` at that + /// same cell is **not** garbage: `ΔH` reads *upwards*, and the cell above is in band. So + /// `abs[i] = abs[i-1] + ΔH_G[i, i-w] - gap_open` seeds the row exactly, from a value the + /// fill already computed, one anti-diagonal's `abs` and one scalar add. + /// + /// That seed doubles as the scrub, and has to: the partial chunk's tail writes garbage to + /// rows *below* the strip, and under a band those rows are reached far later than the + /// unbanded corner seed would have scrubbed them. The row is seeded at the moment it enters + /// the band, which is after the last garbage write to it and before its first real read. + /// # Which mode compiles which flags + /// + /// Every public mode reaches this function through one of three wrappers. `REVERSED` is + /// threaded separately (see `prepare`) and never reaches `run`. + /// + /// | mode | wrapper | TRACE | LAST_ROW | ROW_MAX | ARG_LARGEST | BANDED | score read from | + /// |---|---|---|---|---|---|---|---| + /// | `global`, exact | `align` | y | - | - | - | - | `anchor` | + /// | `global`, banded | `row_pass` | y | - | y | - | y | `abs[qlen]` | + /// | `local_reference_*`, exact | `align` | y | y | - | - | - | `best_score` | + /// | `local_reference_*`, banded | `row_pass` | y | - | y | - | y | `row_max[qlen]` | + /// | `local_end` / `local_start` | `local` | y | - | y | y | opt | reduce over `row_max` | + /// | `split_reference` (two passes) | `row_pass` | y | - | y | - | opt | `row_max_b[k] + row_max[qlen-k]` | + /// + /// **`TRACK_LAST_ROW` and `BANDED` are never compiled together**: `align` is the only user + /// of the former and is never banded, because its `anchor` walk cannot survive a band (see + /// [`U8Probe::local_reference`]). + /// + #[allow(clippy::too_many_arguments)] // the scheme is five scalars, plus the bandwidth + unsafe fn run< + const TRACE: bool, + const TRACK_LAST_ROW: bool, + const TRACK_ROW_MAX: bool, + const ARG_LARGEST: bool, + const BANDED: bool, + >( + &mut self, + qlen: usize, + rlen: usize, + match_score: u8, + mismatch: u8, + gap_open: u8, + gap_extend: u8, + w: usize, + ) -> (i32, usize) { + unsafe { + let m = match_score as i64; + let mm = mismatch as i64; + let go = gap_open as i64 - gap_extend as i64; + let ge = gap_extend as i64; + let c = 2 * (go + ge); // == 2 * gap_open + let upper = (m + c) as u8; + + // The two values of s_G, broadcast once. Only the mismatch side needs the clamp: + // s_G(match) = M + C is exactly `upper`, which `fits_u8_lanes` already bounds and + // which cannot be negative. s_G(mismatch) = C - mismatch **can** be negative (any + // scheme with mismatch > 2*gap_open), and clamping it to zero is lossless because + // A_G only ever maxes it against ΔE'_G/ΔF'_G, which are >= 0. Wrapping instead + // would make the worst substitution score the largest value in the lane, win every + // max, and return a score BETTER than optimal - the signature of a sentinel bug. + let s_g_match = upper; + let s_g_mismatch = (c - mm).max(0) as u8; + let go_u8 = go as u8; + + // ΔV = ΔV_G - gap_open, un-offset, widened to the accumulator's lane type. + let gap_open_vec32 = _mm256_set1_epi32(gap_open as i32); + + let match_vec = _mm256_set1_epi8(s_g_match as i8); + let mismatch_vec = _mm256_set1_epi8(s_g_mismatch as i8); + let go_vec = _mm256_set1_epi8(go_u8 as i8); + // All-ones exactly when `s_G(mismatch)` did NOT clamp, i.e. when `cmpeq(A_G, s_G)` + // means what it says at every cell. Loop-invariant, and a whole-vector constant, so + // the `or` it feeds costs one uop on any port and sits off the critical path + // (`a_g -> de_next/df_next -> store`). See TB_FROM_DIAG. + let unclamped_vec = if c - mm >= 0 { + _mm256_set1_epi8(-1) + } else { + _mm256_setzero_si256() + }; + + let q_ptr = self.q_codes.as_ptr(); + let r_ptr = self.r_rev.as_ptr(); + + // Hoist every buffer to a raw pointer, ONCE, and ping-pong the pointers rather + // than the `Vec`s. + // + // `mem::swap(&mut self.prev, &mut self.cur)` would swap two `Deltas`, i.e. eight + // `Vec` headers per anti-diagonal, which dominates at small `n` where the ramp runs + // more anti-diagonals than chunks of actual DP. Four pointer swaps are register + // moves. + let mut p_dh = self.prev.dh.as_mut_ptr(); + let mut p_dv = self.prev.dv.as_mut_ptr(); + let mut p_de = self.prev.de.as_mut_ptr(); + let mut p_df = self.prev.df.as_mut_ptr(); + let mut c_dh = self.cur.dh.as_mut_ptr(); + let mut c_dv = self.cur.dv.as_mut_ptr(); + let mut c_de = self.cur.de.as_mut_ptr(); + let mut c_df = self.cur.df.as_mut_ptr(); + let tb_ptr = self.tb.as_mut_ptr(); + let tb_base_ptr = self.tb_base.as_mut_ptr(); + let abs_ptr = self.abs.as_mut_ptr(); + let row_max_ptr = self.row_max.as_mut_ptr(); + let nm_ptr = self.nm.as_mut_ptr(); + + // The running absolute value at the anti-diagonal's bottom row: the paper's 64-bit + // large offset, reduced to what a global score actually needs - one i64, not a + // per-lane middle delta. + let mut anchor: i64 = 0; + // `local_reference_end`'s objective, for free. + // + // `anchor` holds the absolute score at the anti-diagonal's *bottom* row, + // hi(p) = min(qlen, p). So for every p >= qlen it is exactly S[qlen][p - qlen] - + // **it already walks the last row**, one cell per anti-diagonal, which is precisely + // the row this mode maximizes. No per-cell absolute values, no extra vector work: + // one scalar compare per anti-diagonal, and only for p >= qlen. + // + // Ties resolve to the smallest `ref_end`, which falls out of updating on a strict + // improvement only, since p (and therefore j) increases monotonically. + let mut best_score = i64::MIN; + let mut best_j = 0usize; + + // Row 0, in closed form. + // + // `H[0][j]` is "consume j reference bases, align no query": 0 at j=0, and + // `-(gap_open + (j-1)*gap_extend) <= 0` beyond it. Both gap costs are unsigned, so + // **row 0's max is always 0**, attained at j=0, and the only question is which `j` + // the tie-break names: + // + // * `!ARG_LARGEST` wants the smallest `j`, and j=0 attains the max, so: 0. + // * `gap_open > 0` makes every j>=1 strictly negative, so: 0. + // * `gap_open == 0 < gap_extend` leaves `H[0][1] == 0` tying with j=0 and + // everything past it negative, so: 1. + // * both zero ties the whole row at 0, so the largest j is `rlen` - and that is + // not a curiosity, it is the free-gap case `local_end` has a test for. + // + // `abs[0]` is deliberately not seeded: nothing reads it. The group loop starts at + // `i_from`, which is `>= 1` always, so row 0 has no lane and never enters it. + // The last case is the one the band touches: with both gap costs zero the whole of + // row 0 ties at 0, and the largest tying `j` is the furthest cell the band still + // reaches, `min(rlen, w)`, not the end of the reference. + let row0_arg: i32 = if !ARG_LARGEST || gap_open > 0 { + 0 + } else if gap_extend > 0 { + 1.min(rlen) as i32 + } else if BANDED { + rlen.min(w) as i32 + } else { + rlen as i32 + }; + if TRACK_ROW_MAX { + *row_max_ptr.add(0) = 0; + } + // The running chunk counter, so `tb_base` is filled by the fill itself rather + // than by a separate O(qlen+rlen) pass. + // + // The flag address is recomputed per chunk rather than carried as a bumped + // pointer: a live `*mut u32` across the chunk loop stops LLVM proving it does not + // alias the delta buffers, so it stops keeping their pointers in registers. + let mut chunk: usize = 0; + + // The interior row range, maintained INCREMENTALLY rather than recomputed. + // + // `i_from = max(1, p - rlen)` and `i_to = min(qlen, p - 1)` are three scalar ops + // per anti-diagonal, charged on nearly every iteration of work at small `n`. + // + // Both are monotone step functions of `p`: `i_to` climbs to `qlen` and stops; + // `i_from` sits at 1 until `p > rlen + 1` and then climbs. So carry them and + // advance by one, with two perfectly-predicted branches (always-taken, then + // never-taken) instead of three ALU ops on the dependency path. + let mut i_from = 1usize; + let mut i_to = 0usize; + + // The band's own two bounds, carried the same way and for the same reason. + // + // `b_lo = ceil((p - w) / 2)` and `b_hi = floor((p + w) / 2)`. Each advances by one + // on alternating anti-diagonals - and, since their parities are opposite, **exactly + // one of the two advances every step**. So the pair costs one branch on a parity + // bit, not two divisions, on a path where the whole anti-diagonal may be only three + // chunks of real work. + // + // These are `band_rows`' two band terms evaluated at `p = 1`, written the way they + // are *defined* rather than the way they are shortest: clippy would have `(1 + w) / + // 2` as `w.div_ceil(2)`, which is the same number and hides that this is + // `floor((p + w) / 2)` at `p = 1`. The `debug_assert` against `band_rows` in the + // loop is what checks the pair, and it can only do that if they stay recognisable. + #[allow(clippy::manual_div_ceil)] + let (mut b_lo, mut b_hi) = if BANDED { + (1usize.saturating_sub(w).div_ceil(2), (1 + w) / 2) + } else { + (0, usize::MAX) + }; + // `(lo, hi)` at `p - 1`, which is what says whether each edge just advanced - and + // therefore whether the cell beyond it is one the last anti-diagonal wrote or one + // the seal has to cover. Seeded at `p = 0`'s empty range. + let (mut prev_lo, mut prev_hi) = (1usize, 0usize); + + for p in 1..=(qlen + rlen) { + debug_assert_eq!(i_from, 1.max(p.saturating_sub(rlen))); + debug_assert_eq!(i_to, qlen.min(p - 1)); + + // The interior, band included. Under `!BANDED` these fold to `i_from`/`i_to` + // and every branch guarded by `BANDED` below disappears with them. + let lo = if BANDED { i_from.max(b_lo) } else { i_from }; + let hi = if BANDED { i_to.min(b_hi) } else { i_to }; + debug_assert_eq!((lo, hi), band_rows::(p, qlen, rlen, w)); + + // Is this anti-diagonal's top (bottom) row sitting exactly ON the band's edge, + // so that the cell beyond it is outside the strip? + // + // Row `i` sits at `i - j = 2i - p`, so the top row is at `+w` exactly when it is + // the band's own bound *and* `p - w` is even - at odd `p - w` the bound rounds + // inwards and the row above is still in the strip. `lo == b_lo` is the "band, + // not the matrix corner, is what is limiting here" half; without it a top row + // pinned by `max(1, p - rlen)` would be sealed against a neighbour that is + // perfectly real. + // + // The two share a parity, and that is not a coincidence: the bounds move at + // half of `p`'s rate on opposite phases, so the strip is `w + 1` rows wide with + // both edges live, then `w` rows wide with neither. Every edge fix below fires + // on the wide ones and none on the narrow ones. + // + // `live` is not pedantry. An anti-diagonal with no interior still *has* band + // bounds, and they can sit anywhere - the strip runs off the end of the query + // long before `p` does whenever `qlen` is short against the band. Without this + // the edge fixes below would write a row nothing computed and, worse, `|=` a + // flag word at `chunk_base` that belongs to the *next* anti-diagonal, because an + // empty one advances `chunk` by zero. + let live = lo <= hi; + let edge_phase = BANDED && live && (p + w).is_multiple_of(2); + let top_sealed = edge_phase && lo == b_lo; + let bottom_sealed = edge_phase && hi == b_hi; + debug_assert!( + !(top_sealed && bottom_sealed) || lo < hi, + "both edges live means a strip {} rows wide, so they cannot be the same row", + w + 1 + ); + // Once the strip leaves the matrix it never comes back, so a live anti-diagonal + // always has a live predecessor to read from. `p <= 2` is the one exception, and + // it is the ordinary one: `p = 1` has no interior in any mode, banded or not, + // and `p = 2`'s reads land on the boundary cells rather than on `p = 1`'s + // interior. + debug_assert!( + !BANDED || !live || p <= 2 || prev_lo <= prev_hi, + "anti-diagonal {p} is live but {} was empty: its reads land on rows that \ + nothing wrote", + p - 1 + ); + + if TRACE { + *tb_base_ptr.add(p) = chunk as u32; + } + // `chunk` is advanced past this anti-diagonal before the TRACK_ROW_MAX pass + // runs, so that pass cannot use it directly - it needs THIS anti-diagonal's + // base, which is what the flags were written against. + let chunk_base = chunk; + + let mut i0 = lo; + while i0 <= hi { + // Two loads from anti-diagonal p-1, one byte apart. This is the trick that + // replaces the paper's shiftl/shiftr - the addressing does it. + let dv_up = _mm256_loadu_si256(p_dv.add(i0 - 1) as *const __m256i); + let de_up = _mm256_loadu_si256(p_de.add(i0 - 1) as *const __m256i); + let dh_left = _mm256_loadu_si256(p_dh.add(i0) as *const __m256i); + let df_left = _mm256_loadu_si256(p_df.add(i0) as *const __m256i); + + // `rlen + i0 - p`, never `rlen - p + i0`: the index is provably + // non-negative on every live lane, but `rlen - p` alone underflows the + // moment the anti-diagonal stops growing and starts sliding right. The + // grouping is load-bearing. + let qv = _mm256_loadu_si256(q_ptr.add(i0 - 1) as *const __m256i); + let rv = _mm256_loadu_si256(r_ptr.add(rlen + i0 - p) as *const __m256i); + // Hoisted out of the blend because the traceback needs it too: on a clamped + // scheme `from_diag` is only trustworthy at match cells. See TB_FROM_DIAG. + let eq_qr = _mm256_cmpeq_epi8(qv, rv); + let s_g = _mm256_blendv_epi8(mismatch_vec, match_vec, eq_qr); + + // Critical path 4: two maxes with nothing feeding them. ΔE'_G and ΔF'_G + // have already absorbed ΔV and ΔH, so A_G reads them directly rather than + // adding first. That is the entire reason Eq 3 exists. + let a_g = _mm256_max_epu8(_mm256_max_epu8(s_g, de_up), df_left); + + let de_next = _mm256_max_epu8(a_g, _mm256_add_epi8(de_up, go_vec)); + let df_next = _mm256_max_epu8(a_g, _mm256_add_epi8(df_left, go_vec)); + + // The traceback flags: equalities on values already in registers, extracted + // a bit per lane. `const TRACE` so the score-only path monomorphizes without + // them entirely. + if TRACE { + let tb = tb_ptr.add((chunk + (i0 - lo) / LANES) * TB_WORDS); + // `from_diag`/`from_e`: which operands achieved A_G's max. `from_f` is + // NOT stored - it is the fallthrough, recovered by the replay from "at + // least one of the three achieved the max". See TB_WORDS for why the + // diagonal and not a gap is the one that must be stored. + // + // `unclamped_vec` is all-ones for every scheme that does not clamp + // `s_G` (the default among them), which makes the `or` a no-op on the + // mask and the `and` an identity - so the correction is paid for only + // where it is needed, without a second monomorphization of this loop. + *tb.add(TB_FROM_DIAG) = _mm256_movemask_epi8(_mm256_and_si256( + _mm256_cmpeq_epi8(a_g, s_g), + _mm256_or_si256(eq_qr, unclamped_vec), + )) as u32; + *tb.add(TB_FROM_E) = + _mm256_movemask_epi8(_mm256_cmpeq_epi8(a_g, de_up)) as u32; + // `e_open`/`f_open`: whether the gap state opened from A_G rather than + // extending. Phrased against the max the fill already computed, so it + // is an equality and needs no unsigned compare (AVX2 has none for u8). + *tb.add(TB_E_OPEN) = + _mm256_movemask_epi8(_mm256_cmpeq_epi8(de_next, a_g)) as u32; + *tb.add(TB_F_OPEN) = + _mm256_movemask_epi8(_mm256_cmpeq_epi8(df_next, a_g)) as u32; + } + + // Plain subs, no saturation: every result is provably in [0, upper], so + // the minuend is always >= the subtrahend. No sentinel exists to cross. + _mm256_storeu_si256(c_dh.add(i0) as *mut __m256i, _mm256_sub_epi8(a_g, dv_up)); + _mm256_storeu_si256( + c_dv.add(i0) as *mut __m256i, + _mm256_sub_epi8(a_g, dh_left), + ); + _mm256_storeu_si256( + c_de.add(i0) as *mut __m256i, + _mm256_sub_epi8(de_next, dh_left), + ); + _mm256_storeu_si256( + c_df.add(i0) as *mut __m256i, + _mm256_sub_epi8(df_next, dv_up), + ); + + i0 += LANES; + } + if TRACE && lo <= hi { + chunk += (hi - lo) / LANES + 1; + } + + // Does row `hi` enter the band here, at its own first in-band column rather + // than at the matrix corner? Then its `abs` has to be seeded rather than + // accumulated - see (3) in this function's docs. + // + // `hi > prev_hi` says the row is new; `hi > w` says it is too far off the + // diagonal to have had a column-0 cell, which is what separates this from the + // corner seed below. The two are exclusive by construction and cover every row + // exactly once. + // `bottom_sealed` IS "row `hi` enters the band here": the band's lower bound + // reaches row `hi` for the first time exactly at `p = 2*hi - w`, which is the + // anti-diagonal where it is the bound and the parity is even. Rows within `w` of + // the origin enter at the matrix corner instead and are seeded there, and + // `bottom_sealed` provably never fires for them - it needs `hi = (p + w) / 2` + // and `hi <= p - 1`, which together force `hi > w`. + let seeds_row = bottom_sealed && TRACK_ROW_MAX; + debug_assert!( + !seeds_row || hi > w, + "a corner-seeded row cannot also band-enter" + ); + // `abs[hi - 1]` **before** the group loop overwrites it: it must be row + // `hi - 1`'s absolute at anti-diagonal `p - 1`, i.e. `S[hi-1, p-hi]`, which is + // the cell diagonally up-left of the one being seeded. One anti-diagonal later + // and it is the wrong cell. + // + // `hi > w >= 1` puts `hi - 1` at 1 or above, so this never reaches row 0 - + // which has no `abs` and needs none. + let seed_carry = if seeds_row { *abs_ptr.add(hi - 1) } else { 0 }; + + // The per-row absolutes and their running maxima, for the interior of this + // anti-diagonal. Deliberately a second pass over the same rows rather than + // work folded into the chunk loop: it needs `cur.dv`, which the chunk loop is + // still writing, and it is 8-lane i32 work against the loop's 32-lane u8 - + // interleaving them would spill registers in the hot loop for the benefit of + // the three modes that use this, and cost the three that do not. + if TRACK_ROW_MAX { + // One 8-lane i32 group. A macro rather than a closure so both call sites + // inline identically - `&mut self` plus intrinsics makes a closure here + // fight the borrow checker for nothing. + macro_rules! group { + ($base:expr) => {{ + let base = $base; + // 8 u8 deltas -> 8 i32, un-offset to true ΔV. + let dv = _mm256_cvtepu8_epi32(_mm_loadl_epi64( + c_dv.add(base) as *const __m128i + )); + let a = _mm256_add_epi32( + _mm256_loadu_si256(abs_ptr.add(base) as *const __m256i), + _mm256_sub_epi32(dv, gap_open_vec32), + ); + _mm256_storeu_si256(abs_ptr.add(base) as *mut __m256i, a); + + // No `j` or `i` is built here: the bit stream records *that* a cell + // took the max, not which, and `arg_for_row` reconstructs the + // coordinates from the bit's `(p, lane)` position afterwards. + let m = _mm256_loadu_si256(row_max_ptr.add(base) as *const __m256i); + + // STAY OFF PORT 5. `blendv_epi8` and the `cvtepu8_epi32` above both + // issue only on the shuffle port, and this block saturates it long + // before it runs out of instructions. So `max_epi32` takes the max + // with no blend at all, and `cmpeq(new, a)` recovers the mask - both + // on p0/p1. A blend is fewer instructions here and slower. + let new_m = _mm256_max_epi32(m, a); + // `cmpeq(max(m,a), a)` is `a >= m` - which IS the largest-j rule. + // The smallest-j rule needs a strict `a > m`. + let take = if ARG_LARGEST { + _mm256_cmpeq_epi32(new_m, a) + } else { + _mm256_cmpgt_epi32(a, m) + }; + _mm256_storeu_si256( + row_max_ptr.add(base) as *mut __m256i, + if ARG_LARGEST { + new_m + } else { + _mm256_max_epi32(m, a) + }, + ); + // One bit per lane instead of a `j` per lane, and the group + // *returns* its 8 bits rather than storing them - the caller packs + // all four into one word and stores once per 32 cells. See the + // store site for why the store count, not the store width, is what + // had to come down. + _mm256_movemask_ps(_mm256_castsi256_ps(take)) as u32 + }}; + } + + // A CONSTANT trip count, so the four groups unroll. Bounding it to skip the + // dead lanes of a partial chunk takes the unroll with it and costs more than + // the lanes are worth. + let mut i0 = lo; + while i0 <= hi { + // Four groups, four 8-bit masks, ONE 32-bit store for the chunk. + // + // One full word per chunk, not four narrow stores: `nm` is written once + // and never re-read, so it streams, and four 4-byte streaming stores per + // chunk miss D1 badly where one 32-bit store does not. + let m0 = group!(i0); + let m1 = group!(i0 + 8); + let m2 = group!(i0 + 16); + let m3 = group!(i0 + 24); + *nm_ptr.add(chunk_base + (i0 - lo) / LANES) = + m0 | (m1 << 8) | (m2 << 16) | (m3 << 24); + i0 += LANES; + } + } + + // The band's seals and its one seed, AFTER both loops and BEFORE the boundary + // cells. Both halves of that ordering are load-bearing: + // + // * after the chunk loop, for the reason the boundary cells are - a 32-lane + // store is unconditional, so the partial last chunk scribbles over the row + // just past the interior, which is exactly the row the bottom seal owns; + // * before the boundary cells, because at the matrix corners the seals aim at + // rows that ARE the boundary (`hi + 1 == p` while the anti-diagonal is still + // growing, `lo - 1 == 0` before it lifts off). The boundary is the truth + // there; letting it land second is cheaper than teaching the seals to duck. + if BANDED && live { + // --- the seals: the two cells the recurrence reaches past the strip --- + // + // `live` is what keeps these in bounds, and it is not a formality: once the + // strip slides off the bottom of the query, `lo` keeps climbing with `p` + // while the buffers stop at `qlen + 1 + LANES`, so `c_de[lo - 1]` walks + // straight off the end of the allocation. Nothing reads those anti-diagonals + // - the strip never re-enters the matrix, which is what the `debug_assert` + // above pins - so there is nothing to seal. (Found as heap corruption, not + // as a wrong answer.) + // + // Given `live`: `1 <= lo <= hi <= qlen`, so `lo - 1` and `hi + 1` are both + // inside every delta buffer. Where they land on a boundary cell that + // legitimately exists - row 0 while the band still reaches it, row `p` while + // the anti-diagonal is still growing - the boundary block below runs second + // and takes it back. + // + // `p = 1` has no interior in any mode and needs no seals: everything `p = 2` + // reads is a boundary cell, which the block below writes. + *c_de.add(lo - 1) = 0; + *c_df.add(hi + 1) = 0; + + // --- the edges: where a seal of `0` is not enough --- + // + // A seal makes the *three-way* max fall the right way: `0` is the smallest a + // lane can hold, so it loses `A_G = max(s_G, ΔE'_G, ΔF'_G)` and the + // out-of-band direction is never taken. + // + // **It does not survive the gap recurrence.** `ΔE'_G[i,j] = max(A_G, + // ΔE'_G[i-1,j] + go)` adds `go` to the seal *before* the max, so a sealed + // `0` arrives as `go` - and `go` beats `A_G` whenever `A_G < go`, which is + // most cells under a strongly affine scheme. The result: a `ΔE'_G` too large + // by `go - A_G`, and an `e_open` bit that says the gap extends from a cell + // outside the band. The score comes out wrong, and the replay walks up out + // of the strip and reads flags the fill never wrote. + // + // Both are fixed by saying what the band means: at the strip's top edge + // `E[i-1,j]` is unreachable, so the E-run **must open here**, and + // `ΔE'_G[i,j] = A_G[i,j] - ΔH_G[i,j-1]` - which is exactly what the chunk + // loop just stored in `ΔV_G[i,j]`. So the correction is a byte copy, not an + // arithmetic redo, and the F edge mirrors it: `ΔF'_G[i,j] = A_G[i,j] - + // ΔV_G[i-1,j]`, which is `ΔH_G[i,j]`. + // + // Each edge copies from the delta whose own inputs are in band - `ΔV_G` reads + // leftwards and the top edge's left neighbour is inside; `ΔH_G` reads upwards + // and the bottom edge's upper neighbour is inside. That is why this works at + // all, and why the copies do not cross: `top_sealed && bottom_sealed` implies + // a strip `w + 1 >= 2` rows wide, so `lo != hi`. + if top_sealed { + *c_de.add(lo) = *c_dv.add(lo); + if TRACE { + // In `St::H`: "H came from E[i-1,j]" - out of band, so clear it. It + // could only ever have been the seal tying an `A_G` of 0. The replay + // then falls through to F, which at this edge leads *inwards*. + *tb_ptr.add(chunk_base * TB_WORDS + TB_FROM_E) &= !1u32; + // In `St::E`: "the run opened here" - it must have. + *tb_ptr.add(chunk_base * TB_WORDS + TB_E_OPEN) |= 1u32; + } + } + if bottom_sealed { + *c_df.add(hi) = *c_dh.add(hi); + if TRACE { + // In `St::F`: "the run opened here". `from_f` itself needs no mask - + // it is the replay's fallthrough, and at this edge one of the two + // flags it falls through *from* is always set. See `replay`. + let lane = hi - lo; + *tb_ptr.add((chunk_base + lane / LANES) * TB_WORDS + TB_F_OPEN) |= + 1u32 << (lane % LANES); + } + } + + // Row `hi` enters the band at column `hi - w`: seed its running absolute + // from the cell above-left, which is in band, via a `ΔH` the chunk loop just + // computed from in-band inputs. See (3) in this function's docs. + // + // `row_max` is *assigned*, not maxed: this is the row's first candidate, and + // whatever is sitting there is the partial chunk's garbage. + if seeds_row { + let s = seed_carry + *c_dh.add(hi) as i32 - gap_open as i32; + *abs_ptr.add(hi) = s; + *row_max_ptr.add(hi) = s; + // No `nm` bit is written for it. `arg_for_row_banded` returns this + // column when its scan finds nothing above, which is precisely when this + // cell is still the row's argmax - so the bit's two possible values give + // the same answer. See that function. + } + } + + // Boundary cells, AFTER the chunk loop - see the hazard note in the module + // docs. The values are Eq 3's initial conditions. + // Under a band the boundary cells are **conditional**, and this is where the + // band's edge and the matrix's edge finally meet. `H[0][j]` sits `j` off the + // diagonal and `H[i][0]` sits `i` off it, so the boundary leaves the strip after + // `w` steps and past that these cells **do not exist**. Writing them anyway + // would hand the recurrence a real, reachable-looking predecessor outside the + // band - which is the one thing the seals are there to prevent, undone one line + // later by the code that runs after them. + // + // Past `w` the seals above are the truth, and they already cover the only two + // rows that get read: `c_de[lo - 1]` reaches row 0 while `lo` is still 1, and + // `c_df[hi + 1]` reaches row `p` while the anti-diagonal is still growing. + // Everything further out is never read at all - by the time `p > w + 1` the + // band no longer reaches either boundary. + if p <= rlen && (!BANDED || p <= w) { + let b = if p == 1 { 0 } else { go_u8 }; + *c_dv.add(0) = b; + *c_de.add(0) = b; + // Row 0's max and argmax are hoisted out of the loop - see the closed form + // above it. `abs[0]` goes with them: it is never read, because `i_from >= 1` + // means the group loop never touches row 0. + } + if p <= qlen && (!BANDED || p <= w) { + let b = if p == 1 { 0 } else { go_u8 }; + *c_dh.add(p) = b; + *c_df.add(p) = b; + // The corner seed, gated with the boundary it reads. Rows past `w` are + // seeded at their real first column instead, by `seeds_row` above; between + // the two, every row is seeded exactly once. + if TRACK_ROW_MAX { + // Row p enters the matrix here, at its column-0 cell. Seeding from + // S[p,0] is load-bearing twice over: H[i][0] - "span query[..i] as one + // insertion, consume no reference" - is a real argmax candidate and for + // a query segment nothing explains it is the only one; and this write + // is what scrubs the garbage the partial chunk tail left on this row at + // earlier anti-diagonals. + // No arg store: `j = 0` is what `arg_for_row` returns when no interior + // cell of the row ever took the max, so the column-0 seed is implicit. + let s_i0 = -(gap_open as i32) - (p as i32 - 1) * gap_extend as i32; + *abs_ptr.add(p) = s_i0; + *row_max_ptr.add(p) = s_i0; + } + } + + // While the anti-diagonal still grows downwards the bottom cell is (p, 0), a + // boundary known in closed form. Once it slides right the bottom row is pinned + // at qlen and each step adds one ΔV, un-offset. At p = qlen + rlen that lands + // exactly on S[qlen][rlen]. + // + // `anchor` is dead in the `TRACK_ROW_MAX` modes - `local` and `row_pass` + // discard `run`'s return value - but it needs no gate: `run` is private, every + // call site is visible under LTO, and the unused return value is enough for LLVM + // to kill the whole chain including the `c_dv[qlen]` load. + anchor = if p <= qlen { + if p == 1 { + -(gap_open as i64) + } else { + anchor - ge + } + } else { + anchor + *c_dv.add(qlen) as i64 - (go + ge) + }; + + if TRACK_LAST_ROW && p >= qlen && anchor > best_score { + best_score = anchor; + best_j = p - qlen; + } + + // Four register moves, not eight Vec headers. + std::mem::swap(&mut p_dh, &mut c_dh); + std::mem::swap(&mut p_dv, &mut c_dv); + std::mem::swap(&mut p_de, &mut c_de); + std::mem::swap(&mut p_df, &mut c_df); + + // Advance the row range for p+1. AFTER the body, not before: the first cut + // incremented on entry and so read min(qlen, p) where it wanted min(qlen, p-1) + // - caught immediately by the debug_assert above, which is why it is there. + if i_to < qlen { + i_to += 1; + } + if p >= rlen + 1 { + i_from += 1; + } + if BANDED { + (prev_lo, prev_hi) = (lo, hi); + // Exactly one of the two band edges advances per anti-diagonal, and which + // one is decided by a single parity bit: `b_hi = floor((p + w) / 2)` steps + // when `p + w` becomes even, `b_lo = ceil((p - w) / 2)` when it becomes odd. + // This is that test for `p + 1`. + // + // `p + 1 > w` is the saturation, carried rather than recomputed. Until the + // band's lower edge lifts off row 0 - `p <= w`, where `ceil((p - w) / 2)` + // is negative and `band_rows`' `saturating_sub` clamps it - it must not + // move at all. Without this the edge climbs from the very first + // anti-diagonal and the strip walks off the top of the matrix; the + // `debug_assert` against the closed form above is what says so. + if (p + 1 + w).is_multiple_of(2) { + b_hi += 1; + } else if p + 1 > w { + b_lo += 1; + } + } + } + + let _ = upper; + // The bound `size_flags` computed, against the chunks this fill actually wrote. + // `chunk` is one past the last one written, so `<=` is the exact condition. Every + // store above is unchecked, so a bound that is too tight is heap corruption rather + // than a wrong score - and the one time this crate got a band bound wrong, that is + // exactly how it showed up. + // + // Read the lengths here rather than hoisting them beside `tb_ptr`: this chunk loop + // is register-starved (see the note beside `chunk`), and down here the whole thing + // is inside the `debug_assert!` so a release build reads nothing at all. + debug_assert!( + !TRACE || chunk * TB_WORDS <= self.tb.len(), + "the fill wrote {chunk} chunks but tb holds {} ({TB_WORDS} words per chunk): \ + the chunk bound is too tight, and every store past it was heap corruption", + self.tb.len() / TB_WORDS, + ); + debug_assert!( + !(TRACE && TRACK_ROW_MAX) || chunk * NM_WORDS <= self.nm.len(), + "the fill wrote {chunk} chunks but nm holds {} - see the tb assert above", + self.nm.len() / NM_WORDS + ); + if TRACK_ROW_MAX { + self.row0_arg = row0_arg; + } + if TRACK_LAST_ROW { + (best_score as i32, best_j) + } else { + (anchor as i32, rlen) + } + } + } +} diff --git a/src/simdaligner/mod.rs b/src/simdaligner/mod.rs new file mode 100644 index 00000000..e91b910a --- /dev/null +++ b/src/simdaligner/mod.rs @@ -0,0 +1,160 @@ +//! Exact pairwise sequence alignment with affine gap costs, SIMD-only. +//! +//! The entry point is [`SimdAligner`]: build one with a [`Scores`] scheme (or +//! [`Scores::default()`]) and call an alignment method on it repeatedly. Construct +//! one per thread and reuse it - the scoring matrix and every scratch buffer live +//! inside it and are reused across calls. +//! +//! Five alignment types are named on a single axis - *which side is local*, i.e. free +//! to stop or start wherever it scores best: +//! +//! - [`SimdAligner::global_alignment`] - nothing is local; both sequences are aligned +//! end to end. +//! - [`SimdAligner::local_reference_end_alignment`] - the query is still spanned in +//! full, but the reference may stop anywhere and its trailing tail is free. +//! - [`SimdAligner::local_reference_start_alignment`] - the mirror: the reference may +//! *begin* anywhere and its leading prefix is free. +//! - [`SimdAligner::local_end_alignment`] - both sequences start at 0 and the +//! alignment ends wherever it scores best; both trailing tails are free, and +//! [`Scores::end_bonus`] nudges it towards covering the whole query. Its score is +//! never negative: the empty alignment is always available. +//! - [`SimdAligner::local_start_alignment`] - the mirror: the alignment *begins* +//! wherever it scores best, and both leading prefixes are free. +//! +//! And one on a second axis - *which side is split* - taking **two** references: +//! +//! - [`SimdAligner::split_reference_alignment`] - one query aligned across a +//! `left_reference` and a `right_reference` with a single jump between them, +//! returning two [`AlignmentResult`]s that partition the query exactly. Built for +//! spotting a tandem duplication that a linear aligner reported as an insertion. +//! +//! ``` +//! use strobealign::simdaligner::{SimdAligner, Scores}; +//! let mut aligner = SimdAligner::new(Scores::default()); +//! let result = aligner.global_alignment(b"ACGTACGT", b"ACGTTACGT", None); +//! assert_eq!(result.cigar.to_string(), "3=1D5="); +//! ``` +//! +//! Every alignment method takes a trailing `bandwidth: Option`: `None` aligns +//! exactly, `Some(w)` confines the fill to `w` diagonals either side of the anchor +//! diagonal. See [`SimdAligner`] for what `w` means per mode. +//! +//! # Alphabet +//! The alphabet is **A, C, G, T and U** (case-insensitive), and nothing else. `U` is RNA's +//! T-equivalent and scores identically to `T`. +//! +//! Any other byte - `N`, an IUPAC ambiguity code, a `-` from a gapped FASTA, a stray `\n` - +//! is **unknown**, and an unknown byte never scores as a match against anything, *including +//! another unknown byte*. Garbage in produces a mismatch rather than a silent false match. +//! +//! `=`/`X` in the CIGAR are **literal byte identity**, deliberately decoupled from how a +//! position scored. So `U` against `T` scores as a match but reports `X`, and an unknown byte +//! against a byte-identical copy of itself scores as a mismatch but reports `=`. Scoring is +//! about biology; `=`/`X` is about the bytes the caller handed us. +//! +//! # Guarantees +//! With `bandwidth: None` the alignment is *exact* at every length: the reported score is the +//! optimal score under the given scheme, and the CIGAR is an alignment achieving it. There is +//! no x-drop and no early exit. With `Some(w)` it is optimal among alignments staying inside +//! the band. Either way every input pair produces an alignment, including ones whose optimal +//! score is negative. +//! +//! `piecewisealigner.rs` passes `None` between anchors and `Some(w)` for the end extensions, so +//! "exact" is not a property of the aligner as strobealign uses it. +//! +//! # The kernel +//! One kernel - u8 lanes, 32 wide, on a difference recurrence - and every call takes it. There +//! is no scalar path and no fallback: every stored value is bounded by the **scoring scheme +//! alone**, with no length term, so no pair of lengths can leave it. AVX2 is required, so this +//! module is compiled only on x86_64; [`SimdAligner::new`] panics without AVX2 at runtime, and +//! [`SimdAligner::is_supported`] lets a caller check first. +//! +//! The scoring scheme *is* bounded - `match + 3*gap_open - gap_extend <= 255`, i.e. +//! `gap_open <= ~84` at `match = 2` - and [`SimdAligner::new`] refuses anything past it rather +//! than silently wrap a lane. + +// `piecewisealigner.rs` calls three of the six modes today; both `local_reference_*` modes and +// `split_reference_alignment` are tested but not yet wired up, and are kept for planned callers. +#![allow(dead_code, unused_imports)] + +#[cfg(target_arch = "x86_64")] +mod aligner; +#[cfg(target_arch = "x86_64")] +mod avx2_u8; + +#[cfg(all(test, target_arch = "x86_64"))] +mod tests; + +#[cfg(target_arch = "x86_64")] +pub use aligner::SimdAligner; + +// Re-exported so the kernel can reach them as `super::Cigar` without a conversion at the +// boundary. +pub use super::{ + aligner::Scores, + cigar::{Cigar, CigarOperation}, +}; + +/// Kernel internals, for tests. **Not part of the alignment API.** +#[doc(hidden)] +#[cfg(target_arch = "x86_64")] +pub mod internal { + pub use super::avx2_u8::{LANES as U8_LANES, U8Probe, fits_u8_lanes}; +} + +/// Result of aligning a query against a reference. +/// +/// `query_start`/`query_end`/`ref_start`/`ref_end` are half-open ranges into the +/// original slices, and the `cigar` covers exactly those ranges and nothing else - a +/// free end gap is **not** represented in the CIGAR. +/// +/// Which of the four coordinates can move depends on the alignment type: +/// +/// - [`SimdAligner::global_alignment`]: both spans cover their inputs in full. +/// - [`SimdAligner::local_reference_end_alignment`]: the query span is still the whole +/// query, but `ref_end` is wherever the alignment chose to stop. +/// - [`SimdAligner::local_reference_start_alignment`]: likewise, but it is `ref_start` +/// that moves and `ref_end` is always `reference.len()`. +/// - [`SimdAligner::local_end_alignment`]: `query_end` and `ref_end` both move. +/// - [`SimdAligner::local_start_alignment`]: `query_start` and `ref_start` both move. +/// +/// `Default` is the **empty alignment**: score 0, all four coordinates at 0, and an empty +/// CIGAR - a real result, returned by [`SimdAligner::local_end_alignment`] when no extension +/// scores better than stopping immediately. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct AlignmentResult { + pub score: i32, + pub query_start: usize, + pub query_end: usize, + pub ref_start: usize, + pub ref_end: usize, + pub cigar: Cigar, +} + +/// Result of [`SimdAligner::split_reference_alignment`]: one query aligned across two +/// references, with a single jump point between them. +/// +/// The two arms **partition the query exactly**: `right.query_start` is 0, +/// `left.query_end` is `query.len()`, and `right.query_end == left.query_start` - the +/// jump point. Each arm's CIGAR covers its own half of the query against its own +/// reference span, and the two CIGARs together consume every query base once. +/// +/// `right` is anchored at the start of `right_reference` (`ref_start == 0`) and may +/// stop anywhere; `left` is anchored at the end of `left_reference` +/// (`ref_end == left_reference.len()`) and may begin anywhere. +/// +/// `score` is `left.score + right.score`, and unlike the local types it **can be +/// negative** - the query must be covered whether or not either reference explains it. +/// An arm may be the empty alignment (score 0, no CIGAR), and an arm may consume no +/// reference at all, reporting its query segment as one pure insertion. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SplitReferenceAlignment { + /// The arm against `left_reference`: covers the query *suffix*, and ends at the + /// end of the left reference. + pub left: AlignmentResult, + /// The arm against `right_reference`: covers the query *prefix*, and starts at the + /// start of the right reference. + pub right: AlignmentResult, + /// `left.score + right.score`. + pub score: i32, +} diff --git a/src/simdaligner/tests.rs b/src/simdaligner/tests.rs new file mode 100644 index 00000000..7092eb10 --- /dev/null +++ b/src/simdaligner/tests.rs @@ -0,0 +1,450 @@ +//! Fast, hardcoded sanity tests for the AVX2 aligner. +//! +//! These are deliberately *not* an oracle suite: there is no scalar reference implementation +//! and nothing sweeps randomized inputs. They are hand-written cases with hand-checked expected +//! values, meant to catch an obvious break quickly. +//! +//! Two invariants are worth more than the individual cases, and every test that produces an +//! alignment routes through them via [`check`]: +//! +//! 1. **The CIGAR re-scores to the reported score.** A wrong traceback under a right score is +//! this kernel's characteristic failure (see `TB_FROM_DIAG` in `avx2_u8.rs`), and comparing +//! scores alone cannot see it. +//! 2. **The CIGAR's consumption matches the reported spans**, so the coordinates and the +//! operations cannot disagree about how much sequence was used. + +use super::{AlignmentResult, Scores, SimdAligner}; + +const SCORES: Scores = Scores { + match_: 2, + mismatch: 8, + gap_open: 12, + gap_extend: 1, + end_bonus: 10, +}; + +fn aligner() -> SimdAligner { + SimdAligner::new(SCORES) +} + +/// The kernel's scoring alphabet: A/C/G/T/U, case-insensitive. Anything else is unknown and +/// never matches, *including another unknown byte*. +fn scores_as_match(q: u8, r: u8) -> bool { + fn code(b: u8) -> Option { + match b.to_ascii_uppercase() { + b'A' => Some(0), + b'C' => Some(1), + b'G' => Some(2), + b'T' | b'U' => Some(3), + _ => None, + } + } + matches!((code(q), code(r)), (Some(a), Some(b)) if a == b) +} + +/// Parse a CIGAR string back into `(len, op)` pairs. +fn parse(cigar: &str) -> Vec<(usize, char)> { + let mut out = vec![]; + let mut len = 0usize; + for ch in cigar.chars() { + if let Some(d) = ch.to_digit(10) { + len = len * 10 + d as usize; + } else { + out.push((len, ch)); + len = 0; + } + } + assert_eq!( + len, 0, + "CIGAR ended with a length and no operation: {cigar}" + ); + out +} + +/// Re-score `result`'s CIGAR against the sequences it claims to align, and assert it agrees +/// with the reported score and spans. `end_bonus` is added by the caller for the modes that +/// earn it, since the CIGAR itself cannot show whether the query was spanned. +fn check(result: &AlignmentResult, query: &[u8], reference: &[u8], end_bonus: i32) { + let (mut qi, mut ri) = (result.query_start, result.ref_start); + let mut score = 0i32; + + for (len, op) in parse(&result.cigar.to_string()) { + match op { + '=' | 'X' | 'M' => { + for _ in 0..len { + let (q, r) = (query[qi], reference[ri]); + score += if scores_as_match(q, r) { + SCORES.match_ as i32 + } else { + -(SCORES.mismatch as i32) + }; + // `=`/`X` are literal byte identity, decoupled from how the cell scored. + if op != 'M' { + let identical = q.to_ascii_uppercase() == r.to_ascii_uppercase(); + assert_eq!( + identical, + op == '=', + "op {op} at query[{qi}]={} reference[{ri}]={} disagrees with byte identity", + q as char, + r as char + ); + } + qi += 1; + ri += 1; + } + } + 'I' => { + score -= SCORES.gap_open as i32 + (len as i32 - 1) * SCORES.gap_extend as i32; + qi += len; + } + 'D' => { + score -= SCORES.gap_open as i32 + (len as i32 - 1) * SCORES.gap_extend as i32; + ri += len; + } + other => panic!("unexpected CIGAR operation {other}"), + } + } + + assert_eq!( + qi, result.query_end, + "CIGAR consumed query up to {qi} but result says {}", + result.query_end + ); + assert_eq!( + ri, result.ref_end, + "CIGAR consumed reference up to {ri} but result says {}", + result.ref_end + ); + assert_eq!( + score + end_bonus, + result.score, + "CIGAR {} re-scores to {} but result says {}", + result.cigar, + score + end_bonus, + result.score + ); +} + +/// Assert a global alignment's score and CIGAR, and re-score it. +fn global_case(query: &[u8], reference: &[u8], score: i32, cigar: &str) { + let result = aligner().global_alignment(query, reference, None); + check(&result, query, reference, 0); + assert_eq!(result.score, score); + assert_eq!(result.cigar.to_string(), cigar); + assert_eq!((result.query_start, result.query_end), (0, query.len())); + assert_eq!((result.ref_start, result.ref_end), (0, reference.len())); +} + +// --------------------------------------------------------------------------------- +// Global. +// --------------------------------------------------------------------------------- + +#[test] +fn global_perfect_match() { + global_case(b"AAATTT", b"AAATTT", 6 * 2, "6="); +} + +#[test] +fn global_complete_mismatch() { + global_case(b"AAA", b"TTT", 3 * -8, "3X"); +} + +#[test] +fn global_single_mismatch() { + global_case(b"AAATAA", b"AAAAAA", 5 * 2 - 8, "3=1X2="); +} + +#[test] +fn global_gap_in_query() { + global_case(b"AAATTT", b"AAAATTTT", 6 * 2 - 12 - 1, "3=2D3="); +} + +#[test] +fn global_gap_in_reference() { + global_case(b"AAAATTTT", b"AAATTT", 6 * 2 - 12 - 1, "3=2I3="); +} + +// The four end-gap cases exercise the boundary cells, which the fill writes after the chunk +// loop rather than before - a hazard called out in the kernel's module docs. + +#[test] +fn global_gap_at_query_start() { + global_case(b"AAATTT", b"TTAAATTT", 6 * 2 - 12 - 1, "2D6="); +} + +#[test] +fn global_gap_at_query_end() { + global_case(b"AAATTT", b"AAATTTAA", 6 * 2 - 12 - 1, "6=2D"); +} + +#[test] +fn global_gap_at_reference_start() { + global_case(b"TTAAATTT", b"AAATTT", 6 * 2 - 12 - 1, "2I6="); +} + +#[test] +fn global_gap_at_reference_end() { + global_case(b"AAATTTAA", b"AAATTT", 6 * 2 - 12 - 1, "6=2I"); +} + +#[test] +fn global_multiple_mismatches() { + global_case( + b"ATATATATAT", + b"AAAAAAAAAA", + 5 * 2 - 5 * 8, + "1=1X1=1X1=1X1=1X1=1X", + ); +} + +#[test] +fn global_complex_alignment() { + global_case( + b"AAACTTAAACCTT", + b"AAAATTGAAATT", + 10 * 2 - 8 - 2 * 12 - 1, + "3=1X2=1D3=2I2=", + ); +} + +#[test] +fn global_empty_query_is_all_deletion() { + let result = aligner().global_alignment(b"", b"ACGT", None); + check(&result, b"", b"ACGT", 0); + assert_eq!(result.score, -(12 + 3)); + assert_eq!(result.cigar.to_string(), "4D"); +} + +#[test] +fn global_both_empty() { + let result = aligner().global_alignment(b"", b"", None); + assert_eq!(result.score, 0); + assert!(result.cigar.is_empty()); +} + +// --------------------------------------------------------------------------------- +// The local modes the piecewise extension path actually calls. +// --------------------------------------------------------------------------------- + +#[test] +fn local_end_stops_before_a_bad_tail() { + // The reference matches for 6 bases then diverges completely. Extending into the tail + // costs more than it earns, so the alignment stops at the divergence. + let (query, reference) = (b"AAATTTGGGGGG", b"AAATTTCCCCCC"); + let result = aligner().local_end_alignment(query, reference, None); + check(&result, query, reference, 0); + assert_eq!(result.score, 6 * 2); + assert_eq!(result.cigar.to_string(), "6="); + assert_eq!((result.query_start, result.query_end), (0, 6)); + assert_eq!((result.ref_start, result.ref_end), (0, 6)); +} + +#[test] +fn local_end_spans_the_query_for_the_end_bonus() { + // A single trailing mismatch is worth crossing: -8 for the mismatch against +10 of + // end_bonus for reaching the end of the query. + let (query, reference) = (b"AAATTTG", b"AAATTTC"); + let result = aligner().local_end_alignment(query, reference, None); + check(&result, query, reference, SCORES.end_bonus as i32); + assert_eq!(result.score, 6 * 2 - 8 + 10); + assert_eq!(result.cigar.to_string(), "6=1X"); + assert_eq!(result.query_end, 7); +} + +#[test] +fn local_end_never_scores_below_zero() { + // Nothing is worth aligning, so the empty alignment wins. + let result = aligner().local_end_alignment(b"GGGG", b"CCCC", None); + assert_eq!(result.score, 0); + assert!(result.cigar.is_empty()); + assert_eq!((result.query_end, result.ref_end), (0, 0)); +} + +#[test] +fn local_start_is_the_mirror_of_local_end() { + // Same shape as `local_end_stops_before_a_bad_tail`, reversed: the junk is a leading + // prefix and the alignment begins after it. + let (query, reference) = (b"GGGGGGAAATTT", b"CCCCCCAAATTT"); + let result = aligner().local_start_alignment(query, reference, None); + check(&result, query, reference, 0); + assert_eq!(result.score, 6 * 2); + assert_eq!(result.cigar.to_string(), "6="); + assert_eq!((result.query_start, result.query_end), (6, 12)); + assert_eq!((result.ref_start, result.ref_end), (6, 12)); +} + +#[test] +fn local_reference_end_spans_the_whole_query() { + // The query is spanned in full; only the reference's trailing tail is free. + let (query, reference) = (b"AAATTT", b"AAATTTGGGGGGGG"); + let result = aligner().local_reference_end_alignment(query, reference, None); + check(&result, query, reference, 0); + assert_eq!(result.score, 6 * 2); + assert_eq!((result.query_start, result.query_end), (0, 6)); + assert_eq!(result.ref_end, 6); +} + +#[test] +fn local_reference_start_spans_the_whole_query() { + let (query, reference) = (b"AAATTT", b"GGGGGGGGAAATTT"); + let result = aligner().local_reference_start_alignment(query, reference, None); + check(&result, query, reference, 0); + assert_eq!(result.score, 6 * 2); + assert_eq!((result.query_start, result.query_end), (0, 6)); + assert_eq!((result.ref_start, result.ref_end), (8, 14)); +} + +// --------------------------------------------------------------------------------- +// Banding. `piecewisealigner.rs` always passes `Some(bandwidth)` for the end extensions and +// lets the aligner decide, so at the default `--bw 1024` a short-read extension takes the exact +// path anyway and only long reads are really banded. These pin the two things that make it safe: +// a band wide enough to reach every cell reproduces the exact answer, and a band too narrow for +// the optimum still returns a consistent alignment. +// --------------------------------------------------------------------------------- + +#[test] +fn a_band_wider_than_the_problem_matches_the_exact_answer() { + let query = b"ACGTACGTAAGGTTCCAACGTTGGCCAATTGGCCAAGGTTCCAA"; + let reference = b"ACGTACGTAAGGTACCAACGTTGGCCAATTGCCCAAGGTTCCAA"; + let mut aligner = aligner(); + + let exact = aligner.global_alignment(query, reference, None); + for w in [64, 128, 1024] { + let banded = aligner.global_alignment(query, reference, Some(w)); + assert_eq!(banded.score, exact.score, "band {w} changed the score"); + assert_eq!( + banded.cigar.to_string(), + exact.cigar.to_string(), + "band {w} changed the CIGAR" + ); + } +} + +#[test] +fn a_narrow_band_still_returns_a_consistent_alignment() { + // A band too narrow to contain the optimum must still produce a self-consistent result: + // the score it reports is the score its own CIGAR achieves. + let query = b"AAAAAAAAAACCCCCCCCCCGGGGGGGGGGTTTTTTTTTT"; + let reference = b"AAAAAAAAAAGGGGGGGGGGTTTTTTTTTT"; + let result = aligner().global_alignment(query, reference, Some(2)); + check(&result, query, reference, 0); +} + +#[test] +fn a_band_is_never_worse_than_the_exact_optimum() { + let query = b"AAAAAAAAAACCCCCCCCCCGGGGGGGGGGTTTTTTTTTT"; + let reference = b"AAAAAAAAAAGGGGGGGGGGTTTTTTTTTT"; + let mut aligner = aligner(); + let exact = aligner.global_alignment(query, reference, None); + let banded = aligner.global_alignment(query, reference, Some(2)); + assert!( + banded.score <= exact.score, + "banded score {} beat the exact optimum {}", + banded.score, + exact.score + ); +} + +// --------------------------------------------------------------------------------- +// Alphabet. Scoring is biological; `=`/`X` is literal byte identity. They disagree on +// purpose, and `check` asserts the byte-identity half on every case above. +// --------------------------------------------------------------------------------- + +#[test] +fn u_scores_as_t_but_reports_a_mismatch_operation() { + let result = aligner().global_alignment(b"ACGU", b"ACGT", None); + assert_eq!(result.score, 4 * 2, "U should score as T"); + assert_eq!( + result.cigar.to_string(), + "3=1X", + "U vs T is not byte-identical, so it reports X" + ); +} + +#[test] +fn case_is_ignored() { + let result = aligner().global_alignment(b"acgt", b"ACGT", None); + assert_eq!(result.score, 4 * 2); +} + +#[test] +fn unknown_bytes_never_match_even_themselves() { + // Two identical `N`s: byte-identical, so the CIGAR says `=`, but they score as mismatches. + let result = aligner().global_alignment(b"ANNT", b"ANNT", None); + assert_eq!( + result.score, + 2 * 2 - 2 * 8, + "N must not match N; got {}", + result.cigar + ); + assert_eq!(result.cigar.to_string(), "4="); +} + +// --------------------------------------------------------------------------------- +// Split reference: one query across two references with a single jump. +// --------------------------------------------------------------------------------- + +#[test] +fn split_reference_partitions_the_query() { + let query = b"AAAATTTTCCCCGGGG"; + let left = b"CCCCGGGG"; + let right = b"AAAATTTT"; + let result = aligner().split_reference_alignment(query, left, right, None); + + check(&result.right, query, right, 0); + check(&result.left, query, left, 0); + + assert_eq!(result.score, result.left.score + result.right.score); + assert_eq!( + result.right.query_start, 0, + "the right arm covers the query prefix" + ); + assert_eq!( + result.left.query_end, + query.len(), + "the left arm covers the query suffix" + ); + assert_eq!( + result.right.query_end, result.left.query_start, + "the two arms must meet at the jump point" + ); + assert_eq!(result.score, 16 * 2); +} + +// --------------------------------------------------------------------------------- +// Scale: enough to cross the 32-lane chunking and the anti-diagonal bookkeeping, still fast. +// --------------------------------------------------------------------------------- + +#[test] +fn a_few_hundred_bases_stay_consistent() { + let unit = b"ACGTTGCAAGGCTTAC"; + let query: Vec = unit.iter().cycle().take(500).copied().collect(); + let mut reference = query.clone(); + reference[100] = b'A'; + reference[101] = b'A'; + reference.remove(300); + + let result = aligner().global_alignment(&query, &reference, None); + check(&result, &query, &reference, 0); + assert_eq!(result.query_end, query.len()); + assert_eq!(result.ref_end, reference.len()); + assert!( + result.score > 400 * 2, + "a near-identical 500-base pair should score well; got {}", + result.score + ); +} + +#[test] +fn identical_long_sequences_are_one_run_of_matches() { + let query: Vec = b"ACGTTGCAAGGCTTAC" + .iter() + .cycle() + .take(300) + .copied() + .collect(); + let result = aligner().global_alignment(&query, &query, None); + check(&result, &query, &query, 0); + assert_eq!(result.score, 300 * 2); + assert_eq!(result.cigar.to_string(), "300="); +}