Skip to main content

fastx/
bgzf.rs

1//! BGZF: the block-compressed gzip variant used across the samtools ecosystem.
2//!
3//! A BGZF file is an ordinary gzip file — any gzip tool can decompress it — made
4//! of independent members of at most 64 KiB each, every one carrying its own
5//! compressed size in a `BC` extra field. Because the blocks are independent, a
6//! reader that knows where they start can seek to any uncompressed position by
7//! jumping to the enclosing block and decompressing just that. This is what makes
8//! a 3 GB bgzipped reference genome randomly accessible.
9//!
10//! [`BgzfReader`] implements [`Read`] and [`Seek`] in *uncompressed* coordinates,
11//! so anything generic over those two traits — including [`crate::IndexedFasta`] —
12//! works on a bgzipped file without knowing it.
13//!
14//! ```no_run
15//! use fastx::bgzf::{BgzfReader, BgzfWriter};
16//! use std::io::{Read, Seek, SeekFrom, Write};
17//!
18//! // Write a seekable gzip file.
19//! let mut writer = BgzfWriter::create("ref.fa.gz")?;
20//! writer.write_all(b">chr1\nACGT\n")?;
21//! writer.finish()?;
22//!
23//! // Read 4 bytes starting at uncompressed offset 6.
24//! let mut reader = BgzfReader::open("ref.fa.gz")?;
25//! reader.seek(SeekFrom::Start(6))?;
26//! let mut buf = [0u8; 4];
27//! reader.read_exact(&mut buf)?;
28//! assert_eq!(&buf, b"ACGT");
29//! # Ok::<(), fastx::Error>(())
30//! ```
31
32use std::fs::File;
33use std::io::{self, BufReader, BufWriter, Read, Seek, SeekFrom, Write};
34use std::path::{Path, PathBuf};
35
36use crate::error::{Error, Result};
37use crate::format::CompressionLevel;
38
39/// Largest uncompressed payload placed in one block.
40///
41/// The spec caps a whole block at 64 KiB; htslib uses 0xff00 for the payload so
42/// that even incompressible data plus headers stays under the limit.
43pub const MAX_BLOCK_PAYLOAD: usize = 0xff00;
44
45/// Fixed part of a BGZF gzip header, up to and including `XLEN`.
46const HEADER_LEN: usize = 12;
47/// The `BC` extra subfield: `SI1`, `SI2`, `SLEN` and `BSIZE`.
48const EXTRA_LEN: usize = 6;
49/// The gzip trailer: CRC32 and ISIZE.
50const TRAILER_LEN: usize = 8;
51
52/// The 28-byte empty block that marks a complete BGZF file.
53///
54/// Its presence is how tools tell a truncated file from a finished one, so
55/// [`BgzfWriter::finish`] always appends it.
56pub const EOF_BLOCK: [u8; 28] = [
57    0x1f, 0x8b, 0x08, 0x04, 0x00, 0x00, 0x00, 0x00, 0x00, 0xff, 0x06, 0x00, 0x42, 0x43, 0x02, 0x00,
58    0x1b, 0x00, 0x03, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
59];
60
61/// Header of one BGZF block, as read from the file.
62#[derive(Debug, Clone, Copy, PartialEq, Eq)]
63struct BlockHeader {
64    /// Total bytes the block occupies in the file, `BSIZE + 1`.
65    compressed_len: usize,
66    /// Bytes of the block that precede the deflate payload.
67    payload_offset: usize,
68}
69
70/// Read and validate one block header from `bytes`.
71fn parse_block_header(bytes: &[u8]) -> Result<BlockHeader> {
72    if bytes.len() < HEADER_LEN {
73        return Err(bgzf_error("block header is truncated"));
74    }
75    if bytes[0] != 0x1f || bytes[1] != 0x8b {
76        return Err(bgzf_error("not a gzip member"));
77    }
78    if bytes[2] != 8 {
79        return Err(bgzf_error("unsupported compression method"));
80    }
81    if bytes[3] & 0x04 == 0 {
82        return Err(bgzf_error(
83            "gzip member has no extra field, so it is gzip but not BGZF",
84        ));
85    }
86    let extra_len = u16::from_le_bytes([bytes[10], bytes[11]]) as usize;
87    if bytes.len() < HEADER_LEN + extra_len {
88        return Err(bgzf_error("extra field is truncated"));
89    }
90    let extra = &bytes[HEADER_LEN..HEADER_LEN + extra_len];
91
92    // Walk the subfields looking for `BC`; other subfields are legal.
93    let mut cursor = 0;
94    while cursor + 4 <= extra.len() {
95        let si1 = extra[cursor];
96        let si2 = extra[cursor + 1];
97        let slen = u16::from_le_bytes([extra[cursor + 2], extra[cursor + 3]]) as usize;
98        let value = cursor + 4;
99        if value + slen > extra.len() {
100            return Err(bgzf_error("extra subfield runs past the extra field"));
101        }
102        if si1 == b'B' && si2 == b'C' {
103            if slen != 2 {
104                return Err(bgzf_error("BC subfield is not two bytes"));
105            }
106            let bsize = u16::from_le_bytes([extra[value], extra[value + 1]]) as usize;
107            let compressed_len = bsize + 1;
108            let overhead = HEADER_LEN + extra_len + TRAILER_LEN;
109            if compressed_len <= overhead {
110                return Err(bgzf_error("BSIZE is smaller than the block overhead"));
111            }
112            return Ok(BlockHeader {
113                compressed_len,
114                payload_offset: HEADER_LEN + extra_len,
115            });
116        }
117        cursor = value + slen;
118    }
119    Err(bgzf_error("gzip member has no BC extra subfield"))
120}
121
122fn bgzf_error(message: &'static str) -> Error {
123    Error::Other(format!("BGZF: {message}"))
124}
125
126/// Inflate a raw deflate stream of known uncompressed size.
127fn inflate(payload: &[u8], expected: usize, out: &mut Vec<u8>) -> Result<()> {
128    out.clear();
129    out.reserve(expected);
130    let mut decoder = flate2::Decompress::new(false);
131    decoder
132        .decompress_vec(payload, out, flate2::FlushDecompress::Finish)
133        .map_err(|e| Error::Other(format!("BGZF: corrupt block: {e}")))?;
134    if out.len() != expected {
135        return Err(bgzf_error("block size does not match its ISIZE field"));
136    }
137    Ok(())
138}
139
140/// The `.gzi` index that makes a BGZF file seekable by uncompressed offset.
141///
142/// The on-disk layout is the one `bgzip --index` writes: a little-endian `u64`
143/// count followed by that many `(compressed_offset, uncompressed_offset)` pairs.
144/// The first block is implicit — it always sits at `(0, 0)` — so a file of *n*
145/// blocks yields *n − 1* entries.
146#[derive(Debug, Clone, Default, PartialEq, Eq)]
147pub struct GziIndex {
148    /// Block starts, including the implicit `(0, 0)` first entry.
149    blocks: Vec<BlockOffset>,
150    /// Total uncompressed size, known only when the index came from scanning a
151    /// file or from the writer that produced it. A `.gzi` on disk records block
152    /// starts and nothing else, so a parsed index cannot know where the data
153    /// ends — and must not guess, or `SeekFrom::End` would land in the middle of
154    /// the last block.
155    total_uncompressed: Option<u64>,
156}
157
158/// Where one block begins, in both coordinate systems.
159#[derive(Debug, Clone, Copy, PartialEq, Eq)]
160pub struct BlockOffset {
161    /// Byte offset of the block in the compressed file.
162    pub compressed: u64,
163    /// Byte offset of the block's first byte in the decompressed stream.
164    pub uncompressed: u64,
165}
166
167impl GziIndex {
168    /// Build an index by walking every block header in the file.
169    ///
170    /// Only the headers are read, not the payloads, so this is fast: it touches
171    /// about 18 bytes per 64 KiB of input.
172    pub fn build<R: Read + Seek>(mut reader: R) -> Result<GziIndex> {
173        reader.seek(SeekFrom::Start(0))?;
174        let mut reader = BufReader::with_capacity(64 * 1024, reader);
175        let mut blocks = Vec::new();
176        let mut compressed = 0u64;
177        let mut uncompressed = 0u64;
178        let mut header = [0u8; HEADER_LEN + 64];
179
180        loop {
181            if !read_exact_or_eof(&mut reader, &mut header[..HEADER_LEN])? {
182                break;
183            }
184            let extra_len = u16::from_le_bytes([header[10], header[11]]) as usize;
185            if HEADER_LEN + extra_len > header.len() {
186                return Err(bgzf_error("extra field is implausibly large"));
187            }
188            reader.read_exact(&mut header[HEADER_LEN..HEADER_LEN + extra_len])?;
189            let block = parse_block_header(&header[..HEADER_LEN + extra_len])?;
190
191            // Skip the payload and CRC, then read ISIZE.
192            let skip = block.compressed_len - block.payload_offset - 4;
193            io::copy(&mut reader.by_ref().take(skip as u64), &mut io::sink())?;
194            let mut isize_bytes = [0u8; 4];
195            reader.read_exact(&mut isize_bytes)?;
196            let payload_len = u32::from_le_bytes(isize_bytes) as u64;
197
198            blocks.push(BlockOffset {
199                compressed,
200                uncompressed,
201            });
202            compressed += block.compressed_len as u64;
203            uncompressed += payload_len;
204
205            // A zero-length block is the EOF marker; nothing follows it.
206            if payload_len == 0 {
207                break;
208            }
209        }
210        Ok(GziIndex {
211            blocks,
212            total_uncompressed: Some(uncompressed),
213        })
214    }
215
216    /// Build an index for a file on disk.
217    pub fn build_from_path<P: AsRef<Path>>(path: P) -> Result<GziIndex> {
218        let path = path.as_ref();
219        GziIndex::build(
220            File::open(path).map_err(|e| {
221                Error::Io(io::Error::new(e.kind(), format!("{}: {e}", path.display())))
222            })?,
223        )
224    }
225
226    /// Parse a `.gzi` file.
227    pub fn parse<R: Read>(mut reader: R) -> Result<GziIndex> {
228        let mut count_bytes = [0u8; 8];
229        reader.read_exact(&mut count_bytes)?;
230        let count = u64::from_le_bytes(count_bytes);
231        // Guard against a corrupt count asking for a terabyte of allocation.
232        if count > 1 << 32 {
233            return Err(bgzf_error("index claims an implausible number of blocks"));
234        }
235        // The first block is implicit.
236        let mut blocks = Vec::with_capacity(count as usize + 1);
237        blocks.push(BlockOffset {
238            compressed: 0,
239            uncompressed: 0,
240        });
241        let mut pair = [0u8; 16];
242        for _ in 0..count {
243            reader.read_exact(&mut pair)?;
244            blocks.push(BlockOffset {
245                compressed: u64::from_le_bytes(pair[..8].try_into().expect("8 bytes")),
246                uncompressed: u64::from_le_bytes(pair[8..].try_into().expect("8 bytes")),
247            });
248        }
249        Ok(GziIndex {
250            blocks,
251            // A `.gzi` does not record the total size.
252            total_uncompressed: None,
253        })
254    }
255
256    /// Parse a `.gzi` file from disk.
257    pub fn from_path<P: AsRef<Path>>(path: P) -> Result<GziIndex> {
258        let path = path.as_ref();
259        GziIndex::parse(BufReader::new(File::open(path).map_err(|e| {
260            Error::Io(io::Error::new(e.kind(), format!("{}: {e}", path.display())))
261        })?))
262    }
263
264    /// Serialise in `bgzip --index` format, omitting the implicit first block.
265    pub fn write<W: Write>(&self, out: &mut W) -> Result<()> {
266        let count = self.blocks.len().saturating_sub(1) as u64;
267        out.write_all(&count.to_le_bytes())?;
268        for block in self.blocks.iter().skip(1) {
269            out.write_all(&block.compressed.to_le_bytes())?;
270            out.write_all(&block.uncompressed.to_le_bytes())?;
271        }
272        Ok(())
273    }
274
275    /// Write the index next to the BGZF file as `<path>.gzi`.
276    pub fn write_to_path<P: AsRef<Path>>(&self, bgzf_path: P) -> Result<PathBuf> {
277        let target = gzi_path(bgzf_path.as_ref());
278        let mut file = BufWriter::new(File::create(&target)?);
279        self.write(&mut file)?;
280        file.flush()?;
281        Ok(target)
282    }
283
284    /// Block starts, in file order, including the implicit first one.
285    pub fn blocks(&self) -> &[BlockOffset] {
286        &self.blocks
287    }
288
289    /// Number of blocks, including the EOF marker if the file has one.
290    pub fn len(&self) -> usize {
291        self.blocks.len()
292    }
293
294    /// True when the index describes no blocks at all.
295    pub fn is_empty(&self) -> bool {
296        self.blocks.is_empty()
297    }
298
299    /// Total uncompressed size.
300    ///
301    /// `None` for an index parsed from a `.gzi` file, which records only where
302    /// blocks start. Build the index with [`GziIndex::build`] if you need this —
303    /// it is a header-only scan, so it is cheap.
304    pub fn uncompressed_len(&self) -> Option<u64> {
305        self.total_uncompressed
306    }
307
308    /// The block containing uncompressed byte `offset`.
309    fn block_for(&self, offset: u64) -> Option<BlockOffset> {
310        // The last block whose uncompressed start is <= offset.
311        match self
312            .blocks
313            .binary_search_by(|b| b.uncompressed.cmp(&offset))
314        {
315            Ok(i) => Some(self.blocks[i]),
316            Err(0) => None,
317            Err(i) => Some(self.blocks[i - 1]),
318        }
319    }
320}
321
322/// The conventional index path for a BGZF file: `<path>.gzi`.
323pub fn gzi_path(bgzf: &Path) -> PathBuf {
324    let mut name = bgzf.as_os_str().to_os_string();
325    name.push(".gzi");
326    PathBuf::from(name)
327}
328
329/// True when the first bytes look like a BGZF block rather than plain gzip.
330///
331/// ```
332/// # use fastx::bgzf::{is_bgzf, EOF_BLOCK};
333/// assert!(is_bgzf(&EOF_BLOCK));
334/// assert!(!is_bgzf(&[0x1f, 0x8b, 0x08, 0x00]));  // gzip without FEXTRA
335/// assert!(!is_bgzf(b"ACGT"));
336/// ```
337pub fn is_bgzf(bytes: &[u8]) -> bool {
338    parse_block_header(bytes).is_ok()
339}
340
341/// A BGZF reader that seeks in uncompressed coordinates.
342///
343/// Sequential reads need no index. [`Seek`] does: build one with
344/// [`GziIndex::build`] (fast, header-only) or load a `.gzi` file.
345pub struct BgzfReader<R: Read + Seek> {
346    inner: R,
347    index: Option<GziIndex>,
348    /// Decompressed contents of the block currently in hand.
349    block: Vec<u8>,
350    /// Read cursor within `block`.
351    block_pos: usize,
352    /// Uncompressed offset at which `block` begins.
353    block_start: u64,
354    /// Compressed offset of the next block to read.
355    next_compressed: u64,
356    eof: bool,
357    /// Scratch buffer for one compressed block.
358    raw: Vec<u8>,
359}
360
361impl BgzfReader<File> {
362    /// Open a BGZF file, loading `<path>.gzi` if it is present.
363    ///
364    /// Without a `.gzi` the reader still works sequentially, and [`Seek`] will
365    /// build an index on first use.
366    pub fn open<P: AsRef<Path>>(path: P) -> Result<BgzfReader<File>> {
367        let path = path.as_ref();
368        let file = File::open(path)
369            .map_err(|e| Error::Io(io::Error::new(e.kind(), format!("{}: {e}", path.display()))))?;
370        let index = match GziIndex::from_path(gzi_path(path)) {
371            Ok(index) => Some(index),
372            Err(Error::Io(e)) if e.kind() == io::ErrorKind::NotFound => None,
373            Err(e) => return Err(e),
374        };
375        let mut reader = BgzfReader::new(file)?;
376        reader.index = index;
377        Ok(reader)
378    }
379}
380
381impl<R: Read + Seek> BgzfReader<R> {
382    /// Wrap a seekable reader, checking that it really is BGZF.
383    pub fn new(mut inner: R) -> Result<BgzfReader<R>> {
384        inner.seek(SeekFrom::Start(0))?;
385        let mut probe = [0u8; HEADER_LEN + 64];
386        let read = read_up_to(&mut inner, &mut probe)?;
387        if read == 0 {
388            return Err(bgzf_error("file is empty"));
389        }
390        parse_block_header(&probe[..read])?;
391        inner.seek(SeekFrom::Start(0))?;
392        Ok(BgzfReader {
393            inner,
394            index: None,
395            block: Vec::new(),
396            block_pos: 0,
397            block_start: 0,
398            next_compressed: 0,
399            eof: false,
400            raw: Vec::new(),
401        })
402    }
403
404    /// Attach an index, enabling [`Seek`].
405    pub fn with_index(mut self, index: GziIndex) -> Self {
406        self.index = Some(index);
407        self
408    }
409
410    /// The index in use, if any.
411    pub fn index(&self) -> Option<&GziIndex> {
412        self.index.as_ref()
413    }
414
415    /// Current uncompressed position.
416    pub fn position(&self) -> u64 {
417        self.block_start + self.block_pos as u64
418    }
419
420    /// Unwrap the underlying reader.
421    pub fn into_inner(self) -> R {
422        self.inner
423    }
424
425    /// Decompress the block at compressed offset `at`, which becomes current.
426    fn load_block_at(&mut self, at: u64, uncompressed_start: u64) -> Result<()> {
427        self.inner.seek(SeekFrom::Start(at))?;
428        self.next_compressed = at;
429        self.block_start = uncompressed_start;
430        self.block_pos = 0;
431        self.block.clear();
432        self.eof = false;
433        self.read_next_block()
434    }
435
436    /// Read and decompress the block at `self.next_compressed`.
437    fn read_next_block(&mut self) -> Result<()> {
438        let mut header = [0u8; HEADER_LEN + 64];
439        if !read_exact_or_eof(&mut self.inner, &mut header[..HEADER_LEN])? {
440            self.eof = true;
441            self.block.clear();
442            self.block_pos = 0;
443            return Ok(());
444        }
445        let extra_len = u16::from_le_bytes([header[10], header[11]]) as usize;
446        if HEADER_LEN + extra_len > header.len() {
447            return Err(bgzf_error("extra field is implausibly large"));
448        }
449        self.inner
450            .read_exact(&mut header[HEADER_LEN..HEADER_LEN + extra_len])?;
451        let block = parse_block_header(&header[..HEADER_LEN + extra_len])?;
452
453        let payload_len = block.compressed_len - block.payload_offset - TRAILER_LEN;
454        self.raw.resize(payload_len, 0);
455        self.inner.read_exact(&mut self.raw)?;
456        let mut trailer = [0u8; TRAILER_LEN];
457        self.inner.read_exact(&mut trailer)?;
458        let expected_crc = u32::from_le_bytes(trailer[..4].try_into().expect("4 bytes"));
459        let expected_len = u32::from_le_bytes(trailer[4..].try_into().expect("4 bytes")) as usize;
460
461        let mut decompressed = std::mem::take(&mut self.block);
462        let result = inflate(&self.raw, expected_len, &mut decompressed);
463        self.block = decompressed;
464        result?;
465
466        let mut crc = flate2::Crc::new();
467        crc.update(&self.block);
468        if crc.sum() != expected_crc {
469            return Err(bgzf_error("block CRC32 does not match"));
470        }
471
472        self.block_pos = 0;
473        self.next_compressed += block.compressed_len as u64;
474        // The EOF marker decompresses to nothing; treat it as end of stream.
475        if self.block.is_empty() {
476            self.eof = true;
477        }
478        Ok(())
479    }
480
481    /// Make sure the current block has unread bytes, or set `eof`.
482    fn fill(&mut self) -> Result<()> {
483        while !self.eof && self.block_pos == self.block.len() {
484            let consumed = self.block.len() as u64;
485            self.block_start += consumed;
486            self.read_next_block()?;
487        }
488        Ok(())
489    }
490
491    /// The index, built on demand if it was not supplied.
492    fn ensure_index(&mut self) -> Result<()> {
493        if self.index.is_none() {
494            let saved = self.inner.stream_position()?;
495            let index = GziIndex::build(&mut self.inner)?;
496            self.inner.seek(SeekFrom::Start(saved))?;
497            self.index = Some(index);
498        }
499        Ok(())
500    }
501}
502
503impl<R: Read + Seek> Read for BgzfReader<R> {
504    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
505        if buf.is_empty() {
506            return Ok(0);
507        }
508        self.fill()?;
509        if self.eof && self.block_pos == self.block.len() {
510            return Ok(0);
511        }
512        let available = &self.block[self.block_pos..];
513        let take = available.len().min(buf.len());
514        buf[..take].copy_from_slice(&available[..take]);
515        self.block_pos += take;
516        Ok(take)
517    }
518}
519
520impl<R: Read + Seek> Seek for BgzfReader<R> {
521    /// Seek in *uncompressed* coordinates.
522    ///
523    /// `SeekFrom::End` needs the total uncompressed size, so it requires an
524    /// index that covers the whole file.
525    fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
526        self.ensure_index()?;
527
528        let target = match pos {
529            SeekFrom::Start(offset) => offset,
530            SeekFrom::Current(delta) => add_signed(self.position(), delta)?,
531            SeekFrom::End(delta) => {
532                let end = self
533                    .index
534                    .as_ref()
535                    .and_then(|index| index.uncompressed_len())
536                    .ok_or_else(|| {
537                        io::Error::new(
538                            io::ErrorKind::InvalidInput,
539                            "BGZF: this index records block starts only, so the end of the \
540                             stream is unknown; rebuild it with GziIndex::build",
541                        )
542                    })?;
543                add_signed(end, delta)?
544            }
545        };
546
547        // Staying inside the current block is the common case for short hops.
548        let within = target.checked_sub(self.block_start);
549        if let Some(within) = within {
550            if !self.block.is_empty() && within <= self.block.len() as u64 {
551                self.block_pos = within as usize;
552                return Ok(target);
553            }
554        }
555
556        // `BlockOffset` is `Copy`, so the borrow of the index ends here and the
557        // load below needs no clone of it.
558        let block = self
559            .index
560            .as_ref()
561            .and_then(|index| index.block_for(target))
562            .ok_or_else(|| {
563                io::Error::new(
564                    io::ErrorKind::InvalidInput,
565                    "BGZF: no block covers that offset",
566                )
567            })?;
568        self.load_block_at(block.compressed, block.uncompressed)?;
569        let within = (target - block.uncompressed) as usize;
570        if within > self.block.len() {
571            // Past the end of the data: leave the cursor at the end.
572            self.block_pos = self.block.len();
573            self.eof = true;
574        } else {
575            self.block_pos = within;
576        }
577        Ok(target)
578    }
579}
580
581fn add_signed(base: u64, delta: i64) -> io::Result<u64> {
582    let result = if delta >= 0 {
583        base.checked_add(delta as u64)
584    } else {
585        base.checked_sub(delta.unsigned_abs())
586    };
587    result.ok_or_else(|| {
588        io::Error::new(
589            io::ErrorKind::InvalidInput,
590            "BGZF: seek would leave the file",
591        )
592    })
593}
594
595/// A writer that emits BGZF blocks.
596///
597/// Output is valid gzip, so any gzip tool can read it, and it is seekable by
598/// anything that understands BGZF. Call [`BgzfWriter::finish`] to append the
599/// EOF marker; without it the file looks truncated to samtools.
600pub struct BgzfWriter<W: Write> {
601    /// `None` once `finish` has handed the writer back. Wrapped in an `Option`
602    /// only because a type with a `Drop` impl cannot give a field away.
603    inner: Option<W>,
604    buffer: Vec<u8>,
605    level: CompressionLevel,
606    /// Block starts, recorded so that an index can be written afterwards.
607    blocks: Vec<BlockOffset>,
608    compressed: u64,
609    uncompressed: u64,
610    finished: bool,
611}
612
613impl BgzfWriter<BufWriter<File>> {
614    /// Create a BGZF file.
615    pub fn create<P: AsRef<Path>>(path: P) -> Result<BgzfWriter<BufWriter<File>>> {
616        let path = path.as_ref();
617        let file = File::create(path)
618            .map_err(|e| Error::Io(io::Error::new(e.kind(), format!("{}: {e}", path.display()))))?;
619        Ok(BgzfWriter::new(BufWriter::with_capacity(128 * 1024, file)))
620    }
621}
622
623impl<W: Write> BgzfWriter<W> {
624    /// Wrap a writer, compressing at the default level.
625    pub fn new(inner: W) -> BgzfWriter<W> {
626        BgzfWriter::with_level(inner, CompressionLevel::default())
627    }
628
629    /// Wrap a writer, compressing at `level`.
630    pub fn with_level(inner: W, level: CompressionLevel) -> BgzfWriter<W> {
631        BgzfWriter {
632            inner: Some(inner),
633            buffer: Vec::with_capacity(MAX_BLOCK_PAYLOAD),
634            level,
635            blocks: vec![BlockOffset {
636                compressed: 0,
637                uncompressed: 0,
638            }],
639            compressed: 0,
640            uncompressed: 0,
641            finished: false,
642        }
643    }
644
645    /// The index describing the blocks written *so far*.
646    ///
647    /// Buffered data that has not been compressed yet is not in it, so calling
648    /// this before [`BgzfWriter::finish_with_index`] gives an index that stops
649    /// short of the end of the file. Prefer `finish_with_index`, which cannot be
650    /// wrong.
651    pub fn index(&self) -> GziIndex {
652        GziIndex {
653            blocks: self.blocks.clone(),
654            total_uncompressed: Some(self.uncompressed),
655        }
656    }
657
658    /// The underlying writer, or an error once `finish` has taken it.
659    fn sink(&mut self) -> io::Result<&mut W> {
660        self.inner.as_mut().ok_or_else(|| {
661            io::Error::new(
662                io::ErrorKind::BrokenPipe,
663                "BGZF: writer was already finished",
664            )
665        })
666    }
667
668    /// Compress and emit whatever is buffered as one block.
669    fn flush_block(&mut self) -> io::Result<()> {
670        if self.buffer.is_empty() {
671            return Ok(());
672        }
673        let payload = deflate_raw(&self.buffer, self.level)?;
674        let block_len = HEADER_LEN + EXTRA_LEN + payload.len() + TRAILER_LEN;
675        if block_len > u16::MAX as usize + 1 {
676            // Cannot happen with MAX_BLOCK_PAYLOAD, but a wrong constant here
677            // would produce files other tools silently misread.
678            return Err(io::Error::new(
679                io::ErrorKind::InvalidData,
680                "BGZF: block would exceed 64 KiB",
681            ));
682        }
683
684        let mut header = [0u8; HEADER_LEN + EXTRA_LEN];
685        header[0] = 0x1f;
686        header[1] = 0x8b;
687        header[2] = 8; // deflate
688        header[3] = 4; // FEXTRA
689                       // MTIME stays zero: reproducible output matters more than a timestamp.
690        header[9] = 0xff; // unknown OS
691        header[10..12].copy_from_slice(&(EXTRA_LEN as u16).to_le_bytes());
692        header[12] = b'B';
693        header[13] = b'C';
694        header[14..16].copy_from_slice(&2u16.to_le_bytes());
695        header[16..18].copy_from_slice(&((block_len - 1) as u16).to_le_bytes());
696
697        let mut crc = flate2::Crc::new();
698        crc.update(&self.buffer);
699        let checksum = crc.sum().to_le_bytes();
700        let payload_len = (self.buffer.len() as u32).to_le_bytes();
701
702        let sink = self.sink()?;
703        sink.write_all(&header)?;
704        sink.write_all(&payload)?;
705        sink.write_all(&checksum)?;
706        sink.write_all(&payload_len)?;
707
708        self.compressed += block_len as u64;
709        self.uncompressed += self.buffer.len() as u64;
710        self.blocks.push(BlockOffset {
711            compressed: self.compressed,
712            uncompressed: self.uncompressed,
713        });
714        self.buffer.clear();
715        Ok(())
716    }
717
718    /// Flush the pending block, append the EOF marker and hand the writer back.
719    ///
720    /// Dropping the writer does the same on a best-effort basis, but only
721    /// `finish` reports a failure.
722    pub fn finish(self) -> Result<W> {
723        self.finish_with_index().map(|(inner, _)| inner)
724    }
725
726    /// Finish, and return the complete index alongside the writer.
727    ///
728    /// This is the safe way to obtain a `.gzi`: every block has been emitted by
729    /// the time the index is taken, so it cannot be missing the tail.
730    ///
731    /// ```no_run
732    /// use fastx::bgzf::BgzfWriter;
733    /// use std::io::Write;
734    ///
735    /// let mut writer = BgzfWriter::create("reads.fq.gz")?;
736    /// writer.write_all(b"@r\nACGT\n+\nIIII\n")?;
737    /// let (_file, index) = writer.finish_with_index()?;
738    /// index.write_to_path("reads.fq.gz")?;   // reads.fq.gz.gzi
739    /// # Ok::<(), fastx::Error>(())
740    /// ```
741    pub fn finish_with_index(mut self) -> Result<(W, GziIndex)> {
742        self.finish_in_place()?;
743        let index = self.index();
744        let inner = self
745            .inner
746            .take()
747            .ok_or_else(|| Error::Other("BGZF: writer was already finished".to_string()))?;
748        Ok((inner, index))
749    }
750
751    fn finish_in_place(&mut self) -> Result<()> {
752        if self.finished || self.inner.is_none() {
753            return Ok(());
754        }
755        self.flush_block()?;
756        let sink = self.sink()?;
757        sink.write_all(&EOF_BLOCK)?;
758        sink.flush()?;
759        self.compressed += EOF_BLOCK.len() as u64;
760        self.finished = true;
761        Ok(())
762    }
763}
764
765impl<W: Write> Write for BgzfWriter<W> {
766    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
767        let room = MAX_BLOCK_PAYLOAD - self.buffer.len();
768        let take = room.min(buf.len());
769        self.buffer.extend_from_slice(&buf[..take]);
770        if self.buffer.len() == MAX_BLOCK_PAYLOAD {
771            self.flush_block()?;
772        }
773        Ok(take)
774    }
775
776    /// Ends the current block, so the next byte starts a fresh one.
777    ///
778    /// This is what makes a flush point seekable, and it costs a little
779    /// compression, so flush at record boundaries rather than per record.
780    fn flush(&mut self) -> io::Result<()> {
781        self.flush_block()?;
782        self.sink()?.flush()
783    }
784}
785
786impl<W: Write> Drop for BgzfWriter<W> {
787    fn drop(&mut self) {
788        // Best effort: a caller who wants to see errors uses finish().
789        let _ = self.finish_in_place();
790    }
791}
792
793/// Deflate with no zlib or gzip wrapper, which is what a gzip member holds.
794fn deflate_raw(data: &[u8], level: CompressionLevel) -> io::Result<Vec<u8>> {
795    use flate2::write::DeflateEncoder;
796    let mut encoder = DeflateEncoder::new(
797        Vec::with_capacity(data.len() / 2 + 64),
798        flate2::Compression::new(level.0.min(9)),
799    );
800    encoder.write_all(data)?;
801    encoder.finish()
802}
803
804/// Read into `buf` fully, or report `false` if the reader was already at EOF.
805fn read_exact_or_eof<R: Read>(reader: &mut R, buf: &mut [u8]) -> Result<bool> {
806    let mut filled = 0;
807    while filled < buf.len() {
808        match reader.read(&mut buf[filled..]) {
809            Ok(0) if filled == 0 => return Ok(false),
810            Ok(0) => return Err(bgzf_error("file ends in the middle of a block")),
811            Ok(n) => filled += n,
812            Err(e) if e.kind() == io::ErrorKind::Interrupted => continue,
813            Err(e) => return Err(Error::Io(e)),
814        }
815    }
816    Ok(true)
817}
818
819/// Read as much as is available, up to `buf.len()`.
820fn read_up_to<R: Read>(reader: &mut R, buf: &mut [u8]) -> Result<usize> {
821    let mut filled = 0;
822    while filled < buf.len() {
823        match reader.read(&mut buf[filled..]) {
824            Ok(0) => break,
825            Ok(n) => filled += n,
826            Err(e) if e.kind() == io::ErrorKind::Interrupted => continue,
827            Err(e) => return Err(Error::Io(e)),
828        }
829    }
830    Ok(filled)
831}
832
833#[cfg(test)]
834mod tests {
835    use super::*;
836    use std::io::Cursor;
837
838    fn compress(data: &[u8]) -> Vec<u8> {
839        let mut writer = BgzfWriter::new(Vec::new());
840        writer.write_all(data).unwrap();
841        writer.finish().unwrap()
842    }
843
844    #[test]
845    fn round_trips_through_our_own_reader() {
846        for size in [
847            0usize,
848            1,
849            100,
850            MAX_BLOCK_PAYLOAD - 1,
851            MAX_BLOCK_PAYLOAD,
852            MAX_BLOCK_PAYLOAD + 1,
853            300_000,
854        ] {
855            let data: Vec<u8> = (0..size).map(|i| b"ACGTN"[i % 5]).collect();
856            let compressed = compress(&data);
857            assert!(is_bgzf(&compressed), "size {size} did not produce BGZF");
858
859            let mut reader = BgzfReader::new(Cursor::new(&compressed)).unwrap();
860            let mut out = Vec::new();
861            reader.read_to_end(&mut out).unwrap();
862            assert_eq!(out, data, "size {size}");
863        }
864    }
865
866    #[test]
867    fn output_is_plain_gzip_too() {
868        // The whole point of BGZF: ordinary gzip tools must still read it.
869        let data: Vec<u8> = (0..200_000).map(|i| b"ACGT"[i % 4]).collect();
870        let compressed = compress(&data);
871        let mut out = Vec::new();
872        flate2::read::MultiGzDecoder::new(&compressed[..])
873            .read_to_end(&mut out)
874            .unwrap();
875        assert_eq!(out, data);
876    }
877
878    #[test]
879    fn ends_with_the_eof_marker() {
880        let compressed = compress(b"ACGT");
881        assert_eq!(
882            &compressed[compressed.len() - EOF_BLOCK.len()..],
883            &EOF_BLOCK
884        );
885        // An empty file is just the marker.
886        assert_eq!(compress(b""), EOF_BLOCK.to_vec());
887    }
888
889    #[test]
890    fn blocks_stay_within_the_size_limit() {
891        // Incompressible data is the case that could overflow a block.
892        let data: Vec<u8> = (0..500_000)
893            .map(|i| ((i * 2_654_435_761u64 as usize) >> 7) as u8)
894            .collect();
895        let compressed = compress(&data);
896        let index = GziIndex::build(Cursor::new(&compressed)).unwrap();
897        for pair in index.blocks().windows(2) {
898            let block_len = pair[1].compressed - pair[0].compressed;
899            assert!(block_len <= 65_536, "block of {block_len} bytes");
900        }
901        let mut out = Vec::new();
902        BgzfReader::new(Cursor::new(&compressed))
903            .unwrap()
904            .read_to_end(&mut out)
905            .unwrap();
906        assert_eq!(out, data);
907    }
908
909    #[test]
910    fn index_from_the_writer_matches_a_rescan() {
911        let data: Vec<u8> = (0..250_000).map(|i| b"ACGTN"[i % 5]).collect();
912        let mut writer = BgzfWriter::new(Vec::new());
913        writer.write_all(&data).unwrap();
914        let from_writer = writer.index();
915        let compressed = writer.finish().unwrap();
916
917        let from_scan = GziIndex::build(Cursor::new(&compressed)).unwrap();
918        // The rescan also sees the EOF block, which the writer's index predates.
919        assert_eq!(
920            &from_scan.blocks()[..from_writer.len()],
921            from_writer.blocks()
922        );
923        assert_eq!(from_scan.uncompressed_len(), Some(data.len() as u64));
924    }
925
926    #[test]
927    fn gzi_round_trips_and_omits_the_first_block() {
928        let data: Vec<u8> = (0..200_000).map(|i| b"ACGT"[i % 4]).collect();
929        let compressed = compress(&data);
930        let index = GziIndex::build(Cursor::new(&compressed)).unwrap();
931
932        let mut text = Vec::new();
933        index.write(&mut text).unwrap();
934        // 8-byte count plus 16 bytes per entry, first block implicit.
935        assert_eq!(text.len(), 8 + 16 * (index.len() - 1));
936        assert_eq!(
937            u64::from_le_bytes(text[..8].try_into().unwrap()),
938            index.len() as u64 - 1
939        );
940
941        let reparsed = GziIndex::parse(&text[..]).unwrap();
942        assert_eq!(reparsed.blocks(), index.blocks());
943        // A parsed index cannot know where the data ends, and must not pretend.
944        assert_eq!(reparsed.uncompressed_len(), None);
945        assert_eq!(index.uncompressed_len(), Some(200_000));
946    }
947
948    #[test]
949    fn seek_from_end_needs_a_scanned_index() {
950        let data: Vec<u8> = (0..200_000).map(|i| b"ACGT"[i % 4]).collect();
951        let compressed = compress(&data);
952        let scanned = GziIndex::build(Cursor::new(&compressed)).unwrap();
953        let mut text = Vec::new();
954        scanned.write(&mut text).unwrap();
955        let parsed = GziIndex::parse(&text[..]).unwrap();
956
957        // With block starts only, End-relative seeks must fail loudly rather
958        // than landing at the start of the last block.
959        let mut reader = BgzfReader::new(Cursor::new(&compressed))
960            .unwrap()
961            .with_index(parsed);
962        assert!(reader.seek(SeekFrom::End(0)).is_err());
963        // Absolute seeks still work.
964        reader.seek(SeekFrom::Start(199_998)).unwrap();
965        let mut tail = Vec::new();
966        reader.read_to_end(&mut tail).unwrap();
967        assert_eq!(tail, &data[199_998..]);
968
969        let mut reader = BgzfReader::new(Cursor::new(&compressed))
970            .unwrap()
971            .with_index(scanned);
972        assert_eq!(reader.seek(SeekFrom::End(0)).unwrap(), 200_000);
973    }
974
975    #[test]
976    fn seeks_to_any_offset() {
977        let data: Vec<u8> = (0..300_000).map(|i| (i % 251) as u8).collect();
978        let compressed = compress(&data);
979        let mut reader = BgzfReader::new(Cursor::new(&compressed)).unwrap();
980
981        // Offsets that land in the first block, deep inside, and on boundaries.
982        for target in [
983            0usize,
984            1,
985            MAX_BLOCK_PAYLOAD - 1,
986            MAX_BLOCK_PAYLOAD,
987            MAX_BLOCK_PAYLOAD + 1,
988            2 * MAX_BLOCK_PAYLOAD,
989            299_999,
990        ] {
991            reader.seek(SeekFrom::Start(target as u64)).unwrap();
992            assert_eq!(reader.position(), target as u64);
993            let mut buf = [0u8; 8];
994            let want = (data.len() - target).min(buf.len());
995            reader.read_exact(&mut buf[..want]).unwrap();
996            assert_eq!(&buf[..want], &data[target..target + want], "at {target}");
997        }
998
999        // Relative and end-relative seeks.
1000        reader.seek(SeekFrom::Start(10)).unwrap();
1001        reader.seek(SeekFrom::Current(5)).unwrap();
1002        assert_eq!(reader.position(), 15);
1003        assert_eq!(reader.seek(SeekFrom::End(0)).unwrap(), data.len() as u64);
1004        let mut rest = Vec::new();
1005        reader.read_to_end(&mut rest).unwrap();
1006        assert!(rest.is_empty());
1007
1008        // Seeking backwards must work as well as forwards.
1009        reader.seek(SeekFrom::Start(7)).unwrap();
1010        let mut buf = [0u8; 4];
1011        reader.read_exact(&mut buf).unwrap();
1012        assert_eq!(&buf, &data[7..11]);
1013    }
1014
1015    #[test]
1016    fn seek_before_the_start_is_an_error() {
1017        let compressed = compress(b"ACGT");
1018        let mut reader = BgzfReader::new(Cursor::new(&compressed)).unwrap();
1019        assert!(reader.seek(SeekFrom::Current(-1)).is_err());
1020        assert!(reader.seek(SeekFrom::End(-100)).is_err());
1021    }
1022
1023    #[test]
1024    fn rejects_plain_gzip_and_garbage() {
1025        // Valid gzip, but no BC extra field: not BGZF.
1026        let mut plain = Vec::new();
1027        {
1028            let mut encoder =
1029                flate2::write::GzEncoder::new(&mut plain, flate2::Compression::default());
1030            encoder.write_all(b"ACGT").unwrap();
1031            encoder.finish().unwrap();
1032        }
1033        assert!(!is_bgzf(&plain));
1034        assert!(BgzfReader::new(Cursor::new(&plain)).is_err());
1035
1036        assert!(BgzfReader::new(Cursor::new(b"not gzip at all".to_vec())).is_err());
1037        assert!(BgzfReader::new(Cursor::new(Vec::new())).is_err());
1038    }
1039
1040    #[test]
1041    fn detects_a_corrupt_block() {
1042        let mut compressed = compress(&vec![b'A'; 5_000]);
1043        // Flip a byte in the deflate payload; the CRC or the inflate must catch it.
1044        let victim = HEADER_LEN + EXTRA_LEN + 5;
1045        compressed[victim] ^= 0xff;
1046        let mut out = Vec::new();
1047        let result = BgzfReader::new(Cursor::new(&compressed))
1048            .unwrap()
1049            .read_to_end(&mut out);
1050        assert!(result.is_err(), "corruption went unnoticed");
1051    }
1052
1053    #[test]
1054    fn truncated_file_is_an_error_not_silent_truncation() {
1055        let compressed = compress(&vec![b'A'; 200_000]);
1056        let cut = compressed.len() / 2;
1057        let mut out = Vec::new();
1058        let result = BgzfReader::new(Cursor::new(compressed[..cut].to_vec()))
1059            .unwrap()
1060            .read_to_end(&mut out);
1061        assert!(result.is_err(), "truncation went unnoticed");
1062    }
1063
1064    #[test]
1065    fn parse_rejects_an_implausible_index() {
1066        let mut bad = u64::MAX.to_le_bytes().to_vec();
1067        bad.extend_from_slice(&[0u8; 16]);
1068        assert!(GziIndex::parse(&bad[..]).is_err());
1069        assert!(GziIndex::parse(&[0u8; 3][..]).is_err());
1070    }
1071
1072    #[test]
1073    fn flush_starts_a_new_block() {
1074        let mut writer = BgzfWriter::new(Vec::new());
1075        writer.write_all(b"first").unwrap();
1076        writer.flush().unwrap();
1077        writer.write_all(b"second").unwrap();
1078        // index() before finishing sees only the flushed block, which is exactly
1079        // why finish_with_index exists.
1080        assert_eq!(writer.index().len(), 2);
1081        let (compressed, index) = writer.finish_with_index().unwrap();
1082
1083        // Two data blocks, each a seek point.
1084        assert_eq!(index.len(), 3); // implicit start + two blocks
1085        assert_eq!(index.blocks()[1].uncompressed, 5);
1086        assert_eq!(index.uncompressed_len(), Some(11));
1087
1088        let mut reader = BgzfReader::new(Cursor::new(&compressed)).unwrap();
1089        reader.seek(SeekFrom::Start(5)).unwrap();
1090        let mut out = Vec::new();
1091        reader.read_to_end(&mut out).unwrap();
1092        assert_eq!(out, b"second");
1093    }
1094}