Skip to main content

slow5lib/
index.rs

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
11/// Binary index file format constants.
12///
13/// Wire layout of the `.blow5.idx` file header:
14///
15/// ```text
16/// Offset  Size  Type      Description
17///      0     9  [u8; 9]   Magic: b"SLOW5IDX\x01"
18///      9     3  [u8; 3]   Version: major, minor, patch
19///  12-63     -  zeros     Padding to offset 64
20///     64+    -  records   Repeated: [read_id_len: u16 le][read_id: N bytes][offset: u64 le][size: u64 le]
21///    end     8  [u8; 8]   EOF marker: b"XDI5WOLS"
22/// ```
23///
24/// `offset` is the byte offset of the record's `record_size` field in the data file.
25/// `size` = 8 + `record_size` (covers the 8-byte length prefix and the compressed data).
26const INDEX_MAGIC: &[u8; 9] = b"SLOW5IDX\x01";
27const INDEX_EOF: &[u8; 8] = b"XDI5WOLS";
28const INDEX_HEADER_SIZE_OFFSET: usize = 64;
29
30/// In-memory index over a `.slow5.idx` or `.blow5.idx` file.
31///
32/// Maps each `read_id` to the byte offset and length of its record in the
33/// data file, enabling random access via [`Slow5IndexedReader::get()`](crate::Slow5IndexedReader::get).
34/// The index is loaded or built automatically by `Slow5IndexedReader::open()`.
35#[derive(Debug)]
36pub struct Index {
37    pub(crate) entries: HashMap<String, IndexEntry>,
38}
39
40#[derive(Debug, Clone, Copy)]
41pub(crate) struct IndexEntry {
42    /// Byte offset of the `record_size` field in the data file.
43    pub offset: u64,
44    /// Total byte count: 8 (record_size field) + record_size (compressed data).
45    pub size: u64,
46}
47
48impl Index {
49    fn new() -> Self {
50        Self {
51            entries: HashMap::new(),
52        }
53    }
54
55    /// Iterator over all read ids in the index.
56    ///
57    /// Iteration order is not defined (backed by a `HashMap`).
58    ///
59    /// # Examples
60    ///
61    /// ```no_run
62    /// use slow5lib::Slow5IndexedReader;
63    ///
64    /// # fn example() -> slow5lib::Result<()> {
65    /// let reader = Slow5IndexedReader::open("reads.blow5")?;
66    /// for id in reader.read_ids() {
67    ///     println!("{id}");
68    /// }
69    /// # Ok(()) }
70    /// ```
71    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    /// Returns the number of reads in the index.
80    ///
81    /// # Examples
82    ///
83    /// ```no_run
84    /// use slow5lib::Slow5IndexedReader;
85    ///
86    /// # fn example() -> slow5lib::Result<()> {
87    /// let reader = Slow5IndexedReader::open("reads.blow5")?;
88    /// println!("{} reads indexed", reader.read_ids().count());
89    /// # Ok(()) }
90    /// ```
91    pub fn len(&self) -> usize {
92        self.entries.len()
93    }
94
95    /// Returns `true` if the index contains no reads.
96    ///
97    /// # Examples
98    ///
99    /// ```no_run
100    /// use slow5lib::Slow5IndexedReader;
101    ///
102    /// # fn example() -> slow5lib::Result<()> {
103    /// let reader = Slow5IndexedReader::open("reads.blow5")?;
104    /// assert!(!reader.read_ids().next().is_none());
105    /// # Ok(()) }
106    /// ```
107    pub fn is_empty(&self) -> bool {
108        self.entries.is_empty()
109    }
110
111    /// Build an index by scanning a BLOW5 file sequentially.
112    ///
113    /// `reader` must be positioned at the first record (after the file header).
114    /// `start_offset` is the corresponding byte offset in the file.
115    ///
116    /// For compressed records, only the first few bytes are decompressed (enough
117    /// to extract the read_id). The signal is never decoded during index building.
118    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(); // reused across records to avoid per-record allocation
125
126        loop {
127            let rec_offset = pos;
128
129            // Read first 5 bytes -- might be the start of the EOF marker
130            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            // Read the full compressed record (must advance the file position regardless).
157            buf.resize(record_size as usize, 0);
158            reader.read_exact(&mut buf)?;
159
160            // Extract read_id with early-stop decompression -- the signal is never decoded.
161            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    /// Build an index by scanning a SLOW5 text file sequentially.
176    ///
177    /// `reader` must be positioned at the first data line (after the file header).
178    /// `start_offset` is the corresponding byte offset in the file.
179    ///
180    /// Each data line's byte range is stored as `IndexEntry { offset: line_start, size: line_len }`.
181    /// The stored `size` includes the terminating `\n` exactly as `read_line` returns it.
182    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    /// Read a `.blow5.idx` binary index file.
224    ///
225    /// Fails if the version embedded in the index does not match `expected_version`.
226    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    /// Write this index to a `.blow5.idx` binary file.
294    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
318// ── Index build helpers ───────────────────────────────────────────────────────
319
320/// Extract the read_id from a (possibly compressed) record buffer.
321///
322/// For compressed formats, decompression stops as soon as the read_id bytes
323/// are available -- the signal is never decoded.
324fn 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
336/// Read the read_id directly from an uncompressed byte slice.
337fn 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
354/// Read the read_id from a streaming decompressor, stopping immediately after.
355///
356/// Dropping the reader here terminates decompression -- the signal bytes are
357/// never processed.
358fn 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}