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, );