1#![allow(clippy::result_large_err)]
2
3use crate::pattern::BarcodeRegex;
4use regex::bytes::Captures;
5
6use seq_io::fastq::{OwnedRecord, Record, RefRecord};
7use std::str;
8
9use crate::error::Error;
10
11const TRANSLATION_TABLE: [u8; 256] = {
13 let mut table = [b'A'; 256];
14
15 table[b'A' as usize] = b'T';
16 table[b'T' as usize] = b'A';
17 table[b'G' as usize] = b'C';
18 table[b'C' as usize] = b'G';
19 table[b'R' as usize] = b'R';
20 table[b'Y' as usize] = b'Y';
21 table[b'S' as usize] = b'S';
22 table[b'W' as usize] = b'W';
23 table[b'K' as usize] = b'K';
24 table[b'M' as usize] = b'M';
25 table[b'B' as usize] = b'B';
26 table[b'D' as usize] = b'D';
27 table[b'H' as usize] = b'H';
28 table[b'V' as usize] = b'V';
29 table[b'N' as usize] = b'N';
30
31 table
32};
33
34pub struct BarcodeParser {
35 barcode_regex: BarcodeRegex,
37
38 skip_trimming: bool,
40
41 rc_barcodes: bool,
43}
44
45impl BarcodeParser {
46 pub fn new(
47 barcode_regex: Option<&BarcodeRegex>,
48 skip_trimming: bool,
49 rc_barcodes: bool,
50 ) -> Option<Self> {
51 Some(BarcodeParser {
52 barcode_regex: barcode_regex?.to_owned(),
53 skip_trimming,
54 rc_barcodes,
55 })
56 }
57
58 pub fn parse_barcodes(&self, record: &RefRecord) -> Option<OwnedRecord> {
59 let read_captures = self.barcode_regex.get_captures(record.seq());
60 let read_seq_rc: Vec<u8>;
61 let read_captures = if read_captures.is_err() && self.rc_barcodes {
62 read_seq_rc = get_reverse_complement(record.seq());
63 self.barcode_regex.get_captures(&read_seq_rc)
64 } else {
65 read_captures
66 };
67 self.create_read(read_captures.map(Some), record)
68 }
69
70 fn create_read(
71 &self,
72 read_captures: Result<Option<Captures>, Error>,
73 record: &RefRecord,
74 ) -> Option<seq_io::fastq::OwnedRecord> {
75 match (read_captures, self.skip_trimming) {
76 (Ok(Some(captures)), true) => {
77 Some(self.create_read_with_new_header(&captures, record).ok()?)
78 }
79 (Ok(Some(captures)), false) => {
80 let new_read = self.create_read_with_new_header(&captures, record).ok()?;
81 Some(trim_adapters(captures, &new_read).ok()?)
82 }
83 (Ok(None), _) => Some(OwnedRecord {
84 head: record.head().to_vec(),
85 seq: record.seq().to_vec(),
86 qual: record.qual().to_vec(),
87 }),
88 (Err(_), _) => None,
89 }
90 }
91
92 fn create_read_with_new_header(
93 &self,
94 captures: &Captures,
95 record: &RefRecord,
96 ) -> Result<OwnedRecord, Error> {
97 let mut head = record.head().to_vec();
98 let seq = record.seq().to_vec();
99 let qual = record.qual().to_vec();
100
101 for barcode in &self.barcode_regex.get_barcode_types() {
102 let barcode_name = barcode.to_string();
103 let (barcode_start, barcode_end) =
104 get_barcode_match_positions(&barcode_name, captures)?;
105 head = add_to_the_header(
106 &barcode_name,
107 &head,
108 &seq,
109 &qual,
110 barcode_start,
111 barcode_end,
112 )?;
113 }
114
115 Ok(OwnedRecord { head, seq, qual })
116 }
117}
118
119fn get_full_match_positions(captures: &Captures) -> Result<(usize, usize), Error> {
120 let full_match = captures
121 .get(0)
122 .ok_or(Error::BarcodeCaptureGroupNotFound("0".to_owned()))?;
123
124 Ok((full_match.start(), full_match.end()))
125}
126
127fn get_barcode_match_positions(
128 barcode_name: &str,
129 captures: &Captures,
130) -> Result<(usize, usize), Error> {
131 let full_match = captures
132 .name(barcode_name)
133 .ok_or(Error::BarcodeCaptureGroupNotFound(barcode_name.to_string()))?;
134
135 Ok((full_match.start(), full_match.end()))
136}
137
138fn trim_adapters(captures: Captures, record: &OwnedRecord) -> Result<OwnedRecord, Error> {
139 let (start, end) = get_full_match_positions(&captures)?;
140 let seq = [&record.seq()[..start], &record.seq()[end..]].concat();
141 let qual = [&record.qual()[..start], &record.qual()[end..]].concat();
142
143 Ok(OwnedRecord {
144 head: record.head().to_vec(),
145 seq,
146 qual,
147 })
148}
149
150fn add_to_the_header(
151 barcode_type: &str,
152 head: &[u8],
153 seq: &[u8],
154 qual: &[u8],
155 start: usize,
156 end: usize,
157) -> Result<Vec<u8>, Error> {
158 let barcode_seq = &seq[start..end];
159 let barcode_qual = &qual[start..end];
160
161 let mut result = Vec::with_capacity(
162 head.len() + barcode_type.len() + barcode_seq.len() + barcode_qual.len() + 3,
163 );
164 result.extend_from_slice(head);
165 result.extend_from_slice(
166 format!(" {}:{}:", barcode_type, std::str::from_utf8(barcode_seq)?).as_bytes(),
167 );
168 result.extend_from_slice(barcode_qual);
169
170 Ok(result)
171}
172
173pub fn get_reverse_complement(sequence: &[u8]) -> Vec<u8> {
174 sequence
175 .iter()
176 .rev()
177 .map(|&base| TRANSLATION_TABLE[base as usize])
178 .collect()
179}
180
181#[cfg(test)]
182mod tests {
183 use rstest::rstest;
184
185 use crate::parse::get_reverse_complement;
186
187 #[rstest]
188 #[case(b"", b"")]
189 #[case(b"GGGCCCAAATTT", b"AAATTTGGGCCC")]
190 #[case(b"ATGCN", b"NGCAT")]
191 #[case(b"AAP", b"ATT")]
192 #[case(b"CCX", b"AGG")]
193 #[case(b"PPP", b"AAA")]
194 fn test_get_reverse_complement(#[case] sequence: &[u8], #[case] rc_sequence: &[u8]) {
195 assert_eq!(get_reverse_complement(sequence), rc_sequence);
196 }
197}