1use std::collections::HashMap;
2use std::io::{BufRead, Read, Write};
3use std::path::Path;
4
5use flate2::read::ZlibDecoder;
6
7use crate::Result;
8use crate::error::SlowError;
9use crate::header::RecordCompression;
10
11const INDEX_MAGIC: &[u8; 9] = b"SLOW5IDX\x01";
27const INDEX_EOF: &[u8; 8] = b"XDI5WOLS";
28const INDEX_HEADER_SIZE_OFFSET: usize = 64;
29
30#[derive(Debug)]
36pub struct Index {
37 pub(crate) entries: HashMap<String, IndexEntry>,
38}
39
40#[derive(Debug, Clone, Copy)]
41pub(crate) struct IndexEntry {
42 pub offset: u64,
44 pub size: u64,
46}
47
48impl Index {
49 fn new() -> Self {
50 Self {
51 entries: HashMap::new(),
52 }
53 }
54
55 pub fn read_ids(&self) -> impl Iterator<Item = &str> {
72 self.entries.keys().map(String::as_str)
73 }
74
75 pub(crate) fn get(&self, read_id: &str) -> Option<IndexEntry> {
76 self.entries.get(read_id).copied()
77 }
78
79 pub fn len(&self) -> usize {
92 self.entries.len()
93 }
94
95 pub fn is_empty(&self) -> bool {
108 self.entries.is_empty()
109 }
110
111 pub(crate) fn build_from_reader<R: Read>(
119 reader: &mut R,
120 record_compression: RecordCompression,
121 mut pos: u64,
122 ) -> Result<Self> {
123 let mut index = Self::new();
124 let mut buf = Vec::new(); loop {
127 let rec_offset = pos;
128
129 let mut prefix = [0u8; 5];
131 match reader.read_exact(&mut prefix) {
132 Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => {
133 return Err(SlowError::InvalidFormat(
134 "unexpected EOF while building index".into(),
135 ));
136 }
137 Err(e) => return Err(SlowError::Io(e)),
138 Ok(()) => {}
139 }
140
141 if &prefix == crate::blow5::EOF_MARKER {
142 break;
143 }
144
145 let mut suffix = [0u8; 3];
146 reader.read_exact(&mut suffix)?;
147
148 let record_size = u64::from_le_bytes([
149 prefix[0], prefix[1], prefix[2], prefix[3], prefix[4], suffix[0], suffix[1],
150 suffix[2],
151 ]);
152
153 let total_size = 8 + record_size;
154 pos += total_size;
155
156 buf.resize(record_size as usize, 0);
158 reader.read_exact(&mut buf)?;
159
160 let read_id = read_id_from_record(&buf, record_compression)?;
162
163 index.entries.insert(
164 read_id,
165 IndexEntry {
166 offset: rec_offset,
167 size: total_size,
168 },
169 );
170 }
171
172 Ok(index)
173 }
174
175 pub(crate) fn build_from_slow5_reader<R: BufRead>(
183 reader: &mut R,
184 mut pos: u64,
185 ) -> Result<Self> {
186 let mut index = Self::new();
187 let mut line = String::new();
188
189 loop {
190 line.clear();
191 let n = reader.read_line(&mut line).map_err(SlowError::Io)?;
192 if n == 0 {
193 break;
194 }
195
196 let line_start = pos;
197 let line_len = n as u64;
198 pos += line_len;
199
200 let t = line.trim_end();
201 if t.is_empty() || t.starts_with('#') || t.starts_with('@') {
202 continue;
203 }
204
205 let read_id = t
206 .split('\t')
207 .next()
208 .ok_or_else(|| SlowError::InvalidFormat("empty data line in SLOW5 file".into()))?
209 .to_string();
210
211 index.entries.insert(
212 read_id,
213 IndexEntry {
214 offset: line_start,
215 size: line_len,
216 },
217 );
218 }
219
220 Ok(index)
221 }
222
223 pub(crate) fn read_from_path(path: &Path, expected_version: (u8, u8, u8)) -> Result<Self> {
227 let mut file = std::fs::File::open(path)?;
228 let mut buf = Vec::new();
229 file.read_to_end(&mut buf)?;
230
231 if buf.len() < INDEX_MAGIC.len() || &buf[..INDEX_MAGIC.len()] != INDEX_MAGIC.as_slice() {
232 return Err(SlowError::Index("not a valid BLOW5 index file".into()));
233 }
234
235 let ver_start = INDEX_MAGIC.len();
236 if buf.len() < ver_start + 3 {
237 return Err(SlowError::Index("index file too short for version".into()));
238 }
239 let idx_version = (buf[ver_start], buf[ver_start + 1], buf[ver_start + 2]);
240 if idx_version != expected_version {
241 return Err(SlowError::Index(format!(
242 "index version {}.{}.{} does not match file version {}.{}.{}; re-index required",
243 idx_version.0,
244 idx_version.1,
245 idx_version.2,
246 expected_version.0,
247 expected_version.1,
248 expected_version.2,
249 )));
250 }
251
252 let mut pos = INDEX_HEADER_SIZE_OFFSET;
253 let mut index = Self::new();
254
255 loop {
256 if pos >= buf.len() {
257 return Err(SlowError::Index("index missing EOF marker".into()));
258 }
259 if buf.len() - pos >= INDEX_EOF.len()
260 && &buf[pos..pos + INDEX_EOF.len()] == INDEX_EOF.as_slice()
261 {
262 break;
263 }
264 if pos + 2 > buf.len() {
265 return Err(SlowError::Index(
266 "index truncated before read_id_len".into(),
267 ));
268 }
269 let read_id_len = u16::from_le_bytes([buf[pos], buf[pos + 1]]) as usize;
270 pos += 2;
271
272 if pos + read_id_len > buf.len() {
273 return Err(SlowError::Index("index truncated in read_id".into()));
274 }
275 let read_id = std::str::from_utf8(&buf[pos..pos + read_id_len])
276 .map_err(|e| SlowError::Index(format!("read_id not valid UTF-8 in index: {e}")))?
277 .to_string();
278 pos += read_id_len;
279
280 if pos + 16 > buf.len() {
281 return Err(SlowError::Index("index truncated in offset/size".into()));
282 }
283 let offset = u64::from_le_bytes(buf[pos..pos + 8].try_into().unwrap());
284 let size = u64::from_le_bytes(buf[pos + 8..pos + 16].try_into().unwrap());
285 pos += 16;
286
287 index.entries.insert(read_id, IndexEntry { offset, size });
288 }
289
290 Ok(index)
291 }
292
293 pub(crate) fn write_to_path(&self, path: &Path, version: (u8, u8, u8)) -> Result<()> {
295 let mut file = std::fs::File::create(path)?;
296
297 file.write_all(INDEX_MAGIC)?;
298 file.write_all(&[version.0, version.1, version.2])?;
299
300 let padding_len = INDEX_HEADER_SIZE_OFFSET - INDEX_MAGIC.len() - 3;
301 file.write_all(&vec![0u8; padding_len])?;
302
303 for (read_id, entry) in &self.entries {
304 let id_bytes = read_id.as_bytes();
305 let id_len = id_bytes.len() as u16;
306 file.write_all(&id_len.to_le_bytes())?;
307 file.write_all(id_bytes)?;
308 file.write_all(&entry.offset.to_le_bytes())?;
309 file.write_all(&entry.size.to_le_bytes())?;
310 }
311
312 file.write_all(INDEX_EOF)?;
313
314 Ok(())
315 }
316}
317
318fn read_id_from_record(record: &[u8], compression: RecordCompression) -> Result<String> {
325 match compression {
326 RecordCompression::None => read_id_from_bytes(record),
327 RecordCompression::Zstd => {
328 let decoder =
329 zstd::Decoder::new(record).map_err(|e| SlowError::Decompression(e.to_string()))?;
330 read_id_from_stream(decoder)
331 }
332 RecordCompression::Zlib => read_id_from_stream(ZlibDecoder::new(record)),
333 }
334}
335
336fn read_id_from_bytes(buf: &[u8]) -> Result<String> {
338 if buf.len() < 2 {
339 return Err(SlowError::InvalidFormat(
340 "record too short for read_id_len during index build".into(),
341 ));
342 }
343 let n = u16::from_le_bytes([buf[0], buf[1]]) as usize;
344 if buf.len() < 2 + n {
345 return Err(SlowError::InvalidFormat(
346 "record too short for read_id during index build".into(),
347 ));
348 }
349 std::str::from_utf8(&buf[2..2 + n])
350 .map_err(|e| SlowError::InvalidFormat(format!("read_id not valid UTF-8: {e}")))
351 .map(str::to_string)
352}
353
354fn read_id_from_stream<R: Read>(mut r: R) -> Result<String> {
359 let mut len_buf = [0u8; 2];
360 r.read_exact(&mut len_buf)
361 .map_err(|e| SlowError::InvalidFormat(format!("read_id_len read failed: {e}")))?;
362 let n = u16::from_le_bytes(len_buf) as usize;
363
364 let mut id_buf = vec![0u8; n];
365 r.read_exact(&mut id_buf)
366 .map_err(|e| SlowError::InvalidFormat(format!("read_id read failed: {e}")))?;
367
368 std::str::from_utf8(&id_buf)
369 .map_err(|e| SlowError::InvalidFormat(format!("read_id not valid UTF-8: {e}")))
370 .map(str::to_string)
371}