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}