cuttlefish_rs/
input.rs

1//! FASTA/FASTQ input discovery and zero-copy fragment parsing.
2//!
3//! Parsers split records at non-ACGT symbols. Borrowed callbacks are preferred
4//! by the production partitioner so sequence data does not need to be copied.
5
6use crate::dna::is_dna_ascii;
7use crate::params::BuildParams;
8use flate2::read::MultiGzDecoder;
9use std::fs;
10use std::io::{BufRead, BufReader, Read, Seek};
11use std::path::{Path, PathBuf};
12
13#[derive(Debug, Clone, PartialEq, Eq)]
14pub struct SequenceFragment {
15    pub source_id: u32,
16    pub record_id: u64,
17    pub offset: usize,
18    pub seq: Vec<u8>,
19}
20
21#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22pub struct BorrowedSequenceFragment<'a> {
23    pub source_id: u32,
24    pub record_id: u64,
25    pub offset: usize,
26    pub seq: &'a [u8],
27}
28
29pub fn expand_input_paths(params: &BuildParams) -> Result<Vec<PathBuf>, InputError> {
30    let mut paths = Vec::new();
31
32    for path in &params.seqs {
33        paths.push(PathBuf::from(path));
34    }
35
36    for list in &params.lists {
37        let file = fs::File::open(list).map_err(|source| InputError::Io {
38            path: PathBuf::from(list),
39            source,
40        })?;
41        for line in BufReader::new(file).lines() {
42            let line = line.map_err(|source| InputError::Io {
43                path: PathBuf::from(list),
44                source,
45            })?;
46            let trimmed = line.trim();
47            if !trimmed.is_empty() {
48                paths.push(PathBuf::from(trimmed));
49            }
50        }
51    }
52
53    for dir in &params.dirs {
54        let mut entries = fs::read_dir(dir)
55            .map_err(|source| InputError::Io {
56                path: PathBuf::from(dir),
57                source,
58            })?
59            .collect::<Result<Vec<_>, _>>()
60            .map_err(|source| InputError::Io {
61                path: PathBuf::from(dir),
62                source,
63            })?;
64        entries.sort_by_key(|entry| entry.path());
65        for entry in entries {
66            let path = entry.path();
67            if path.is_file() {
68                paths.push(path);
69            }
70        }
71    }
72
73    if paths.is_empty() {
74        return Err(InputError::NoInput);
75    }
76
77    Ok(paths)
78}
79
80pub fn parse_fragments<P, F>(
81    path: P,
82    source_id: u32,
83    min_len: usize,
84    mut on_fragment: F,
85) -> Result<u64, InputError>
86where
87    P: AsRef<Path>,
88    F: FnMut(SequenceFragment) -> Result<(), InputError>,
89{
90    parse_fragments_borrowed(path, source_id, min_len, |fragment| {
91        on_fragment(SequenceFragment {
92            source_id: fragment.source_id,
93            record_id: fragment.record_id,
94            offset: fragment.offset,
95            seq: normalized_fragment_seq(fragment.seq),
96        })
97    })
98}
99
100pub fn parse_fragments_borrowed<P, F>(
101    path: P,
102    source_id: u32,
103    min_len: usize,
104    on_fragment: F,
105) -> Result<u64, InputError>
106where
107    P: AsRef<Path>,
108    F: for<'a> FnMut(BorrowedSequenceFragment<'a>) -> Result<(), InputError>,
109{
110    parse_fragments_borrowed_with(path, source_id, min_len, 1, on_fragment)
111}
112
113/// As [`parse_fragments_borrowed`], but permitted to use `inflate_workers`
114/// threads to decompress block-structured (BGZF) gzip input.
115///
116/// The decompressed byte stream is identical either way; only the work of
117/// producing it is shared out. Plain gzip ignores the budget.
118pub fn parse_fragments_borrowed_with<P, F>(
119    path: P,
120    source_id: u32,
121    min_len: usize,
122    inflate_workers: usize,
123    mut on_fragment: F,
124) -> Result<u64, InputError>
125where
126    P: AsRef<Path>,
127    F: for<'a> FnMut(BorrowedSequenceFragment<'a>) -> Result<(), InputError>,
128{
129    let path = path.as_ref();
130    let mut file = fs::File::open(path).map_err(|source| InputError::Io {
131        path: path.to_path_buf(),
132        source,
133    })?;
134    let input: Box<dyn Read> = if path.extension().is_some_and(|ext| ext == "gz") {
135        // BGZF concatenates independent gzip members, so its blocks can be
136        // inflated concurrently. A plain member cannot be split, and falls back
137        // to the serial decoder.
138        let mut head = [0u8; crate::bgzf::PROBE_BYTES];
139        let probed = read_probe(&mut file, &mut head).map_err(|source| InputError::Io {
140            path: path.to_path_buf(),
141            source,
142        })?;
143        file.rewind().map_err(|source| InputError::Io {
144            path: path.to_path_buf(),
145            source,
146        })?;
147        if inflate_workers > 1 && crate::bgzf::is_bgzf(&head[..probed]) {
148            Box::new(crate::bgzf::ParallelBgzfReader::new(file, inflate_workers))
149        } else {
150            Box::new(MultiGzDecoder::new(file))
151        }
152    } else {
153        Box::new(file)
154    };
155    let mut reader = BufReader::with_capacity(1024 * 1024, input);
156    let first_line = next_non_empty_line(&mut reader, path)?;
157    match first_line.first().copied() {
158        Some(b'>') => parse_fasta_reader(first_line, reader, source_id, min_len, &mut on_fragment),
159        Some(b'@') => parse_fastq_reader(first_line, reader, source_id, min_len, &mut on_fragment),
160        Some(_) if first_line.iter().copied().all(is_dna_ascii) => {
161            parse_plain_sequence_reader(first_line, reader, source_id, min_len, &mut on_fragment)
162        }
163        Some(_) => Err(InputError::UnknownFormat(path.to_path_buf())),
164        None => Err(InputError::EmptyFile(path.to_path_buf())),
165    }
166}
167
168/// Fills as much of `head` as the file provides, tolerating short reads.
169fn read_probe<R: Read>(source: &mut R, head: &mut [u8]) -> std::io::Result<usize> {
170    let mut filled = 0;
171    while filled < head.len() {
172        match source.read(&mut head[filled..]) {
173            Ok(0) => break,
174            Ok(n) => filled += n,
175            Err(error) if error.kind() == std::io::ErrorKind::Interrupted => continue,
176            Err(error) => return Err(error),
177        }
178    }
179    Ok(filled)
180}
181
182/// Appends a payload line, dropping embedded whitespace.
183///
184/// The filtering form is a per-byte predicate and push, which cannot vectorize.
185/// Sequence lines essentially never contain interior whitespace, so probe first
186/// with a scan the compiler can vectorize and fall back only when needed. The
187/// slow path is kept so that interior whitespace still joins a record rather
188/// than splitting it into fragments.
189#[inline]
190fn append_sequence_line(seq: &mut Vec<u8>, line: &[u8]) {
191    if line.iter().any(u8::is_ascii_whitespace) {
192        seq.extend(line.iter().copied().filter(|b| !b.is_ascii_whitespace()));
193    } else {
194        seq.extend_from_slice(line);
195    }
196}
197
198fn parse_plain_sequence_reader<R, F>(
199    first_line: Vec<u8>,
200    mut reader: R,
201    source_id: u32,
202    min_len: usize,
203    on_fragment: &mut F,
204) -> Result<u64, InputError>
205where
206    R: BufRead,
207    F: for<'a> FnMut(BorrowedSequenceFragment<'a>) -> Result<(), InputError>,
208{
209    let mut seq = first_line;
210    let mut line = Vec::new();
211    loop {
212        line.clear();
213        if reader
214            .read_until(b'\n', &mut line)
215            .map_err(|source| InputError::Read { source })?
216            == 0
217        {
218            break;
219        }
220        trim_ascii_line_in_place(&mut line);
221        append_sequence_line(&mut seq, &line);
222    }
223    emit_actg_fragments(source_id, 1, &seq, min_len, on_fragment)?;
224    Ok(1)
225}
226
227fn next_non_empty_line<R: BufRead>(reader: &mut R, path: &Path) -> Result<Vec<u8>, InputError> {
228    let mut line = Vec::new();
229    loop {
230        line.clear();
231        let n = reader
232            .read_until(b'\n', &mut line)
233            .map_err(|source| InputError::Io {
234                path: path.to_path_buf(),
235                source,
236            })?;
237        if n == 0 {
238            return Err(InputError::EmptyFile(path.to_path_buf()));
239        }
240        trim_ascii_line_in_place(&mut line);
241        if !line.is_empty() {
242            return Ok(line);
243        }
244    }
245}
246
247fn parse_fasta_reader<R, F>(
248    first_header: Vec<u8>,
249    mut reader: R,
250    source_id: u32,
251    min_len: usize,
252    on_fragment: &mut F,
253) -> Result<u64, InputError>
254where
255    R: BufRead,
256    F: for<'a> FnMut(BorrowedSequenceFragment<'a>) -> Result<(), InputError>,
257{
258    debug_assert!(first_header.starts_with(b">"));
259    let mut records = 0u64;
260    let mut record_id = 1u64;
261    let mut seq = Vec::new();
262    let mut line = Vec::new();
263
264    loop {
265        line.clear();
266        let n = reader
267            .read_until(b'\n', &mut line)
268            .map_err(|source| InputError::Read { source })?;
269        if n == 0 {
270            break;
271        }
272        trim_ascii_line_in_place(&mut line);
273        if line.starts_with(b">") {
274            records += 1;
275            emit_actg_fragments(source_id, record_id, &seq, min_len, on_fragment)?;
276            seq.clear();
277            record_id += 1;
278        } else if !line.is_empty() {
279            append_sequence_line(&mut seq, &line);
280        }
281    }
282
283    records += 1;
284    emit_actg_fragments(source_id, record_id, &seq, min_len, on_fragment)?;
285
286    Ok(records)
287}
288
289fn parse_fastq_reader<R, F>(
290    first_header: Vec<u8>,
291    mut reader: R,
292    source_id: u32,
293    min_len: usize,
294    on_fragment: &mut F,
295) -> Result<u64, InputError>
296where
297    R: BufRead,
298    F: for<'a> FnMut(BorrowedSequenceFragment<'a>) -> Result<(), InputError>,
299{
300    debug_assert!(first_header.starts_with(b"@"));
301    let mut record_id = 0u64;
302    let mut header = first_header;
303    let mut seq = Vec::new();
304    let mut line = Vec::new();
305
306    loop {
307        record_id += 1;
308        if !header.starts_with(b"@") {
309            return Err(InputError::MalformedFastq(record_id));
310        }
311
312        seq.clear();
313        loop {
314            line.clear();
315            if reader
316                .read_until(b'\n', &mut line)
317                .map_err(|source| InputError::Read { source })?
318                == 0
319            {
320                return Err(InputError::MalformedFastq(record_id));
321            }
322            trim_ascii_line_in_place(&mut line);
323            if line.starts_with(b"+") {
324                break;
325            }
326            append_sequence_line(&mut seq, &line);
327        }
328
329        let mut qual_len = 0usize;
330        while qual_len < seq.len() {
331            line.clear();
332            if reader
333                .read_until(b'\n', &mut line)
334                .map_err(|source| InputError::Read { source })?
335                == 0
336            {
337                return Err(InputError::MalformedFastq(record_id));
338            }
339            trim_ascii_line_in_place(&mut line);
340            qual_len += line.len();
341        }
342
343        emit_actg_fragments(source_id, record_id, &seq, min_len, on_fragment)?;
344
345        header.clear();
346        let bytes = reader
347            .read_until(b'\n', &mut header)
348            .map_err(|source| InputError::Read { source })?;
349        if bytes == 0 {
350            break;
351        }
352        trim_ascii_line_in_place(&mut header);
353        if header.is_empty() {
354            return Err(InputError::MalformedFastq(record_id + 1));
355        }
356    }
357
358    Ok(record_id)
359}
360
361#[cfg(test)]
362fn parse_fasta_bytes<F>(
363    bytes: &[u8],
364    source_id: u32,
365    min_len: usize,
366    on_fragment: &mut F,
367) -> Result<u64, InputError>
368where
369    F: FnMut(SequenceFragment) -> Result<(), InputError>,
370{
371    let mut reader = BufReader::new(bytes);
372    let first = next_non_empty_line(&mut reader, Path::new("<memory>"))?;
373    let mut on_borrowed = |fragment: BorrowedSequenceFragment<'_>| {
374        on_fragment(SequenceFragment {
375            source_id: fragment.source_id,
376            record_id: fragment.record_id,
377            offset: fragment.offset,
378            seq: normalized_fragment_seq(fragment.seq),
379        })
380    };
381    match first.first().copied() {
382        Some(b'>') => parse_fasta_reader(first, reader, source_id, min_len, &mut on_borrowed),
383        Some(_) => Err(InputError::UnknownFormat(PathBuf::from("<memory>"))),
384        None => Err(InputError::EmptyFile(PathBuf::from("<memory>"))),
385    }
386}
387
388#[cfg(test)]
389fn parse_fastq_bytes<F>(
390    bytes: &[u8],
391    source_id: u32,
392    min_len: usize,
393    on_fragment: &mut F,
394) -> Result<u64, InputError>
395where
396    F: FnMut(SequenceFragment) -> Result<(), InputError>,
397{
398    let mut reader = BufReader::new(bytes);
399    let first = next_non_empty_line(&mut reader, Path::new("<memory>"))?;
400    let mut on_borrowed = |fragment: BorrowedSequenceFragment<'_>| {
401        on_fragment(SequenceFragment {
402            source_id: fragment.source_id,
403            record_id: fragment.record_id,
404            offset: fragment.offset,
405            seq: normalized_fragment_seq(fragment.seq),
406        })
407    };
408    match first.first().copied() {
409        Some(b'@') => parse_fastq_reader(first, reader, source_id, min_len, &mut on_borrowed),
410        Some(_) => Err(InputError::UnknownFormat(PathBuf::from("<memory>"))),
411        None => Err(InputError::EmptyFile(PathBuf::from("<memory>"))),
412    }
413}
414
415fn emit_actg_fragments<F>(
416    source_id: u32,
417    record_id: u64,
418    seq: &[u8],
419    min_len: usize,
420    on_fragment: &mut F,
421) -> Result<(), InputError>
422where
423    F: for<'a> FnMut(BorrowedSequenceFragment<'a>) -> Result<(), InputError>,
424{
425    let mut start = None;
426    for (idx, &base) in seq.iter().enumerate() {
427        if is_dna_ascii(base) {
428            start.get_or_insert(idx);
429        } else if let Some(beg) = start.take() {
430            emit_fragment(
431                source_id,
432                record_id,
433                beg,
434                &seq[beg..idx],
435                min_len,
436                on_fragment,
437            )?;
438        }
439    }
440
441    if let Some(beg) = start {
442        emit_fragment(source_id, record_id, beg, &seq[beg..], min_len, on_fragment)?;
443    }
444
445    Ok(())
446}
447
448fn emit_fragment<F>(
449    source_id: u32,
450    record_id: u64,
451    offset: usize,
452    seq: &[u8],
453    min_len: usize,
454    on_fragment: &mut F,
455) -> Result<(), InputError>
456where
457    F: for<'a> FnMut(BorrowedSequenceFragment<'a>) -> Result<(), InputError>,
458{
459    if seq.len() >= min_len {
460        on_fragment(BorrowedSequenceFragment {
461            source_id,
462            record_id,
463            offset,
464            seq,
465        })?;
466    }
467    Ok(())
468}
469
470pub fn normalized_fragment_seq(seq: &[u8]) -> Vec<u8> {
471    if seq.iter().all(|&b| matches!(b, b'A' | b'C' | b'G' | b'T')) {
472        seq.to_vec()
473    } else {
474        seq.iter().map(|b| b.to_ascii_uppercase()).collect()
475    }
476}
477
478#[inline]
479fn trim_ascii_line_in_place(line: &mut Vec<u8>) {
480    while line.last().is_some_and(|b| b.is_ascii_whitespace()) {
481        line.pop();
482    }
483}
484
485#[derive(Debug)]
486pub enum InputError {
487    NoInput,
488    EmptyFile(PathBuf),
489    UnknownFormat(PathBuf),
490    MalformedFastq(u64),
491    Io {
492        path: PathBuf,
493        source: std::io::Error,
494    },
495    Read {
496        source: std::io::Error,
497    },
498    Partition(crate::partition::PartitionError),
499    Bucket(crate::buckets::BucketError),
500}
501
502impl std::fmt::Display for InputError {
503    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
504        match self {
505            Self::NoInput => write!(f, "no input files resolved"),
506            Self::EmptyFile(path) => write!(f, "input file is empty: {}", path.display()),
507            Self::UnknownFormat(path) => {
508                write!(f, "unknown FASTA/FASTQ format: {}", path.display())
509            }
510            Self::MalformedFastq(record) => write!(f, "malformed FASTQ record {record}"),
511            Self::Io { path, source } => write!(f, "{}: {source}", path.display()),
512            Self::Read { source } => write!(f, "{source}"),
513            Self::Partition(err) => write!(f, "{err}"),
514            Self::Bucket(err) => write!(f, "{err}"),
515        }
516    }
517}
518
519impl std::error::Error for InputError {}
520
521impl From<crate::buckets::BucketError> for InputError {
522    fn from(value: crate::buckets::BucketError) -> Self {
523        Self::Bucket(value)
524    }
525}
526
527#[cfg(test)]
528mod tests {
529    use super::*;
530    use flate2::Compression;
531    use flate2::write::GzEncoder;
532    use std::io::Write;
533
534    #[test]
535    fn fasta_splits_on_non_actg() {
536        let mut fragments = Vec::new();
537        let records = parse_fasta_bytes(b">r1\nAACNNttg\n>r2\nCC\n", 1, 2, &mut |frag| {
538            fragments.push(frag);
539            Ok(())
540        })
541        .unwrap();
542
543        assert_eq!(records, 2);
544        assert_eq!(
545            fragments
546                .iter()
547                .map(|f| f.seq.as_slice())
548                .collect::<Vec<_>>(),
549            vec![b"AAC".as_slice(), b"TTG".as_slice(), b"CC".as_slice()]
550        );
551        assert_eq!(fragments[1].offset, 5);
552    }
553
554    #[test]
555    fn fastq_splits_and_counts_records() {
556        let mut fragments = Vec::new();
557        let records = parse_fastq_bytes(b"@r1\nACNTA\n+\nIIIII\n", 7, 2, &mut |frag| {
558            fragments.push(frag);
559            Ok(())
560        })
561        .unwrap();
562
563        assert_eq!(records, 1);
564        assert_eq!(fragments.len(), 2);
565        assert_eq!(fragments[0].source_id, 7);
566        assert_eq!(fragments[0].seq, b"AC");
567        assert_eq!(fragments[1].seq, b"TA");
568    }
569
570    #[test]
571    fn parses_gzipped_fastq() {
572        let path =
573            std::env::temp_dir().join(format!("cf3rs-input-{}.fastq.gz", std::process::id()));
574        let mut encoder = GzEncoder::new(Vec::new(), Compression::default());
575        encoder.write_all(b"@r1\nACGTNNTA\n+\nIIIIIIII\n").unwrap();
576        fs::write(&path, encoder.finish().unwrap()).unwrap();
577
578        let mut fragments = Vec::new();
579        let records = parse_fragments(&path, 3, 2, |fragment| {
580            fragments.push(fragment);
581            Ok(())
582        })
583        .unwrap();
584
585        assert_eq!(records, 1);
586        assert_eq!(
587            fragments
588                .iter()
589                .map(|fragment| fragment.seq.as_slice())
590                .collect::<Vec<_>>(),
591            vec![b"ACGT".as_slice(), b"TA".as_slice()]
592        );
593
594        let _ = fs::remove_file(path);
595    }
596
597    #[test]
598    fn parses_headerless_wrapped_sequence() {
599        let path =
600            std::env::temp_dir().join(format!("cf3rs-input-{}-plain.fna.gz", std::process::id()));
601        let mut encoder = GzEncoder::new(Vec::new(), Compression::default());
602        encoder.write_all(b"\nAACGT\nTNNACGT\n").unwrap();
603        fs::write(&path, encoder.finish().unwrap()).unwrap();
604
605        let mut fragments = Vec::new();
606        let records = parse_fragments(&path, 9, 3, |fragment| {
607            fragments.push(fragment);
608            Ok(())
609        })
610        .unwrap();
611        assert_eq!(records, 1);
612        assert_eq!(fragments[0].seq, b"AACGTT");
613        assert_eq!(fragments[1].seq, b"ACGT");
614        assert_eq!(fragments[1].offset, 8);
615
616        let _ = fs::remove_file(path);
617    }
618}
cuttlefish_rs/input.rs

cuttlefish_rs/
input.rs
