perf(align): memoize the junction-position scan inside stitchWindowAligns

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<JunctionScanKey, (i32, SpliceMotif, i32, u32, u32)>,
+}
+
+impl JunctionScanCache {
+    pub fn new() -> Self {
+        Self::default()
+    }
+}
+
 /// Splice junction motif types
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 pub enum SpliceMotif {

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<WorkingTranscript> {
     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<WorkingTranscript> = 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,
     );