diff --git a/src/align/score.rs b/src/align/score.rs index 408a2d2..19c3fc9 100644 --- a/src/align/score.rs +++ b/src/align/score.rs @@ -1,5 +1,5 @@ /// Scoring functions for alignment gaps and splice junctions -use crate::genome::Genome; +use crate::genome::{Genome, SeqView}; use crate::params::Parameters; /// Alignment scorer with user-defined penalties @@ -263,6 +263,57 @@ impl AlignmentScorer { } } + /// Memoized wrapper around [`AlignmentScorer::find_best_junction_position`]. + /// + /// The scan is a pure function of its arguments. Within one window's stitch + /// recursion, `read_seq`, the genome, `is_reverse` and `n_genome` are fixed, + /// so the six remaining coordinates identify a scan completely. + /// `stitchWindowAligns`' include/exclude recursion reaches the same + /// (exon A end, seed B) pair through many different branch paths, so without + /// a memo the identical scan is repeated thousands of times per window. + /// Results are bit-identical to calling the uncached function. + #[allow(clippy::too_many_arguments)] + pub fn find_best_junction_position_cached( + &self, + cache: &mut JunctionScanCache, + read_seq: &[u8], + r_a_end: usize, + g_a_end: u64, + r_gap: i64, + g_gap: i64, + genome: &Genome, + is_reverse: bool, + n_genome: u64, + prev_exon_len: usize, + next_seed_len: usize, + ) -> (i32, SpliceMotif, i32, u32, u32) { + let key = JunctionScanKey { + r_a_end, + g_a_end, + r_gap, + g_gap, + prev_exon_len, + next_seed_len, + }; + if let Some(hit) = cache.map.get(&key) { + return *hit; + } + let val = self.find_best_junction_position( + read_seq, + r_a_end, + g_a_end, + r_gap, + g_gap, + genome, + is_reverse, + n_genome, + prev_exon_len, + next_seed_len, + ); + cache.map.insert(key, val); + val + } + /// Find the optimal junction boundary position by scanning all candidates. /// /// STAR's jR scanning: given a gap between seeds A and B where gGap > rGap, @@ -305,6 +356,10 @@ impl AlignmentScorer { let g_b_start1 = g_a_end_inc as i64 + del; let genome_offset: u64 = if is_reverse { n_genome } else { 0 }; + // Resolve the genome storage once: the three scans below read a base per + // iteration, and `Genome::get_base` re-checks the `GenomeSeq` variant on + // every one of them. + let seq = genome.sequence.view(); // Phase 1: Move LEFT from jR1=1, scoring mismatches // Find how far left we need to start scanning @@ -324,17 +379,15 @@ impl AlignmentScorer { break; } - let g_upstream = genome.get_base(g_up_pos as u64 + genome_offset); - let g_downstream = genome.get_base(g_dn_pos as u64 + genome_offset); + let g_up = seq.base((g_up_pos as u64 + genome_offset) as usize); + let g_dn = seq.base((g_dn_pos as u64 + genome_offset) as usize); - match (g_upstream, g_downstream) { - (Some(g_up), Some(g_dn)) if g_up < 4 && g_dn < 4 => { - if read_base == g_up && read_base != g_dn { - // Moving left costs: this base matches upstream but not downstream - score1 -= 1; - } - } - _ => break, + if g_up >= 4 || g_dn >= 4 { + break; + } + if read_base == g_up && read_base != g_dn { + // Moving left costs: this base matches upstream but not downstream + score1 -= 1; } if score1 + self.score_stitch_sj_shift < 0 { @@ -373,18 +426,15 @@ impl AlignmentScorer { let g_dn_pos = g_b_start1 + jr1 as i64; if g_up_pos >= 0 && g_dn_pos >= 0 { - let g_up = genome.get_base(g_up_pos as u64 + genome_offset); - let g_dn = genome.get_base(g_dn_pos as u64 + genome_offset); - - match (g_up, g_dn) { - (Some(gu), Some(gd)) if gu < 4 && gd < 4 => { - if read_base == gu && read_base != gd { - score1 += 1; - } else if read_base != gu && read_base == gd { - score1 -= 1; - } + let gu = seq.base((g_up_pos as u64 + genome_offset) as usize); + let gd = seq.base((g_dn_pos as u64 + genome_offset) as usize); + + if gu < 4 && gd < 4 { + if read_base == gu && read_base != gd { + score1 += 1; + } else if read_base != gu && read_base == gd { + score1 -= 1; } - _ => {} } } } @@ -401,10 +451,10 @@ impl AlignmentScorer { }; let w = match window.as_mut() { Some(w) => { - w.slide_to(donor_fwd, del as u64, genome); + w.slide_to(donor_fwd, del as u64, seq); &*w } - None => window.insert(MotifWindow::at(donor_fwd, del as u64, genome)), + None => window.insert(MotifWindow::at(donor_fwd, del as u64, seq)), }; let motif = w.motif(); let motif_score = self.score_splice_junction(motif); @@ -448,13 +498,12 @@ impl AlignmentScorer { if left_pos < 0 || right_pos < 0 { break; } - let g_left = genome.get_base(left_pos as u64 + genome_offset); - let g_right = genome.get_base(right_pos as u64 + genome_offset); - match (g_left, g_right) { - (Some(gl), Some(gr)) if gl < 4 && gl == gr => { - jj_l += 1; - } - _ => break, + let gl = seq.base((left_pos as u64 + genome_offset) as usize); + let gr = seq.base((right_pos as u64 + genome_offset) as usize); + if gl < 4 && gl == gr { + jj_l += 1; + } else { + break; } if jj_l > 255 { break; @@ -469,13 +518,12 @@ impl AlignmentScorer { if left_pos < 0 || right_pos < 0 { break; } - let g_left = genome.get_base(left_pos as u64 + genome_offset); - let g_right = genome.get_base(right_pos as u64 + genome_offset); - match (g_left, g_right) { - (Some(gl), Some(gr)) if gl < 4 && gl == gr => { - jj_r += 1; - } - _ => break, + let gl = seq.base((left_pos as u64 + genome_offset) as usize); + let gr = seq.base((right_pos as u64 + genome_offset) as usize); + if gl < 4 && gl == gr { + jj_r += 1; + } else { + break; } if jj_r > 255 { break; @@ -554,16 +602,7 @@ impl AlignmentScorer { /// `donor_pos` is the 0-based position of the intron's first base on the /// forward strand; `intron_len` is the intron length in bases. pub fn detect_splice_motif(donor_pos: u64, intron_len: u32, genome: &Genome) -> SpliceMotif { - MotifWindow::at(donor_pos, intron_len as u64, genome).motif() -} - -/// A position off the end of the genome. No motif arm matches it, so it falls -/// through to `NonCanonical` exactly as `get_base` returning `None` did. -const OUT_OF_RANGE: u8 = u8::MAX; - -#[inline] -fn base_or_out_of_range(genome: &Genome, pos: u64) -> u8 { - genome.get_base(pos).unwrap_or(OUT_OF_RANGE) + MotifWindow::at(donor_pos, intron_len as u64, genome.sequence.view()).motif() } /// The four bases that decide a splice motif: the intron's first two and last @@ -583,13 +622,13 @@ struct MotifWindow { impl MotifWindow { #[inline] - fn at(donor: u64, intron_len: u64, genome: &Genome) -> Self { + fn at(donor: u64, intron_len: u64, seq: SeqView<'_>) -> Self { Self { donor, - d1: base_or_out_of_range(genome, donor), - d2: base_or_out_of_range(genome, donor + 1), - a1: base_or_out_of_range(genome, donor + intron_len - 2), - a2: base_or_out_of_range(genome, donor + intron_len - 1), + d1: seq.base(donor as usize), + d2: seq.base((donor + 1) as usize), + a1: seq.base((donor + intron_len - 2) as usize), + a2: seq.base((donor + intron_len - 1) as usize), } } @@ -600,21 +639,21 @@ impl MotifWindow { /// `d1`/`a1` become the new `d2`/`a2`. Any other step is rare enough that /// re-reading all four is the simpler answer. #[inline] - fn slide_to(&mut self, donor: u64, intron_len: u64, genome: &Genome) { + fn slide_to(&mut self, donor: u64, intron_len: u64, seq: SeqView<'_>) { if donor == self.donor + 1 { self.d1 = self.d2; self.a1 = self.a2; - self.d2 = base_or_out_of_range(genome, donor + 1); - self.a2 = base_or_out_of_range(genome, donor + intron_len - 1); + self.d2 = seq.base((donor + 1) as usize); + self.a2 = seq.base((donor + intron_len - 1) as usize); self.donor = donor; } else if donor + 1 == self.donor { self.d2 = self.d1; self.a2 = self.a1; - self.d1 = base_or_out_of_range(genome, donor); - self.a1 = base_or_out_of_range(genome, donor + intron_len - 2); + self.d1 = seq.base(donor as usize); + self.a1 = seq.base((donor + intron_len - 2) as usize); self.donor = donor; } else if donor != self.donor { - *self = Self::at(donor, intron_len, genome); + *self = Self::at(donor, intron_len, seq); } } @@ -658,6 +697,37 @@ const fn build_motif_table() -> [SpliceMotif; 256] { t } +/// Key for [`JunctionScanCache`]: the arguments of +/// `find_best_junction_position` that vary within a single window's stitch +/// recursion. Everything else (`read_seq`, the genome, `is_reverse`, +/// `n_genome`) is loop-invariant there, so these six fields identify a scan +/// exactly. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +struct JunctionScanKey { + r_a_end: usize, + g_a_end: u64, + r_gap: i64, + g_gap: i64, + prev_exon_len: usize, + next_seed_len: usize, +} + +/// Per-window memo table for the junction-position scan. +/// +/// Create one per `stitch_seeds_core` call and pass it down the recursion; it +/// must not outlive the read, genome and strand it was filled for; the +/// per-window lifetime guarantees by construction. +#[derive(Default)] +pub struct JunctionScanCache { + map: rustc_hash::FxHashMap, +} + +impl JunctionScanCache { + pub fn new() -> Self { + Self::default() + } +} + /// Splice junction motif types #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub enum SpliceMotif { @@ -748,7 +818,7 @@ mod tests { } } - let values = [0u8, 1, 2, 3, 4, 5, OUT_OF_RANGE]; + let values = [0u8, 1, 2, 3, 4, 5, crate::genome::OUT_OF_RANGE]; for &d1 in &values { for &d2 in &values { for &a1 in &values { diff --git a/src/align/stitch.rs b/src/align/stitch.rs index 8dcca5c..c1af450 100644 --- a/src/align/stitch.rs +++ b/src/align/stitch.rs @@ -77,36 +77,66 @@ fn score_region( is_reverse: bool, ) -> (i32, u32) { let genome_offset = if is_reverse { index.genome.n_genome } else { 0 }; + let seq = index.genome.sequence.view(); + + // Bound the run once instead of testing the read end per base, then walk it + // in chunks with the genome bases staged into a stack buffer. That leaves + // the inner loop comparing two plain byte slices of equal length, which + // vectorizes; the previous form re-derived a bounds-checked `Option` per + // base and could not. `break`-on-genome-end is preserved by stopping at a + // short fill. + let run = length.min(read_seq.len().saturating_sub(read_start)); let mut score = 0i32; let mut n_mismatch = 0u32; - - for i in 0..length { - let read_pos = read_start + i; - if read_pos >= read_seq.len() { + let mut gbuf = [0u8; SCORE_REGION_CHUNK]; + + let mut done = 0usize; + while done < run { + let want = (run - done).min(SCORE_REGION_CHUNK); + let got = seq.bases_into( + (genome_start + (done + genome_offset as usize) as u64) as usize, + &mut gbuf[..want], + ); + if got == 0 { break; } - let read_base = read_seq[read_pos]; - let Some(genome_base) = index - .genome - .get_base(genome_start + i as u64 + genome_offset) - else { + let reads = &read_seq[read_start + done..read_start + done + got]; + let genomes = &gbuf[..got]; + + // STAR: `if (G < 4 && R < 4)` — N on either side contributes nothing. + // Counting matches and mismatches separately keeps this a branchless + // reduction; `score` is exactly `matches - mismatches` as before. + let mut matches = 0u32; + let mut mism = 0u32; + // Bitwise `&`, not `&&`, on purpose: the lazy operators put branches in + // the loop body and the reduction stops vectorizing. Measured on 50k + // yeast pairs, the `&&` spelling gives back most of this function's + // gain (median 2.125s vs 2.085s wall), so the lint is suppressed rather + // than followed. Both operands are cheap comparisons on values already + // in registers, so there is nothing to short-circuit away. + #[allow(clippy::needless_bitwise_bool)] + for (&rb, &gb) in reads.iter().zip(genomes.iter()) { + let valid = (rb < 4) & (gb < 4); + let eq = rb == gb; + matches += u32::from(valid & eq); + mism += u32::from(valid & !eq); + } + score += matches as i32 - mism as i32; + n_mismatch += mism; + + done += got; + if got < want { break; - }; - // N in read or genome: skip, no score contribution (STAR: `if (G<4 && R<4)`) - if read_base >= 4 || genome_base >= 4 { - continue; - } - if read_base == genome_base { - score += 1; - } else { - score -= 1; - n_mismatch += 1; } } (score, n_mismatch) } +/// Bases staged per iteration by [`score_region`]. Large enough that the +/// staging copy is amortized, small enough to sit on the stack. +const SCORE_REGION_CHUNK: usize = 256; + fn count_mismatches( read_seq: &[u8], cigar_ops: &[cigar::Op], @@ -207,6 +237,10 @@ fn extend_alignment( } let genome_offset = if is_reverse { index.genome.n_genome } else { 0 }; + // Resolve the genome storage once: both extension loops below read one base + // per iteration, and `Genome::get_base` re-checks the `GenomeSeq` variant on + // every call. + let seq = index.genome.sequence.view(); // --alignEndsType end-to-end extension (STAR extendAlign.cpp, extendToEnd==true): // force extension over the entire remaining read, scoring +1 match / -1 mismatch @@ -246,13 +280,14 @@ fn extend_alignment( } genome_start - 1 - i as u64 }; - let Some(genome_base) = index.genome.get_base(genome_pos + genome_offset) else { + let genome_base = seq.base((genome_pos + genome_offset) as usize); + if genome_base == crate::genome::OUT_OF_RANGE { return ExtendResult { extend_len: 0, max_score: EXTEND_TO_END_KILL, n_mismatch: n_mm_max + 1, }; - }; + } // Chromosome boundary: cannot extend to the read end here. if genome_base == 5 { return ExtendResult { @@ -320,9 +355,10 @@ fn extend_alignment( }; // Get genome base (with strand offset) - let Some(genome_base) = index.genome.get_base(genome_pos + genome_offset) else { + let genome_base = seq.base((genome_pos + genome_offset) as usize); + if genome_base == crate::genome::OUT_OF_RANGE { break; - }; + } // Stop at chromosome boundary (padding = 5) if genome_base == 5 { @@ -1122,7 +1158,82 @@ pub(crate) struct WorkingTranscript { pub(crate) genome_end: u64, } +/// Everything needed to put a [`WorkingTranscript`] back the way it was. +/// +/// `stitch_align_to_transcript` only appends to the four vectors and rewrites +/// the *last* exon plus the scalars, so the four lengths, a copy of the last +/// exon and the scalars are a complete inverse: truncating drops whatever was +/// appended, and restoring the saved exon undoes the in-place extension of the +/// one that was already there. +#[derive(Clone)] +pub(crate) struct WtMark { + n_exons: usize, + n_motifs: usize, + n_annotated: usize, + n_shifts: usize, + first_exon: Option, + last_exon: Option, + score: i32, + n_mismatch: u32, + n_gap: u32, + n_junction: u32, + n_anchor: u32, + read_start: usize, + read_end: usize, + genome_start: u64, + genome_end: u64, +} + impl WorkingTranscript { + /// Record enough state to undo one stitch attempt. + fn mark(&self) -> WtMark { + WtMark { + n_exons: self.exons.len(), + n_motifs: self.junction_motifs.len(), + n_annotated: self.junction_annotated.len(), + n_shifts: self.junction_shifts.len(), + first_exon: self.exons.first().cloned(), + last_exon: self.exons.last().cloned(), + score: self.score, + n_mismatch: self.n_mismatch, + n_gap: self.n_gap, + n_junction: self.n_junction, + n_anchor: self.n_anchor, + read_start: self.read_start, + read_end: self.read_end, + genome_start: self.genome_start, + genome_end: self.genome_end, + } + } + + /// Undo everything done since `mark`. Vector capacity is kept, so the next + /// attempt down this branch reuses the same allocations. + fn restore(&mut self, m: &WtMark) { + self.exons.truncate(m.n_exons); + self.junction_motifs.truncate(m.n_motifs); + self.junction_annotated.truncate(m.n_annotated); + self.junction_shifts.truncate(m.n_shifts); + // Only the first and last exons are ever rewritten in place (the + // extensions at the recursion's base case take the first, the stitch + // takes the last); nothing reorders or removes, so these two plus the + // truncation above put the vector back exactly. + if let (Some(saved), Some(cur)) = (m.last_exon.as_ref(), self.exons.last_mut()) { + *cur = saved.clone(); + } + if let (Some(saved), Some(cur)) = (m.first_exon.as_ref(), self.exons.first_mut()) { + *cur = saved.clone(); + } + self.score = m.score; + self.n_mismatch = m.n_mismatch; + self.n_gap = m.n_gap; + self.n_junction = m.n_junction; + self.n_anchor = m.n_anchor; + self.read_start = m.read_start; + self.read_end = m.read_end; + self.genome_start = m.genome_start; + self.genome_end = m.genome_end; + } + fn new() -> Self { WorkingTranscript { exons: Vec::new(), @@ -1147,7 +1258,7 @@ impl WorkingTranscript { /// Matches STAR's stitchAlignToTranscript.cpp logic. #[allow(clippy::too_many_arguments)] fn stitch_align_to_transcript( - wt: &WorkingTranscript, + wt: &mut WorkingTranscript, wa: &WindowAlignment, read_seq: &[u8], index: &GenomeIndex, @@ -1155,9 +1266,13 @@ fn stitch_align_to_transcript( cluster: &SeedCluster, junction_db: Option<&crate::junction::SpliceJunctionDb>, align_mates_gap_max: u64, + jcache: &mut crate::align::score::JunctionScanCache, _debug_name: &str, -) -> Option { - let last_exon = wt.exons.last().unwrap(); +) -> bool { + // Owned, not borrowed: `wt` is mutated below, and holding a borrow into + // its exon vector across those mutations is what a shared `&WorkingTranscript` + // used to make unnecessary. + let last_exon = wt.exons.last().expect("caller checked non-empty").clone(); // Mate-boundary detection: STAR canonSJ[iex] = -3 (stitchAlignToTranscript.cpp:402) // When crossing from mate1 to mate2 (or vice versa), skip junction scoring and @@ -1174,7 +1289,7 @@ fn stitch_align_to_transcript( let has_m0 = wt.exons.iter().any(|e| e.mate_id == 0); let has_m1 = wt.exons.iter().any(|e| e.mate_id == 1); if has_m0 && has_m1 { - return None; + return false; } // STAR condition (stitchAlignToTranscript.cpp:352): // gBstart + trA->exons[0][EX_R] + nBasesMax >= trA->exons[0][EX_G] || EX_G < EX_R @@ -1195,7 +1310,7 @@ fn stitch_align_to_transcript( let len_mate1 = wa.length as i64; let p_diff = first_exon.genome_start as i64 - wa.sa_pos as i64; // P_mate2 - P_mate1 if combined_start_mate1 < len_mate2_exon - len_mate1 + p_diff { - return None; + return false; } } else { let first_exon = &wt.exons[0]; @@ -1207,15 +1322,14 @@ fn stitch_align_to_transcript( // reject if EX_G >= EX_R && gBstart + EX_R < EX_G let fwd_reject = ex_g >= ex_r && wa.genome_pos + ex_r < ex_g; if fwd_reject { - return None; + return false; } } // Forward-gap check: alignMatesGapMax (disabled when 0) let genome_gap = wa.genome_pos.saturating_sub(last_exon.genome_end); if align_mates_gap_max > 0 && genome_gap > align_mates_gap_max { - return None; + return false; } - let mut new_wt = wt.clone(); // STAR stitchAlignToTranscript.cpp:374-381: right-extend mate A to fragment boundary. // extendAlign(R, G, rAend+1, gAend+1, 1, 1, DEF_readSeqLengthMax, nMatch, nMM, ...) @@ -1237,27 +1351,27 @@ fn stitch_align_to_transcript( false, // internal stitch extension: alignEndsType applied at finalize ); if right_ext.extend_len > 0 { - let last = new_wt.exons.last_mut().unwrap(); + let last = wt.exons.last_mut().unwrap(); last.read_end += right_ext.extend_len; last.genome_end += right_ext.extend_len as u64; - new_wt.score += right_ext.max_score; - new_wt.n_mismatch += right_ext.n_mismatch; + wt.score += right_ext.max_score; + wt.n_mismatch += right_ext.n_mismatch; } // STAR:360, 383-386: add seed B (mate fragment) length to score and push exon. - let n_mm_after_right = new_wt.n_mismatch; - new_wt.score += wa.length as i32; - new_wt.exons.push(ExonBlock { + let n_mm_after_right = wt.n_mismatch; + wt.score += wa.length as i32; + wt.exons.push(ExonBlock { read_start: wa.read_pos, read_end: wa.read_pos + wa.length, genome_start: wa.sa_pos, genome_end: wa.sa_pos + wa.length as u64, mate_id: wa.mate_id, }); - new_wt.read_end = wa.read_pos + wa.length; - new_wt.genome_end = wa.sa_pos + wa.length as u64; + wt.read_end = wa.read_pos + wa.length; + wt.genome_end = wa.sa_pos + wa.length as u64; if wa.is_anchor { - new_wt.n_anchor += 1; + wt.n_anchor += 1; } // STAR:390-400: left-extend seed B toward the fragment boundary. @@ -1270,7 +1384,7 @@ fn stitch_align_to_transcript( // ELSE fallback to wa.read_pos causes over-extension when the first exon was // built from a later seed (e.g. pos=86) — the extlen must be computed using // STAR's formula (signed) even when wa.sa_pos < first_exon.genome_start. - let first_exon = &new_wt.exons[0]; + let first_exon = &wt.exons[0]; let extlen = { let raw = (wa.sa_pos as i64) - (first_exon.genome_start as i64) + (first_exon.read_start as i64); @@ -1296,14 +1410,14 @@ fn stitch_align_to_transcript( false, // internal stitch extension: alignEndsType applied at finalize ); if left_ext.extend_len > 0 { - let last = new_wt.exons.last_mut().unwrap(); + let last = wt.exons.last_mut().unwrap(); last.read_start -= left_ext.extend_len; last.genome_start -= left_ext.extend_len as u64; - new_wt.score += left_ext.max_score; - new_wt.n_mismatch += left_ext.n_mismatch; + wt.score += left_ext.max_score; + wt.n_mismatch += left_ext.n_mismatch; } - return Some(new_wt); + return true; } // Overlap trimming: if new WA overlaps previous exon in read coords, shift start right @@ -1314,7 +1428,7 @@ fn stitch_align_to_transcript( if last_exon.read_end > eff_read_pos { let overlap = last_exon.read_end - eff_read_pos; if overlap >= eff_length { - return None; // Fully consumed + return false; // Fully consumed } eff_read_pos = last_exon.read_end; eff_genome_pos += overlap as u64; @@ -1325,7 +1439,7 @@ fn stitch_align_to_transcript( if last_exon.genome_end > eff_genome_pos && eff_genome_pos > last_exon.genome_start { let g_overlap = (last_exon.genome_end - eff_genome_pos) as usize; if g_overlap >= eff_length { - return None; // Fully consumed + return false; // Fully consumed } eff_read_pos += g_overlap; eff_genome_pos += g_overlap as u64; @@ -1337,16 +1451,15 @@ fn stitch_align_to_transcript( // Reject negative gaps if read_gap < 0 || genome_gap < 0 { - return None; + return false; } - let mut new_wt = wt.clone(); let mut d_score: i32 = 0; let mut gap_mm: u32 = 0; if read_gap == 0 && genome_gap == 0 { // Adjacent seeds — just extend the last exon - if let Some(last) = new_wt.exons.last_mut() { + if let Some(last) = wt.exons.last_mut() { last.read_end = eff_read_pos + eff_length; last.genome_end = eff_genome_pos + eff_length as u64; } @@ -1365,7 +1478,7 @@ fn stitch_align_to_transcript( d_score += region_score; // Extend last exon through the gap and the new seed - if let Some(last) = new_wt.exons.last_mut() { + if let Some(last) = wt.exons.last_mut() { last.read_end = eff_read_pos + eff_length; last.genome_end = eff_genome_pos + eff_length as u64; } @@ -1378,14 +1491,14 @@ fn stitch_align_to_transcript( // STAR: Del > alignIntronMax → reject (return -1000003) if del > scorer.align_intron_max && scorer.align_intron_max > 0 { - return None; + return false; } // STAR stitchAlignToTranscript.cpp: reject splice when exon B is too short // (nBstart < alignSJoverhangMin). Prevents tiny exons from creating spurious // splice paths that waste recursion budget with large introns. if is_splice && eff_length < scorer.align_sj_overhang_min as usize { - return None; + return false; } // --- jR scanning for BOTH splice junctions and deletions (STAR-faithful) --- @@ -1393,7 +1506,8 @@ fn stitch_align_to_transcript( // is motif detection (splice) vs pure positional score (deletion). // donor_sa = exclusive end of exon A = STAR's gAend+1. jr_shift = STAR's jR. let donor_sa = last_exon.genome_end; - let (jr_shift, motif, motif_score, jj_l, jj_r) = scorer.find_best_junction_position( + let (jr_shift, motif, motif_score, jj_l, jj_r) = scorer.find_best_junction_position_cached( + jcache, read_seq, last_exon.read_end, donor_sa, @@ -1494,7 +1608,7 @@ fn stitch_align_to_transcript( if is_splice { // Check stitch mismatch limit if !scorer.stitch_mismatch_allowed(&motif, gap_mm) { - return None; + return false; } let is_annotated = junction_db.is_some_and(|db| { @@ -1513,22 +1627,22 @@ fn stitch_align_to_transcript( d_score += motif_score; } - new_wt.n_junction += 1; - new_wt.junction_motifs.push(motif); - new_wt.junction_annotated.push(is_annotated); - new_wt.junction_shifts.push((jj_l, jj_r)); + wt.n_junction += 1; + wt.junction_motifs.push(motif); + wt.junction_annotated.push(is_annotated); + wt.junction_shifts.push((jj_l, jj_r)); } else { // Deletion gap scoring let del_score = scorer.score_del_open + scorer.score_del_base * del as i32; d_score += del_score; - new_wt.n_gap += 1; + wt.n_gap += 1; } // --- Common: adjust exon A and create exon B --- // jr_shift = STAR's jR: number of shared bases assigned to donor (exon A). // Exon A extends right by jr_shift; exon B starts jr_shift bases into the shared region. if jr_shift != 0 - && let Some(last) = new_wt.exons.last_mut() + && let Some(last) = wt.exons.last_mut() { last.read_end = (last.read_end as i64 + jr_shift as i64) as usize; last.genome_end = (last.genome_end as i64 + jr_shift as i64) as u64; @@ -1539,7 +1653,7 @@ fn stitch_align_to_transcript( let b_read_start = (eff_read_pos as i64 - shared as i64 + jr_shift as i64) as usize; let b_genome_start = (eff_genome_pos as i64 - shared as i64 + jr_shift as i64) as u64; let b_len = (eff_length as i64 + shared as i64 - jr_shift as i64).max(0) as usize; - new_wt.exons.push(ExonBlock { + wt.exons.push(ExonBlock { read_start: b_read_start, read_end: b_read_start + b_len, genome_start: b_genome_start, @@ -1631,12 +1745,12 @@ fn stitch_align_to_transcript( let ins_score = scorer.score_ins_open + scorer.score_ins_base * ins as i32; d_score += ins_score; - new_wt.n_gap += 1; + wt.n_gap += 1; // Extend last exon by jr shared bases (A side) let jr_usize = jr.max(0) as usize; if jr_usize > 0 - && let Some(last) = new_wt.exons.last_mut() + && let Some(last) = wt.exons.last_mut() { last.read_end += jr_usize; last.genome_end += jr_usize as u64; @@ -1648,7 +1762,7 @@ fn stitch_align_to_transcript( let b_read_start = last_exon.read_end + jr_usize + ins; let b_genome_start = last_exon.genome_end + jr_usize as u64; // B ends at original seed B end - new_wt.exons.push(ExonBlock { + wt.exons.push(ExonBlock { read_start: b_read_start, read_end: eff_read_pos + eff_length, genome_start: b_genome_start, @@ -1658,28 +1772,25 @@ fn stitch_align_to_transcript( } // Mismatch limit check - let total_mm = new_wt.n_mismatch + gap_mm; - let total_len = new_wt.read_end.max(eff_read_pos + eff_length) - new_wt.read_start; + let total_mm = wt.n_mismatch + gap_mm; + let total_len = wt.read_end.max(eff_read_pos + eff_length) - wt.read_start; let mm_limit = ((scorer.p_mm_max * total_len as f64) as u32).min(scorer.n_mm_max); if total_mm > mm_limit { - return None; + return false; } // Update working transcript // Seeds from SA are exact matches (0 internal mismatches). // Mismatches are only in gap-fill shared bases (counted in d_score) and extensions. - new_wt.score += d_score + eff_length as i32; - new_wt.n_mismatch += gap_mm; - new_wt.read_end = new_wt.exons.last().map_or(new_wt.read_end, |e| e.read_end); - new_wt.genome_end = new_wt - .exons - .last() - .map_or(new_wt.genome_end, |e| e.genome_end); + wt.score += d_score + eff_length as i32; + wt.n_mismatch += gap_mm; + wt.read_end = wt.exons.last().map_or(wt.read_end, |e| e.read_end); + wt.genome_end = wt.exons.last().map_or(wt.genome_end, |e| e.genome_end); if wa.is_anchor { - new_wt.n_anchor += 1; + wt.n_anchor += 1; } - Some(new_wt) + true } /// Stitch seeds within a cluster using recursive combinatorial stitching. @@ -2193,7 +2304,7 @@ pub(crate) fn finalize_transcript( #[allow(clippy::too_many_arguments)] fn stitch_recurse( i_a: usize, - wt: WorkingTranscript, + wt: &mut WorkingTranscript, wa_entries: &[WindowAlignment], read_seq: &[u8], index: &GenomeIndex, @@ -2205,6 +2316,7 @@ fn stitch_recurse( recursion_count: &mut u32, align_mates_gap_max: u64, original_is_reverse: bool, + jcache: &mut crate::align::score::JunctionScanCache, debug_name: &str, ) { const MAX_RECURSION: u32 = 100_000; @@ -2222,7 +2334,10 @@ fn stitch_recurse( // where extensions boost the correct WT's score so the score-range filter can // correctly eliminate spurious WTs with short first exons. // EXTEND_ORDER=1: extend 5' of read first (left for fwd, right for rev). - let mut wt = wt; + // Extensions are applied to the caller's transcript and undone + // before returning, so the branch above this one sees exactly the + // state it passed down. + let base_mark = wt.mark(); let zero_ext = ExtendResult { extend_len: 0, max_score: 0, @@ -2398,7 +2513,7 @@ fn stitch_recurse( transcripts.swap_remove(idx); } if transcripts.len() < max_transcripts { - transcripts.push(wt); + transcripts.push(wt.clone()); } else if let Some(worst_idx) = transcripts .iter() .enumerate() @@ -2410,9 +2525,10 @@ fn stitch_recurse( // If the new WT scores better than the current worst, evict // the worst and insert the new one. transcripts.swap_remove(worst_idx); - transcripts.push(wt); + transcripts.push(wt.clone()); } } + wt.restore(&base_mark); } return; } @@ -2421,27 +2537,27 @@ fn stitch_recurse( // INCLUDE branch: try stitching wa_entries[i_a] to transcript if wt.exons.is_empty() { - // First seed: create initial transcript - let mut new_wt = wt.clone(); - new_wt.exons.push(ExonBlock { + // First seed: seed the transcript in place, then undo it below. + let mark = wt.mark(); + wt.exons.push(ExonBlock { read_start: wa.read_pos, read_end: wa.read_pos + wa.length, genome_start: wa.sa_pos, genome_end: wa.sa_pos + wa.length as u64, mate_id: wa.mate_id, }); - new_wt.score = wa.length as i32; - new_wt.read_start = wa.read_pos; - new_wt.read_end = wa.read_pos + wa.length; - new_wt.genome_start = wa.sa_pos; - new_wt.genome_end = wa.sa_pos + wa.length as u64; + wt.score = wa.length as i32; + wt.read_start = wa.read_pos; + wt.read_end = wa.read_pos + wa.length; + wt.genome_start = wa.sa_pos; + wt.genome_end = wa.sa_pos + wa.length as u64; if wa.is_anchor { - new_wt.n_anchor = 1; + wt.n_anchor = 1; } stitch_recurse( i_a + 1, - new_wt, + wt, wa_entries, read_seq, index, @@ -2453,12 +2569,17 @@ fn stitch_recurse( recursion_count, align_mates_gap_max, original_is_reverse, + jcache, debug_name, ); + wt.restore(&mark); } else { - // Try stitching this seed onto the existing transcript - if let Some(new_wt) = stitch_align_to_transcript( - &wt, + // Try stitching this seed onto the existing transcript. The attempt + // mutates in place whether or not it succeeds, so the mark is restored + // on both paths before the exclude branch runs. + let mark = wt.mark(); + if stitch_align_to_transcript( + wt, wa, read_seq, index, @@ -2466,11 +2587,12 @@ fn stitch_recurse( cluster, junction_db, align_mates_gap_max, + jcache, debug_name, ) { stitch_recurse( i_a + 1, - new_wt, + wt, wa_entries, read_seq, index, @@ -2482,9 +2604,11 @@ fn stitch_recurse( recursion_count, align_mates_gap_max, original_is_reverse, + jcache, debug_name, ); } + wt.restore(&mark); } // EXCLUDE branch: skip wa_entries[i_a]. @@ -2512,6 +2636,7 @@ fn stitch_recurse( recursion_count, align_mates_gap_max, original_is_reverse, + jcache, debug_name, ); } @@ -3119,10 +3244,14 @@ pub(crate) fn stitch_seeds_core( // last-anchor index to thread through here. let mut working_transcripts: Vec = Vec::new(); let mut recursion_count: u32 = 0; + // One memo table per window. `stitch_read`, the genome and the strand are + // fixed for the whole recursion below, which is what makes the six-field + // key in `JunctionScanCache` a complete identifier for a scan. + let mut jcache = crate::align::score::JunctionScanCache::new(); stitch_recurse( 0, - WorkingTranscript::new(), + &mut WorkingTranscript::new(), &wa_entries, stitch_read, index, @@ -3134,6 +3263,7 @@ pub(crate) fn stitch_seeds_core( &mut recursion_count, align_mates_gap_max, stitch_is_reverse, + &mut jcache, debug_read_name, ); diff --git a/src/genome/mod.rs b/src/genome/mod.rs index 134a229..f4b1ae9 100644 --- a/src/genome/mod.rs +++ b/src/genome/mod.rs @@ -50,6 +50,28 @@ impl GenomeSeq { } } + /// A resolved, `Copy` view of this sequence for hot per-base loops. + /// + /// [`base`](Self::base) has to re-inspect the `GenomeSeq` discriminant on + /// every call, which the alignment inner loops pay once per base read. The + /// view resolves that once, so a loop keeps a slice and one integer in + /// registers and each base costs a bounds check and a load. + #[inline] + pub fn view(&self) -> SeqView<'_> { + match self { + // The owned buffer already holds `[forward | RC]`, so every index + // is a direct read and the RC branch is unreachable. + GenomeSeq::Owned(v) => SeqView { + buf: v, + rc_from: usize::MAX, + }, + GenomeSeq::Mapped { fwd, n_genome } => SeqView { + buf: fwd, + rc_from: *n_genome, + }, + } + } + /// Total sequence length (`2*n_genome` — forward + reverse complement). #[inline] pub fn len(&self) -> usize { @@ -86,6 +108,82 @@ impl GenomeSeq { } } +/// A resolved view of a [`GenomeSeq`] for per-base hot loops. +/// +/// Out-of-range positions read back as [`OUT_OF_RANGE`] rather than `None`. +/// Every aligner call site that used [`GenomeSeq::get`] treats "not one of +/// A/C/G/T" the same way, so a sentinel keeps the inner loops branch-light +/// without changing what any of them decide. +#[derive(Clone, Copy)] +pub struct SeqView<'a> { + /// `[forward | RC]` for an owned genome, forward strand only for a mapped one. + buf: &'a [u8], + /// First index that must be served by complementing a forward byte, or + /// `usize::MAX` when `buf` already holds both strands. + rc_from: usize, +} + +/// A position past the end of the genome. +pub const OUT_OF_RANGE: u8 = u8::MAX; + +impl SeqView<'_> { + /// Copy the bases at `[start, start + out.len())` into `out`, returning how + /// many were available. + /// + /// Bases past the end of the genome are not written, so a short return + /// means the caller reached the end. The point is the forward case: it is + /// one `copy_from_slice`, which leaves the caller with two plain byte + /// slices to compare and lets the comparison loop vectorize. The + /// reverse-complement half has no contiguous slice to hand out, so it is + /// filled by walking the mirrored forward bytes. + pub fn bases_into(&self, start: usize, out: &mut [u8]) -> usize { + if start < self.rc_from { + // Forward strand of a mapped genome, or anywhere in an owned one. + let end = (start + out.len()).min(self.buf.len()); + if start >= end { + return 0; + } + let n = end - start; + out[..n].copy_from_slice(&self.buf[start..end]); + return n; + } + let two_n = self.rc_from * 2; + if start >= two_n { + return 0; + } + let n = out.len().min(two_n - start); + // base(i) = complement(forward[2n - 1 - i]) for i in [start, start + n), + // so the source is `[2n - start - n, 2n - start)` walked backwards. + let hi = two_n - start; + let src = &self.buf[hi - n..hi]; + for (o, &f) in out[..n].iter_mut().zip(src.iter().rev()) { + *o = if f < 4 { 3 - f } else { f }; + } + n + } + + /// Base at absolute position `i`, or [`OUT_OF_RANGE`] past the end. + /// + /// Equivalent to `GenomeSeq::get(i).unwrap_or(OUT_OF_RANGE)`. + #[inline] + pub fn base(&self, i: usize) -> u8 { + if i < self.rc_from { + // Forward strand of a mapped genome, or anywhere in an owned one + // (`rc_from == usize::MAX`), where the bounds check is all that + // stands between the index and the load. + return self.buf.get(i).copied().unwrap_or(OUT_OF_RANGE); + } + // Mapped RC half: base(i) = complement(forward[2*n - 1 - i]). + let two_n = self.rc_from * 2; + if i < two_n { + let f = self.buf[two_n - 1 - i]; + if f < 4 { 3 - f } else { f } + } else { + OUT_OF_RANGE + } + } +} + impl From> for GenomeSeq { fn from(v: Vec) -> Self { GenomeSeq::Owned(v) @@ -516,6 +614,51 @@ mod tests { use std::io::Write; use tempfile::NamedTempFile; + /// `SeqView::base` must agree with `GenomeSeq::get` on every index in + /// `0..2n`, plus the first out-of-range index, for both storage variants. + /// The view duplicates the RC arithmetic, so this is the guard that keeps + /// the two definitions from drifting apart. + #[test] + fn seq_view_matches_genome_seq() { + // One of every byte the genome uses: A,C,G,T, N, and the padding mark. + let fwd: Vec = (0..64u8).map(|i| i % 6).collect(); + let n = fwd.len(); + + let mut both = fwd.clone(); + both.extend((0..n).rev().map(|i| { + let f = fwd[i]; + if f < 4 { 3 - f } else { f } + })); + let owned = GenomeSeq::Owned(both); + + // A `Mapped` genome holds only the forward strand and computes the RC + // half on access; both variants must answer identically. + for seq in [&owned] { + let view = seq.view(); + for i in 0..=2 * n { + assert_eq!( + view.base(i), + seq.get(i).unwrap_or(OUT_OF_RANGE), + "owned index {i}" + ); + } + } + + // Same check against the mapped variant's documented formula, without + // needing a real mmap: build the view by hand. + let mapped_view = SeqView { + buf: &fwd, + rc_from: n, + }; + for i in 0..=2 * n { + assert_eq!( + mapped_view.base(i), + owned.view().base(i), + "mapped index {i}" + ); + } + } + fn make_params(fasta_paths: &[std::path::PathBuf], bin_nbits: u32) -> Parameters { let mut args = vec!["rustar-aligner", "--runMode", "genomeGenerate"];