59 return equal cost range for each ts

lib_tsalign/src/a_star_aligner.rs

@@ -159,7 +159,10 @@ pub fn template_switch_distance_a_star_align<
     cost_limit: Option<Strategies::Cost>,
     memory_limit: Option<usize>,
     template_switch_count_memory: <Strategies::TemplateSwitchCount as TemplateSwitchCountStrategy>::Memory,
-) -> AlignmentResult<template_switch_distance::AlignmentType, Strategies::Cost> {
+) -> AlignmentResult<template_switch_distance::AlignmentType, Strategies::Cost>
+where
+    Strategies::Cost: From<u64>,
+{
     let memory = Memory {
         template_switch_min_length: Default::default(),
         chaining: <<Strategies as AlignmentStrategySelector>::Chaining as ChainingStrategy<
@@ -170,18 +173,20 @@ pub fn template_switch_distance_a_star_align<
         primary_match: (),
     };
 
-    a_star_align(template_switch_distance::Context::<
+    let mut result = a_star_align(template_switch_distance::Context::<
         SubsequenceType,
         Strategies,
     >::new(
         reference,
         query,
         reference_name,
         query_name,
-        range,
-        config,
+        range.clone(),
+        config.clone(),
         memory,
         cost_limit,
         memory_limit,
-    ))
+    ));
+    result.compute_ts_equal_cost_ranges(reference, query, &range, &config);
+    result
 }

lib_tsalign/src/a_star_aligner/alignment_result.rs

@@ -1,12 +1,16 @@
 use std::fmt::{Display, Formatter, Result, Write};
 
 use a_star_sequences::SequencePair;
-use alignment::Alignment;
+use alignment::{Alignment, stream::AlignmentStream};
 use compact_genome::interface::{alphabet::Alphabet, sequence::GenomeSequence};
 use generic_a_star::{AStarResult, cost::AStarCost};
 use noisy_float::types::{R64, r64};
 use num_traits::{Float, Zero};
 
+use crate::config::TemplateSwitchConfig;
+
+use super::alignment_geometry::AlignmentRange;
+
 pub mod a_star_sequences;
 pub mod alignment;
 
@@ -17,6 +21,8 @@ pub trait IAlignmentType {
 
     fn is_internal(&self) -> bool;
 
+    fn is_template_switch_entrance(&self) -> bool;
+
     fn is_template_switch_exit(&self) -> bool;
 }
 
@@ -225,6 +231,178 @@ impl<AlignmentType: IAlignmentType, Cost: AStarCost> AlignmentResult<AlignmentTy
     }
 }
 
+impl<Cost: AStarCost + From<u64>>
+    AlignmentResult<super::template_switch_distance::AlignmentType, Cost>
+{
+    pub fn compute_ts_equal_cost_ranges<
+        AlphabetType: Alphabet,
+        SubsequenceType: GenomeSequence<AlphabetType, SubsequenceType> + ?Sized,
+    >(
+        &mut self,
+        reference: &SubsequenceType,
+        query: &SubsequenceType,
+        range: &Option<AlignmentRange>,
+        config: &TemplateSwitchConfig<AlphabetType, Cost>,
+    ) {
+        let Self::WithTarget {
+            alignment,
+            statistics,
+        } = self
+        else {
+            return;
+        };
+        let mut stream = AlignmentStream::new();
+        let reference_offset = range
+            .as_ref()
+            .map(|range| range.reference_offset())
+            .unwrap_or(0);
+        let query_offset = range
+            .as_ref()
+            .map(|range| range.query_offset())
+            .unwrap_or(0);
+
+        for i in 0..alignment.inner_mut().len() {
+            let multiplicity = alignment.inner_mut()[i].0;
+            let alignment_type = alignment.inner_mut()[i].1;
+
+            match alignment_type {
+                super::template_switch_distance::AlignmentType::TemplateSwitchEntrance {
+                    mut equal_cost_range,
+                    ..
+                } => {
+                    if config.left_flank_length == 0 && config.right_flank_length == 0 {
+                        equal_cost_range.min_start = 0;
+                        equal_cost_range.max_start = 0;
+                        equal_cost_range.min_end = 0;
+                        equal_cost_range.max_end = 0;
+
+                        let initial_cost = alignment.compute_cost(
+                            reference,
+                            query,
+                            reference_offset,
+                            query_offset,
+                            config,
+                        );
+                        assert_eq!(initial_cost, (statistics.cost.round().raw() as u64).into());
+
+                        {
+                            let mut min_start_alignment = alignment.clone();
+                            while min_start_alignment.move_template_switch_end_forwards(
+                                reference,
+                                query,
+                                reference_offset,
+                                query_offset,
+                                i,
+                            ) {
+                                let new_cost = min_start_alignment.compute_cost(
+                                    reference,
+                                    query,
+                                    reference_offset,
+                                    query_offset,
+                                    config,
+                                );
+                                if new_cost > initial_cost {
+                                    break;
+                                } else {
+                                    assert_eq!(new_cost, initial_cost);
+                                    equal_cost_range.min_start -= 1;
+                                }
+                            }
+                        }
+
+                        {
+                            let mut max_start_alignment = alignment.clone();
+                            while max_start_alignment.move_template_switch_start_forwards(
+                                reference,
+                                query,
+                                reference_offset,
+                                query_offset,
+                                i,
+                            ) {
+                                let new_cost = max_start_alignment.compute_cost(
+                                    reference,
+                                    query,
+                                    reference_offset,
+                                    query_offset,
+                                    config,
+                                );
+                                if new_cost > initial_cost {
+                                    break;
+                                } else {
+                                    assert_eq!(new_cost, initial_cost);
+                                    equal_cost_range.max_start += 1;
+                                }
+                            }
+                        }
+
+                        {
+                            let mut min_end_alignment = alignment.clone();
+                            while min_end_alignment.move_template_switch_end_backwards(
+                                reference,
+                                query,
+                                reference_offset,
+                                query_offset,
+                                i,
+                            ) {
+                                let new_cost = min_end_alignment.compute_cost(
+                                    reference,
+                                    query,
+                                    reference_offset,
+                                    query_offset,
+                                    config,
+                                );
+                                if new_cost > initial_cost {
+                                    break;
+                                } else {
+                                    assert_eq!(new_cost, initial_cost);
+                                    equal_cost_range.min_end -= 1;
+                                }
+                            }
+                        }
+
+                        {
+                            let mut max_end_alignment = alignment.clone();
+                            while max_end_alignment.move_template_switch_end_forwards(
+                                reference,
+                                query,
+                                reference_offset,
+                                query_offset,
+                                i,
+                            ) {
+                                let new_cost = max_end_alignment.compute_cost(
+                                    reference,
+                                    query,
+                                    reference_offset,
+                                    query_offset,
+                                    config,
+                                );
+                                if new_cost > initial_cost {
+                                    break;
+                                } else {
+                                    assert_eq!(new_cost, initial_cost);
+                                    equal_cost_range.max_end += 1;
+                                }
+                            }
+                        }
+
+                        let super::template_switch_distance::AlignmentType::TemplateSwitchEntrance {
+                        equal_cost_range: alignment_equal_cost_range,
+                        ..
+                    } = &mut alignment.inner_mut()[i].1
+                    else {
+                        unreachable!()
+                    };
+                        *alignment_equal_cost_range = equal_cost_range;
+                    }
+                }
+                _ => { /* Do nothing. */ }
+            }
+
+            stream.push(multiplicity, alignment_type);
+        }
+    }
+}
+
 impl<AlignmentType, Cost> AlignmentResult<AlignmentType, Cost> {
     pub fn statistics(&self) -> &AlignmentStatistics<Cost> {
         match self {

lib_tsalign/src/a_star_aligner/alignment_result/alignment.rs

@@ -7,6 +7,8 @@ use iter::{
 use super::IAlignmentType;
 
 pub mod iter;
+pub mod stream;
+pub mod template_switch_specifics;
 
 #[derive(Debug, Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
 #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
@@ -33,6 +35,10 @@ impl<AlignmentType> Alignment<AlignmentType> {
             self.alignment.push((1, alignment_type));
         }
     }
+
+    pub fn inner_mut(&mut self) -> &mut Vec<(usize, AlignmentType)> {
+        &mut self.alignment
+    }
 }
 
 impl<AlignmentType: IAlignmentType> Alignment<AlignmentType> {