Skip to main content

barkit_extract/
parse.rs

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
11/// https://www.bioinformatics.org/sms/iupac.html
12const 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    /// Prepared regex pattern to parse barcodes
36    barcode_regex: BarcodeRegex,
37
38    /// If `true`, all captured patterns will not be trimmed
39    skip_trimming: bool,
40
41    /// If `true`, the barcode pattern will also be matched in the reverse complement sequence.
42    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}