jellyfish-reader 0.1.0

use crate::mer::{MerDna, encode_base};

/// Iterator that extracts k-mers from a DNA sequence string.
///
/// Slides a window of size `k` across the input sequence, yielding each
/// valid k-mer. Invalid characters (non-ACGT) cause the window to reset.
///
/// # Examples
///
/// ```
/// use jellyfish_reader::StringMers;
///
/// let mers: Vec<_> = StringMers::new("ACGTACGT", 4)
///     .map(|m| m.to_string())
///     .collect();
/// assert_eq!(mers, vec!["ACGT", "CGTA", "GTAC", "TACG", "ACGT"]);
/// ```
pub struct StringMers<'a> {
    seq: &'a [u8],
    k: usize,
    pos: usize,
    valid_count: usize,
    current: MerDna,
    canonical: bool,
}

impl<'a> StringMers<'a> {
    /// Create a new k-mer iterator over the given DNA sequence.
    pub fn new(seq: &'a str, k: usize) -> Self {
        Self {
            seq: seq.as_bytes(),
            k,
            pos: 0,
            valid_count: 0,
            current: MerDna::new(k),
            canonical: false,
        }
    }

    /// Create a new canonical k-mer iterator.
    ///
    /// Each yielded k-mer is the lexicographically smaller of itself
    /// and its reverse complement.
    pub fn canonicals(seq: &'a str, k: usize) -> Self {
        Self {
            seq: seq.as_bytes(),
            k,
            pos: 0,
            valid_count: 0,
            current: MerDna::new(k),
            canonical: true,
        }
    }
}

impl<'a> Iterator for StringMers<'a> {
    type Item = MerDna;

    fn next(&mut self) -> Option<MerDna> {
        while self.pos < self.seq.len() {
            let ch = self.seq[self.pos];
            self.pos += 1;

            match encode_base(ch) {
                Some(_code) => {
                    self.current.shift_left(ch);
                    self.valid_count += 1;

                    if self.valid_count >= self.k {
                        return if self.canonical {
                            Some(self.current.get_canonical())
                        } else {
                            Some(self.current.clone())
                        };
                    }
                }
                None => {
                    // Invalid character, reset window
                    self.valid_count = 0;
                    self.current = MerDna::new(self.k);
                }
            }
        }
        None
    }
}

/// Convenience function to create a k-mer iterator.
pub fn string_mers(seq: &str, k: usize) -> StringMers<'_> {
    StringMers::new(seq, k)
}

/// Convenience function to create a canonical k-mer iterator.
pub fn string_canonicals(seq: &str, k: usize) -> StringMers<'_> {
    StringMers::canonicals(seq, k)
}

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

    #[test]
    fn test_basic_extraction() {
        let mers: Vec<String> = StringMers::new("ACGTACGT", 4)
            .map(|m| m.to_string())
            .collect();
        assert_eq!(mers, vec!["ACGT", "CGTA", "GTAC", "TACG", "ACGT"]);
    }

    #[test]
    fn test_exact_length() {
        let mers: Vec<String> = StringMers::new("ACGT", 4).map(|m| m.to_string()).collect();
        assert_eq!(mers, vec!["ACGT"]);
    }

    #[test]
    fn test_shorter_than_k() {
        let mers: Vec<String> = StringMers::new("ACG", 4).map(|m| m.to_string()).collect();
        assert!(mers.is_empty());
    }

    #[test]
    fn test_empty_string() {
        let mers: Vec<String> = StringMers::new("", 4).map(|m| m.to_string()).collect();
        assert!(mers.is_empty());
    }

    #[test]
    fn test_skip_invalid_characters() {
        // N in the middle should reset the window
        let mers: Vec<String> = StringMers::new("ACGTNACGT", 4)
            .map(|m| m.to_string())
            .collect();
        // First 4-mer: "ACGT", then N resets, then we need 4 more valid bases
        assert_eq!(mers, vec!["ACGT", "ACGT"]);
    }

    #[test]
    fn test_invalid_at_start() {
        let mers: Vec<String> = StringMers::new("NACGT", 4).map(|m| m.to_string()).collect();
        assert_eq!(mers, vec!["ACGT"]);
    }

    #[test]
    fn test_all_invalid() {
        let mers: Vec<String> = StringMers::new("NNNN", 4).map(|m| m.to_string()).collect();
        assert!(mers.is_empty());
    }

    #[test]
    fn test_k_equals_1() {
        let mers: Vec<String> = StringMers::new("ACGT", 1).map(|m| m.to_string()).collect();
        assert_eq!(mers, vec!["A", "C", "G", "T"]);
    }

    #[test]
    fn test_lowercase() {
        let mers: Vec<String> = StringMers::new("acgt", 4).map(|m| m.to_string()).collect();
        assert_eq!(mers, vec!["ACGT"]);
    }

    #[test]
    fn test_canonical_mode() {
        // AAAA and TTTT are reversecomplement pairs; canonical picks AAAA
        let mers: Vec<String> = StringMers::canonicals("TTTT", 4)
            .map(|m| m.to_string())
            .collect();
        assert_eq!(mers, vec!["AAAA"]);
    }

    #[test]
    fn test_canonical_palindrome() {
        let mers: Vec<String> = StringMers::canonicals("ACGT", 4)
            .map(|m| m.to_string())
            .collect();
        assert_eq!(mers, vec!["ACGT"]);
    }

    #[test]
    fn test_canonical_various() {
        // For each k-mer, canonical should be the lexicographically smaller of
        // itself and its reverse complement
        let mers: Vec<String> = StringMers::canonicals("ACGTACGT", 4)
            .map(|m| m.to_string())
            .collect();

        // Verify each is canonical (equal to its own canonical)
        for mer_str in &mers {
            let mer: MerDna = mer_str.parse().unwrap();
            assert_eq!(mer.get_canonical().to_string(), *mer_str);
        }
    }

    #[test]
    fn test_convenience_functions() {
        let mers1: Vec<String> = string_mers("ACGT", 4).map(|m| m.to_string()).collect();
        let mers2: Vec<String> = StringMers::new("ACGT", 4).map(|m| m.to_string()).collect();
        assert_eq!(mers1, mers2);

        let can1: Vec<String> = string_canonicals("ACGT", 4)
            .map(|m| m.to_string())
            .collect();
        let can2: Vec<String> = StringMers::canonicals("ACGT", 4)
            .map(|m| m.to_string())
            .collect();
        assert_eq!(can1, can2);
    }

    #[test]
    fn test_multiple_invalid_regions() {
        let mers: Vec<String> = StringMers::new("ACGTNNTACG", 3)
            .map(|m| m.to_string())
            .collect();
        // "ACG" then "T" then NN resets, then "TAC" then "ACG"
        assert_eq!(mers, vec!["ACG", "CGT", "TAC", "ACG"]);
    }

    #[test]
    fn test_long_sequence() {
        let seq = "ACGT".repeat(100);
        let mers: Vec<String> = StringMers::new(&seq, 25).map(|m| m.to_string()).collect();
        assert_eq!(mers.len(), seq.len() - 24);
    }

    #[test]
    fn test_homopolymer_run() {
        let mers: Vec<String> = StringMers::new("AAAAA", 3).map(|m| m.to_string()).collect();
        assert_eq!(mers, vec!["AAA", "AAA", "AAA"]);
    }
}