Skip to main content

dino_seq/fastq/
chunk.rs

1use crate::Result;
2
3use super::FastqVisitRecord;
4
5/// Caller-provided sink for single-pass FASTQ record processing.
6///
7/// This trait is the allocation-control counterpart to
8/// [`FastqReader::next_batch`](super::FastqReader::next_batch). Use it when a
9/// downstream tool wants to fill its own output buffers directly and does not
10/// need Dino Seq to build a batch side table.
11///
12/// If `record` returns an error, [`FastqReader::next_chunk_with_sink`](super::FastqReader::next_chunk_with_sink)
13/// returns that error immediately. The sink may already have received records
14/// from the current chunk, and the reader should be dropped rather than reused
15/// after a sink error.
16pub trait FastqRecordSink {
17    /// Consume one borrowed FASTQ record.
18    fn record(&mut self, record: FastqVisitRecord<'_>) -> Result<()>;
19}
20
21impl<F> FastqRecordSink for F
22where
23    F: for<'a> FnMut(FastqVisitRecord<'a>) -> Result<()>,
24{
25    fn record(&mut self, record: FastqVisitRecord<'_>) -> Result<()> {
26        self(record)
27    }
28}
29
30/// Optional extension point for sinks that can preallocate per chunk.
31///
32/// Dino Seq calls [`FastqRecordSink::record`] for correctness. Downstream tools
33/// that own a growable output buffer can also implement this trait locally and
34/// call [`FastqChunkConfig::estimated_records`] before invoking
35/// [`FastqReader::next_chunk_with_sink`](super::FastqReader::next_chunk_with_sink).
36pub trait FastqChunkSinkExt: FastqRecordSink {
37    /// Reserve space for approximately `records` upcoming records.
38    fn reserve_records(&mut self, _records: usize) {}
39}
40
41/// Chunking policy for [`FastqReader::next_chunk_with_sink`](super::FastqReader::next_chunk_with_sink).
42#[non_exhaustive]
43#[derive(Debug, Clone, Copy, PartialEq, Eq)]
44pub struct FastqChunkConfig {
45    /// Target number of sequence bases to emit before returning a chunk.
46    ///
47    /// A value of zero disables the target limit.
48    pub target_bases: u64,
49    /// Minimum number of records required before the target-base limit may stop
50    /// a chunk.
51    ///
52    /// Values below one are raised to one.
53    pub min_records: usize,
54    /// Hard maximum number of sequence bases to emit before returning a chunk.
55    ///
56    /// A value of zero disables the hard limit. Unlike [`target_bases`](Self::target_bases),
57    /// this limit is not gated by [`min_records`](Self::min_records).
58    pub max_bases: u64,
59}
60
61impl FastqChunkConfig {
62    /// Create a chunking policy with a target base count.
63    pub fn new(target_bases: u64) -> Self {
64        Self {
65            target_bases,
66            min_records: 1,
67            max_bases: 0,
68        }
69    }
70
71    /// Return a copy with a minimum record count.
72    pub fn min_records(mut self, min_records: usize) -> Self {
73        self.min_records = min_records.max(1);
74        self
75    }
76
77    /// Return a copy with a hard maximum base count.
78    pub fn max_bases(mut self, max_bases: u64) -> Self {
79        self.max_bases = max_bases;
80        self
81    }
82
83    /// Target bases per chunk. Zero means no target-base limit.
84    pub fn target_bases(self) -> u64 {
85        self.target_bases
86    }
87
88    /// Minimum records required before the target-base limit may stop a chunk.
89    pub fn min_records_value(self) -> usize {
90        self.min_records
91    }
92
93    /// Hard maximum bases per chunk. Zero means no hard limit.
94    pub fn max_bases_value(self) -> u64 {
95        self.max_bases
96    }
97
98    /// Estimate records per chunk for a fixed or representative read length.
99    pub fn estimated_records(self, read_len: usize) -> usize {
100        let read_len = read_len.max(1) as u64;
101        let target = if self.max_bases != 0 {
102            self.max_bases
103        } else {
104            self.target_bases
105        };
106        if target == 0 {
107            return self.min_records.max(1);
108        }
109        target
110            .div_ceil(read_len)
111            .max(self.min_records.max(1) as u64)
112            .min(usize::MAX as u64) as usize
113    }
114
115    pub(crate) fn normalized(self) -> Self {
116        Self {
117            min_records: self.min_records.max(1),
118            ..self
119        }
120    }
121
122    pub(crate) fn should_stop(self, records: u64, bases: u64) -> bool {
123        if self.max_bases != 0 && bases >= self.max_bases {
124            return true;
125        }
126        self.target_bases != 0 && bases >= self.target_bases && records >= self.min_records as u64
127    }
128}
129
130/// Summary for one chunk emitted by [`FastqReader::next_chunk_with_sink`](super::FastqReader::next_chunk_with_sink).
131#[non_exhaustive]
132#[derive(Debug, Clone, Copy, PartialEq, Eq)]
133pub struct FastqChunkStats {
134    /// Zero-based index of the first record in this chunk.
135    pub first_record_index: u64,
136    /// Number of records emitted into the sink.
137    pub records: u64,
138    /// Number of sequence bases emitted into the sink.
139    pub bases: u64,
140}
141
142impl FastqChunkStats {
143    pub(crate) fn new(first_record_index: u64) -> Self {
144        Self {
145            first_record_index,
146            records: 0,
147            bases: 0,
148        }
149    }
150
151    /// Zero-based index of the first record in this chunk.
152    pub fn first_record_index(self) -> u64 {
153        self.first_record_index
154    }
155
156    /// Number of records emitted into the sink.
157    pub fn records(self) -> u64 {
158        self.records
159    }
160
161    /// Number of sequence bases emitted into the sink.
162    pub fn bases(self) -> u64 {
163        self.bases
164    }
165
166    /// Whether the chunk contains no records.
167    pub fn is_empty(self) -> bool {
168        self.records == 0
169    }
170}