mtnoc-rs 0.1.0

//! Canonical k-mer set for mtDNA read prefiltering (fast mode).
//! Builds a set of canonical 21-mers from a (circular-wrapped) rCRS reference; a read is a chrM
//! candidate if it shares >=1 canonical 21-mer with the set. Pure Rust, no external tools.

use rustc_hash::FxHashSet;

pub const K: usize = 21;

#[inline]
fn base2(b: u8) -> Option<u64> {
    match b {
        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,
    }
}

/// Fold a sequence into its canonical 21-mers, calling `f` on each. Resets on any non-ACGT base.
/// Forward and reverse-complement codes are maintained incrementally (O(1) per base) — no per-k-mer
/// reverse-complement recomputation.
fn each_kmer(seq: &[u8], mut f: impl FnMut(u64)) {
    let mask: u64 = (1u64 << (2 * K)) - 1;
    let shift: u64 = 2 * (K as u64 - 1);
    let mut fwd = 0u64;
    let mut rev = 0u64;
    let mut len = 0usize;
    for &b in seq {
        match base2(b) {
            Some(v) => {
                fwd = ((fwd << 2) | v) & mask;
                rev = (rev >> 2) | ((3 - v) << shift);
                len += 1;
                if len >= K {
                    f(if fwd < rev { fwd } else { rev });
                }
            }
            None => {
                len = 0;
                fwd = 0;
                rev = 0;
            }
        }
    }
}

/// Forward (non-canonical) 21-mer codes with their 0-based start positions.
pub fn kmers_fwd(seq: &[u8]) -> Vec<(usize, u64)> {
    let mask: u64 = (1u64 << (2 * K)) - 1;
    let mut fwd = 0u64;
    let mut len = 0usize;
    let mut out = Vec::new();
    for (i, &b) in seq.iter().enumerate() {
        match base2(b) {
            Some(v) => {
                fwd = ((fwd << 2) | v) & mask;
                len += 1;
                if len >= K {
                    out.push((i + 1 - K, fwd));
                }
            }
            None => {
                len = 0;
                fwd = 0;
            }
        }
    }
    out
}

/// Reverse-complement of a nucleotide sequence.
pub fn revcomp_seq(seq: &[u8]) -> Vec<u8> {
    seq.iter()
        .rev()
        .map(|&b| match b {
            b'A' | b'a' => b'T',
            b'C' | b'c' => b'G',
            b'G' | b'g' => b'C',
            b'T' | b't' => b'A',
            _ => b'N',
        })
        .collect()
}

/// Number of bits (log2) in the pre-screen bit-array. 2^20 bits = 128 KiB, which stays
/// L2-resident. With ~16.5k reference k-mers the false-positive rate is ~1.6%; false positives
/// only cost an extra (rare) hashset confirm, never a missed match — so the filter is LOSSLESS.
const BLOOM_BITS: u32 = 20;
const BLOOM_MASK: u64 = (1u64 << BLOOM_BITS) - 1;
/// Fibonacci-hashing multiplier; using the high bits of `c * PHI` mixes the 42-bit k-mer code well.
const PHI: u64 = 0x9E37_79B9_7F4A_7C15;

#[inline(always)]
fn bloom_idx(c: u64) -> usize {
    // Multiplicative (Fibonacci) hash; take the top BLOOM_BITS bits for good low-collision spread.
    (c.wrapping_mul(PHI) >> (64 - BLOOM_BITS)) as usize
}

pub struct KmerSet {
    set: FxHashSet<u64>,
    /// One bit per hashed k-mer slot; a cheap, cache-resident necessary condition for membership.
    bits: Vec<u64>,
}

impl KmerSet {
    pub fn from_seq(seq: &[u8]) -> Self {
        let mut set = FxHashSet::default();
        let mut bits = vec![0u64; ((BLOOM_MASK + 1) / 64) as usize];
        each_kmer(seq, |c| {
            set.insert(c);
            let i = bloom_idx(c);
            bits[i >> 6] |= 1u64 << (i & 63);
        });
        KmerSet { set, bits }
    }

    pub fn len(&self) -> usize {
        self.set.len()
    }

    /// Membership test with a cache-friendly bit-array pre-screen. The bit-array has no false
    /// negatives, so this returns exactly the same answer as `set.contains(&c)` (lossless).
    #[inline(always)]
    fn contains(&self, c: u64) -> bool {
        let i = bloom_idx(c);
        // SAFETY: i is masked to BLOOM_BITS, so i>>6 is always in range.
        let word = unsafe { *self.bits.get_unchecked(i >> 6) };
        (word >> (i & 63)) & 1 != 0 && self.set.contains(&c)
    }

    /// True if the read shares at least `min` canonical 21-mers with the set.
    pub fn matches(&self, seq: &[u8], min: usize) -> bool {
        if min <= 1 {
            let mut hit = false;
            each_kmer(seq, |c| {
                if !hit && self.contains(c) {
                    hit = true;
                }
            });
            hit
        } else {
            let mut n = 0usize;
            each_kmer(seq, |c| {
                if self.contains(c) {
                    n += 1;
                }
            });
            n >= min
        }
    }
}