use std::fs::File;
use std::io::{self, BufReader, BufWriter, Read, Seek, SeekFrom, Write};
use std::path::{Path, PathBuf};
use crate::error::{Error, Result};
use crate::format::CompressionLevel;
pub const MAX_BLOCK_PAYLOAD: usize = 0xff00;
const HEADER_LEN: usize = 12;
const EXTRA_LEN: usize = 6;
const TRAILER_LEN: usize = 8;
pub const EOF_BLOCK: [u8; 28] = [
0x1f, 0x8b, 0x08, 0x04, 0x00, 0x00, 0x00, 0x00, 0x00, 0xff, 0x06, 0x00, 0x42, 0x43, 0x02, 0x00,
0x1b, 0x00, 0x03, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
];
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct BlockHeader {
compressed_len: usize,
payload_offset: usize,
}
fn parse_block_header(bytes: &[u8]) -> Result<BlockHeader> {
if bytes.len() < HEADER_LEN {
return Err(bgzf_error("block header is truncated"));
}
if bytes[0] != 0x1f || bytes[1] != 0x8b {
return Err(bgzf_error("not a gzip member"));
}
if bytes[2] != 8 {
return Err(bgzf_error("unsupported compression method"));
}
if bytes[3] & 0x04 == 0 {
return Err(bgzf_error(
"gzip member has no extra field, so it is gzip but not BGZF",
));
}
let extra_len = u16::from_le_bytes([bytes[10], bytes[11]]) as usize;
if bytes.len() < HEADER_LEN + extra_len {
return Err(bgzf_error("extra field is truncated"));
}
let extra = &bytes[HEADER_LEN..HEADER_LEN + extra_len];
let mut cursor = 0;
while cursor + 4 <= extra.len() {
let si1 = extra[cursor];
let si2 = extra[cursor + 1];
let slen = u16::from_le_bytes([extra[cursor + 2], extra[cursor + 3]]) as usize;
let value = cursor + 4;
if value + slen > extra.len() {
return Err(bgzf_error("extra subfield runs past the extra field"));
}
if si1 == b'B' && si2 == b'C' {
if slen != 2 {
return Err(bgzf_error("BC subfield is not two bytes"));
}
let bsize = u16::from_le_bytes([extra[value], extra[value + 1]]) as usize;
let compressed_len = bsize + 1;
let overhead = HEADER_LEN + extra_len + TRAILER_LEN;
if compressed_len <= overhead {
return Err(bgzf_error("BSIZE is smaller than the block overhead"));
}
return Ok(BlockHeader {
compressed_len,
payload_offset: HEADER_LEN + extra_len,
});
}
cursor = value + slen;
}
Err(bgzf_error("gzip member has no BC extra subfield"))
}
fn bgzf_error(message: &'static str) -> Error {
Error::Other(format!("BGZF: {message}"))
}
struct Deflater {
#[cfg(not(feature = "libdeflate"))]
level: CompressionLevel,
#[cfg(feature = "libdeflate")]
compressor: libdeflater::Compressor,
}
impl Deflater {
fn new(level: CompressionLevel) -> Deflater {
Deflater {
#[cfg(not(feature = "libdeflate"))]
level,
#[cfg(feature = "libdeflate")]
compressor: libdeflater::Compressor::new(libdeflate_level(level)),
}
}
#[cfg(feature = "libdeflate")]
fn deflate(&mut self, data: &[u8]) -> io::Result<Vec<u8>> {
let mut out = vec![0u8; self.compressor.deflate_compress_bound(data.len())];
let written = self
.compressor
.deflate_compress(data, &mut out)
.map_err(|e| io::Error::other(format!("libdeflate: {e}")))?;
out.truncate(written);
Ok(out)
}
#[cfg(not(feature = "libdeflate"))]
fn deflate(&mut self, data: &[u8]) -> io::Result<Vec<u8>> {
use flate2::write::DeflateEncoder;
let mut encoder = DeflateEncoder::new(
Vec::with_capacity(data.len() / 2 + 64),
flate2::Compression::new(self.level.0.min(9)),
);
encoder.write_all(data)?;
encoder.finish()
}
}
#[cfg(feature = "libdeflate")]
fn libdeflate_level(level: CompressionLevel) -> libdeflater::CompressionLvl {
let mapped = match level.0 {
0 => 0,
level => ((level.min(9) as i32 - 1) * 11 / 8 + 1).min(12),
};
libdeflater::CompressionLvl::new(mapped)
.unwrap_or_else(|_| libdeflater::CompressionLvl::default())
}
struct Inflater {
#[cfg(feature = "libdeflate")]
decompressor: libdeflater::Decompressor,
}
impl Inflater {
fn new() -> Inflater {
Inflater {
#[cfg(feature = "libdeflate")]
decompressor: libdeflater::Decompressor::new(),
}
}
#[cfg(feature = "libdeflate")]
fn inflate(&mut self, payload: &[u8], expected: usize, out: &mut Vec<u8>) -> Result<()> {
out.clear();
if expected == 0 {
return Ok(());
}
out.resize(expected, 0);
let written = self
.decompressor
.deflate_decompress(payload, out)
.map_err(|e| Error::Other(format!("BGZF: corrupt block: {e}")))?;
if written != expected {
return Err(bgzf_error("block size does not match its ISIZE field"));
}
Ok(())
}
#[cfg(not(feature = "libdeflate"))]
fn inflate(&mut self, payload: &[u8], expected: usize, out: &mut Vec<u8>) -> Result<()> {
out.clear();
out.reserve(expected);
let mut decoder = flate2::Decompress::new(false);
decoder
.decompress_vec(payload, out, flate2::FlushDecompress::Finish)
.map_err(|e| Error::Other(format!("BGZF: corrupt block: {e}")))?;
if out.len() != expected {
return Err(bgzf_error("block size does not match its ISIZE field"));
}
Ok(())
}
}
fn inflate(payload: &[u8], expected: usize, out: &mut Vec<u8>) -> Result<()> {
Inflater::new().inflate(payload, expected, out)
}
fn crc32(data: &[u8]) -> u32 {
#[cfg(feature = "libdeflate")]
{
let mut crc = libdeflater::Crc::new();
crc.update(data);
crc.sum()
}
#[cfg(not(feature = "libdeflate"))]
{
let mut crc = flate2::Crc::new();
crc.update(data);
crc.sum()
}
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct GziIndex {
blocks: Vec<BlockOffset>,
total_uncompressed: Option<u64>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct BlockOffset {
pub compressed: u64,
pub uncompressed: u64,
}
impl GziIndex {
pub fn build<R: Read + Seek>(mut reader: R) -> Result<GziIndex> {
reader.seek(SeekFrom::Start(0))?;
let mut reader = BufReader::with_capacity(64 * 1024, reader);
let mut blocks = Vec::new();
let mut compressed = 0u64;
let mut uncompressed = 0u64;
let mut header = [0u8; HEADER_LEN + 64];
loop {
if !read_exact_or_eof(&mut reader, &mut header[..HEADER_LEN])? {
break;
}
let extra_len = u16::from_le_bytes([header[10], header[11]]) as usize;
if HEADER_LEN + extra_len > header.len() {
return Err(bgzf_error("extra field is implausibly large"));
}
reader.read_exact(&mut header[HEADER_LEN..HEADER_LEN + extra_len])?;
let block = parse_block_header(&header[..HEADER_LEN + extra_len])?;
let skip = block.compressed_len - block.payload_offset - 4;
io::copy(&mut reader.by_ref().take(skip as u64), &mut io::sink())?;
let mut isize_bytes = [0u8; 4];
reader.read_exact(&mut isize_bytes)?;
let payload_len = u32::from_le_bytes(isize_bytes) as u64;
blocks.push(BlockOffset {
compressed,
uncompressed,
});
compressed += block.compressed_len as u64;
uncompressed += payload_len;
if payload_len == 0 {
break;
}
}
Ok(GziIndex {
blocks,
total_uncompressed: Some(uncompressed),
})
}
pub fn build_from_path<P: AsRef<Path>>(path: P) -> Result<GziIndex> {
let path = path.as_ref();
GziIndex::build(
File::open(path).map_err(|e| {
Error::Io(io::Error::new(e.kind(), format!("{}: {e}", path.display())))
})?,
)
}
pub fn parse<R: Read>(mut reader: R) -> Result<GziIndex> {
let mut count_bytes = [0u8; 8];
reader.read_exact(&mut count_bytes)?;
let count = u64::from_le_bytes(count_bytes);
if count > 1 << 32 {
return Err(bgzf_error("index claims an implausible number of blocks"));
}
let mut blocks = Vec::with_capacity(count as usize + 1);
blocks.push(BlockOffset {
compressed: 0,
uncompressed: 0,
});
let mut pair = [0u8; 16];
for _ in 0..count {
reader.read_exact(&mut pair)?;
blocks.push(BlockOffset {
compressed: u64::from_le_bytes(pair[..8].try_into().expect("8 bytes")),
uncompressed: u64::from_le_bytes(pair[8..].try_into().expect("8 bytes")),
});
}
Ok(GziIndex {
blocks,
total_uncompressed: None,
})
}
pub fn from_path<P: AsRef<Path>>(path: P) -> Result<GziIndex> {
let path = path.as_ref();
GziIndex::parse(BufReader::new(File::open(path).map_err(|e| {
Error::Io(io::Error::new(e.kind(), format!("{}: {e}", path.display())))
})?))
}
pub fn write<W: Write>(&self, out: &mut W) -> Result<()> {
let count = self.blocks.len().saturating_sub(1) as u64;
out.write_all(&count.to_le_bytes())?;
for block in self.blocks.iter().skip(1) {
out.write_all(&block.compressed.to_le_bytes())?;
out.write_all(&block.uncompressed.to_le_bytes())?;
}
Ok(())
}
pub fn write_to_path<P: AsRef<Path>>(&self, bgzf_path: P) -> Result<PathBuf> {
let target = gzi_path(bgzf_path.as_ref());
let mut file = BufWriter::new(File::create(&target)?);
self.write(&mut file)?;
file.flush()?;
Ok(target)
}
pub fn blocks(&self) -> &[BlockOffset] {
&self.blocks
}
pub fn len(&self) -> usize {
self.blocks.len()
}
pub fn is_empty(&self) -> bool {
self.blocks.is_empty()
}
pub fn uncompressed_len(&self) -> Option<u64> {
self.total_uncompressed
}
fn block_for(&self, offset: u64) -> Option<BlockOffset> {
match self
.blocks
.binary_search_by(|b| b.uncompressed.cmp(&offset))
{
Ok(i) => Some(self.blocks[i]),
Err(0) => None,
Err(i) => Some(self.blocks[i - 1]),
}
}
}
pub fn gzi_path(bgzf: &Path) -> PathBuf {
let mut name = bgzf.as_os_str().to_os_string();
name.push(".gzi");
PathBuf::from(name)
}
pub fn is_bgzf(bytes: &[u8]) -> bool {
parse_block_header(bytes).is_ok()
}
pub struct BgzfReader<R: Read + Seek> {
inner: R,
index: Option<GziIndex>,
block: Vec<u8>,
block_pos: usize,
block_start: u64,
next_compressed: u64,
eof: bool,
raw: Vec<u8>,
}
impl BgzfReader<File> {
pub fn open<P: AsRef<Path>>(path: P) -> Result<BgzfReader<File>> {
let path = path.as_ref();
let file = File::open(path)
.map_err(|e| Error::Io(io::Error::new(e.kind(), format!("{}: {e}", path.display()))))?;
let index = match GziIndex::from_path(gzi_path(path)) {
Ok(index) => Some(index),
Err(Error::Io(e)) if e.kind() == io::ErrorKind::NotFound => None,
Err(e) => return Err(e),
};
let mut reader = BgzfReader::new(file)?;
reader.index = index;
Ok(reader)
}
}
impl<R: Read + Seek> BgzfReader<R> {
pub fn new(mut inner: R) -> Result<BgzfReader<R>> {
inner.seek(SeekFrom::Start(0))?;
let mut probe = [0u8; HEADER_LEN + 64];
let read = read_up_to(&mut inner, &mut probe)?;
if read == 0 {
return Err(bgzf_error("file is empty"));
}
parse_block_header(&probe[..read])?;
inner.seek(SeekFrom::Start(0))?;
Ok(BgzfReader {
inner,
index: None,
block: Vec::new(),
block_pos: 0,
block_start: 0,
next_compressed: 0,
eof: false,
raw: Vec::new(),
})
}
pub fn with_index(mut self, index: GziIndex) -> Self {
self.index = Some(index);
self
}
pub fn index(&self) -> Option<&GziIndex> {
self.index.as_ref()
}
pub fn position(&self) -> u64 {
self.block_start + self.block_pos as u64
}
pub fn into_inner(self) -> R {
self.inner
}
fn load_block_at(&mut self, at: u64, uncompressed_start: u64) -> Result<()> {
self.inner.seek(SeekFrom::Start(at))?;
self.next_compressed = at;
self.block_start = uncompressed_start;
self.block_pos = 0;
self.block.clear();
self.eof = false;
self.read_next_block()
}
fn read_next_block(&mut self) -> Result<()> {
let mut header = [0u8; HEADER_LEN + 64];
if !read_exact_or_eof(&mut self.inner, &mut header[..HEADER_LEN])? {
self.eof = true;
self.block.clear();
self.block_pos = 0;
return Ok(());
}
let extra_len = u16::from_le_bytes([header[10], header[11]]) as usize;
if HEADER_LEN + extra_len > header.len() {
return Err(bgzf_error("extra field is implausibly large"));
}
self.inner
.read_exact(&mut header[HEADER_LEN..HEADER_LEN + extra_len])?;
let block = parse_block_header(&header[..HEADER_LEN + extra_len])?;
let payload_len = block.compressed_len - block.payload_offset - TRAILER_LEN;
self.raw.resize(payload_len, 0);
self.inner.read_exact(&mut self.raw)?;
let mut trailer = [0u8; TRAILER_LEN];
self.inner.read_exact(&mut trailer)?;
let expected_crc = u32::from_le_bytes(trailer[..4].try_into().expect("4 bytes"));
let expected_len = u32::from_le_bytes(trailer[4..].try_into().expect("4 bytes")) as usize;
let mut decompressed = std::mem::take(&mut self.block);
let result = inflate(&self.raw, expected_len, &mut decompressed);
self.block = decompressed;
result?;
if crc32(&self.block) != expected_crc {
return Err(bgzf_error("block CRC32 does not match"));
}
self.block_pos = 0;
self.next_compressed += block.compressed_len as u64;
if self.block.is_empty() {
self.eof = true;
}
Ok(())
}
fn fill(&mut self) -> Result<()> {
while !self.eof && self.block_pos == self.block.len() {
let consumed = self.block.len() as u64;
self.block_start += consumed;
self.read_next_block()?;
}
Ok(())
}
fn ensure_index(&mut self) -> Result<()> {
if self.index.is_none() {
let saved = self.inner.stream_position()?;
let index = GziIndex::build(&mut self.inner)?;
self.inner.seek(SeekFrom::Start(saved))?;
self.index = Some(index);
}
Ok(())
}
}
impl<R: Read + Seek> Read for BgzfReader<R> {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
if buf.is_empty() {
return Ok(0);
}
self.fill()?;
if self.eof && self.block_pos == self.block.len() {
return Ok(0);
}
let available = &self.block[self.block_pos..];
let take = available.len().min(buf.len());
buf[..take].copy_from_slice(&available[..take]);
self.block_pos += take;
Ok(take)
}
}
impl<R: Read + Seek> Seek for BgzfReader<R> {
fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
self.ensure_index()?;
let target = match pos {
SeekFrom::Start(offset) => offset,
SeekFrom::Current(delta) => add_signed(self.position(), delta)?,
SeekFrom::End(delta) => {
let end = self
.index
.as_ref()
.and_then(|index| index.uncompressed_len())
.ok_or_else(|| {
io::Error::new(
io::ErrorKind::InvalidInput,
"BGZF: this index records block starts only, so the end of the \
stream is unknown; rebuild it with GziIndex::build",
)
})?;
add_signed(end, delta)?
}
};
let within = target.checked_sub(self.block_start);
if let Some(within) = within {
if !self.block.is_empty() && within <= self.block.len() as u64 {
self.block_pos = within as usize;
return Ok(target);
}
}
let block = self
.index
.as_ref()
.and_then(|index| index.block_for(target))
.ok_or_else(|| {
io::Error::new(
io::ErrorKind::InvalidInput,
"BGZF: no block covers that offset",
)
})?;
self.load_block_at(block.compressed, block.uncompressed)?;
let within = (target - block.uncompressed) as usize;
if within > self.block.len() {
self.block_pos = self.block.len();
self.eof = true;
} else {
self.block_pos = within;
}
Ok(target)
}
}
fn add_signed(base: u64, delta: i64) -> io::Result<u64> {
let result = if delta >= 0 {
base.checked_add(delta as u64)
} else {
base.checked_sub(delta.unsigned_abs())
};
result.ok_or_else(|| {
io::Error::new(
io::ErrorKind::InvalidInput,
"BGZF: seek would leave the file",
)
})
}
#[cfg(feature = "parallel")]
struct RawBlock {
payload: Vec<u8>,
crc: u32,
uncompressed_len: usize,
}
#[cfg(feature = "parallel")]
fn read_raw_block<R: Read>(reader: &mut R) -> Result<Option<RawBlock>> {
let mut header = [0u8; HEADER_LEN + 64];
if !read_exact_or_eof(reader, &mut header[..HEADER_LEN])? {
return Ok(None);
}
let extra_len = u16::from_le_bytes([header[10], header[11]]) as usize;
if HEADER_LEN + extra_len > header.len() {
return Err(bgzf_error("extra field is implausibly large"));
}
reader.read_exact(&mut header[HEADER_LEN..HEADER_LEN + extra_len])?;
let block = parse_block_header(&header[..HEADER_LEN + extra_len])?;
let payload_len = block.compressed_len - block.payload_offset - TRAILER_LEN;
let mut payload = vec![0u8; payload_len];
reader.read_exact(&mut payload)?;
let mut trailer = [0u8; TRAILER_LEN];
reader.read_exact(&mut trailer)?;
Ok(Some(RawBlock {
payload,
crc: u32::from_le_bytes(trailer[..4].try_into().expect("4 bytes")),
uncompressed_len: u32::from_le_bytes(trailer[4..].try_into().expect("4 bytes")) as usize,
}))
}
#[cfg(feature = "parallel")]
fn inflate_checked(inflater: &mut Inflater, raw: &RawBlock) -> Result<Vec<u8>> {
let mut out = Vec::new();
inflater.inflate(&raw.payload, raw.uncompressed_len, &mut out)?;
if crc32(&out) != raw.crc {
return Err(bgzf_error("block CRC32 does not match"));
}
Ok(out)
}
#[cfg(feature = "parallel")]
pub struct ParallelBgzfReader<R: Read> {
inner: R,
batch: usize,
buffer: Vec<u8>,
pos: usize,
eof: bool,
}
#[cfg(feature = "parallel")]
impl<R: Read> ParallelBgzfReader<R> {
pub fn new(inner: R) -> ParallelBgzfReader<R> {
ParallelBgzfReader {
inner,
batch: default_batch().max(2),
buffer: Vec::new(),
pos: 0,
eof: false,
}
}
pub fn blocks_per_batch(mut self, blocks: usize) -> Self {
self.batch = blocks.max(1);
self
}
pub fn into_inner(self) -> R {
self.inner
}
fn fill(&mut self) -> Result<()> {
use rayon::prelude::*;
self.buffer.clear();
self.pos = 0;
let mut raws: Vec<RawBlock> = Vec::with_capacity(self.batch);
while raws.len() < self.batch {
match read_raw_block(&mut self.inner)? {
None => {
self.eof = true;
break;
}
Some(raw) if raw.uncompressed_len == 0 => {
self.eof = true;
break;
}
Some(raw) => raws.push(raw),
}
}
if raws.is_empty() {
return Ok(());
}
let blocks: Vec<Vec<u8>> = if raws.len() > 1 {
raws.par_iter()
.map_init(Inflater::new, inflate_checked)
.collect::<Result<_>>()?
} else {
vec![inflate_checked(&mut Inflater::new(), &raws[0])?]
};
self.buffer.reserve(blocks.iter().map(Vec::len).sum());
for block in &blocks {
self.buffer.extend_from_slice(block);
}
Ok(())
}
}
#[cfg(feature = "parallel")]
impl<R: Read> Read for ParallelBgzfReader<R> {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
if buf.is_empty() {
return Ok(0);
}
while self.pos == self.buffer.len() {
if self.eof {
return Ok(0);
}
self.fill()?;
}
let available = &self.buffer[self.pos..];
let take = available.len().min(buf.len());
buf[..take].copy_from_slice(&available[..take]);
self.pos += take;
Ok(take)
}
}
pub struct BgzfWriter<W: Write> {
inner: Option<W>,
buffer: Vec<u8>,
pending: Vec<Vec<u8>>,
recycled: Vec<Vec<u8>>,
batch: usize,
level: CompressionLevel,
blocks: Vec<BlockOffset>,
compressed: u64,
uncompressed: u64,
finished: bool,
}
struct EncodedBlock {
header: [u8; HEADER_LEN + EXTRA_LEN],
payload: Vec<u8>,
trailer: [u8; TRAILER_LEN],
block_len: usize,
uncompressed_len: usize,
}
fn encode_block(deflater: &mut Deflater, data: &[u8]) -> io::Result<EncodedBlock> {
let payload = deflater.deflate(data)?;
let block_len = HEADER_LEN + EXTRA_LEN + payload.len() + TRAILER_LEN;
if block_len > u16::MAX as usize + 1 {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"BGZF: block would exceed 64 KiB",
));
}
let mut header = [0u8; HEADER_LEN + EXTRA_LEN];
header[0] = 0x1f;
header[1] = 0x8b;
header[2] = 8; header[3] = 4; header[9] = 0xff; header[10..12].copy_from_slice(&(EXTRA_LEN as u16).to_le_bytes());
header[12] = b'B';
header[13] = b'C';
header[14..16].copy_from_slice(&2u16.to_le_bytes());
header[16..18].copy_from_slice(&((block_len - 1) as u16).to_le_bytes());
let mut trailer = [0u8; TRAILER_LEN];
trailer[..4].copy_from_slice(&crc32(data).to_le_bytes());
trailer[4..].copy_from_slice(&(data.len() as u32).to_le_bytes());
Ok(EncodedBlock {
header,
payload,
trailer,
block_len,
uncompressed_len: data.len(),
})
}
fn encode_batch(payloads: &[Vec<u8>], level: CompressionLevel) -> io::Result<Vec<EncodedBlock>> {
#[cfg(feature = "parallel")]
{
if payloads.len() > 1 {
use rayon::prelude::*;
return payloads
.par_iter()
.map_init(
|| Deflater::new(level),
|deflater, payload| encode_block(deflater, payload),
)
.collect();
}
}
let mut deflater = Deflater::new(level);
payloads
.iter()
.map(|payload| encode_block(&mut deflater, payload))
.collect()
}
#[cfg(feature = "parallel")]
const BATCHES_PER_CORE: usize = 8;
fn default_batch() -> usize {
#[cfg(feature = "parallel")]
{
std::thread::available_parallelism().map_or(1, |n| n.get() * BATCHES_PER_CORE)
}
#[cfg(not(feature = "parallel"))]
{
1
}
}
impl BgzfWriter<BufWriter<File>> {
pub fn create<P: AsRef<Path>>(path: P) -> Result<BgzfWriter<BufWriter<File>>> {
let path = path.as_ref();
let file = File::create(path)
.map_err(|e| Error::Io(io::Error::new(e.kind(), format!("{}: {e}", path.display()))))?;
Ok(BgzfWriter::new(BufWriter::with_capacity(128 * 1024, file)))
}
}
impl<W: Write> BgzfWriter<W> {
pub fn new(inner: W) -> BgzfWriter<W> {
BgzfWriter::with_level(inner, CompressionLevel::default())
}
pub fn with_level(inner: W, level: CompressionLevel) -> BgzfWriter<W> {
BgzfWriter {
inner: Some(inner),
buffer: Vec::with_capacity(MAX_BLOCK_PAYLOAD),
pending: Vec::new(),
recycled: Vec::new(),
batch: default_batch(),
level,
blocks: vec![BlockOffset {
compressed: 0,
uncompressed: 0,
}],
compressed: 0,
uncompressed: 0,
finished: false,
}
}
pub fn get_ref(&self) -> Option<&W> {
self.inner.as_ref()
}
pub fn blocks_per_batch(mut self, blocks: usize) -> Self {
self.batch = blocks.max(1);
self
}
pub fn index(&self) -> GziIndex {
GziIndex {
blocks: self.blocks.clone(),
total_uncompressed: Some(self.uncompressed),
}
}
fn sink(&mut self) -> io::Result<&mut W> {
self.inner.as_mut().ok_or_else(|| {
io::Error::new(
io::ErrorKind::BrokenPipe,
"BGZF: writer was already finished",
)
})
}
fn seal_buffer(&mut self) -> io::Result<()> {
if self.buffer.is_empty() {
return Ok(());
}
let mut fresh = self.recycled.pop().unwrap_or_else(|| {
let mut buffer = Vec::new();
buffer.reserve_exact(MAX_BLOCK_PAYLOAD);
buffer
});
std::mem::swap(&mut self.buffer, &mut fresh);
self.pending.push(fresh);
if self.pending.len() >= self.batch {
self.compress_pending()?;
}
Ok(())
}
fn compress_pending(&mut self) -> io::Result<()> {
if self.pending.is_empty() {
return Ok(());
}
let level = self.level;
let encoded = encode_batch(&self.pending, level)?;
for block in encoded {
self.emit(block)?;
}
for mut buffer in self.pending.drain(..) {
buffer.clear();
if self.recycled.len() < self.batch {
self.recycled.push(buffer);
}
}
Ok(())
}
fn emit(&mut self, block: EncodedBlock) -> io::Result<()> {
let sink = self.sink()?;
sink.write_all(&block.header)?;
sink.write_all(&block.payload)?;
sink.write_all(&block.trailer)?;
self.compressed += block.block_len as u64;
self.uncompressed += block.uncompressed_len as u64;
self.blocks.push(BlockOffset {
compressed: self.compressed,
uncompressed: self.uncompressed,
});
Ok(())
}
fn flush_block(&mut self) -> io::Result<()> {
self.seal_buffer()?;
self.compress_pending()
}
pub fn finish(self) -> Result<W> {
self.finish_with_index().map(|(inner, _)| inner)
}
pub fn finish_with_index(mut self) -> Result<(W, GziIndex)> {
self.finish_in_place()?;
let index = self.index();
let inner = self
.inner
.take()
.ok_or_else(|| Error::Other("BGZF: writer was already finished".to_string()))?;
Ok((inner, index))
}
fn finish_in_place(&mut self) -> Result<()> {
if self.finished || self.inner.is_none() {
return Ok(());
}
self.flush_block()?;
let sink = self.sink()?;
sink.write_all(&EOF_BLOCK)?;
sink.flush()?;
self.compressed += EOF_BLOCK.len() as u64;
self.finished = true;
Ok(())
}
}
impl<W: Write> Write for BgzfWriter<W> {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
let room = MAX_BLOCK_PAYLOAD - self.buffer.len();
let take = room.min(buf.len());
self.buffer.extend_from_slice(&buf[..take]);
if self.buffer.len() == MAX_BLOCK_PAYLOAD {
self.seal_buffer()?;
}
Ok(take)
}
fn flush(&mut self) -> io::Result<()> {
self.flush_block()?;
self.sink()?.flush()
}
}
impl<W: Write> Drop for BgzfWriter<W> {
fn drop(&mut self) {
let _ = self.finish_in_place();
}
}
fn read_exact_or_eof<R: Read>(reader: &mut R, buf: &mut [u8]) -> Result<bool> {
let mut filled = 0;
while filled < buf.len() {
match reader.read(&mut buf[filled..]) {
Ok(0) if filled == 0 => return Ok(false),
Ok(0) => return Err(bgzf_error("file ends in the middle of a block")),
Ok(n) => filled += n,
Err(e) if e.kind() == io::ErrorKind::Interrupted => continue,
Err(e) => return Err(Error::Io(e)),
}
}
Ok(true)
}
fn read_up_to<R: Read>(reader: &mut R, buf: &mut [u8]) -> Result<usize> {
let mut filled = 0;
while filled < buf.len() {
match reader.read(&mut buf[filled..]) {
Ok(0) => break,
Ok(n) => filled += n,
Err(e) if e.kind() == io::ErrorKind::Interrupted => continue,
Err(e) => return Err(Error::Io(e)),
}
}
Ok(filled)
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::Cursor;
fn compress(data: &[u8]) -> Vec<u8> {
let mut writer = BgzfWriter::new(Vec::new());
writer.write_all(data).unwrap();
writer.finish().unwrap()
}
#[test]
fn round_trips_through_our_own_reader() {
for size in [
0usize,
1,
100,
MAX_BLOCK_PAYLOAD - 1,
MAX_BLOCK_PAYLOAD,
MAX_BLOCK_PAYLOAD + 1,
300_000,
] {
let data: Vec<u8> = (0..size).map(|i| b"ACGTN"[i % 5]).collect();
let compressed = compress(&data);
assert!(is_bgzf(&compressed), "size {size} did not produce BGZF");
let mut reader = BgzfReader::new(Cursor::new(&compressed)).unwrap();
let mut out = Vec::new();
reader.read_to_end(&mut out).unwrap();
assert_eq!(out, data, "size {size}");
}
}
#[test]
fn output_is_plain_gzip_too() {
let data: Vec<u8> = (0..200_000).map(|i| b"ACGT"[i % 4]).collect();
let compressed = compress(&data);
let mut out = Vec::new();
flate2::read::MultiGzDecoder::new(&compressed[..])
.read_to_end(&mut out)
.unwrap();
assert_eq!(out, data);
}
#[test]
fn ends_with_the_eof_marker() {
let compressed = compress(b"ACGT");
assert_eq!(
&compressed[compressed.len() - EOF_BLOCK.len()..],
&EOF_BLOCK
);
assert_eq!(compress(b""), EOF_BLOCK.to_vec());
}
#[test]
fn blocks_stay_within_the_size_limit() {
let data: Vec<u8> = (0..500_000)
.map(|i| ((i * 2_654_435_761u64 as usize) >> 7) as u8)
.collect();
let compressed = compress(&data);
let index = GziIndex::build(Cursor::new(&compressed)).unwrap();
for pair in index.blocks().windows(2) {
let block_len = pair[1].compressed - pair[0].compressed;
assert!(block_len <= 65_536, "block of {block_len} bytes");
}
let mut out = Vec::new();
BgzfReader::new(Cursor::new(&compressed))
.unwrap()
.read_to_end(&mut out)
.unwrap();
assert_eq!(out, data);
}
#[test]
fn index_from_the_writer_matches_a_rescan() {
let data: Vec<u8> = (0..250_000).map(|i| b"ACGTN"[i % 5]).collect();
let mut writer = BgzfWriter::new(Vec::new());
writer.write_all(&data).unwrap();
let from_writer = writer.index();
let compressed = writer.finish().unwrap();
let from_scan = GziIndex::build(Cursor::new(&compressed)).unwrap();
assert_eq!(
&from_scan.blocks()[..from_writer.len()],
from_writer.blocks()
);
assert_eq!(from_scan.uncompressed_len(), Some(data.len() as u64));
}
#[test]
fn gzi_round_trips_and_omits_the_first_block() {
let data: Vec<u8> = (0..200_000).map(|i| b"ACGT"[i % 4]).collect();
let compressed = compress(&data);
let index = GziIndex::build(Cursor::new(&compressed)).unwrap();
let mut text = Vec::new();
index.write(&mut text).unwrap();
assert_eq!(text.len(), 8 + 16 * (index.len() - 1));
assert_eq!(
u64::from_le_bytes(text[..8].try_into().unwrap()),
index.len() as u64 - 1
);
let reparsed = GziIndex::parse(&text[..]).unwrap();
assert_eq!(reparsed.blocks(), index.blocks());
assert_eq!(reparsed.uncompressed_len(), None);
assert_eq!(index.uncompressed_len(), Some(200_000));
}
#[test]
fn seek_from_end_needs_a_scanned_index() {
let data: Vec<u8> = (0..200_000).map(|i| b"ACGT"[i % 4]).collect();
let compressed = compress(&data);
let scanned = GziIndex::build(Cursor::new(&compressed)).unwrap();
let mut text = Vec::new();
scanned.write(&mut text).unwrap();
let parsed = GziIndex::parse(&text[..]).unwrap();
let mut reader = BgzfReader::new(Cursor::new(&compressed))
.unwrap()
.with_index(parsed);
assert!(reader.seek(SeekFrom::End(0)).is_err());
reader.seek(SeekFrom::Start(199_998)).unwrap();
let mut tail = Vec::new();
reader.read_to_end(&mut tail).unwrap();
assert_eq!(tail, &data[199_998..]);
let mut reader = BgzfReader::new(Cursor::new(&compressed))
.unwrap()
.with_index(scanned);
assert_eq!(reader.seek(SeekFrom::End(0)).unwrap(), 200_000);
}
#[test]
fn seeks_to_any_offset() {
let data: Vec<u8> = (0..300_000).map(|i| (i % 251) as u8).collect();
let compressed = compress(&data);
let mut reader = BgzfReader::new(Cursor::new(&compressed)).unwrap();
for target in [
0usize,
1,
MAX_BLOCK_PAYLOAD - 1,
MAX_BLOCK_PAYLOAD,
MAX_BLOCK_PAYLOAD + 1,
2 * MAX_BLOCK_PAYLOAD,
299_999,
] {
reader.seek(SeekFrom::Start(target as u64)).unwrap();
assert_eq!(reader.position(), target as u64);
let mut buf = [0u8; 8];
let want = (data.len() - target).min(buf.len());
reader.read_exact(&mut buf[..want]).unwrap();
assert_eq!(&buf[..want], &data[target..target + want], "at {target}");
}
reader.seek(SeekFrom::Start(10)).unwrap();
reader.seek(SeekFrom::Current(5)).unwrap();
assert_eq!(reader.position(), 15);
assert_eq!(reader.seek(SeekFrom::End(0)).unwrap(), data.len() as u64);
let mut rest = Vec::new();
reader.read_to_end(&mut rest).unwrap();
assert!(rest.is_empty());
reader.seek(SeekFrom::Start(7)).unwrap();
let mut buf = [0u8; 4];
reader.read_exact(&mut buf).unwrap();
assert_eq!(&buf, &data[7..11]);
}
#[test]
fn seek_before_the_start_is_an_error() {
let compressed = compress(b"ACGT");
let mut reader = BgzfReader::new(Cursor::new(&compressed)).unwrap();
assert!(reader.seek(SeekFrom::Current(-1)).is_err());
assert!(reader.seek(SeekFrom::End(-100)).is_err());
}
#[test]
fn rejects_plain_gzip_and_garbage() {
let mut plain = Vec::new();
{
let mut encoder =
flate2::write::GzEncoder::new(&mut plain, flate2::Compression::default());
encoder.write_all(b"ACGT").unwrap();
encoder.finish().unwrap();
}
assert!(!is_bgzf(&plain));
assert!(BgzfReader::new(Cursor::new(&plain)).is_err());
assert!(BgzfReader::new(Cursor::new(b"not gzip at all".to_vec())).is_err());
assert!(BgzfReader::new(Cursor::new(Vec::new())).is_err());
}
#[test]
fn detects_a_corrupt_block() {
let mut compressed = compress(&vec![b'A'; 5_000]);
let victim = HEADER_LEN + EXTRA_LEN + 5;
compressed[victim] ^= 0xff;
let mut out = Vec::new();
let result = BgzfReader::new(Cursor::new(&compressed))
.unwrap()
.read_to_end(&mut out);
assert!(result.is_err(), "corruption went unnoticed");
}
#[test]
fn truncated_file_is_an_error_not_silent_truncation() {
let compressed = compress(&vec![b'A'; 200_000]);
let cut = compressed.len() / 2;
let mut out = Vec::new();
let result = BgzfReader::new(Cursor::new(compressed[..cut].to_vec()))
.unwrap()
.read_to_end(&mut out);
assert!(result.is_err(), "truncation went unnoticed");
}
#[test]
fn parse_rejects_an_implausible_index() {
let mut bad = u64::MAX.to_le_bytes().to_vec();
bad.extend_from_slice(&[0u8; 16]);
assert!(GziIndex::parse(&bad[..]).is_err());
assert!(GziIndex::parse(&[0u8; 3][..]).is_err());
}
#[test]
fn batch_size_never_changes_the_bytes() {
let data: Vec<u8> = (0..400_000).map(|i| b"ACGTN"[i % 5]).collect();
let mut reference = None;
for blocks in [1usize, 2, 3, 7, 64, 1024] {
let mut writer = BgzfWriter::new(Vec::new()).blocks_per_batch(blocks);
writer.write_all(&data).unwrap();
let (bytes, index) = writer.finish_with_index().unwrap();
match &reference {
None => reference = Some((bytes, index)),
Some((expected_bytes, expected_index)) => {
assert_eq!(&bytes, expected_bytes, "batch of {blocks} differs");
assert_eq!(index.blocks(), expected_index.blocks(), "index differs");
}
}
}
let (bytes, _) = reference.unwrap();
let mut out = Vec::new();
BgzfReader::new(Cursor::new(&bytes))
.unwrap()
.read_to_end(&mut out)
.unwrap();
assert_eq!(out, data);
}
#[test]
fn write_flush_boundaries_survive_batching() {
let mut writer = BgzfWriter::new(Vec::new()).blocks_per_batch(16);
writer.write_all(b"first").unwrap();
writer.flush().unwrap();
assert!(
!writer.get_ref().expect("still open").is_empty(),
"flush left the batch uncompressed"
);
writer.write_all(b"second").unwrap();
let (bytes, index) = writer.finish_with_index().unwrap();
assert_eq!(index.len(), 3); let mut out = Vec::new();
BgzfReader::new(Cursor::new(&bytes))
.unwrap()
.read_to_end(&mut out)
.unwrap();
assert_eq!(out, b"firstsecond");
}
#[cfg(feature = "parallel")]
#[test]
fn parallel_reader_agrees_with_the_serial_one() {
for size in [
0usize,
1,
MAX_BLOCK_PAYLOAD - 1,
MAX_BLOCK_PAYLOAD,
MAX_BLOCK_PAYLOAD + 1,
500_000,
] {
let data: Vec<u8> = (0..size).map(|i| (i % 251) as u8).collect();
let compressed = compress(&data);
for blocks in [1usize, 2, 5, 64] {
let mut parallel =
ParallelBgzfReader::new(Cursor::new(&compressed)).blocks_per_batch(blocks);
let mut out = Vec::new();
parallel.read_to_end(&mut out).unwrap();
assert_eq!(out, data, "size {size}, batch {blocks}");
}
let mut parallel = ParallelBgzfReader::new(Cursor::new(&compressed));
let mut out = Vec::new();
let mut byte = [0u8; 1];
while parallel.read(&mut byte).unwrap() == 1 {
out.push(byte[0]);
}
assert_eq!(out, data, "size {size}, byte at a time");
}
}
#[cfg(feature = "parallel")]
#[test]
fn parallel_reader_still_catches_corruption() {
let mut compressed = compress(&vec![b'A'; 300_000]);
let victim = compressed.len() / 2;
compressed[victim] ^= 0xff;
let mut out = Vec::new();
let result = ParallelBgzfReader::new(Cursor::new(&compressed)).read_to_end(&mut out);
assert!(result.is_err(), "corruption went unnoticed");
let whole = compress(&vec![b'C'; 300_000]);
let mut out = Vec::new();
let result = ParallelBgzfReader::new(Cursor::new(whole[..whole.len() / 2].to_vec()))
.read_to_end(&mut out);
assert!(result.is_err(), "truncation went unnoticed");
}
#[test]
fn flush_starts_a_new_block() {
let mut writer = BgzfWriter::new(Vec::new());
writer.write_all(b"first").unwrap();
writer.flush().unwrap();
writer.write_all(b"second").unwrap();
assert_eq!(writer.index().len(), 2);
let (compressed, index) = writer.finish_with_index().unwrap();
assert_eq!(index.len(), 3); assert_eq!(index.blocks()[1].uncompressed, 5);
assert_eq!(index.uncompressed_len(), Some(11));
let mut reader = BgzfReader::new(Cursor::new(&compressed)).unwrap();
reader.seek(SeekFrom::Start(5)).unwrap();
let mut out = Vec::new();
reader.read_to_end(&mut out).unwrap();
assert_eq!(out, b"second");
}
}