1use 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
39pub const MAX_BLOCK_PAYLOAD: usize = 0xff00;
44
45const HEADER_LEN: usize = 12;
47const EXTRA_LEN: usize = 6;
49const TRAILER_LEN: usize = 8;
51
52pub 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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
63struct BlockHeader {
64 compressed_len: usize,
66 payload_offset: usize,
68}
69
70fn 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 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
126struct Deflater {
132 #[cfg(not(feature = "libdeflate"))]
135 level: CompressionLevel,
136 #[cfg(feature = "libdeflate")]
137 compressor: libdeflater::Compressor,
138}
139
140impl Deflater {
141 fn new(level: CompressionLevel) -> Deflater {
142 Deflater {
143 #[cfg(not(feature = "libdeflate"))]
144 level,
145 #[cfg(feature = "libdeflate")]
146 compressor: libdeflater::Compressor::new(libdeflate_level(level)),
147 }
148 }
149
150 #[cfg(feature = "libdeflate")]
152 fn deflate(&mut self, data: &[u8]) -> io::Result<Vec<u8>> {
153 let mut out = vec![0u8; self.compressor.deflate_compress_bound(data.len())];
154 let written = self
155 .compressor
156 .deflate_compress(data, &mut out)
157 .map_err(|e| io::Error::other(format!("libdeflate: {e}")))?;
158 out.truncate(written);
159 Ok(out)
160 }
161
162 #[cfg(not(feature = "libdeflate"))]
163 fn deflate(&mut self, data: &[u8]) -> io::Result<Vec<u8>> {
164 use flate2::write::DeflateEncoder;
165 let mut encoder = DeflateEncoder::new(
166 Vec::with_capacity(data.len() / 2 + 64),
167 flate2::Compression::new(self.level.0.min(9)),
168 );
169 encoder.write_all(data)?;
170 encoder.finish()
171 }
172}
173
174#[cfg(feature = "libdeflate")]
176fn libdeflate_level(level: CompressionLevel) -> libdeflater::CompressionLvl {
177 let mapped = match level.0 {
178 0 => 0,
179 level => ((level.min(9) as i32 - 1) * 11 / 8 + 1).min(12),
180 };
181 libdeflater::CompressionLvl::new(mapped)
182 .unwrap_or_else(|_| libdeflater::CompressionLvl::default())
183}
184
185struct Inflater {
187 #[cfg(feature = "libdeflate")]
188 decompressor: libdeflater::Decompressor,
189}
190
191impl Inflater {
192 fn new() -> Inflater {
193 Inflater {
194 #[cfg(feature = "libdeflate")]
195 decompressor: libdeflater::Decompressor::new(),
196 }
197 }
198
199 #[cfg(feature = "libdeflate")]
201 fn inflate(&mut self, payload: &[u8], expected: usize, out: &mut Vec<u8>) -> Result<()> {
202 out.clear();
203 if expected == 0 {
204 return Ok(());
207 }
208 out.resize(expected, 0);
209 let written = self
210 .decompressor
211 .deflate_decompress(payload, out)
212 .map_err(|e| Error::Other(format!("BGZF: corrupt block: {e}")))?;
213 if written != expected {
214 return Err(bgzf_error("block size does not match its ISIZE field"));
215 }
216 Ok(())
217 }
218
219 #[cfg(not(feature = "libdeflate"))]
220 fn inflate(&mut self, payload: &[u8], expected: usize, out: &mut Vec<u8>) -> Result<()> {
221 out.clear();
222 out.reserve(expected);
223 let mut decoder = flate2::Decompress::new(false);
224 decoder
225 .decompress_vec(payload, out, flate2::FlushDecompress::Finish)
226 .map_err(|e| Error::Other(format!("BGZF: corrupt block: {e}")))?;
227 if out.len() != expected {
228 return Err(bgzf_error("block size does not match its ISIZE field"));
229 }
230 Ok(())
231 }
232}
233
234fn inflate(payload: &[u8], expected: usize, out: &mut Vec<u8>) -> Result<()> {
236 Inflater::new().inflate(payload, expected, out)
237}
238
239fn crc32(data: &[u8]) -> u32 {
244 #[cfg(feature = "libdeflate")]
245 {
246 let mut crc = libdeflater::Crc::new();
247 crc.update(data);
248 crc.sum()
249 }
250 #[cfg(not(feature = "libdeflate"))]
251 {
252 let mut crc = flate2::Crc::new();
253 crc.update(data);
254 crc.sum()
255 }
256}
257
258#[derive(Debug, Clone, Default, PartialEq, Eq)]
265pub struct GziIndex {
266 blocks: Vec<BlockOffset>,
268 total_uncompressed: Option<u64>,
274}
275
276#[derive(Debug, Clone, Copy, PartialEq, Eq)]
278pub struct BlockOffset {
279 pub compressed: u64,
281 pub uncompressed: u64,
283}
284
285impl GziIndex {
286 pub fn build<R: Read + Seek>(mut reader: R) -> Result<GziIndex> {
291 reader.seek(SeekFrom::Start(0))?;
292 let mut reader = BufReader::with_capacity(64 * 1024, reader);
293 let mut blocks = Vec::new();
294 let mut compressed = 0u64;
295 let mut uncompressed = 0u64;
296 let mut header = [0u8; HEADER_LEN + 64];
297
298 loop {
299 if !read_exact_or_eof(&mut reader, &mut header[..HEADER_LEN])? {
300 break;
301 }
302 let extra_len = u16::from_le_bytes([header[10], header[11]]) as usize;
303 if HEADER_LEN + extra_len > header.len() {
304 return Err(bgzf_error("extra field is implausibly large"));
305 }
306 reader.read_exact(&mut header[HEADER_LEN..HEADER_LEN + extra_len])?;
307 let block = parse_block_header(&header[..HEADER_LEN + extra_len])?;
308
309 let skip = block.compressed_len - block.payload_offset - 4;
311 io::copy(&mut reader.by_ref().take(skip as u64), &mut io::sink())?;
312 let mut isize_bytes = [0u8; 4];
313 reader.read_exact(&mut isize_bytes)?;
314 let payload_len = u32::from_le_bytes(isize_bytes) as u64;
315
316 blocks.push(BlockOffset {
317 compressed,
318 uncompressed,
319 });
320 compressed += block.compressed_len as u64;
321 uncompressed += payload_len;
322
323 if payload_len == 0 {
325 break;
326 }
327 }
328 Ok(GziIndex {
329 blocks,
330 total_uncompressed: Some(uncompressed),
331 })
332 }
333
334 pub fn build_from_path<P: AsRef<Path>>(path: P) -> Result<GziIndex> {
336 let path = path.as_ref();
337 GziIndex::build(
338 File::open(path).map_err(|e| {
339 Error::Io(io::Error::new(e.kind(), format!("{}: {e}", path.display())))
340 })?,
341 )
342 }
343
344 pub fn parse<R: Read>(mut reader: R) -> Result<GziIndex> {
346 let mut count_bytes = [0u8; 8];
347 reader.read_exact(&mut count_bytes)?;
348 let count = u64::from_le_bytes(count_bytes);
349 if count > 1 << 32 {
351 return Err(bgzf_error("index claims an implausible number of blocks"));
352 }
353 let mut blocks = Vec::with_capacity(count as usize + 1);
355 blocks.push(BlockOffset {
356 compressed: 0,
357 uncompressed: 0,
358 });
359 let mut pair = [0u8; 16];
360 for _ in 0..count {
361 reader.read_exact(&mut pair)?;
362 blocks.push(BlockOffset {
363 compressed: u64::from_le_bytes(pair[..8].try_into().expect("8 bytes")),
364 uncompressed: u64::from_le_bytes(pair[8..].try_into().expect("8 bytes")),
365 });
366 }
367 Ok(GziIndex {
368 blocks,
369 total_uncompressed: None,
371 })
372 }
373
374 pub fn from_path<P: AsRef<Path>>(path: P) -> Result<GziIndex> {
376 let path = path.as_ref();
377 GziIndex::parse(BufReader::new(File::open(path).map_err(|e| {
378 Error::Io(io::Error::new(e.kind(), format!("{}: {e}", path.display())))
379 })?))
380 }
381
382 pub fn write<W: Write>(&self, out: &mut W) -> Result<()> {
384 let count = self.blocks.len().saturating_sub(1) as u64;
385 out.write_all(&count.to_le_bytes())?;
386 for block in self.blocks.iter().skip(1) {
387 out.write_all(&block.compressed.to_le_bytes())?;
388 out.write_all(&block.uncompressed.to_le_bytes())?;
389 }
390 Ok(())
391 }
392
393 pub fn write_to_path<P: AsRef<Path>>(&self, bgzf_path: P) -> Result<PathBuf> {
395 let target = gzi_path(bgzf_path.as_ref());
396 let mut file = BufWriter::new(File::create(&target)?);
397 self.write(&mut file)?;
398 file.flush()?;
399 Ok(target)
400 }
401
402 pub fn blocks(&self) -> &[BlockOffset] {
404 &self.blocks
405 }
406
407 pub fn len(&self) -> usize {
409 self.blocks.len()
410 }
411
412 pub fn is_empty(&self) -> bool {
414 self.blocks.is_empty()
415 }
416
417 pub fn uncompressed_len(&self) -> Option<u64> {
423 self.total_uncompressed
424 }
425
426 fn block_for(&self, offset: u64) -> Option<BlockOffset> {
428 match self
430 .blocks
431 .binary_search_by(|b| b.uncompressed.cmp(&offset))
432 {
433 Ok(i) => Some(self.blocks[i]),
434 Err(0) => None,
435 Err(i) => Some(self.blocks[i - 1]),
436 }
437 }
438}
439
440pub fn gzi_path(bgzf: &Path) -> PathBuf {
442 let mut name = bgzf.as_os_str().to_os_string();
443 name.push(".gzi");
444 PathBuf::from(name)
445}
446
447pub fn is_bgzf(bytes: &[u8]) -> bool {
456 parse_block_header(bytes).is_ok()
457}
458
459pub struct BgzfReader<R: Read + Seek> {
464 inner: R,
465 index: Option<GziIndex>,
466 block: Vec<u8>,
468 block_pos: usize,
470 block_start: u64,
472 next_compressed: u64,
474 eof: bool,
475 raw: Vec<u8>,
477}
478
479impl BgzfReader<File> {
480 pub fn open<P: AsRef<Path>>(path: P) -> Result<BgzfReader<File>> {
485 let path = path.as_ref();
486 let file = File::open(path)
487 .map_err(|e| Error::Io(io::Error::new(e.kind(), format!("{}: {e}", path.display()))))?;
488 let index = match GziIndex::from_path(gzi_path(path)) {
489 Ok(index) => Some(index),
490 Err(Error::Io(e)) if e.kind() == io::ErrorKind::NotFound => None,
491 Err(e) => return Err(e),
492 };
493 let mut reader = BgzfReader::new(file)?;
494 reader.index = index;
495 Ok(reader)
496 }
497}
498
499impl<R: Read + Seek> BgzfReader<R> {
500 pub fn new(mut inner: R) -> Result<BgzfReader<R>> {
502 inner.seek(SeekFrom::Start(0))?;
503 let mut probe = [0u8; HEADER_LEN + 64];
504 let read = read_up_to(&mut inner, &mut probe)?;
505 if read == 0 {
506 return Err(bgzf_error("file is empty"));
507 }
508 parse_block_header(&probe[..read])?;
509 inner.seek(SeekFrom::Start(0))?;
510 Ok(BgzfReader {
511 inner,
512 index: None,
513 block: Vec::new(),
514 block_pos: 0,
515 block_start: 0,
516 next_compressed: 0,
517 eof: false,
518 raw: Vec::new(),
519 })
520 }
521
522 pub fn with_index(mut self, index: GziIndex) -> Self {
524 self.index = Some(index);
525 self
526 }
527
528 pub fn index(&self) -> Option<&GziIndex> {
530 self.index.as_ref()
531 }
532
533 pub fn position(&self) -> u64 {
535 self.block_start + self.block_pos as u64
536 }
537
538 pub fn into_inner(self) -> R {
540 self.inner
541 }
542
543 fn load_block_at(&mut self, at: u64, uncompressed_start: u64) -> Result<()> {
545 self.inner.seek(SeekFrom::Start(at))?;
546 self.next_compressed = at;
547 self.block_start = uncompressed_start;
548 self.block_pos = 0;
549 self.block.clear();
550 self.eof = false;
551 self.read_next_block()
552 }
553
554 fn read_next_block(&mut self) -> Result<()> {
556 let mut header = [0u8; HEADER_LEN + 64];
557 if !read_exact_or_eof(&mut self.inner, &mut header[..HEADER_LEN])? {
558 self.eof = true;
559 self.block.clear();
560 self.block_pos = 0;
561 return Ok(());
562 }
563 let extra_len = u16::from_le_bytes([header[10], header[11]]) as usize;
564 if HEADER_LEN + extra_len > header.len() {
565 return Err(bgzf_error("extra field is implausibly large"));
566 }
567 self.inner
568 .read_exact(&mut header[HEADER_LEN..HEADER_LEN + extra_len])?;
569 let block = parse_block_header(&header[..HEADER_LEN + extra_len])?;
570
571 let payload_len = block.compressed_len - block.payload_offset - TRAILER_LEN;
572 self.raw.resize(payload_len, 0);
573 self.inner.read_exact(&mut self.raw)?;
574 let mut trailer = [0u8; TRAILER_LEN];
575 self.inner.read_exact(&mut trailer)?;
576 let expected_crc = u32::from_le_bytes(trailer[..4].try_into().expect("4 bytes"));
577 let expected_len = u32::from_le_bytes(trailer[4..].try_into().expect("4 bytes")) as usize;
578
579 let mut decompressed = std::mem::take(&mut self.block);
580 let result = inflate(&self.raw, expected_len, &mut decompressed);
581 self.block = decompressed;
582 result?;
583
584 if crc32(&self.block) != expected_crc {
585 return Err(bgzf_error("block CRC32 does not match"));
586 }
587
588 self.block_pos = 0;
589 self.next_compressed += block.compressed_len as u64;
590 if self.block.is_empty() {
592 self.eof = true;
593 }
594 Ok(())
595 }
596
597 fn fill(&mut self) -> Result<()> {
599 while !self.eof && self.block_pos == self.block.len() {
600 let consumed = self.block.len() as u64;
601 self.block_start += consumed;
602 self.read_next_block()?;
603 }
604 Ok(())
605 }
606
607 fn ensure_index(&mut self) -> Result<()> {
609 if self.index.is_none() {
610 let saved = self.inner.stream_position()?;
611 let index = GziIndex::build(&mut self.inner)?;
612 self.inner.seek(SeekFrom::Start(saved))?;
613 self.index = Some(index);
614 }
615 Ok(())
616 }
617}
618
619impl<R: Read + Seek> Read for BgzfReader<R> {
620 fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
621 if buf.is_empty() {
622 return Ok(0);
623 }
624 self.fill()?;
625 if self.eof && self.block_pos == self.block.len() {
626 return Ok(0);
627 }
628 let available = &self.block[self.block_pos..];
629 let take = available.len().min(buf.len());
630 buf[..take].copy_from_slice(&available[..take]);
631 self.block_pos += take;
632 Ok(take)
633 }
634}
635
636impl<R: Read + Seek> Seek for BgzfReader<R> {
637 fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
642 self.ensure_index()?;
643
644 let target = match pos {
645 SeekFrom::Start(offset) => offset,
646 SeekFrom::Current(delta) => add_signed(self.position(), delta)?,
647 SeekFrom::End(delta) => {
648 let end = self
649 .index
650 .as_ref()
651 .and_then(|index| index.uncompressed_len())
652 .ok_or_else(|| {
653 io::Error::new(
654 io::ErrorKind::InvalidInput,
655 "BGZF: this index records block starts only, so the end of the \
656 stream is unknown; rebuild it with GziIndex::build",
657 )
658 })?;
659 add_signed(end, delta)?
660 }
661 };
662
663 let within = target.checked_sub(self.block_start);
665 if let Some(within) = within {
666 if !self.block.is_empty() && within <= self.block.len() as u64 {
667 self.block_pos = within as usize;
668 return Ok(target);
669 }
670 }
671
672 let block = self
675 .index
676 .as_ref()
677 .and_then(|index| index.block_for(target))
678 .ok_or_else(|| {
679 io::Error::new(
680 io::ErrorKind::InvalidInput,
681 "BGZF: no block covers that offset",
682 )
683 })?;
684 self.load_block_at(block.compressed, block.uncompressed)?;
685 let within = (target - block.uncompressed) as usize;
686 if within > self.block.len() {
687 self.block_pos = self.block.len();
689 self.eof = true;
690 } else {
691 self.block_pos = within;
692 }
693 Ok(target)
694 }
695}
696
697fn add_signed(base: u64, delta: i64) -> io::Result<u64> {
698 let result = if delta >= 0 {
699 base.checked_add(delta as u64)
700 } else {
701 base.checked_sub(delta.unsigned_abs())
702 };
703 result.ok_or_else(|| {
704 io::Error::new(
705 io::ErrorKind::InvalidInput,
706 "BGZF: seek would leave the file",
707 )
708 })
709}
710
711#[cfg(feature = "parallel")]
717struct RawBlock {
718 payload: Vec<u8>,
719 crc: u32,
720 uncompressed_len: usize,
721}
722
723#[cfg(feature = "parallel")]
725fn read_raw_block<R: Read>(reader: &mut R) -> Result<Option<RawBlock>> {
726 let mut header = [0u8; HEADER_LEN + 64];
727 if !read_exact_or_eof(reader, &mut header[..HEADER_LEN])? {
728 return Ok(None);
729 }
730 let extra_len = u16::from_le_bytes([header[10], header[11]]) as usize;
731 if HEADER_LEN + extra_len > header.len() {
732 return Err(bgzf_error("extra field is implausibly large"));
733 }
734 reader.read_exact(&mut header[HEADER_LEN..HEADER_LEN + extra_len])?;
735 let block = parse_block_header(&header[..HEADER_LEN + extra_len])?;
736
737 let payload_len = block.compressed_len - block.payload_offset - TRAILER_LEN;
738 let mut payload = vec![0u8; payload_len];
739 reader.read_exact(&mut payload)?;
740 let mut trailer = [0u8; TRAILER_LEN];
741 reader.read_exact(&mut trailer)?;
742
743 Ok(Some(RawBlock {
744 payload,
745 crc: u32::from_le_bytes(trailer[..4].try_into().expect("4 bytes")),
746 uncompressed_len: u32::from_le_bytes(trailer[4..].try_into().expect("4 bytes")) as usize,
747 }))
748}
749
750#[cfg(feature = "parallel")]
752fn inflate_checked(inflater: &mut Inflater, raw: &RawBlock) -> Result<Vec<u8>> {
753 let mut out = Vec::new();
754 inflater.inflate(&raw.payload, raw.uncompressed_len, &mut out)?;
755 if crc32(&out) != raw.crc {
756 return Err(bgzf_error("block CRC32 does not match"));
757 }
758 Ok(out)
759}
760
761#[cfg(feature = "parallel")]
782pub struct ParallelBgzfReader<R: Read> {
783 inner: R,
784 batch: usize,
785 buffer: Vec<u8>,
787 pos: usize,
788 eof: bool,
789}
790
791#[cfg(feature = "parallel")]
792impl<R: Read> ParallelBgzfReader<R> {
793 pub fn new(inner: R) -> ParallelBgzfReader<R> {
795 ParallelBgzfReader {
796 inner,
797 batch: default_batch().max(2),
798 buffer: Vec::new(),
799 pos: 0,
800 eof: false,
801 }
802 }
803
804 pub fn blocks_per_batch(mut self, blocks: usize) -> Self {
807 self.batch = blocks.max(1);
808 self
809 }
810
811 pub fn into_inner(self) -> R {
813 self.inner
814 }
815
816 fn fill(&mut self) -> Result<()> {
818 use rayon::prelude::*;
819
820 self.buffer.clear();
821 self.pos = 0;
822
823 let mut raws: Vec<RawBlock> = Vec::with_capacity(self.batch);
824 while raws.len() < self.batch {
825 match read_raw_block(&mut self.inner)? {
826 None => {
827 self.eof = true;
828 break;
829 }
830 Some(raw) if raw.uncompressed_len == 0 => {
832 self.eof = true;
833 break;
834 }
835 Some(raw) => raws.push(raw),
836 }
837 }
838 if raws.is_empty() {
839 return Ok(());
840 }
841
842 let blocks: Vec<Vec<u8>> = if raws.len() > 1 {
843 raws.par_iter()
845 .map_init(Inflater::new, inflate_checked)
846 .collect::<Result<_>>()?
847 } else {
848 vec![inflate_checked(&mut Inflater::new(), &raws[0])?]
849 };
850 self.buffer.reserve(blocks.iter().map(Vec::len).sum());
851 for block in &blocks {
852 self.buffer.extend_from_slice(block);
853 }
854 Ok(())
855 }
856}
857
858#[cfg(feature = "parallel")]
859impl<R: Read> Read for ParallelBgzfReader<R> {
860 fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
861 if buf.is_empty() {
862 return Ok(0);
863 }
864 while self.pos == self.buffer.len() {
865 if self.eof {
866 return Ok(0);
867 }
868 self.fill()?;
869 }
870 let available = &self.buffer[self.pos..];
871 let take = available.len().min(buf.len());
872 buf[..take].copy_from_slice(&available[..take]);
873 self.pos += take;
874 Ok(take)
875 }
876}
877
878pub struct BgzfWriter<W: Write> {
884 inner: Option<W>,
887 buffer: Vec<u8>,
889 pending: Vec<Vec<u8>>,
891 recycled: Vec<Vec<u8>>,
893 batch: usize,
897 level: CompressionLevel,
898 blocks: Vec<BlockOffset>,
900 compressed: u64,
901 uncompressed: u64,
902 finished: bool,
903}
904
905struct EncodedBlock {
910 header: [u8; HEADER_LEN + EXTRA_LEN],
911 payload: Vec<u8>,
912 trailer: [u8; TRAILER_LEN],
913 block_len: usize,
914 uncompressed_len: usize,
915}
916
917fn encode_block(deflater: &mut Deflater, data: &[u8]) -> io::Result<EncodedBlock> {
919 let payload = deflater.deflate(data)?;
920 let block_len = HEADER_LEN + EXTRA_LEN + payload.len() + TRAILER_LEN;
921 if block_len > u16::MAX as usize + 1 {
922 return Err(io::Error::new(
925 io::ErrorKind::InvalidData,
926 "BGZF: block would exceed 64 KiB",
927 ));
928 }
929
930 let mut header = [0u8; HEADER_LEN + EXTRA_LEN];
931 header[0] = 0x1f;
932 header[1] = 0x8b;
933 header[2] = 8; header[3] = 4; header[9] = 0xff; header[10..12].copy_from_slice(&(EXTRA_LEN as u16).to_le_bytes());
938 header[12] = b'B';
939 header[13] = b'C';
940 header[14..16].copy_from_slice(&2u16.to_le_bytes());
941 header[16..18].copy_from_slice(&((block_len - 1) as u16).to_le_bytes());
942
943 let mut trailer = [0u8; TRAILER_LEN];
944 trailer[..4].copy_from_slice(&crc32(data).to_le_bytes());
945 trailer[4..].copy_from_slice(&(data.len() as u32).to_le_bytes());
946
947 Ok(EncodedBlock {
948 header,
949 payload,
950 trailer,
951 block_len,
952 uncompressed_len: data.len(),
953 })
954}
955
956fn encode_batch(payloads: &[Vec<u8>], level: CompressionLevel) -> io::Result<Vec<EncodedBlock>> {
961 #[cfg(feature = "parallel")]
962 {
963 if payloads.len() > 1 {
964 use rayon::prelude::*;
965 return payloads
969 .par_iter()
970 .map_init(
971 || Deflater::new(level),
972 |deflater, payload| encode_block(deflater, payload),
973 )
974 .collect();
975 }
976 }
977 let mut deflater = Deflater::new(level);
978 payloads
979 .iter()
980 .map(|payload| encode_block(&mut deflater, payload))
981 .collect()
982}
983
984#[cfg(feature = "parallel")]
992const BATCHES_PER_CORE: usize = 8;
993
994fn default_batch() -> usize {
996 #[cfg(feature = "parallel")]
997 {
998 std::thread::available_parallelism().map_or(1, |n| n.get() * BATCHES_PER_CORE)
999 }
1000 #[cfg(not(feature = "parallel"))]
1001 {
1002 1
1003 }
1004}
1005
1006impl BgzfWriter<BufWriter<File>> {
1007 pub fn create<P: AsRef<Path>>(path: P) -> Result<BgzfWriter<BufWriter<File>>> {
1009 let path = path.as_ref();
1010 let file = File::create(path)
1011 .map_err(|e| Error::Io(io::Error::new(e.kind(), format!("{}: {e}", path.display()))))?;
1012 Ok(BgzfWriter::new(BufWriter::with_capacity(128 * 1024, file)))
1013 }
1014}
1015
1016impl<W: Write> BgzfWriter<W> {
1017 pub fn new(inner: W) -> BgzfWriter<W> {
1019 BgzfWriter::with_level(inner, CompressionLevel::default())
1020 }
1021
1022 pub fn with_level(inner: W, level: CompressionLevel) -> BgzfWriter<W> {
1024 BgzfWriter {
1025 inner: Some(inner),
1026 buffer: Vec::with_capacity(MAX_BLOCK_PAYLOAD),
1027 pending: Vec::new(),
1028 recycled: Vec::new(),
1029 batch: default_batch(),
1030 level,
1031 blocks: vec![BlockOffset {
1032 compressed: 0,
1033 uncompressed: 0,
1034 }],
1035 compressed: 0,
1036 uncompressed: 0,
1037 finished: false,
1038 }
1039 }
1040
1041 pub fn get_ref(&self) -> Option<&W> {
1043 self.inner.as_ref()
1044 }
1045
1046 pub fn blocks_per_batch(mut self, blocks: usize) -> Self {
1056 self.batch = blocks.max(1);
1057 self
1058 }
1059
1060 pub fn index(&self) -> GziIndex {
1067 GziIndex {
1068 blocks: self.blocks.clone(),
1069 total_uncompressed: Some(self.uncompressed),
1070 }
1071 }
1072
1073 fn sink(&mut self) -> io::Result<&mut W> {
1075 self.inner.as_mut().ok_or_else(|| {
1076 io::Error::new(
1077 io::ErrorKind::BrokenPipe,
1078 "BGZF: writer was already finished",
1079 )
1080 })
1081 }
1082
1083 fn seal_buffer(&mut self) -> io::Result<()> {
1087 if self.buffer.is_empty() {
1088 return Ok(());
1089 }
1090 let mut fresh = self.recycled.pop().unwrap_or_else(|| {
1091 let mut buffer = Vec::new();
1092 buffer.reserve_exact(MAX_BLOCK_PAYLOAD);
1093 buffer
1094 });
1095 std::mem::swap(&mut self.buffer, &mut fresh);
1096 self.pending.push(fresh);
1097 if self.pending.len() >= self.batch {
1098 self.compress_pending()?;
1099 }
1100 Ok(())
1101 }
1102
1103 fn compress_pending(&mut self) -> io::Result<()> {
1105 if self.pending.is_empty() {
1106 return Ok(());
1107 }
1108 let level = self.level;
1109 let encoded = encode_batch(&self.pending, level)?;
1110 for block in encoded {
1111 self.emit(block)?;
1112 }
1113 for mut buffer in self.pending.drain(..) {
1115 buffer.clear();
1116 if self.recycled.len() < self.batch {
1117 self.recycled.push(buffer);
1118 }
1119 }
1120 Ok(())
1121 }
1122
1123 fn emit(&mut self, block: EncodedBlock) -> io::Result<()> {
1125 let sink = self.sink()?;
1126 sink.write_all(&block.header)?;
1127 sink.write_all(&block.payload)?;
1128 sink.write_all(&block.trailer)?;
1129
1130 self.compressed += block.block_len as u64;
1131 self.uncompressed += block.uncompressed_len as u64;
1132 self.blocks.push(BlockOffset {
1133 compressed: self.compressed,
1134 uncompressed: self.uncompressed,
1135 });
1136 Ok(())
1137 }
1138
1139 fn flush_block(&mut self) -> io::Result<()> {
1141 self.seal_buffer()?;
1142 self.compress_pending()
1143 }
1144
1145 pub fn finish(self) -> Result<W> {
1150 self.finish_with_index().map(|(inner, _)| inner)
1151 }
1152
1153 pub fn finish_with_index(mut self) -> Result<(W, GziIndex)> {
1169 self.finish_in_place()?;
1170 let index = self.index();
1171 let inner = self
1172 .inner
1173 .take()
1174 .ok_or_else(|| Error::Other("BGZF: writer was already finished".to_string()))?;
1175 Ok((inner, index))
1176 }
1177
1178 fn finish_in_place(&mut self) -> Result<()> {
1179 if self.finished || self.inner.is_none() {
1180 return Ok(());
1181 }
1182 self.flush_block()?;
1183 let sink = self.sink()?;
1184 sink.write_all(&EOF_BLOCK)?;
1185 sink.flush()?;
1186 self.compressed += EOF_BLOCK.len() as u64;
1187 self.finished = true;
1188 Ok(())
1189 }
1190}
1191
1192impl<W: Write> Write for BgzfWriter<W> {
1193 fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
1194 let room = MAX_BLOCK_PAYLOAD - self.buffer.len();
1195 let take = room.min(buf.len());
1196 self.buffer.extend_from_slice(&buf[..take]);
1197 if self.buffer.len() == MAX_BLOCK_PAYLOAD {
1198 self.seal_buffer()?;
1200 }
1201 Ok(take)
1202 }
1203
1204 fn flush(&mut self) -> io::Result<()> {
1209 self.flush_block()?;
1210 self.sink()?.flush()
1211 }
1212}
1213
1214impl<W: Write> Drop for BgzfWriter<W> {
1215 fn drop(&mut self) {
1216 let _ = self.finish_in_place();
1218 }
1219}
1220
1221fn read_exact_or_eof<R: Read>(reader: &mut R, buf: &mut [u8]) -> Result<bool> {
1223 let mut filled = 0;
1224 while filled < buf.len() {
1225 match reader.read(&mut buf[filled..]) {
1226 Ok(0) if filled == 0 => return Ok(false),
1227 Ok(0) => return Err(bgzf_error("file ends in the middle of a block")),
1228 Ok(n) => filled += n,
1229 Err(e) if e.kind() == io::ErrorKind::Interrupted => continue,
1230 Err(e) => return Err(Error::Io(e)),
1231 }
1232 }
1233 Ok(true)
1234}
1235
1236fn read_up_to<R: Read>(reader: &mut R, buf: &mut [u8]) -> Result<usize> {
1238 let mut filled = 0;
1239 while filled < buf.len() {
1240 match reader.read(&mut buf[filled..]) {
1241 Ok(0) => break,
1242 Ok(n) => filled += n,
1243 Err(e) if e.kind() == io::ErrorKind::Interrupted => continue,
1244 Err(e) => return Err(Error::Io(e)),
1245 }
1246 }
1247 Ok(filled)
1248}
1249
1250#[cfg(test)]
1251mod tests {
1252 use super::*;
1253 use std::io::Cursor;
1254
1255 fn compress(data: &[u8]) -> Vec<u8> {
1256 let mut writer = BgzfWriter::new(Vec::new());
1257 writer.write_all(data).unwrap();
1258 writer.finish().unwrap()
1259 }
1260
1261 #[test]
1262 fn round_trips_through_our_own_reader() {
1263 for size in [
1264 0usize,
1265 1,
1266 100,
1267 MAX_BLOCK_PAYLOAD - 1,
1268 MAX_BLOCK_PAYLOAD,
1269 MAX_BLOCK_PAYLOAD + 1,
1270 300_000,
1271 ] {
1272 let data: Vec<u8> = (0..size).map(|i| b"ACGTN"[i % 5]).collect();
1273 let compressed = compress(&data);
1274 assert!(is_bgzf(&compressed), "size {size} did not produce BGZF");
1275
1276 let mut reader = BgzfReader::new(Cursor::new(&compressed)).unwrap();
1277 let mut out = Vec::new();
1278 reader.read_to_end(&mut out).unwrap();
1279 assert_eq!(out, data, "size {size}");
1280 }
1281 }
1282
1283 #[test]
1284 fn output_is_plain_gzip_too() {
1285 let data: Vec<u8> = (0..200_000).map(|i| b"ACGT"[i % 4]).collect();
1287 let compressed = compress(&data);
1288 let mut out = Vec::new();
1289 flate2::read::MultiGzDecoder::new(&compressed[..])
1290 .read_to_end(&mut out)
1291 .unwrap();
1292 assert_eq!(out, data);
1293 }
1294
1295 #[test]
1296 fn ends_with_the_eof_marker() {
1297 let compressed = compress(b"ACGT");
1298 assert_eq!(
1299 &compressed[compressed.len() - EOF_BLOCK.len()..],
1300 &EOF_BLOCK
1301 );
1302 assert_eq!(compress(b""), EOF_BLOCK.to_vec());
1304 }
1305
1306 #[test]
1307 fn blocks_stay_within_the_size_limit() {
1308 let data: Vec<u8> = (0..500_000)
1310 .map(|i| ((i * 2_654_435_761u64 as usize) >> 7) as u8)
1311 .collect();
1312 let compressed = compress(&data);
1313 let index = GziIndex::build(Cursor::new(&compressed)).unwrap();
1314 for pair in index.blocks().windows(2) {
1315 let block_len = pair[1].compressed - pair[0].compressed;
1316 assert!(block_len <= 65_536, "block of {block_len} bytes");
1317 }
1318 let mut out = Vec::new();
1319 BgzfReader::new(Cursor::new(&compressed))
1320 .unwrap()
1321 .read_to_end(&mut out)
1322 .unwrap();
1323 assert_eq!(out, data);
1324 }
1325
1326 #[test]
1327 fn index_from_the_writer_matches_a_rescan() {
1328 let data: Vec<u8> = (0..250_000).map(|i| b"ACGTN"[i % 5]).collect();
1329 let mut writer = BgzfWriter::new(Vec::new());
1330 writer.write_all(&data).unwrap();
1331 let from_writer = writer.index();
1332 let compressed = writer.finish().unwrap();
1333
1334 let from_scan = GziIndex::build(Cursor::new(&compressed)).unwrap();
1335 assert_eq!(
1337 &from_scan.blocks()[..from_writer.len()],
1338 from_writer.blocks()
1339 );
1340 assert_eq!(from_scan.uncompressed_len(), Some(data.len() as u64));
1341 }
1342
1343 #[test]
1344 fn gzi_round_trips_and_omits_the_first_block() {
1345 let data: Vec<u8> = (0..200_000).map(|i| b"ACGT"[i % 4]).collect();
1346 let compressed = compress(&data);
1347 let index = GziIndex::build(Cursor::new(&compressed)).unwrap();
1348
1349 let mut text = Vec::new();
1350 index.write(&mut text).unwrap();
1351 assert_eq!(text.len(), 8 + 16 * (index.len() - 1));
1353 assert_eq!(
1354 u64::from_le_bytes(text[..8].try_into().unwrap()),
1355 index.len() as u64 - 1
1356 );
1357
1358 let reparsed = GziIndex::parse(&text[..]).unwrap();
1359 assert_eq!(reparsed.blocks(), index.blocks());
1360 assert_eq!(reparsed.uncompressed_len(), None);
1362 assert_eq!(index.uncompressed_len(), Some(200_000));
1363 }
1364
1365 #[test]
1366 fn seek_from_end_needs_a_scanned_index() {
1367 let data: Vec<u8> = (0..200_000).map(|i| b"ACGT"[i % 4]).collect();
1368 let compressed = compress(&data);
1369 let scanned = GziIndex::build(Cursor::new(&compressed)).unwrap();
1370 let mut text = Vec::new();
1371 scanned.write(&mut text).unwrap();
1372 let parsed = GziIndex::parse(&text[..]).unwrap();
1373
1374 let mut reader = BgzfReader::new(Cursor::new(&compressed))
1377 .unwrap()
1378 .with_index(parsed);
1379 assert!(reader.seek(SeekFrom::End(0)).is_err());
1380 reader.seek(SeekFrom::Start(199_998)).unwrap();
1382 let mut tail = Vec::new();
1383 reader.read_to_end(&mut tail).unwrap();
1384 assert_eq!(tail, &data[199_998..]);
1385
1386 let mut reader = BgzfReader::new(Cursor::new(&compressed))
1387 .unwrap()
1388 .with_index(scanned);
1389 assert_eq!(reader.seek(SeekFrom::End(0)).unwrap(), 200_000);
1390 }
1391
1392 #[test]
1393 fn seeks_to_any_offset() {
1394 let data: Vec<u8> = (0..300_000).map(|i| (i % 251) as u8).collect();
1395 let compressed = compress(&data);
1396 let mut reader = BgzfReader::new(Cursor::new(&compressed)).unwrap();
1397
1398 for target in [
1400 0usize,
1401 1,
1402 MAX_BLOCK_PAYLOAD - 1,
1403 MAX_BLOCK_PAYLOAD,
1404 MAX_BLOCK_PAYLOAD + 1,
1405 2 * MAX_BLOCK_PAYLOAD,
1406 299_999,
1407 ] {
1408 reader.seek(SeekFrom::Start(target as u64)).unwrap();
1409 assert_eq!(reader.position(), target as u64);
1410 let mut buf = [0u8; 8];
1411 let want = (data.len() - target).min(buf.len());
1412 reader.read_exact(&mut buf[..want]).unwrap();
1413 assert_eq!(&buf[..want], &data[target..target + want], "at {target}");
1414 }
1415
1416 reader.seek(SeekFrom::Start(10)).unwrap();
1418 reader.seek(SeekFrom::Current(5)).unwrap();
1419 assert_eq!(reader.position(), 15);
1420 assert_eq!(reader.seek(SeekFrom::End(0)).unwrap(), data.len() as u64);
1421 let mut rest = Vec::new();
1422 reader.read_to_end(&mut rest).unwrap();
1423 assert!(rest.is_empty());
1424
1425 reader.seek(SeekFrom::Start(7)).unwrap();
1427 let mut buf = [0u8; 4];
1428 reader.read_exact(&mut buf).unwrap();
1429 assert_eq!(&buf, &data[7..11]);
1430 }
1431
1432 #[test]
1433 fn seek_before_the_start_is_an_error() {
1434 let compressed = compress(b"ACGT");
1435 let mut reader = BgzfReader::new(Cursor::new(&compressed)).unwrap();
1436 assert!(reader.seek(SeekFrom::Current(-1)).is_err());
1437 assert!(reader.seek(SeekFrom::End(-100)).is_err());
1438 }
1439
1440 #[test]
1441 fn rejects_plain_gzip_and_garbage() {
1442 let mut plain = Vec::new();
1444 {
1445 let mut encoder =
1446 flate2::write::GzEncoder::new(&mut plain, flate2::Compression::default());
1447 encoder.write_all(b"ACGT").unwrap();
1448 encoder.finish().unwrap();
1449 }
1450 assert!(!is_bgzf(&plain));
1451 assert!(BgzfReader::new(Cursor::new(&plain)).is_err());
1452
1453 assert!(BgzfReader::new(Cursor::new(b"not gzip at all".to_vec())).is_err());
1454 assert!(BgzfReader::new(Cursor::new(Vec::new())).is_err());
1455 }
1456
1457 #[test]
1458 fn detects_a_corrupt_block() {
1459 let mut compressed = compress(&vec![b'A'; 5_000]);
1460 let victim = HEADER_LEN + EXTRA_LEN + 5;
1462 compressed[victim] ^= 0xff;
1463 let mut out = Vec::new();
1464 let result = BgzfReader::new(Cursor::new(&compressed))
1465 .unwrap()
1466 .read_to_end(&mut out);
1467 assert!(result.is_err(), "corruption went unnoticed");
1468 }
1469
1470 #[test]
1471 fn truncated_file_is_an_error_not_silent_truncation() {
1472 let compressed = compress(&vec![b'A'; 200_000]);
1473 let cut = compressed.len() / 2;
1474 let mut out = Vec::new();
1475 let result = BgzfReader::new(Cursor::new(compressed[..cut].to_vec()))
1476 .unwrap()
1477 .read_to_end(&mut out);
1478 assert!(result.is_err(), "truncation went unnoticed");
1479 }
1480
1481 #[test]
1482 fn parse_rejects_an_implausible_index() {
1483 let mut bad = u64::MAX.to_le_bytes().to_vec();
1484 bad.extend_from_slice(&[0u8; 16]);
1485 assert!(GziIndex::parse(&bad[..]).is_err());
1486 assert!(GziIndex::parse(&[0u8; 3][..]).is_err());
1487 }
1488
1489 #[test]
1490 fn batch_size_never_changes_the_bytes() {
1491 let data: Vec<u8> = (0..400_000).map(|i| b"ACGTN"[i % 5]).collect();
1495 let mut reference = None;
1496 for blocks in [1usize, 2, 3, 7, 64, 1024] {
1497 let mut writer = BgzfWriter::new(Vec::new()).blocks_per_batch(blocks);
1498 writer.write_all(&data).unwrap();
1499 let (bytes, index) = writer.finish_with_index().unwrap();
1500
1501 match &reference {
1502 None => reference = Some((bytes, index)),
1503 Some((expected_bytes, expected_index)) => {
1504 assert_eq!(&bytes, expected_bytes, "batch of {blocks} differs");
1505 assert_eq!(index.blocks(), expected_index.blocks(), "index differs");
1506 }
1507 }
1508 }
1509 let (bytes, _) = reference.unwrap();
1511 let mut out = Vec::new();
1512 BgzfReader::new(Cursor::new(&bytes))
1513 .unwrap()
1514 .read_to_end(&mut out)
1515 .unwrap();
1516 assert_eq!(out, data);
1517 }
1518
1519 #[test]
1520 fn write_flush_boundaries_survive_batching() {
1521 let mut writer = BgzfWriter::new(Vec::new()).blocks_per_batch(16);
1524 writer.write_all(b"first").unwrap();
1525 writer.flush().unwrap();
1526 assert!(
1527 !writer.get_ref().expect("still open").is_empty(),
1528 "flush left the batch uncompressed"
1529 );
1530 writer.write_all(b"second").unwrap();
1531 let (bytes, index) = writer.finish_with_index().unwrap();
1532
1533 assert_eq!(index.len(), 3); let mut out = Vec::new();
1535 BgzfReader::new(Cursor::new(&bytes))
1536 .unwrap()
1537 .read_to_end(&mut out)
1538 .unwrap();
1539 assert_eq!(out, b"firstsecond");
1540 }
1541
1542 #[cfg(feature = "parallel")]
1543 #[test]
1544 fn parallel_reader_agrees_with_the_serial_one() {
1545 for size in [
1546 0usize,
1547 1,
1548 MAX_BLOCK_PAYLOAD - 1,
1549 MAX_BLOCK_PAYLOAD,
1550 MAX_BLOCK_PAYLOAD + 1,
1551 500_000,
1552 ] {
1553 let data: Vec<u8> = (0..size).map(|i| (i % 251) as u8).collect();
1554 let compressed = compress(&data);
1555
1556 for blocks in [1usize, 2, 5, 64] {
1557 let mut parallel =
1558 ParallelBgzfReader::new(Cursor::new(&compressed)).blocks_per_batch(blocks);
1559 let mut out = Vec::new();
1560 parallel.read_to_end(&mut out).unwrap();
1561 assert_eq!(out, data, "size {size}, batch {blocks}");
1562 }
1563
1564 let mut parallel = ParallelBgzfReader::new(Cursor::new(&compressed));
1566 let mut out = Vec::new();
1567 let mut byte = [0u8; 1];
1568 while parallel.read(&mut byte).unwrap() == 1 {
1569 out.push(byte[0]);
1570 }
1571 assert_eq!(out, data, "size {size}, byte at a time");
1572 }
1573 }
1574
1575 #[cfg(feature = "parallel")]
1576 #[test]
1577 fn parallel_reader_still_catches_corruption() {
1578 let mut compressed = compress(&vec![b'A'; 300_000]);
1579 let victim = compressed.len() / 2;
1581 compressed[victim] ^= 0xff;
1582 let mut out = Vec::new();
1583 let result = ParallelBgzfReader::new(Cursor::new(&compressed)).read_to_end(&mut out);
1584 assert!(result.is_err(), "corruption went unnoticed");
1585
1586 let whole = compress(&vec![b'C'; 300_000]);
1588 let mut out = Vec::new();
1589 let result = ParallelBgzfReader::new(Cursor::new(whole[..whole.len() / 2].to_vec()))
1590 .read_to_end(&mut out);
1591 assert!(result.is_err(), "truncation went unnoticed");
1592 }
1593
1594 #[test]
1595 fn flush_starts_a_new_block() {
1596 let mut writer = BgzfWriter::new(Vec::new());
1597 writer.write_all(b"first").unwrap();
1598 writer.flush().unwrap();
1599 writer.write_all(b"second").unwrap();
1600 assert_eq!(writer.index().len(), 2);
1603 let (compressed, index) = writer.finish_with_index().unwrap();
1604
1605 assert_eq!(index.len(), 3); assert_eq!(index.blocks()[1].uncompressed, 5);
1608 assert_eq!(index.uncompressed_len(), Some(11));
1609
1610 let mut reader = BgzfReader::new(Cursor::new(&compressed)).unwrap();
1611 reader.seek(SeekFrom::Start(5)).unwrap();
1612 let mut out = Vec::new();
1613 reader.read_to_end(&mut out).unwrap();
1614 assert_eq!(out, b"second");
1615 }
1616}