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)]
269pub enum Alphabet {
270 Dna,
272 Rna,
274 Iupac,
276 Protein,
278 Any,
280}
281
282impl Alphabet {
283 pub fn contains(self, byte: u8) -> bool {
285 let upper = byte.to_ascii_uppercase();
286 match self {
287 Alphabet::Dna => matches!(upper, b'A' | b'C' | b'G' | b'T' | b'N'),
288 Alphabet::Rna => matches!(upper, b'A' | b'C' | b'G' | b'U' | b'N'),
289 Alphabet::Iupac => {
290 matches!(
291 upper,
292 b'A' | b'C'
293 | b'G'
294 | b'T'
295 | b'U'
296 | b'R'
297 | b'Y'
298 | b'S'
299 | b'W'
300 | b'K'
301 | b'M'
302 | b'B'
303 | b'D'
304 | b'H'
305 | b'V'
306 | b'N'
307 | b'-'
308 | b'.'
309 )
310 }
311 Alphabet::Protein => upper.is_ascii_uppercase() || matches!(upper, b'*' | b'-' | b'.'),
312 Alphabet::Any => byte.is_ascii_graphic(),
313 }
314 }
315
316 pub fn validate(self, seq: &[u8]) -> Result<()> {
318 self.validate_named(seq, "")
319 }
320
321 pub fn validate_named(self, seq: &[u8], id: &str) -> Result<()> {
323 match seq.iter().position(|&b| !self.contains(b)) {
324 None => Ok(()),
325 Some(pos) => Err(Error::InvalidByte {
326 id: id.to_string(),
327 pos,
328 byte: seq[pos],
329 }),
330 }
331 }
332}
333
334pub fn make_uppercase(seq: &mut [u8]) {
336 seq.make_ascii_uppercase();
337}
338
339pub fn rna_to_dna(seq: &mut [u8]) {
341 for b in seq.iter_mut() {
342 match *b {
343 b'U' => *b = b'T',
344 b'u' => *b = b't',
345 _ => {}
346 }
347 }
348}
349
350const CODON_TABLE: &[u8; 64] = b"FFLLSSSSYY**CC*WLLLLPPPPHHQQRRRRIIIMTTTTNNKKSSRRVVVVAAAADDEEGGGG";
352
353const NOT_A_BASE: u8 = u8::MAX;
355
356static CODON_INDEX: [u8; 256] = build_codon_index_table();
358
359const fn build_codon_index_table() -> [u8; 256] {
360 let mut table = [NOT_A_BASE; 256];
361 table[b'T' as usize] = 0;
362 table[b't' as usize] = 0;
363 table[b'U' as usize] = 0;
364 table[b'u' as usize] = 0;
365 table[b'C' as usize] = 1;
366 table[b'c' as usize] = 1;
367 table[b'A' as usize] = 2;
368 table[b'a' as usize] = 2;
369 table[b'G' as usize] = 3;
370 table[b'g' as usize] = 3;
371 table
372}
373
374#[inline]
375fn base_index(base: u8) -> Option<usize> {
376 match CODON_INDEX[base as usize] {
377 NOT_A_BASE => None,
378 index => Some(index as usize),
379 }
380}
381
382pub fn translate_codon(codon: &[u8]) -> u8 {
393 if codon.len() < 3 {
394 return b'X';
395 }
396 match (
397 base_index(codon[0]),
398 base_index(codon[1]),
399 base_index(codon[2]),
400 ) {
401 (Some(a), Some(b), Some(c)) => CODON_TABLE[a * 16 + b * 4 + c],
402 _ => b'X',
403 }
404}
405
406pub fn translate(seq: &[u8], frame: usize, stop_at_stop: bool) -> Vec<u8> {
419 let seq = if frame < seq.len() {
420 &seq[frame..]
421 } else {
422 &[][..]
423 };
424 let mut out = Vec::with_capacity(seq.len() / 3);
425 for codon in seq.chunks_exact(3) {
426 let aa = translate_codon(codon);
427 if stop_at_stop && aa == b'*' {
428 break;
429 }
430 out.push(aa);
431 }
432 out
433}
434
435pub fn kmers(seq: &[u8], k: usize) -> impl Iterator<Item = &[u8]> {
444 let n = if k == 0 || seq.len() < k {
445 0
446 } else {
447 seq.len() - k + 1
448 };
449 (0..n).map(move |i| &seq[i..i + k])
450}
451
452pub fn canonical_kmer(kmer: &[u8]) -> Vec<u8> {
455 let rc = reverse_complement(kmer);
456 if rc.as_slice() < kmer {
457 rc
458 } else {
459 kmer.to_vec()
460 }
461}
462
463pub fn hamming_distance(a: &[u8], b: &[u8]) -> Option<usize> {
467 if a.len() != b.len() {
468 return None;
469 }
470 Some(
471 a.iter()
472 .zip(b)
473 .filter(|(x, y)| !x.eq_ignore_ascii_case(y))
474 .count(),
475 )
476}
477
478pub fn n50(lengths: &mut [u64]) -> Option<u64> {
487 nx(lengths, 0.5)
488}
489
490pub fn nx(lengths: &mut [u64], fraction: f64) -> Option<u64> {
494 if lengths.is_empty() {
495 return None;
496 }
497 let total: u64 = lengths.iter().sum();
498 if total == 0 {
499 return Some(0);
500 }
501 lengths.sort_unstable_by(|a, b| b.cmp(a));
502 let target = total as f64 * fraction.clamp(0.0, 1.0);
503 let mut acc = 0u64;
504 for &len in lengths.iter() {
505 acc += len;
506 if acc as f64 >= target {
507 return Some(len);
508 }
509 }
510 lengths.last().copied()
511}
512
513#[cfg(test)]
514mod tests {
515 use super::*;
516
517 #[test]
518 fn complements_iupac_and_preserves_case() {
519 assert_eq!(reverse_complement(b"ACGTacgt"), b"acgtACGT");
520 assert_eq!(reverse_complement(b"RYKM"), b"KMRY");
521 assert_eq!(reverse_complement(b""), b"");
522 assert_eq!(reverse_complement(b"AC-GT"), b"AC-GT");
524 }
525
526 #[test]
527 fn in_place_matches_allocating() {
528 for s in [&b""[..], b"A", b"AC", b"ACG", b"ACGTN", b"acgtRYn"] {
529 let mut owned = s.to_vec();
530 reverse_complement_in_place(&mut owned);
531 assert_eq!(
532 owned,
533 reverse_complement(s),
534 "{:?}",
535 String::from_utf8_lossy(s)
536 );
537 }
538 }
539
540 #[test]
541 fn counts_bases() {
542 let c = BaseCounts::of(b"AACCGGTTNNxx");
543 assert_eq!((c.a, c.c, c.g, c.t, c.n, c.other), (2, 2, 2, 2, 2, 2));
544 assert_eq!(c.total(), 12);
545 assert_eq!(c.acgt(), 8);
546 assert_eq!(c.gc_content(), Some(0.5));
547 }
548
549 #[test]
550 fn validates_alphabets() {
551 assert!(Alphabet::Dna.validate(b"ACGTN").is_ok());
552 assert!(Alphabet::Dna.validate(b"ACGU").is_err());
553 assert!(Alphabet::Rna.validate(b"ACGU").is_ok());
554 assert!(Alphabet::Iupac.validate(b"ACGTRYKM-").is_ok());
555 assert!(Alphabet::Protein.validate(b"MEEPQSDPSV*").is_ok());
556 assert!(Alphabet::Any.validate(b"anything!").is_ok());
557 assert!(Alphabet::Any.validate(b"tab\there").is_err());
558
559 match Alphabet::Dna.validate_named(b"ACG!T", "read1") {
560 Err(Error::InvalidByte { id, pos, byte }) => {
561 assert_eq!((id.as_str(), pos, byte), ("read1", 3, b'!'));
562 }
563 other => panic!("expected InvalidByte, got {other:?}"),
564 }
565 }
566
567 #[test]
568 fn translates_all_codons() {
569 assert_eq!(translate(b"TTTTTCTTATTG", 0, false), b"FFLL");
571 assert_eq!(translate(b"ATGCATTAA", 0, false), b"MH*");
572 assert_eq!(translate(b"AUGCAU", 0, false), b"MH"); assert_eq!(translate(b"ATGCA", 0, false), b"M");
575 assert_eq!(translate(b"AT", 0, false), b"");
576 assert_eq!(translate(b"ATG", 5, false), b"");
577 }
578
579 #[test]
580 fn kmers_and_canonical() {
581 assert_eq!(kmers(b"AAAA", 4).count(), 1);
582 assert_eq!(kmers(b"AAAA", 0).count(), 0);
583 assert_eq!(canonical_kmer(b"TTT"), b"AAA");
584 assert_eq!(canonical_kmer(b"AAA"), b"AAA");
585 }
586
587 #[test]
588 fn hamming() {
589 assert_eq!(hamming_distance(b"ACGT", b"acgt"), Some(0));
590 assert_eq!(hamming_distance(b"ACGT", b"ACGA"), Some(1));
591 assert_eq!(hamming_distance(b"ACGT", b"ACG"), None);
592 }
593
594 #[test]
595 fn n_statistics() {
596 assert_eq!(n50(&mut [5, 50, 30, 15]), Some(50));
598 assert_eq!(nx(&mut [5, 50, 30, 15], 0.9), Some(15));
599 assert_eq!(nx(&mut [0, 0], 0.5), Some(0));
600 }
601
602 #[test]
603 fn rna_dna_conversion() {
604 let mut s = b"ACGUacgu".to_vec();
605 rna_to_dna(&mut s);
606 assert_eq!(s, b"ACGTacgt");
607 make_uppercase(&mut s);
608 assert_eq!(s, b"ACGTACGT");
609 }
610}