cuttlefish_rs/
kmer.rs

1//! Packed canonical k-mer primitives.
2//!
3//! [`Kmer`] stores up to 63 bases in two words and uses the same one-word
4//! wyhash specialization as the C++ implementation for `K <= 31`.
5
6use crate::dna::Base;
7use crate::hash::{hash_bytes, wyhash_u64};
8use std::hash::{Hash, Hasher};
9
10const fn ascii_quad_table() -> [u32; 256] {
11    let mut table = [0u32; 256];
12    let mut byte = 0usize;
13    while byte < table.len() {
14        let mut encoded = 0u32;
15        let mut base = 0usize;
16        while base < 4 {
17            let bits = (byte >> (2 * (3 - base))) & 0b11;
18            let ascii = match bits {
19                0 => b'A',
20                1 => b'C',
21                2 => b'G',
22                _ => b'T',
23            };
24            encoded |= (ascii as u32) << (8 * base);
25            base += 1;
26        }
27        table[byte] = encoded;
28        byte += 1;
29    }
30    table
31}
32
33const ASCII_QUADS: [u32; 256] = ascii_quad_table();
34
35/// A two-bit packed DNA k-mer of compile-time length `K`.
36///
37/// Bases are ordered from most significant to least significant bits in the
38/// logical value. The representation is internal; use the provided conversion
39/// and orientation methods rather than depending on word layout.
40#[repr(C)]
41#[derive(Clone, Copy, Debug, PartialEq, Eq)]
42pub struct Kmer<const K: usize> {
43    words: [u64; 2],
44}
45
46impl<const K: usize> Kmer<K> {
47    #[inline]
48    pub fn zero() -> Self {
49        Self { words: [0, 0] }
50    }
51
52    pub fn from_ascii(seq: &[u8]) -> Result<Self, KmerError> {
53        Self::check_k()?;
54        if seq.len() < K {
55            return Err(KmerError::TooShort {
56                expected: K,
57                got: seq.len(),
58            });
59        }
60
61        let mut value = 0u128;
62        for &ch in &seq[..K] {
63            let base = Base::from_ascii(ch);
64            if !base.is_dna() {
65                return Err(KmerError::InvalidBase(ch));
66            }
67            value = (value << 2) | base.bits() as u128;
68        }
69
70        Ok(Self::from_u128(value))
71    }
72
73    #[inline]
74    pub fn as_u128(self) -> u128 {
75        self.words[0] as u128 | ((self.words[1] as u128) << 64)
76    }
77
78    #[inline]
79    pub fn words(self) -> [u64; 2] {
80        self.words
81    }
82
83    #[inline]
84    pub fn get(self, idx: usize) -> Base {
85        assert!(idx < K);
86        let shift = 2 * (K - 1 - idx);
87        match ((self.as_u128() >> shift) & 0b11) as u8 {
88            0 => Base::A,
89            1 => Base::C,
90            2 => Base::G,
91            3 => Base::T,
92            _ => unreachable!(),
93        }
94    }
95
96    #[inline]
97    pub fn front(self) -> Base {
98        self.get(0)
99    }
100
101    #[inline]
102    pub fn back(self) -> Base {
103        self.get(K - 1)
104    }
105
106    #[inline]
107    pub fn reverse_complement(self) -> Self {
108        if K <= 32 {
109            let mut x = self.words[0];
110            x = ((x & 0x3333_3333_3333_3333) << 2) | ((x >> 2) & 0x3333_3333_3333_3333);
111            x = ((x & 0x0f0f_0f0f_0f0f_0f0f) << 4) | ((x >> 4) & 0x0f0f_0f0f_0f0f_0f0f);
112            x = x.swap_bytes();
113            let shift = 64 - 2 * K;
114            let mask = if K == 32 {
115                u64::MAX
116            } else {
117                (1u64 << (2 * K)) - 1
118            };
119            return Self::from_u128(((!x >> shift) & mask) as u128);
120        }
121
122        let mut x = self.as_u128();
123        x = ((x & 0x3333_3333_3333_3333_3333_3333_3333_3333) << 2)
124            | ((x >> 2) & 0x3333_3333_3333_3333_3333_3333_3333_3333);
125        x = ((x & 0x0f0f_0f0f_0f0f_0f0f_0f0f_0f0f_0f0f_0f0f) << 4)
126            | ((x >> 4) & 0x0f0f_0f0f_0f0f_0f0f_0f0f_0f0f_0f0f_0f0f);
127        x = x.swap_bytes();
128        let shift = 128 - 2 * K;
129        let mask = (1u128 << (2 * K)) - 1;
130        Self::from_u128((!x >> shift) & mask)
131    }
132
133    #[inline]
134    pub fn canonical(self) -> Self {
135        self.min(self.reverse_complement())
136    }
137
138    #[inline]
139    pub fn is_canonical(self) -> bool {
140        self <= self.reverse_complement()
141    }
142
143    #[inline]
144    pub fn roll_forward(self, next: Base) -> Self {
145        assert!(next.is_dna());
146        if K <= 32 {
147            let mask = if K == 32 {
148                u64::MAX
149            } else {
150                (1u64 << (2 * K)) - 1
151            };
152            return Self::from_u128(
153                (((self.words[0] << 2) | u64::from(next.bits())) & mask) as u128,
154            );
155        }
156        let mask = if K == 64 {
157            u128::MAX
158        } else {
159            (1u128 << (2 * K)) - 1
160        };
161        Self::from_u128(((self.as_u128() << 2) | next.bits() as u128) & mask)
162    }
163
164    #[inline]
165    pub fn roll_backward(self, prev: Base) -> Self {
166        assert!(prev.is_dna());
167        if K <= 32 {
168            let high = u64::from(prev.bits()) << (2 * (K - 1));
169            return Self::from_u128((high | (self.words[0] >> 2)) as u128);
170        }
171        let high = (prev.bits() as u128) << (2 * (K - 1));
172        Self::from_u128(high | (self.as_u128() >> 2))
173    }
174
175    pub fn to_ascii_string(self) -> String {
176        let mut s = Vec::with_capacity(K);
177        self.append_ascii(&mut s);
178        String::from_utf8(s).unwrap()
179    }
180
181    #[inline]
182    pub(crate) fn append_ascii(self, output: &mut Vec<u8>) {
183        if K <= 32 {
184            output.reserve(K);
185            let leading = K & 3;
186            for index in 0..leading {
187                output.push(self.get(index).to_ascii());
188            }
189            let quads = K / 4;
190            for index in 0..quads {
191                let shift = 8 * (quads - 1 - index);
192                let byte = ((self.words[0] >> shift) & 0xff) as usize;
193                output.extend_from_slice(&ASCII_QUADS[byte].to_le_bytes());
194            }
195            return;
196        }
197        output.reserve(K);
198        for index in 0..K {
199            output.push(self.get(index).to_ascii());
200        }
201    }
202
203    #[inline(always)]
204    pub fn hash64(self, seed: u64) -> u64 {
205        let byte_len = (2 * K).div_ceil(8);
206        if byte_len <= 8 {
207            return wyhash_u64(self.words[0], seed);
208        }
209        let bytes = self.as_u128().to_le_bytes();
210        hash_bytes(&bytes[..byte_len], seed)
211    }
212
213    #[inline]
214    fn from_u128(value: u128) -> Self {
215        Self {
216            words: [value as u64, (value >> 64) as u64],
217        }
218    }
219
220    #[inline]
221    pub(crate) fn from_bits(value: u128) -> Self {
222        Self::from_u128(value)
223    }
224
225    #[inline]
226    fn check_k() -> Result<(), KmerError> {
227        if K == 0 || K > 63 {
228            Err(KmerError::UnsupportedK(K))
229        } else {
230            Ok(())
231        }
232    }
233}
234
235impl<const K: usize> Ord for Kmer<K> {
236    #[inline]
237    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
238        if K <= 32 {
239            self.words[0].cmp(&other.words[0])
240        } else {
241            self.words[1]
242                .cmp(&other.words[1])
243                .then_with(|| self.words[0].cmp(&other.words[0]))
244        }
245    }
246}
247
248impl<const K: usize> PartialOrd for Kmer<K> {
249    #[inline]
250    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
251        Some(self.cmp(other))
252    }
253}
254
255impl<const K: usize> Hash for Kmer<K> {
256    #[inline]
257    fn hash<H: Hasher>(&self, state: &mut H) {
258        state.write_u64(self.words[0]);
259        state.write_u64(self.words[1]);
260    }
261}
262
263#[derive(Debug, Clone, Copy, PartialEq, Eq)]
264pub enum KmerError {
265    UnsupportedK(usize),
266    TooShort { expected: usize, got: usize },
267    InvalidBase(u8),
268}
269
270impl std::fmt::Display for KmerError {
271    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
272        match self {
273            Self::UnsupportedK(k) => write!(f, "unsupported k-mer length {k}; expected 1..=63"),
274            Self::TooShort { expected, got } => {
275                write!(
276                    f,
277                    "sequence too short for k-mer: expected {expected}, got {got}"
278                )
279            }
280            Self::InvalidBase(b) => write!(f, "invalid DNA base '{}'", *b as char),
281        }
282    }
283}
284
285impl std::error::Error for KmerError {}
286
287#[cfg(test)]
288mod tests {
289    use super::*;
290
291    #[test]
292    fn round_trips_ascii() {
293        let k = Kmer::<31>::from_ascii(b"ACGTACGTACGTACGTACGTACGTACGTACG").unwrap();
294        assert_eq!(k.to_ascii_string(), "ACGTACGTACGTACGTACGTACGTACGTACG");
295        assert_eq!(k.front(), Base::A);
296        assert_eq!(k.back(), Base::G);
297    }
298
299    #[test]
300    fn reverse_complement_round_trips() {
301        let k = Kmer::<7>::from_ascii(b"ACGTAGC").unwrap();
302        assert_eq!(k.reverse_complement().to_ascii_string(), "GCTACGT");
303        assert_eq!(k.reverse_complement().reverse_complement(), k);
304    }
305
306    fn slow_reverse_complement<const K: usize>(kmer: Kmer<K>) -> Kmer<K> {
307        let mut value = 0u128;
308        let mut input = kmer.as_u128();
309        for _ in 0..K {
310            value = (value << 2) | ((!input) & 0b11);
311            input >>= 2;
312        }
313        Kmer::from_bits(value)
314    }
315
316    #[test]
317    fn packed_reverse_complement_matches_basewise_reference() {
318        fn check<const K: usize>() {
319            let mask = (1u128 << (2 * K)) - 1;
320            for value in [
321                0,
322                1,
323                mask,
324                0x0123_4567_89ab_cdef_7654_3210_fedc_ba98 & mask,
325                0xaaaa_5555_3333_cccc_0f0f_f0f0_9696_6969 & mask,
326            ] {
327                let kmer = Kmer::<K>::from_bits(value);
328                assert_eq!(kmer.reverse_complement(), slow_reverse_complement(kmer));
329            }
330        }
331
332        check::<1>();
333        check::<4>();
334        check::<31>();
335        check::<32>();
336        check::<33>();
337        check::<63>();
338    }
339
340    #[test]
341    fn rolling_matches_reparse() {
342        let k = Kmer::<5>::from_ascii(b"ACGTA").unwrap();
343        assert_eq!(
344            k.roll_forward(Base::C),
345            Kmer::<5>::from_ascii(b"CGTAC").unwrap()
346        );
347        assert_eq!(
348            k.roll_backward(Base::T),
349            Kmer::<5>::from_ascii(b"TACGT").unwrap()
350        );
351    }
352
353    #[test]
354    fn one_word_hash_matches_cpp_wyhash() {
355        let k = Kmer::<31>::from_ascii(b"ACGTACGTACGTACGTACGTACGTACGTACG").unwrap();
356        for seed in [0, 1, 17, u64::MAX] {
357            assert_eq!(k.hash64(seed), wyhash_u64(k.words()[0], seed));
358        }
359    }
360}
cuttlefish_rs/kmer.rs

cuttlefish_rs/
kmer.rs
