jellyfish_reader/

mer.rs

use std::cmp::Ordering;
use std::fmt;
use std::hash::{Hash, Hasher};
use std::str::FromStr;

use crate::error::Error;

/// Number of bases that fit in a single u64 word.
const BASES_PER_WORD: usize = 32;

/// Encode a DNA base character to its 2-bit representation.
///
/// A=0, C=1, G=2, T=3. Returns None for invalid characters.
#[inline]
pub fn encode_base(c: u8) -> Option<u8> {
    match c {
        b'A' | b'a' => Some(0),
        b'C' | b'c' => Some(1),
        b'G' | b'g' => Some(2),
        b'T' | b't' => Some(3),
        _ => None,
    }
}

/// Decode a 2-bit encoding back to a DNA base character.
#[inline]
pub fn decode_base(code: u8) -> u8 {
    match code & 0x3 {
        0 => b'A',
        1 => b'C',
        2 => b'G',
        3 => b'T',
        _ => unreachable!(),
    }
}

/// Complement a 2-bit encoded base (A↔T, C↔G).
#[inline]
pub fn complement_code(code: u8) -> u8 {
    code ^ 0x3
}

/// Compute the number of u64 words needed for k bases.
#[inline]
pub fn words_for_k(k: usize) -> usize {
    k.div_ceil(BASES_PER_WORD)
}

/// Reverse complement of a single u64 word containing packed 2-bit bases.
///
/// Reverses the order of 2-bit pairs and complements each.
fn word_reverse_complement(mut word: u64) -> u64 {
    // Complement all bits (swaps A↔T, C↔G)
    word = !word;
    // Reverse 2-bit pairs using byte-swap and bit manipulation
    // Swap adjacent 2-bit groups
    word = ((word >> 2) & 0x3333_3333_3333_3333) | ((word & 0x3333_3333_3333_3333) << 2);
    // Swap adjacent 4-bit groups
    word = ((word >> 4) & 0x0F0F_0F0F_0F0F_0F0F) | ((word & 0x0F0F_0F0F_0F0F_0F0F) << 4);
    // Reverse bytes
    word.swap_bytes()
}

/// A DNA k-mer stored as 2-bit packed encoding in u64 words.
///
/// This matches the Jellyfish MerDNA representation:
/// - Each base is 2 bits: A=0, C=1, G=2, T=3
/// - Bases are packed from LSB to MSB within each word
/// - Words are ordered from least significant to most significant
///
/// # Examples
///
/// ```
/// use jellyfish_reader::MerDna;
///
/// let mer: MerDna = "ACGT".parse().unwrap();
/// assert_eq!(mer.to_string(), "ACGT");
/// assert_eq!(mer.k(), 4);
///
/// let rc = mer.get_reverse_complement();
/// assert_eq!(rc.to_string(), "ACGT"); // ACGT is its own reverse complement
/// ```
#[derive(Clone)]
pub struct MerDna {
    /// Packed 2-bit encoded bases.
    words: Vec<u64>,
    /// K-mer length (number of bases).
    k: usize,
}

impl MerDna {
    /// Create a new k-mer of the given length, initialized to all A's.
    pub fn new(k: usize) -> Self {
        Self {
            words: vec![0u64; words_for_k(k)],
            k,
        }
    }

    /// Create a MerDna from raw word data and k-mer length.
    ///
    /// The words should contain 2-bit packed bases matching Jellyfish's encoding.
    pub fn from_words(words: Vec<u64>, k: usize) -> Self {
        debug_assert!(words.len() == words_for_k(k));
        let mut mer = Self { words, k };
        mer.clean_high_bits();
        mer
    }

    /// Create a MerDna by reading packed bytes (as stored in Jellyfish binary files).
    ///
    /// Bytes are read in order and packed into u64 words in little-endian byte order.
    pub fn from_bytes(bytes: &[u8], k: usize) -> Self {
        let n_words = words_for_k(k);
        let mut words = vec![0u64; n_words];

        for (i, &byte) in bytes.iter().enumerate() {
            let word_idx = i / 8;
            let byte_idx = i % 8;
            if word_idx < n_words {
                words[word_idx] |= (byte as u64) << (byte_idx * 8);
            }
        }

        let mut mer = Self { words, k };
        mer.clean_high_bits();
        mer
    }

    /// The k-mer length (number of bases).
    #[inline]
    pub fn k(&self) -> usize {
        self.k
    }

    /// Access the raw u64 words.
    #[inline]
    pub fn words(&self) -> &[u64] {
        &self.words
    }

    /// Get the base at position `i` (0-indexed from the right/LSB end).
    ///
    /// # Panics
    /// Panics if `i >= k`.
    pub fn get_base(&self, i: usize) -> u8 {
        assert!(i < self.k, "base index {i} out of range for k={}", self.k);
        let word_idx = i / BASES_PER_WORD;
        let bit_offset = (i % BASES_PER_WORD) * 2;
        ((self.words[word_idx] >> bit_offset) & 0x3) as u8
    }

    /// Set the base at position `i` (0-indexed from the right/LSB end).
    ///
    /// # Panics
    /// Panics if `i >= k` or if `base_code` is not in 0..4.
    pub fn set_base(&mut self, i: usize, base_code: u8) {
        assert!(i < self.k, "base index {i} out of range for k={}", self.k);
        assert!(base_code < 4, "invalid base code: {base_code}");
        let word_idx = i / BASES_PER_WORD;
        let bit_offset = (i % BASES_PER_WORD) * 2;
        self.words[word_idx] &= !(0x3u64 << bit_offset);
        self.words[word_idx] |= (base_code as u64) << bit_offset;
    }

    /// Shift the k-mer left by one position, inserting `base` at the right end.
    ///
    /// Returns the base character that was shifted out from the left end.
    pub fn shift_left(&mut self, base: u8) -> Option<u8> {
        let code = encode_base(base)?;
        let old_high = self.get_base(self.k - 1);

        // Shift each word left by 2 bits, propagating carries
        let n = self.words.len();
        for i in (1..n).rev() {
            self.words[i] = (self.words[i] << 2) | (self.words[i - 1] >> 62);
        }
        self.words[0] = (self.words[0] << 2) | (code as u64);
        self.clean_high_bits();

        Some(decode_base(old_high))
    }

    /// Shift the k-mer right by one position, inserting `base` at the left end.
    ///
    /// Returns the base character that was shifted out from the right end.
    pub fn shift_right(&mut self, base: u8) -> Option<u8> {
        let code = encode_base(base)?;
        let old_low = self.get_base(0);

        // Shift each word right by 2 bits, propagating carries
        let n = self.words.len();
        for i in 0..n - 1 {
            self.words[i] = (self.words[i] >> 2) | (self.words[i + 1] << 62);
        }
        self.words[n - 1] >>= 2;

        // Insert new base at the high end
        let high_pos = self.k - 1;
        let word_idx = high_pos / BASES_PER_WORD;
        let bit_offset = (high_pos % BASES_PER_WORD) * 2;
        self.words[word_idx] |= (code as u64) << bit_offset;

        Some(decode_base(old_low))
    }

    /// Compute the reverse complement of this k-mer.
    pub fn get_reverse_complement(&self) -> MerDna {
        let n = self.words.len();

        if n == 1 {
            let mut result = vec![0u64; 1];
            result[0] = word_reverse_complement(self.words[0]) >> (64 - self.k * 2);
            let mut mer = MerDna {
                words: result,
                k: self.k,
            };
            mer.clean_high_bits();
            return mer;
        }

        // For multi-word k-mers, use base-by-base approach for correctness.
        // Position i in self maps to position (k-1-i) in result, with complement.
        let mut result = MerDna::new(self.k);
        for i in 0..self.k {
            let base = self.get_base(i);
            result.set_base(self.k - 1 - i, complement_code(base));
        }
        result
    }

    /// Modify this k-mer in place to its reverse complement.
    pub fn reverse_complement(&mut self) {
        *self = self.get_reverse_complement();
    }

    /// Get the canonical form (lexicographically smaller of self and reverse complement).
    pub fn get_canonical(&self) -> MerDna {
        let rc = self.get_reverse_complement();
        if *self <= rc { self.clone() } else { rc }
    }

    /// Modify this k-mer in place to its canonical form.
    pub fn canonicalize(&mut self) {
        let rc = self.get_reverse_complement();
        if rc < *self {
            *self = rc;
        }
    }

    /// Check if this k-mer is a homopolymer (all same base).
    pub fn is_homopolymer(&self) -> bool {
        if self.k == 0 {
            return true;
        }
        let base = self.get_base(0);
        (1..self.k).all(|i| self.get_base(i) == base)
    }

    /// Set all bases to A.
    pub fn poly_a(&mut self) {
        self.words.fill(0);
    }

    /// Set all bases to C.
    pub fn poly_c(&mut self) {
        self.fill_with_code(1);
    }

    /// Set all bases to G.
    pub fn poly_g(&mut self) {
        self.fill_with_code(2);
    }

    /// Set all bases to T.
    pub fn poly_t(&mut self) {
        self.fill_with_code(3);
    }

    /// Fill all bases with the given 2-bit code.
    fn fill_with_code(&mut self, code: u8) {
        let pattern = match code {
            0 => 0x0000_0000_0000_0000u64,
            1 => 0x5555_5555_5555_5555u64,
            2 => 0xAAAA_AAAA_AAAA_AAAAu64,
            3 => 0xFFFF_FFFF_FFFF_FFFFu64,
            _ => unreachable!(),
        };
        self.words.fill(pattern);
        self.clean_high_bits();
    }

    /// Zero out bits above the k-mer length in the highest word.
    fn clean_high_bits(&mut self) {
        if self.k == 0 {
            return;
        }
        let used_bits = self.k * 2;
        let total_bits = self.words.len() * 64;
        if used_bits < total_bits {
            let last = self.words.len() - 1;
            let bits_in_last = used_bits - last * 64;
            self.words[last] &= (1u64 << bits_in_last) - 1;
        }
    }
}

impl fmt::Debug for MerDna {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "MerDna(\"{}\")", self)
    }
}

impl fmt::Display for MerDna {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        for i in (0..self.k).rev() {
            let code = self.get_base(i);
            f.write_str(std::str::from_utf8(&[decode_base(code)]).unwrap())?;
        }
        Ok(())
    }
}

impl FromStr for MerDna {
    type Err = Error;

    fn from_str(s: &str) -> Result<Self, Error> {
        let k = s.len();
        if k == 0 {
            return Err(Error::InvalidKmer("empty k-mer string".to_string()));
        }

        let mut mer = MerDna::new(k);
        let bytes = s.as_bytes();

        for (i, &ch) in bytes.iter().enumerate() {
            let code = encode_base(ch).ok_or_else(|| {
                Error::InvalidKmer(format!("invalid base '{}' at position {i}", ch as char))
            })?;
            // String is stored with first character at highest position
            let pos = k - 1 - i;
            mer.set_base(pos, code);
        }

        Ok(mer)
    }
}

impl PartialEq for MerDna {
    fn eq(&self, other: &Self) -> bool {
        self.k == other.k && self.words == other.words
    }
}

impl Eq for MerDna {}

impl PartialOrd for MerDna {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for MerDna {
    fn cmp(&self, other: &Self) -> Ordering {
        // Compare from most significant word to least
        assert_eq!(
            self.k, other.k,
            "cannot compare k-mers of different lengths"
        );
        for i in (0..self.words.len()).rev() {
            match self.words[i].cmp(&other.words[i]) {
                Ordering::Equal => continue,
                ord => return ord,
            }
        }
        Ordering::Equal
    }
}

impl Hash for MerDna {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.k.hash(state);
        self.words.hash(state);
    }
}

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

    #[test]
    fn test_encode_decode_bases() {
        for (ch, code) in [(b'A', 0), (b'C', 1), (b'G', 2), (b'T', 3)] {
            assert_eq!(encode_base(ch), Some(code));
            assert_eq!(decode_base(code), ch);
        }
        // Lowercase
        for (ch, code) in [(b'a', 0), (b'c', 1), (b'g', 2), (b't', 3)] {
            assert_eq!(encode_base(ch), Some(code));
        }
        assert_eq!(encode_base(b'N'), None);
        assert_eq!(encode_base(b'X'), None);
    }

    #[test]
    fn test_complement_code() {
        assert_eq!(complement_code(0), 3); // A -> T
        assert_eq!(complement_code(1), 2); // C -> G
        assert_eq!(complement_code(2), 1); // G -> C
        assert_eq!(complement_code(3), 0); // T -> A
    }

    #[test]
    fn test_new_mer() {
        let mer = MerDna::new(4);
        assert_eq!(mer.k(), 4);
        assert_eq!(mer.to_string(), "AAAA");
    }

    #[test]
    fn test_from_str_basic() {
        let mer: MerDna = "ACGT".parse().unwrap();
        assert_eq!(mer.k(), 4);
        assert_eq!(mer.to_string(), "ACGT");
    }

    #[test]
    fn test_from_str_lowercase() {
        let mer: MerDna = "acgt".parse().unwrap();
        assert_eq!(mer.to_string(), "ACGT");
    }

    #[test]
    fn test_from_str_single_base() {
        for (ch, expected) in [("A", "A"), ("C", "C"), ("G", "G"), ("T", "T")] {
            let mer: MerDna = ch.parse().unwrap();
            assert_eq!(mer.to_string(), expected);
        }
    }

    #[test]
    fn test_from_str_invalid() {
        assert!("ACGN".parse::<MerDna>().is_err());
        assert!("".parse::<MerDna>().is_err());
        assert!("ACGX".parse::<MerDna>().is_err());
    }

    #[test]
    fn test_roundtrip_various_lengths() {
        let seqs = [
            "A",
            "AC",
            "ACG",
            "ACGT",
            "ACGTACGT",
            "ACGTACGTACGTACGTACGTACGTACGTACGT", // 32 bases = 1 word exactly
            "ACGTACGTACGTACGTACGTACGTACGTACGTA", // 33 bases = 2 words
        ];
        for seq in seqs {
            let mer: MerDna = seq.parse().unwrap();
            assert_eq!(mer.to_string(), seq, "roundtrip failed for {seq}");
        }
    }

    #[test]
    fn test_get_set_base() {
        let mut mer: MerDna = "ACGT".parse().unwrap();
        // String "ACGT": A is at position 3, C at 2, G at 1, T at 0
        assert_eq!(mer.get_base(0), 3); // T
        assert_eq!(mer.get_base(1), 2); // G
        assert_eq!(mer.get_base(2), 1); // C
        assert_eq!(mer.get_base(3), 0); // A

        mer.set_base(0, 0); // T -> A
        assert_eq!(mer.to_string(), "ACGA");
    }

    #[test]
    fn test_reverse_complement_palindrome() {
        // ACGT is its own reverse complement
        let mer: MerDna = "ACGT".parse().unwrap();
        let rc = mer.get_reverse_complement();
        assert_eq!(rc.to_string(), "ACGT");
    }

    #[test]
    fn test_reverse_complement_simple() {
        let mer: MerDna = "AAAA".parse().unwrap();
        let rc = mer.get_reverse_complement();
        assert_eq!(rc.to_string(), "TTTT");
    }

    #[test]
    fn test_reverse_complement_asymmetric() {
        let mer: MerDna = "AACG".parse().unwrap();
        let rc = mer.get_reverse_complement();
        assert_eq!(rc.to_string(), "CGTT");
    }

    #[test]
    fn test_reverse_complement_involution() {
        // RC(RC(x)) == x
        let seqs = ["ACGT", "AAAA", "GCTA", "AACG", "TTTCCCGGGAAA"];
        for seq in seqs {
            let mer: MerDna = seq.parse().unwrap();
            let rc2 = mer.get_reverse_complement().get_reverse_complement();
            assert_eq!(mer, rc2, "RC involution failed for {seq}");
        }
    }

    #[test]
    fn test_canonical_already_canonical() {
        let mer: MerDna = "AAAA".parse().unwrap();
        let canonical = mer.get_canonical();
        assert_eq!(canonical.to_string(), "AAAA"); // AAAA < TTTT
    }

    #[test]
    fn test_canonical_needs_rc() {
        let mer: MerDna = "TTTT".parse().unwrap();
        let canonical = mer.get_canonical();
        assert_eq!(canonical.to_string(), "AAAA"); // AAAA < TTTT
    }

    #[test]
    fn test_canonical_palindrome() {
        let mer: MerDna = "ACGT".parse().unwrap();
        let canonical = mer.get_canonical();
        assert_eq!(canonical.to_string(), "ACGT");
    }

    #[test]
    fn test_canonical_idempotent() {
        let seqs = ["ACGT", "TGCA", "AAAA", "CCCC", "AACG"];
        for seq in seqs {
            let mer: MerDna = seq.parse().unwrap();
            let c1 = mer.get_canonical();
            let c2 = c1.get_canonical();
            assert_eq!(c1, c2, "canonical not idempotent for {seq}");
        }
    }

    #[test]
    fn test_canonicalize_in_place() {
        let mut mer: MerDna = "TTTT".parse().unwrap();
        mer.canonicalize();
        assert_eq!(mer.to_string(), "AAAA");
    }

    #[test]
    fn test_ordering() {
        let a: MerDna = "AAAA".parse().unwrap();
        let c: MerDna = "CCCC".parse().unwrap();
        let g: MerDna = "GGGG".parse().unwrap();
        let t: MerDna = "TTTT".parse().unwrap();
        assert!(a < c);
        assert!(c < g);
        assert!(g < t);
    }

    #[test]
    fn test_hash_consistency() {
        use std::collections::HashMap;
        let mer1: MerDna = "ACGT".parse().unwrap();
        let mer2: MerDna = "ACGT".parse().unwrap();
        let mut map = HashMap::new();
        map.insert(mer1, 42);
        assert_eq!(map.get(&mer2), Some(&42));
    }

    #[test]
    fn test_shift_left() {
        let mut mer: MerDna = "ACGT".parse().unwrap();
        let out = mer.shift_left(b'A');
        assert_eq!(out, Some(b'A'));
        assert_eq!(mer.to_string(), "CGTA");
    }

    #[test]
    fn test_shift_right() {
        let mut mer: MerDna = "ACGT".parse().unwrap();
        let out = mer.shift_right(b'A');
        assert_eq!(out, Some(b'T'));
        assert_eq!(mer.to_string(), "AACG");
    }

    #[test]
    fn test_shift_invalid_base() {
        let mut mer: MerDna = "ACGT".parse().unwrap();
        assert_eq!(mer.shift_left(b'N'), None);
        assert_eq!(mer.to_string(), "ACGT"); // unchanged
    }

    #[test]
    fn test_homopolymer() {
        let aaaa: MerDna = "AAAA".parse().unwrap();
        assert!(aaaa.is_homopolymer());

        let cccc: MerDna = "CCCC".parse().unwrap();
        assert!(cccc.is_homopolymer());

        let acgt: MerDna = "ACGT".parse().unwrap();
        assert!(!acgt.is_homopolymer());
    }

    #[test]
    fn test_poly_constructors() {
        let mut mer = MerDna::new(4);
        mer.poly_a();
        assert_eq!(mer.to_string(), "AAAA");

        mer.poly_c();
        assert_eq!(mer.to_string(), "CCCC");

        mer.poly_g();
        assert_eq!(mer.to_string(), "GGGG");

        mer.poly_t();
        assert_eq!(mer.to_string(), "TTTT");
    }

    #[test]
    fn test_equality() {
        let a: MerDna = "ACGT".parse().unwrap();
        let b: MerDna = "ACGT".parse().unwrap();
        let c: MerDna = "ACGA".parse().unwrap();
        assert_eq!(a, b);
        assert_ne!(a, c);
    }

    #[test]
    fn test_from_bytes() {
        // For k=4 (8 bits), "ACGT" in 2-bit encoding from LSB:
        // Position 0 (T) = 11, Position 1 (G) = 10, Position 2 (C) = 01, Position 3 (A) = 00
        // = 0b00_01_10_11 = 0x1B
        let mer = MerDna::from_bytes(&[0x1B], 4);
        assert_eq!(mer.to_string(), "ACGT");
    }

    #[test]
    fn test_long_kmer() {
        // Test with k=33 (requires 2 words)
        let seq = "ACGTACGTACGTACGTACGTACGTACGTACGTA";
        let mer: MerDna = seq.parse().unwrap();
        assert_eq!(mer.k(), 33);
        assert_eq!(mer.to_string(), seq);

        // Test reverse complement involution
        let rc2 = mer.get_reverse_complement().get_reverse_complement();
        assert_eq!(mer, rc2);
    }

    #[test]
    fn test_word_reverse_complement_basic() {
        // All A's (0x0000) -> All T's (0xFFFF...)
        assert_eq!(word_reverse_complement(0), u64::MAX);
    }
}