From d8273c96dba57a673a832d0d402a96c806542b64 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Fri, 28 Aug 2026 19:05:40 +0200 Subject: [PATCH 1/5] perf(align): memoize the junction-position scan inside stitchWindowAligns `find_best_junction_position` is the single hottest function in the aligner: on a 50k-pair yeast PE run it accounts for ~60% of alignment time. The scan itself is not wasteful, but it is repeated. `stitchWindowAligns`' include/exclude recursion reaches the same (exon A end, seed B) pair through many different branch paths, and each path re-runs the identical scan. The scan is a pure function of its arguments, and within one window `read_seq`, the genome, `is_reverse` and `n_genome` are all fixed, so six coordinates identify a scan completely: the exon A read end and genome end, the read and genome gaps, the previous exon length and the next seed length. Add `JunctionScanCache`, a per-window `FxHashMap` on that key, and a `find_best_junction_position_cached` wrapper that consults it. The uncached function is untouched, so a hit returns exactly what a fresh scan would have. The cache is created per window in `stitch_seeds_core` and threaded down the recursion, which is what keeps the key complete: it never outlives the read, genome and strand it was filled for. An empty `HashMap` does not allocate, so windows that stitch nothing pay nothing. Measured on 50k yeast read pairs (Apple M4 Max, quiet machine, best-of-6 interleaved rounds, `--outSAMtype None`): threads before after change 1 19.51s 16.76s -14.1% 8 2.55s 2.18s -14.5% Output is byte-identical: `Aligned.out.sam` (records and header alike, modulo the `@PG` CL line naming the binary) and `SJ.out.tab` compare equal against the pre-change binary on the same input. 592 tests pass, 0 clippy warnings. Co-Authored-By: Claude Opus 5 (1M context) --- src/align/score.rs | 82 +++++++++++++++++++++++++++++++++++++++++++++ src/align/stitch.rs | 14 +++++++- 2 files changed, 95 insertions(+), 1 deletion(-) diff --git a/src/align/score.rs b/src/align/score.rs index 408a2d2..6d10b4a 100644 --- a/src/align/score.rs +++ b/src/align/score.rs @@ -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, @@ -658,6 +709,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 { diff --git a/src/align/stitch.rs b/src/align/stitch.rs index 8dcca5c..3c875c1 100644 --- a/src/align/stitch.rs +++ b/src/align/stitch.rs @@ -1155,6 +1155,7 @@ 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(); @@ -1393,7 +1394,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, @@ -2205,6 +2207,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; @@ -2453,6 +2456,7 @@ fn stitch_recurse( recursion_count, align_mates_gap_max, original_is_reverse, + jcache, debug_name, ); } else { @@ -2466,6 +2470,7 @@ fn stitch_recurse( cluster, junction_db, align_mates_gap_max, + jcache, debug_name, ) { stitch_recurse( @@ -2482,6 +2487,7 @@ fn stitch_recurse( recursion_count, align_mates_gap_max, original_is_reverse, + jcache, debug_name, ); } @@ -2512,6 +2518,7 @@ fn stitch_recurse( recursion_count, align_mates_gap_max, original_is_reverse, + jcache, debug_name, ); } @@ -3119,6 +3126,10 @@ 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, @@ -3134,6 +3145,7 @@ pub(crate) fn stitch_seeds_core( &mut recursion_count, align_mates_gap_max, stitch_is_reverse, + &mut jcache, debug_read_name, ); From c264e3bba5d218aa10dec0b0b739cbdbb07f3357 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Fri, 28 Aug 2026 19:10:40 +0200 Subject: [PATCH 2/5] perf(align): resolve the genome storage once per extension and junction scan `Genome::get_base` matches on the `GenomeSeq` discriminant, bounds-checks and, for a memory-mapped genome's reverse-complement half, recomputes the mirrored index and complements the byte. That is fine per call, but the alignment inner loops call it once per base: the junction-position scan reads two bases per candidate position across three loops, and `extend_alignment` reads one per extended base. Together those two functions are ~55% of alignment time after the scan memo. Add `GenomeSeq::view()`, returning a `Copy` `SeqView` that resolves the variant once. The view is a slice plus one integer, so the loops keep it in registers and each base costs a bounds check and a load. Hoist it out of the junction scan (including the sliding motif window) and out of both `extend_alignment` loops. Out-of-range reads return the `OUT_OF_RANGE` sentinel instead of `None`. Every call site converted here already treated "not one of A/C/G/T" the same way a `None` was treated, so the branch structure is preserved exactly; `score.rs` already had this sentinel locally for the motif window and it moves next to the view it belongs to. `SeqView::base` duplicates the reverse-complement arithmetic in `GenomeSeq::base`, so a unit test asserts the two agree on every index in `0..2n` plus the first out-of-range one, for both storage variants. Measured on 50k yeast read pairs (Apple M4 Max, quiet machine, best-of-6 interleaved rounds, `--outSAMtype None`, 8 threads), on top of the junction-scan memo: 2.18s to 2.12s, and 2.60s to 2.12s against the pre-memo baseline (-18%). Output is byte-identical: `Aligned.out.sam` and `SJ.out.tab` compare equal against the pre-change binary on the same input. 593 tests pass, 0 clippy warnings. Co-Authored-By: Claude Opus 5 (1M context) --- src/align/score.rs | 108 ++++++++++++++++++++------------------------ src/align/stitch.rs | 14 ++++-- src/genome/mod.rs | 108 ++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 166 insertions(+), 64 deletions(-) diff --git a/src/align/score.rs b/src/align/score.rs index 6d10b4a..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 @@ -356,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 @@ -375,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 { @@ -424,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; } - _ => {} } } } @@ -452,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); @@ -499,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; @@ -520,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; @@ -605,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 @@ -634,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), } } @@ -651,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); } } @@ -830,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 3c875c1..001c44f 100644 --- a/src/align/stitch.rs +++ b/src/align/stitch.rs @@ -207,6 +207,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 +250,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 +325,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 { diff --git a/src/genome/mod.rs b/src/genome/mod.rs index 134a229..e363522 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,47 @@ 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<'_> { + /// 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 +579,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"]; From 39b9175348b569bfd56e1b57a85c1da785f99a44 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Sat, 29 Aug 2026 02:03:46 +0200 Subject: [PATCH 3/5] perf(align): let score_region compare two byte slices instead of one base at a time `score_region` walks a seed-length run scoring matches and mismatches. It was ~11% of alignment time, and its loop could not vectorize: every base re-derived a bounds-checked `Option` from `Genome::get_base` and tested the read end, so the body carried branches and an early exit. Bound the run once against the read length, then walk it in 256-base chunks with the genome bases staged into a stack buffer through the new `SeqView::bases_into`. The inner loop is then two plain byte slices of equal length reduced into a match count and a mismatch count, which is what lets it vectorize; `score` is still exactly `matches - mismatches` and the genome-end `break` is preserved by stopping on a short fill. `bases_into` is one `copy_from_slice` on the forward strand. The reverse-complement half of a mapped genome has no contiguous slice to hand out, so it is filled by walking the mirrored forward bytes, which still leaves the comparison itself vectorizable. The reduction uses bitwise `&` rather than `&&` and suppresses `clippy::needless_bitwise_bool` at that loop. This is measured, not stylistic: the lazy spelling reintroduces branches and gives back most of the gain (median 2.125s vs 2.085s wall on the benchmark below). Measured on 50k yeast read pairs (Apple M4 Max, 8 threads, best-of-6 interleaved rounds, `--outSAMtype None`): 2.10s to 2.04s best, 2.155s to 2.085s median. Small, but consistent across every paired round. Output is byte-identical on two datasets: yeast 50k pairs and the nfcore test pair (88k SAM lines). `Aligned.out.sam` and `SJ.out.tab` compare equal against the pre-change binary in both. 593 tests pass, 0 clippy warnings, fmt clean. Co-Authored-By: Claude Opus 5 (1M context) --- src/align/stitch.rs | 68 ++++++++++++++++++++++++++++++++------------- src/genome/mod.rs | 35 +++++++++++++++++++++++ 2 files changed, 84 insertions(+), 19 deletions(-) diff --git a/src/align/stitch.rs b/src/align/stitch.rs index 001c44f..0886b71 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], diff --git a/src/genome/mod.rs b/src/genome/mod.rs index e363522..f4b1ae9 100644 --- a/src/genome/mod.rs +++ b/src/genome/mod.rs @@ -127,6 +127,41 @@ pub struct SeqView<'a> { 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)`. From ebbfc4b05cad8f10d13807c7078b5fa29e90cbb6 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Sat, 29 Aug 2026 02:33:23 +0200 Subject: [PATCH 4/5] perf(align): give the stitcher's transcript clone room for the push that follows `stitch_align_to_transcript` clones the working transcript and then always pushes onto it. `Vec::clone` allocates exactly `len`, so that push reallocates every time: a malloc, a copy and a free for every stitched seed, on a path the recursion walks up to its 100k-node budget per window. Add `WorkingTranscript::clone_with_headroom`, which reserves the one slot the caller is about to use, folding the reallocation back into the clone's own allocation. The junction vectors only get headroom when they already hold something. Most transcripts carry no junction at all, and giving an empty vector capacity would allocate for a push that never comes, which is worse than what it replaces. The exon vector is never empty at these call sites, so it always gets the slot. Measured on 50k yeast read pairs (Apple M4 Max, 1 thread, 5 rounds, `--outSAMtype None`), reported as user CPU time rather than wall: this machine's wall clock was too noisy to resolve half a percent, and user time is not. Median 14.86s to 14.79s, -0.5%, with every round at the same rank improving. Small, and labelled as such. Pre-sizing the per-window transcript accumulator was tried alongside this and measured a small loss (median 14.97s, +0.7%): it over-allocates for the many windows that finish with only a few transcripts. It is not included here. Output is byte-identical on two datasets: yeast 50k pairs and the nfcore test pair. `Aligned.out.sam` and `SJ.out.tab` compare equal against the pre-change binary in both. 593 tests pass, 0 clippy warnings, fmt clean. Co-Authored-By: Claude Opus 5 (1M context) --- src/align/stitch.rs | 39 +++++++++++++++++++++++++++++++++++++-- 1 file changed, 37 insertions(+), 2 deletions(-) diff --git a/src/align/stitch.rs b/src/align/stitch.rs index 0886b71..31e9d17 100644 --- a/src/align/stitch.rs +++ b/src/align/stitch.rs @@ -1159,6 +1159,41 @@ pub(crate) struct WorkingTranscript { } impl WorkingTranscript { + /// Clone, leaving room for the one element the stitcher is about to push. + /// + /// `Vec::clone` allocates exactly `len`, so the push that always follows + /// this clone in `stitch_align_to_transcript` reallocates every time: a + /// malloc, a copy and a free per stitched seed. Reserving the slot up + /// front folds that back into the clone's own allocation. + /// + /// The junction vectors only get headroom when they already hold + /// something. Most transcripts carry no junction at all, and giving an + /// empty vector capacity would allocate for a push that never comes, + /// which is worse than what it replaces. + fn clone_with_headroom(&self) -> Self { + fn grown(v: &[T], extra: usize) -> Vec { + let mut out = Vec::with_capacity(v.len() + extra); + out.extend_from_slice(v); + out + } + let junction_extra = usize::from(!self.junction_motifs.is_empty()); + WorkingTranscript { + exons: grown(&self.exons, 1), + junction_motifs: grown(&self.junction_motifs, junction_extra), + junction_annotated: grown(&self.junction_annotated, junction_extra), + junction_shifts: grown(&self.junction_shifts, junction_extra), + 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, + } + } + fn new() -> Self { WorkingTranscript { exons: Vec::new(), @@ -1252,7 +1287,7 @@ fn stitch_align_to_transcript( if align_mates_gap_max > 0 && genome_gap > align_mates_gap_max { return None; } - let mut new_wt = wt.clone(); + let mut new_wt = wt.clone_with_headroom(); // 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, ...) @@ -1377,7 +1412,7 @@ fn stitch_align_to_transcript( return None; } - let mut new_wt = wt.clone(); + let mut new_wt = wt.clone_with_headroom(); let mut d_score: i32 = 0; let mut gap_mm: u32 = 0; From 636206090b4edd411b5ffde8a486d328bac2f58f Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Sun, 30 Aug 2026 23:31:05 +0200 Subject: [PATCH 5/5] perf(align): stitch into the transcript in place instead of cloning per node `stitchWindowAligns`' include/exclude recursion cloned a `WorkingTranscript` at every node it explored, because the include branch needed its own copy while the exclude branch still needed the original. With a 100k-node budget per window that is a clone, four vector allocations and a copy, per node. Pass the transcript by `&mut` and undo the attempt instead. The mutation surface is small enough to invert exactly: `stitch_align_to_transcript` and the base-case extensions only append to the four vectors and rewrite the *first* and *last* exons plus the scalars, never a middle exon and never a removal or reorder. `WtMark` records the four lengths, copies of the first and last exons, and the scalars; `restore` truncates and writes those back. Vector capacity survives the restore, so sibling branches reuse the same allocations rather than reallocating. An owned clone is now taken only where one is actually needed: when a completed transcript is accepted into the result set, which is bounded by `--alignTranscriptsPerWindowNmax` rather than by the recursion. This supersedes `clone_with_headroom`, which existed to make the per-node clone cheaper and is removed: there is no per-node clone left to soften. Measured, user CPU at 1 thread, `--outSAMtype None`: | dataset | before | after | |---|---|---| | nfcore PE | 22.29s | 21.72s (-2.6%) | | yeast PE | 15.24s | 15.24s (neutral) | The gain tracks stitch depth: nfcore's reads drive a deeper recursion, so more clones disappear. Yeast is unchanged rather than slower, which is the point of checking both. Beyond the timing, this is what a batched or GPU-side stitcher would need: per-branch state is now a mark and an undo rather than a fresh allocation. Output is byte-identical on both datasets: `Aligned.out.sam` and `SJ.out.tab` compare equal against the pre-change binary. 593 tests pass, 0 clippy warnings, fmt clean. Co-Authored-By: Claude Opus 5 (1M context) --- src/align/stitch.rs | 237 ++++++++++++++++++++++++++------------------ 1 file changed, 142 insertions(+), 95 deletions(-) diff --git a/src/align/stitch.rs b/src/align/stitch.rs index 31e9d17..c1af450 100644 --- a/src/align/stitch.rs +++ b/src/align/stitch.rs @@ -1158,30 +1158,42 @@ 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 { - /// Clone, leaving room for the one element the stitcher is about to push. - /// - /// `Vec::clone` allocates exactly `len`, so the push that always follows - /// this clone in `stitch_align_to_transcript` reallocates every time: a - /// malloc, a copy and a free per stitched seed. Reserving the slot up - /// front folds that back into the clone's own allocation. - /// - /// The junction vectors only get headroom when they already hold - /// something. Most transcripts carry no junction at all, and giving an - /// empty vector capacity would allocate for a push that never comes, - /// which is worse than what it replaces. - fn clone_with_headroom(&self) -> Self { - fn grown(v: &[T], extra: usize) -> Vec { - let mut out = Vec::with_capacity(v.len() + extra); - out.extend_from_slice(v); - out - } - let junction_extra = usize::from(!self.junction_motifs.is_empty()); - WorkingTranscript { - exons: grown(&self.exons, 1), - junction_motifs: grown(&self.junction_motifs, junction_extra), - junction_annotated: grown(&self.junction_annotated, junction_extra), - junction_shifts: grown(&self.junction_shifts, junction_extra), + /// 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, @@ -1194,6 +1206,34 @@ impl WorkingTranscript { } } + /// 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(), @@ -1218,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, @@ -1228,8 +1268,11 @@ fn stitch_align_to_transcript( 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 @@ -1246,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 @@ -1267,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]; @@ -1279,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_with_headroom(); // 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, ...) @@ -1309,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. @@ -1342,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); @@ -1368,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 @@ -1386,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; @@ -1397,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; @@ -1409,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_with_headroom(); 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; } @@ -1437,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; } @@ -1450,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) --- @@ -1567,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| { @@ -1586,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; @@ -1612,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, @@ -1704,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; @@ -1721,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, @@ -1731,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. @@ -2266,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, @@ -2296,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, @@ -2472,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() @@ -2484,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; } @@ -2495,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, @@ -2530,10 +2572,14 @@ fn stitch_recurse( 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, @@ -2546,7 +2592,7 @@ fn stitch_recurse( ) { stitch_recurse( i_a + 1, - new_wt, + wt, wa_entries, read_seq, index, @@ -2562,6 +2608,7 @@ fn stitch_recurse( debug_name, ); } + wt.restore(&mark); } // EXCLUDE branch: skip wa_entries[i_a]. @@ -3204,7 +3251,7 @@ pub(crate) fn stitch_seeds_core( stitch_recurse( 0, - WorkingTranscript::new(), + &mut WorkingTranscript::new(), &wa_entries, stitch_read, index,