1use crate::error::{Error, Result};
7
8const AMBIGUITY_PAIRS: &[(u8, u8)] = &[
10 (b'Y', b'R'),
11 (b'R', b'Y'),
12 (b'S', b'S'),
13 (b'W', b'W'),
14 (b'K', b'M'),
15 (b'M', b'K'),
16 (b'B', b'V'),
17 (b'V', b'B'),
18 (b'D', b'H'),
19 (b'H', b'D'),
20 (b'N', b'N'),
21];
22
23static COMPLEMENT: [u8; 256] = build_complement_table(&[
25 (b'A', b'T'),
26 (b'T', b'A'),
27 (b'U', b'A'),
28 (b'G', b'C'),
29 (b'C', b'G'),
30]);
31
32static COMPLEMENT_RNA: [u8; 256] = build_complement_table(&[
35 (b'A', b'U'),
36 (b'U', b'A'),
37 (b'T', b'A'),
38 (b'G', b'C'),
39 (b'C', b'G'),
40]);
41
42const fn build_complement_table(bases: &[(u8, u8)]) -> [u8; 256] {
43 let mut table = [0u8; 256];
45 let mut i = 0;
46 while i < 256 {
47 table[i] = i as u8;
48 i += 1;
49 }
50 let mut p = 0;
51 while p < bases.len() {
52 let (from, to) = bases[p];
53 table[from as usize] = to;
54 table[(from + 32) as usize] = to + 32; p += 1;
56 }
57 let mut p = 0;
58 while p < AMBIGUITY_PAIRS.len() {
59 let (from, to) = AMBIGUITY_PAIRS[p];
60 table[from as usize] = to;
61 table[(from + 32) as usize] = to + 32;
62 p += 1;
63 }
64 table
65}
66
67#[inline]
86pub fn complement(base: u8) -> u8 {
87 COMPLEMENT[base as usize]
88}
89
90#[inline]
100pub fn complement_rna(base: u8) -> u8 {
101 COMPLEMENT_RNA[base as usize]
102}
103
104pub fn reverse_complement(seq: &[u8]) -> Vec<u8> {
112 seq.iter().rev().map(|&b| complement(b)).collect()
113}
114
115pub fn reverse_complement_rna(seq: &[u8]) -> Vec<u8> {
123 seq.iter().rev().map(|&b| complement_rna(b)).collect()
124}
125
126pub fn reverse_complement_into(seq: &[u8], out: &mut Vec<u8>) {
130 out.clear();
131 out.reserve(seq.len());
132 out.extend(seq.iter().rev().map(|&b| complement(b)));
133}
134
135pub fn reverse_complement_in_place(seq: &mut [u8]) {
137 let n = seq.len();
138 for i in 0..n / 2 {
139 let j = n - 1 - i;
140 let a = complement(seq[i]);
141 let b = complement(seq[j]);
142 seq[i] = b;
143 seq[j] = a;
144 }
145 if n % 2 == 1 {
146 seq[n / 2] = complement(seq[n / 2]);
147 }
148}
149
150#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
152pub struct BaseCounts {
153 pub a: u64,
155 pub c: u64,
157 pub g: u64,
159 pub t: u64,
161 pub n: u64,
163 pub other: u64,
165}
166
167const CLASS_A: u8 = 0;
169const CLASS_C: u8 = 1;
170const CLASS_G: u8 = 2;
171const CLASS_T: u8 = 3;
172const CLASS_AMBIGUOUS: u8 = 4;
173const CLASS_OTHER: u8 = 5;
174
175static BASE_CLASS: [u8; 256] = build_base_class_table();
178
179const fn build_base_class_table() -> [u8; 256] {
180 let mut table = [CLASS_OTHER; 256];
181 table[b'A' as usize] = CLASS_A;
182 table[b'a' as usize] = CLASS_A;
183 table[b'C' as usize] = CLASS_C;
184 table[b'c' as usize] = CLASS_C;
185 table[b'G' as usize] = CLASS_G;
186 table[b'g' as usize] = CLASS_G;
187 table[b'T' as usize] = CLASS_T;
188 table[b't' as usize] = CLASS_T;
189 table[b'U' as usize] = CLASS_T;
190 table[b'u' as usize] = CLASS_T;
191
192 let ambiguous = b"NRYSWKMBDHV";
193 let mut i = 0;
194 while i < ambiguous.len() {
195 table[ambiguous[i] as usize] = CLASS_AMBIGUOUS;
196 table[(ambiguous[i] + 32) as usize] = CLASS_AMBIGUOUS;
197 i += 1;
198 }
199 table
200}
201
202impl BaseCounts {
203 pub fn of(seq: &[u8]) -> BaseCounts {
205 let mut counts = [0u64; 8];
208 for &b in seq {
209 counts[(BASE_CLASS[b as usize] & 7) as usize] += 1;
210 }
211 BaseCounts {
212 a: counts[CLASS_A as usize],
213 c: counts[CLASS_C as usize],
214 g: counts[CLASS_G as usize],
215 t: counts[CLASS_T as usize],
216 n: counts[CLASS_AMBIGUOUS as usize],
217 other: counts[CLASS_OTHER as usize],
218 }
219 }
220
221 pub fn total(&self) -> u64 {
223 self.a + self.c + self.g + self.t + self.n + self.other
224 }
225
226 pub fn acgt(&self) -> u64 {
228 self.a + self.c + self.g + self.t
229 }
230
231 pub fn gc_content(&self) -> Option<f64> {
233 let acgt = self.acgt();
234 if acgt == 0 {
235 None
236 } else {
237 Some((self.g + self.c) as f64 / acgt as f64)
238 }
239 }
240
241 pub fn merge(&mut self, other: &BaseCounts) {
243 self.a += other.a;
244 self.c += other.c;
245 self.g += other.g;
246 self.t += other.t;
247 self.n += other.n;
248 self.other += other.other;
249 }
250}
251
252pub fn gc_content(seq: &[u8]) -> Option<f64> {
264 BaseCounts::of(seq).gc_content()
265}
266
267#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
273#[non_exhaustive]
274pub enum Alphabet {
275 Dna,
277 Rna,
279 Iupac,
281 Protein,
283 Any,
285}
286
287impl Alphabet {
288 pub fn contains(self, byte: u8) -> bool {
290 let upper = byte.to_ascii_uppercase();
291 match self {
292 Alphabet::Dna => matches!(upper, b'A' | b'C' | b'G' | b'T' | b'N'),
293 Alphabet::Rna => matches!(upper, b'A' | b'C' | b'G' | b'U' | b'N'),
294 Alphabet::Iupac => {
295 matches!(
296 upper,
297 b'A' | b'C'
298 | b'G'
299 | b'T'
300 | b'U'
301 | b'R'
302 | b'Y'
303 | b'S'
304 | b'W'
305 | b'K'
306 | b'M'
307 | b'B'
308 | b'D'
309 | b'H'
310 | b'V'
311 | b'N'
312 | b'-'
313 | b'.'
314 )
315 }
316 Alphabet::Protein => upper.is_ascii_uppercase() || matches!(upper, b'*' | b'-' | b'.'),
317 Alphabet::Any => byte.is_ascii_graphic(),
318 }
319 }
320
321 pub fn validate(self, seq: &[u8]) -> Result<()> {
323 self.validate_named(seq, "")
324 }
325
326 pub fn validate_named(self, seq: &[u8], id: &str) -> Result<()> {
328 match seq.iter().position(|&b| !self.contains(b)) {
329 None => Ok(()),
330 Some(pos) => Err(Error::InvalidByte {
331 id: id.to_string(),
332 pos,
333 byte: seq[pos],
334 }),
335 }
336 }
337}
338
339pub fn make_uppercase(seq: &mut [u8]) {
341 seq.make_ascii_uppercase();
342}
343
344pub fn rna_to_dna(seq: &mut [u8]) {
346 for b in seq.iter_mut() {
347 match *b {
348 b'U' => *b = b'T',
349 b'u' => *b = b't',
350 _ => {}
351 }
352 }
353}
354
355const CODON_TABLE: &[u8; 64] = b"FFLLSSSSYY**CC*WLLLLPPPPHHQQRRRRIIIMTTTTNNKKSSRRVVVVAAAADDEEGGGG";
357
358const NOT_A_BASE: u8 = u8::MAX;
360
361static CODON_INDEX: [u8; 256] = build_codon_index_table();
363
364const fn build_codon_index_table() -> [u8; 256] {
365 let mut table = [NOT_A_BASE; 256];
366 table[b'T' as usize] = 0;
367 table[b't' as usize] = 0;
368 table[b'U' as usize] = 0;
369 table[b'u' as usize] = 0;
370 table[b'C' as usize] = 1;
371 table[b'c' as usize] = 1;
372 table[b'A' as usize] = 2;
373 table[b'a' as usize] = 2;
374 table[b'G' as usize] = 3;
375 table[b'g' as usize] = 3;
376 table
377}
378
379#[inline]
380fn base_index(base: u8) -> Option<usize> {
381 match CODON_INDEX[base as usize] {
382 NOT_A_BASE => None,
383 index => Some(index as usize),
384 }
385}
386
387pub fn translate_codon(codon: &[u8]) -> u8 {
398 if codon.len() < 3 {
399 return b'X';
400 }
401 match (
402 base_index(codon[0]),
403 base_index(codon[1]),
404 base_index(codon[2]),
405 ) {
406 (Some(a), Some(b), Some(c)) => CODON_TABLE[a * 16 + b * 4 + c],
407 _ => b'X',
408 }
409}
410
411pub fn translate(seq: &[u8], frame: usize, stop_at_stop: bool) -> Vec<u8> {
424 let seq = if frame < seq.len() {
425 &seq[frame..]
426 } else {
427 &[][..]
428 };
429 let mut out = Vec::with_capacity(seq.len() / 3);
430 for codon in seq.chunks_exact(3) {
431 let aa = translate_codon(codon);
432 if stop_at_stop && aa == b'*' {
433 break;
434 }
435 out.push(aa);
436 }
437 out
438}
439
440pub fn kmers(seq: &[u8], k: usize) -> impl Iterator<Item = &[u8]> {
449 let n = if k == 0 || seq.len() < k {
450 0
451 } else {
452 seq.len() - k + 1
453 };
454 (0..n).map(move |i| &seq[i..i + k])
455}
456
457pub fn canonical_kmer(kmer: &[u8]) -> Vec<u8> {
460 let rc = reverse_complement(kmer);
461 if rc.as_slice() < kmer {
462 rc
463 } else {
464 kmer.to_vec()
465 }
466}
467
468pub fn hamming_distance(a: &[u8], b: &[u8]) -> Option<usize> {
472 if a.len() != b.len() {
473 return None;
474 }
475 Some(
476 a.iter()
477 .zip(b)
478 .filter(|(x, y)| !x.eq_ignore_ascii_case(y))
479 .count(),
480 )
481}
482
483pub fn n50(lengths: &mut [u64]) -> Option<u64> {
492 nx(lengths, 0.5)
493}
494
495pub fn nx(lengths: &mut [u64], fraction: f64) -> Option<u64> {
499 if lengths.is_empty() {
500 return None;
501 }
502 let total: u64 = lengths.iter().sum();
503 if total == 0 {
504 return Some(0);
505 }
506 lengths.sort_unstable_by(|a, b| b.cmp(a));
507 let target = total as f64 * fraction.clamp(0.0, 1.0);
508 let mut acc = 0u64;
509 for &len in lengths.iter() {
510 acc += len;
511 if acc as f64 >= target {
512 return Some(len);
513 }
514 }
515 lengths.last().copied()
516}
517
518#[cfg(test)]
519mod tests {
520 use super::*;
521
522 #[test]
523 fn complements_iupac_and_preserves_case() {
524 assert_eq!(reverse_complement(b"ACGTacgt"), b"acgtACGT");
525 assert_eq!(reverse_complement(b"RYKM"), b"KMRY");
526 assert_eq!(reverse_complement(b""), b"");
527 assert_eq!(reverse_complement(b"AC-GT"), b"AC-GT");
529 }
530
531 #[test]
532 fn in_place_matches_allocating() {
533 for s in [&b""[..], b"A", b"AC", b"ACG", b"ACGTN", b"acgtRYn"] {
534 let mut owned = s.to_vec();
535 reverse_complement_in_place(&mut owned);
536 assert_eq!(
537 owned,
538 reverse_complement(s),
539 "{:?}",
540 String::from_utf8_lossy(s)
541 );
542 }
543 }
544
545 #[test]
546 fn counts_bases() {
547 let c = BaseCounts::of(b"AACCGGTTNNxx");
548 assert_eq!((c.a, c.c, c.g, c.t, c.n, c.other), (2, 2, 2, 2, 2, 2));
549 assert_eq!(c.total(), 12);
550 assert_eq!(c.acgt(), 8);
551 assert_eq!(c.gc_content(), Some(0.5));
552 }
553
554 #[test]
555 fn validates_alphabets() {
556 assert!(Alphabet::Dna.validate(b"ACGTN").is_ok());
557 assert!(Alphabet::Dna.validate(b"ACGU").is_err());
558 assert!(Alphabet::Rna.validate(b"ACGU").is_ok());
559 assert!(Alphabet::Iupac.validate(b"ACGTRYKM-").is_ok());
560 assert!(Alphabet::Protein.validate(b"MEEPQSDPSV*").is_ok());
561 assert!(Alphabet::Any.validate(b"anything!").is_ok());
562 assert!(Alphabet::Any.validate(b"tab\there").is_err());
563
564 match Alphabet::Dna.validate_named(b"ACG!T", "read1") {
565 Err(Error::InvalidByte { id, pos, byte }) => {
566 assert_eq!((id.as_str(), pos, byte), ("read1", 3, b'!'));
567 }
568 other => panic!("expected InvalidByte, got {other:?}"),
569 }
570 }
571
572 #[test]
573 fn translates_all_codons() {
574 assert_eq!(translate(b"TTTTTCTTATTG", 0, false), b"FFLL");
576 assert_eq!(translate(b"ATGCATTAA", 0, false), b"MH*");
577 assert_eq!(translate(b"AUGCAU", 0, false), b"MH"); assert_eq!(translate(b"ATGCA", 0, false), b"M");
580 assert_eq!(translate(b"AT", 0, false), b"");
581 assert_eq!(translate(b"ATG", 5, false), b"");
582 }
583
584 #[test]
585 fn kmers_and_canonical() {
586 assert_eq!(kmers(b"AAAA", 4).count(), 1);
587 assert_eq!(kmers(b"AAAA", 0).count(), 0);
588 assert_eq!(canonical_kmer(b"TTT"), b"AAA");
589 assert_eq!(canonical_kmer(b"AAA"), b"AAA");
590 }
591
592 #[test]
593 fn hamming() {
594 assert_eq!(hamming_distance(b"ACGT", b"acgt"), Some(0));
595 assert_eq!(hamming_distance(b"ACGT", b"ACGA"), Some(1));
596 assert_eq!(hamming_distance(b"ACGT", b"ACG"), None);
597 }
598
599 #[test]
600 fn n_statistics() {
601 assert_eq!(n50(&mut [5, 50, 30, 15]), Some(50));
603 assert_eq!(nx(&mut [5, 50, 30, 15], 0.9), Some(15));
604 assert_eq!(nx(&mut [0, 0], 0.5), Some(0));
605 }
606
607 #[test]
608 fn rna_dna_conversion() {
609 let mut s = b"ACGUacgu".to_vec();
610 rna_to_dna(&mut s);
611 assert_eq!(s, b"ACGTacgt");
612 make_uppercase(&mut s);
613 assert_eq!(s, b"ACGTACGT");
614 }
615}