rustalign_aligner/

lib.rs

//! RustAlign Alignment Algorithms
//!
//! This crate implements the seed-based alignment and Smith-Waterman
//! dynamic programming algorithms used by RustAlign.

#![warn(missing_docs)]
#![warn(clippy::all)]

mod constraint;
mod paired_end;
mod result;
mod seed;
mod smith_waterman;
mod sw_sse;

#[cfg(test)]
mod proptest;

pub use constraint::{Constraint, SeedPolicy, SeedType};
pub use paired_end::{PairedAlignment, PeAlignmentType, PeConstraints, PePolicy};
pub use result::{Alignment, AlignmentResult};
pub use seed::{Seed, SeedAligner, SeedResults};
pub use smith_waterman::{SwAligner, SwParams};
pub use sw_sse::{QueryProfile, SwSseAligner};

use rustalign_common::{Nuc, Score, Strand};
use rustalign_fmindex::{EbwtWithRef, FmIndex};

/// Main aligner that orchestrates seed search and DP extension
pub struct Aligner<I: FmIndex> {
    /// FM-index for forward strand
    index_fw: I,

    /// FM-index for reverse strand
    #[allow(dead_code)]
    index_rc: I,

    /// Alignment parameters
    params: AlignerParams,
}

/// Aligner with reference sequence for proper SW scoring
///
/// This aligner extracts reference sequences from the index and uses them
/// for Smith-Waterman alignment, giving accurate scores for matches/mismatches.
pub struct AlignerWithRef {
    /// Combined FM-index and reference
    index_with_ref: EbwtWithRef,

    /// Alignment parameters
    params: AlignerParams,
}

/// Parameters for the aligner
#[derive(Debug, Clone)]
pub struct AlignerParams {
    /// Maximum number of alignments per read
    pub max_alignments: usize,

    /// Seed policy
    pub seed_policy: SeedPolicy,

    /// Smith-Waterman parameters
    pub sw_params: SwParams,

    /// Minimum score threshold
    pub min_score: Score,
}

impl Default for AlignerParams {
    fn default() -> Self {
        Self {
            max_alignments: 50, // Extend many seeds to find best alignment across chromosomes
            seed_policy: SeedPolicy::default(),
            sw_params: SwParams::default(),
            min_score: 30,
        }
    }
}

impl<I: FmIndex> Aligner<I> {
    /// Create a new aligner with the given FM-indexes
    pub fn new(index_fw: I, index_rc: I, params: AlignerParams) -> Self {
        Self {
            index_fw,
            index_rc,
            params,
        }
    }

    /// Align a single read
    ///
    /// Searches both forward and reverse complement strands and returns
    /// the best alignment found.
    #[allow(clippy::collapsible_if)]
    pub fn align(&self, read: &[Nuc], read_rc: &[Nuc]) -> AlignmentResult {
        use rustalign_common::Strand;

        let mut all_alignments = Vec::new();

        // 1. Search forward strand (using the original read)
        if let Ok(seed_results) = self.search_seeds(read) {
            if !seed_results.hits.is_empty() {
                for seed_hit in seed_results.hits.iter().take(self.params.max_alignments) {
                    if let Ok(mut aln) = self.extend_seed(read, seed_hit) {
                        aln.strand = Strand::Forward;
                        if aln.score >= self.params.min_score {
                            all_alignments.push(aln);
                        }
                    }
                }
            }
        }

        // 2. Search reverse complement strand (using the RC of the read)
        if let Ok(seed_results) = self.search_seeds(read_rc) {
            if !seed_results.hits.is_empty() {
                for seed_hit in seed_results.hits.iter().take(self.params.max_alignments) {
                    if let Ok(mut aln) = self.extend_seed(read_rc, seed_hit) {
                        aln.strand = Strand::Reverse;
                        if aln.score >= self.params.min_score {
                            all_alignments.push(aln);
                        }
                    }
                }
            }
        }

        // 3. Return results
        if all_alignments.is_empty() {
            return AlignmentResult {
                success: false,
                error: Some("No seed matches found on either strand".to_string()),
                ..Default::default()
            };
        }

        // Sort by score (best first)
        all_alignments.sort_by(|a, b| b.score.cmp(&a.score));

        AlignmentResult {
            success: true,
            alignments: all_alignments,
            ..Default::default()
        }
    }

    /// Search for seed matches
    fn search_seeds(&self, read: &[Nuc]) -> rustalign_common::Result<SeedResults> {
        let mut seed_aligner = SeedAligner::new(self.params.seed_policy.clone());
        seed_aligner.search(&self.index_fw, read)
    }

    /// Extend a seed with Smith-Waterman
    fn extend_seed(&self, read: &[Nuc], seed: &Seed) -> rustalign_common::Result<Alignment> {
        let sw = SwAligner::new(self.params.sw_params);
        // TODO: Extract reference sequence around seed position
        // For now, just use the read itself as placeholder
        let mut aln = sw.align(read, read)?;
        // Store the BWT row for later position resolution
        aln.bwt_row = seed.bwt_row;
        // Store the seed's position in the read for position calculation
        aln.seed_offset = seed.read_pos;
        // Store how many positions this seed matched (for tiebreaking)
        aln.seed_hit_count = seed.hit_count;
        Ok(aln)
    }
}

impl AlignerWithRef {
    /// Create a new aligner with reference
    pub fn new(index_with_ref: EbwtWithRef, params: AlignerParams) -> Self {
        Self {
            index_with_ref,
            params,
        }
    }

    /// Align a single read using actual reference sequences
    ///
    /// Searches both forward and reverse complement strands and returns
    /// all valid alignments with scores computed against actual reference.
    #[allow(clippy::collapsible_if)]
    pub fn align(&self, read: &[Nuc], read_rc: &[Nuc]) -> AlignmentResult {
        let mut all_alignments = Vec::new();

        // 1. Search forward strand
        if let Ok(seed_results) = self.search_seeds(read) {
            if !seed_results.hits.is_empty() {
                for seed_hit in seed_results.hits.iter().take(self.params.max_alignments) {
                    if let Ok(aln) = self.extend_seed_with_ref(read, seed_hit, Strand::Forward) {
                        if aln.score >= self.params.min_score {
                            all_alignments.push(aln);
                        }
                    }
                }
            }
        }

        // 2. Search reverse complement strand
        if let Ok(seed_results) = self.search_seeds(read_rc) {
            if !seed_results.hits.is_empty() {
                for seed_hit in seed_results.hits.iter().take(self.params.max_alignments) {
                    if let Ok(aln) = self.extend_seed_with_ref(read_rc, seed_hit, Strand::Reverse) {
                        if aln.score >= self.params.min_score {
                            all_alignments.push(aln);
                        }
                    }
                }
            }
        }

        // 3. Return results
        if all_alignments.is_empty() {
            return AlignmentResult {
                success: false,
                error: Some("No seed matches found on either strand".to_string()),
                ..Default::default()
            };
        }

        // Sort by score (best first)
        all_alignments.sort_by(|a, b| b.score.cmp(&a.score));

        AlignmentResult {
            success: true,
            alignments: all_alignments,
            ..Default::default()
        }
    }

    /// Search for seed matches
    fn search_seeds(&self, read: &[Nuc]) -> rustalign_common::Result<SeedResults> {
        let mut seed_aligner = SeedAligner::new(self.params.seed_policy.clone());
        seed_aligner.search(self.index_with_ref.ebwt(), read)
    }

    /// Extend a seed with Smith-Waterman using actual reference sequence
    #[allow(clippy::implicit_saturating_sub)]
    fn extend_seed_with_ref(
        &self,
        read: &[Nuc],
        seed: &Seed,
        strand: Strand,
    ) -> rustalign_common::Result<Alignment> {
        let read_len = read.len() as u64;

        // Step 1: Resolve position from BWT row
        let (chr_idx, pos_in_chr) = {
            let genome_off = self.index_with_ref.ebwt().get_offset(seed.bwt_row)?;

            // Get chromosome index and position
            // We need to use joined_to_text_off, which requires the internal method
            let (chr_idx, text_off, _chr_len) = self
                .index_with_ref
                .ebwt()
                .joined_to_text_off(read_len, genome_off, true)?;

            // Adjust for seed offset within the read
            let adjusted_pos = if text_off > seed.read_pos as u64 {
                text_off - seed.read_pos as u64
            } else {
                0
            };

            (chr_idx, adjusted_pos)
        };

        // Step 2: Extract reference sequence
        // Extract enough reference for the entire read plus some padding
        let ref_len = read.len() + 20; // Extra padding for potential indels
        let reference = self
            .index_with_ref
            .get_reference(chr_idx, pos_in_chr, ref_len)?;

        // Step 3: Perform Smith-Waterman alignment against actual reference
        let sw = SwAligner::new(self.params.sw_params);
        let mut aln = if reference.len() >= read.len() {
            sw.align(read, &reference)?
        } else {
            // Reference too short (near end of chromosome), use read as fallback
            sw.align(read, read)?
        };

        // Store metadata
        aln.bwt_row = seed.bwt_row;
        aln.seed_offset = seed.read_pos;
        aln.seed_hit_count = seed.hit_count;
        aln.strand = strand;

        Ok(aln)
    }

    /// Get the underlying EbwtWithRef
    pub fn index(&self) -> &EbwtWithRef {
        &self.index_with_ref
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_params_default() {
        let params = AlignerParams::default();
        assert_eq!(params.max_alignments, 50);
        assert_eq!(params.min_score, 30);
    }
}