lbzip2/src/chunk.rs
//! Chunk-level parallel bzip2 decoder.
//!
//! Designed for zero-copy operation: the caller reads ~100 MB of
//! compressed data into a chunk buffer, then passes the `&[u8]`
//! directly to the decoder — no per-block copy.
//!
//! The decoder scans for block boundaries, parallel-decodes all
//! complete blocks, and returns:
//! - `decompressed`: the concatenated output (~800–1000 MB)
//! - `consumed`: how many bytes were fully decoded
//!
//! The caller carries `data[consumed..]` into the next slot.
use crate::bitreader::BitReader;
use crate::block;
use crate::block_scan;
use crate::limits::OutputBudget;
use crate::{BLOCK_MAGIC, FINAL_MAGIC};
/// Which inverse-BWT / RLE1 chase strategy a segment decode uses.
///
/// GATING: [`DecodeMode::Interleaved`] runs several blocks' inverse-BWT pointer
/// chases in lockstep to overlap their dependent-load latency (memory-level
/// parallelism, see [`block::decode_blocks_interleaved`]). This is a large win on a
/// SINGLE core but neutral-to-harmful at high core counts (the socket is already
/// memory-bandwidth-bound), so the CLI selects it ONLY when decoding with one
/// worker. [`DecodeMode::Serial`] is the block-by-block baseline — the default and
/// the exact path every multi-worker and library decode takes (byte-identical to
/// before this option existed).
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum DecodeMode {
/// Block-by-block chase directly into the output (baseline; all multi-worker
/// and library decodes).
Serial,
/// Interleaved multi-block chase (single-worker CLI decode only).
Interleaved,
}
/// Result of splitting a chunk into segment boundaries.
pub struct ChunkSplit {
/// Segment start boundaries (one per segment found).
pub segment_starts: Vec<block_scan::BlockBoundary>,
/// Number of segments to decode (all if is_last, else n-1).
pub decode_segments: usize,
/// Bytes consumed from the data (for carry computation).
pub consumed: usize,
}
/// Split a chunk of compressed data into segment boundaries.
///
/// `n_segments` is typically n_cores (one segment per worker).
/// Uses parallel scanning (scoped-thread `par_map`) for speed.
///
/// Returns `None` if no block boundary found in data.
pub fn split_chunk(
data: &[u8],
n_segments: usize,
max_blocksize: u32,
is_last: bool,
) -> Option<ChunkSplit> {
let first_block = block_scan::find_next_block(data, 0)?;
let splits = block_scan::split_boundaries_parallel(data, n_segments, max_blocksize);
let mut segment_starts = Vec::with_capacity(n_segments);
segment_starts.push(first_block);
for s in &splits {
if segment_starts
.last()
.map_or(true, |prev: &block_scan::BlockBoundary| {
prev.bit_offset != s.bit_offset
})
{
segment_starts.push(*s);
}
}
let n = segment_starts.len();
let decode_segments = if is_last {
n
} else if n > 1 {
n - 1
} else {
return None;
};
let consumed = if decode_segments < n {
segment_starts[decode_segments].byte_offset()
} else {
data.len()
};
Some(ChunkSplit {
segment_starts,
decode_segments,
consumed,
})
}
/// A decoded segment: the concatenated block output plus the ordered list of
/// each decoded block's (verified) CRC-32.
///
/// The per-block CRCs are folded by the caller into the running combined stream
/// CRC so the whole-stream checksum can be verified after a parallel decode.
pub struct SegmentDecoded {
pub data: Vec<u8>,
/// One entry per decoded block, in stream order (already CRC-verified).
pub block_crcs: Vec<u32>,
}
/// Decode a single segment of compressed bzip2 data.
///
/// Reads blocks from `start_bit` (after the 48-bit BLOCK_MAGIC) up to `end_bit`.
/// Handles pbzip2 concatenated streams (FINAL_MAGIC boundaries).
///
/// Zero per-block allocation: decodes directly into the output buffer
/// using `decode_block_into()`. Only the segment Vec grows (amortised).
///
/// This is the **lenient** entry point kept for backward compatibility: on a
/// decode failure it returns the bytes decoded so far (never panics). Callers
/// that must REJECT a corrupt/tampered stream (the parallel library API and the
/// `lbunzip2` CLI) use [`decode_segment_checked`], which surfaces the per-block
/// CRC failure as an error instead of yielding silently-truncated output.
pub fn decode_segment(data: &[u8], start_bit: u64, end_bit: u64, max_blocksize: u32) -> Vec<u8> {
let mut output = Vec::with_capacity(seg_cap_hint(start_bit, end_bit));
let mut block_crcs: Vec<u32> = Vec::new();
decode_segment_inner(
data,
start_bit,
end_bit,
max_blocksize,
&mut output,
&mut block_crcs,
DecodeMode::Serial,
&OutputBudget::unlimited(),
);
output
}
/// Integrity-checked variant of [`decode_segment`].
///
/// SECURITY: a block whose stored CRC-32 does not match its decompressed bytes
/// (a truncated or tampered stream) is returned as a [`block::BlockError`] — it
/// is **never** silently accepted as partial output. This is the parallel-path
/// counterpart of the per-block check the sequential [`crate::stream::decompress`]
/// already performs. Structural framing markers (end-of-segment, FINAL_MAGIC,
/// EOF) still terminate the segment cleanly.
///
/// Returns the decompressed output plus the ordered per-block CRC list (which the
/// caller folds into the combined whole-stream CRC).
pub fn decode_segment_checked(
data: &[u8],
start_bit: u64,
end_bit: u64,
max_blocksize: u32,
) -> Result<SegmentDecoded, block::BlockError> {
decode_segment_checked_budgeted(
data,
start_bit,
end_bit,
max_blocksize,
&OutputBudget::unlimited(),
)
}
/// [`decode_segment_checked`] against a caller-supplied output budget.
///
/// The budget is **shared across every worker decoding the same stream** — see
/// [`crate::limits::OutputBudget`] for why a per-segment share cannot work — and
/// is charged once per decoded block. When it runs out the segment fails with
/// [`crate::limits::OUTPUT_TOO_LARGE`], so the whole decode stops at roughly the
/// ceiling plus one block per worker rather than continuing to the size the
/// stream claimed.
pub fn decode_segment_checked_budgeted(
data: &[u8],
start_bit: u64,
end_bit: u64,
max_blocksize: u32,
budget: &OutputBudget,
) -> Result<SegmentDecoded, block::BlockError> {
let mut output = Vec::with_capacity(seg_cap_hint(start_bit, end_bit).min(budget.remaining()));
let mut block_crcs: Vec<u32> = Vec::new();
match decode_segment_inner(
data,
start_bit,
end_bit,
max_blocksize,
&mut output,
&mut block_crcs,
DecodeMode::Serial,
budget,
) {
None => Ok(SegmentDecoded {
data: output,
block_crcs,
}),
Some(e) => Err(e),
}
}
/// Shared decode loop. Decodes as far as it can and returns the accumulated
/// output plus the first decode error encountered (if any). The lenient wrapper
/// discards the error (returning the partial output, preserving the historical
/// behaviour); the checked wrapper turns a present error into `Err`.
/// Estimated decoded size of a segment from its compressed bit-span (×4: bzip2 text
/// decompresses ~4–8×). Used to presize the output so `decode_block_into` appends
/// without doubling reallocs (each of which memcpys the whole buffer).
#[inline]
fn seg_cap_hint(start_bit: u64, end_bit: u64) -> usize {
let comp_span = (end_bit.saturating_sub(start_bit) / 8) as usize;
comp_span.saturating_mul(4).min(256 * 1024 * 1024)
}
/// Decode the segment into the caller-supplied `output`/`block_crcs` buffers,
/// appending. The caller owns (and may recycle) both — this fn never allocates them,
/// so a warmed buffer pool decodes with ZERO per-segment allocation. Returns the first
/// decode error (if any); on error the partial output is left in `output`.
fn decode_segment_inner(
data: &[u8],
start_bit: u64,
end_bit: u64,
max_blocksize: u32,
output: &mut Vec<u8>,
block_crcs: &mut Vec<u32>,
mode: DecodeMode,
budget: &OutputBudget,
) -> Option<block::BlockError> {
match mode {
// Single-worker CLI decode: interleave N blocks' inverse-BWT chases (MLP).
DecodeMode::Interleaved => decode_segment_interleaved(
data,
start_bit,
end_bit,
max_blocksize,
output,
block_crcs,
budget,
),
// Everything else: the block-by-block baseline (byte-identical to before).
DecodeMode::Serial => decode_segment_serial(
data,
start_bit,
end_bit,
max_blocksize,
output,
block_crcs,
budget,
),
}
}
/// Baseline block-by-block segment decode: each block is chased directly onto the
/// tail of `output`. This is the exact path all multi-worker and library decodes
/// take — unchanged by the single-thread MLP work.
fn decode_segment_serial(
data: &[u8],
start_bit: u64,
end_bit: u64,
max_blocksize: u32,
output: &mut Vec<u8>,
block_crcs: &mut Vec<u32>,
budget: &OutputBudget,
) -> Option<block::BlockError> {
let total_bits = data.len() as u64 * 8;
// Charge the shared budget for what a block actually appended. Once per
// block, so the check is free next to the inverse-BWT chase that produced
// the bytes; the overshoot before a refusal is at most one block.
let charge = |output: &Vec<u8>, before: usize| budget.take(output.len() - before);
let mut reader = BitReader::from_bit_offset(data, (start_bit + 48) as usize);
// The segment's first block MUST decode: a CRC mismatch or malformed block
// here is corruption, not a framing boundary.
let before = output.len();
match decode_block_into_vec(output, &mut reader, max_blocksize) {
Ok(crc) => block_crcs.push(crc),
Err(e) => return Some(e),
}
if let Err(e) = charge(output, before) {
return Some(e);
}
loop {
let pos = reader.position() as u64;
if pos + 48 > total_bits || pos >= end_bit {
break;
}
let magic = match reader.read_u64(48) {
Some(v) => v,
None => break,
};
if magic == BLOCK_MAGIC {
// A block magic we chose to decode: any decode failure (including a
// CRC mismatch) is corruption, recorded and surfaced by the checked
// wrapper rather than silently swallowed.
let before = output.len();
match decode_block_into_vec(output, &mut reader, max_blocksize) {
Ok(crc) => block_crcs.push(crc),
Err(e) => return Some(e),
}
if let Err(e) = charge(output, before) {
return Some(e);
}
} else if magic == FINAL_MAGIC {
if reader.read_u32(32).is_none() {
break;
}
let p = reader.position();
let pad = (8 - (p % 8)) % 8;
if pad > 0 {
BitReader::skip(&mut reader, pad);
}
match reader.read_u32(32) {
Some(h) => {
let b = h.to_be_bytes();
if &b[..3] != b"BZh" {
break;
}
}
None => break,
}
} else {
break;
}
}
None
}
/// Interleave factor for the single-thread memory-level-parallel decode: how many
/// blocks' inverse-BWT pointer-chases run in lockstep. The chains are independent,
/// so their dependent-load latencies overlap (see `block::decode_blocks_interleaved`).
/// N=4 measured best on this microarchitecture (single-thread cycles 78 B → 47 B);
/// N=2 leaves ~half the win on the table, N>4 adds working-set pressure for little
/// further latency hiding.
const MLP_N: usize = 4;
/// Single-thread MLP segment decode: gather up to `MLP_N` block cores (advancing the
/// bitstream sequentially), chase them interleaved to overlap their dependent-load
/// latency, and append their bytes in stream order. Produces byte-identical output to
/// [`decode_segment_serial`]; only the chase scheduling differs.
fn decode_segment_interleaved(
data: &[u8],
start_bit: u64,
end_bit: u64,
max_blocksize: u32,
output: &mut Vec<u8>,
block_crcs: &mut Vec<u32>,
budget: &OutputBudget,
) -> Option<block::BlockError> {
let total_bits = data.len() as u64 * 8;
let mut reader = BitReader::from_bit_offset(data, (start_bit + 48) as usize);
let mut first = true;
let mut cores: Vec<block::BlockCore> = Vec::with_capacity(MLP_N);
loop {
// ── Gather up to MLP_N block cores (header+Huffman+MTF+RLE2+inverse-BWT),
// advancing the bitstream sequentially. ────────────────────────────────
cores.clear();
let mut ended = false;
let mut gather_err: Option<block::BlockError> = None;
while cores.len() < MLP_N {
match next_block_core(&mut reader, end_bit, total_bits, max_blocksize, &mut first) {
Ok(Some(core)) => cores.push(core),
Ok(None) => {
ended = true;
break;
}
Err(e) => {
gather_err = Some(e);
break;
}
}
}
// ── Chase the gathered cores INTERLEAVED (the MLP hot loop) and append
// their bytes in stream order, CRC-checking each. ────────────────────
if !cores.is_empty() {
let before = output.len();
if let Err(e) = block::decode_blocks_interleaved(&mut cores, output, block_crcs) {
return Some(e);
}
// Charged per MLP batch rather than per block — the chase appends all
// `cores` in one call — so the overshoot here is at most MLP_N blocks.
if let Err(e) = budget.take(output.len() - before) {
return Some(e);
}
}
// A gather error surfaces only AFTER the already-parsed (earlier-in-stream)
// cores have been chased+appended, matching the block-by-block order.
if let Some(e) = gather_err {
return Some(e);
}
if ended {
break;
}
}
None
}
/// Advance the reader to the next block and decode its core (header → inverse-BWT),
/// WITHOUT the final RLE1 pointer-chase. Returns `Ok(None)` at a clean framing
/// boundary (end-of-segment, EOF, FINAL_MAGIC with no following stream, or a non-
/// magic word), `Ok(Some(core))` for a block to chase, or `Err` on a malformed block.
///
/// Encapsulates the exact framing state machine [`decode_segment_serial`] uses: the
/// segment's first block MUST decode; subsequent blocks are gated on BLOCK_MAGIC, and
/// FINAL_MAGIC + a valid following `BZh` header continues into a concatenated (pbzip2)
/// stream.
#[inline]
fn next_block_core(
reader: &mut BitReader<'_>,
end_bit: u64,
total_bits: u64,
max_blocksize: u32,
first: &mut bool,
) -> Result<Option<block::BlockCore>, block::BlockError> {
if *first {
// The segment's first block MUST decode: a CRC mismatch or malformed block
// here is corruption, not a framing boundary.
*first = false;
return block::decode_block_core(reader, max_blocksize).map(Some);
}
loop {
// Inherent `BitReader::position` (fully qualified: `&mut BitReader` also
// implements `Iterator`, whose `position` would otherwise shadow it).
let pos = BitReader::position(reader) as u64;
if pos + 48 > total_bits || pos >= end_bit {
return Ok(None);
}
let magic = match reader.read_u64(48) {
Some(v) => v,
None => return Ok(None),
};
if magic == BLOCK_MAGIC {
return block::decode_block_core(reader, max_blocksize).map(Some);
} else if magic == FINAL_MAGIC {
if reader.read_u32(32).is_none() {
return Ok(None);
}
let p = BitReader::position(reader);
let pad = (8 - (p % 8)) % 8;
if pad > 0 {
BitReader::skip(reader, pad);
}
match reader.read_u32(32) {
Some(h) => {
let b = h.to_be_bytes();
if &b[..3] != b"BZh" {
return Ok(None);
}
}
None => return Ok(None),
}
// Concatenated stream: loop to read the next stream's first block magic.
} else {
return Ok(None);
}
}
}
// Per-worker scratch for the CRC list the streaming `_into` path does not return
// (each block is CRC-checked inline; the whole-stream fold is the library path's job).
// Thread-local + reused, so the zero-alloc decode stays zero-alloc.
thread_local! {
static CRC_SCRATCH: std::cell::RefCell<Vec<u32>> = const { std::cell::RefCell::new(Vec::new()) };
}
/// Zero-alloc, CRC-checked segment decode into a caller-owned **reused** buffer.
///
/// Identical integrity guarantees to [`decode_segment_checked`] (a per-block CRC
/// mismatch → `Err`), but the output buffer is supplied by the caller (a recycled
/// slot) and the per-block CRC list lands in a thread-local scratch instead of a
/// fresh `Vec`. A warmed pool of `out` buffers therefore decodes with no allocation.
pub fn decode_segment_checked_into(
data: &[u8],
start_bit: u64,
end_bit: u64,
max_blocksize: u32,
out: &mut Vec<u8>,
mode: DecodeMode,
) -> Result<(), block::BlockError> {
out.clear();
out.reserve(seg_cap_hint(start_bit, end_bit));
CRC_SCRATCH.with(|c| {
let mut crcs = c.borrow_mut();
crcs.clear();
match decode_segment_inner(
data,
start_bit,
end_bit,
max_blocksize,
out,
&mut crcs,
mode,
&OutputBudget::unlimited(),
) {
None => Ok(()),
Some(e) => Err(e),
}
})
}
/// Decode one block directly onto the tail of `output`, avoiding a temporary Vec.
/// `output` grows to fit the block's actual (data-dependent, possibly ≫
/// max_blocksize) RLE1-expanded size; on error it is rolled back to its prior
/// length by [`block::decode_block_into`]. Returns the block's verified CRC-32.
#[inline]
fn decode_block_into_vec(
output: &mut Vec<u8>,
reader: &mut BitReader<'_>,
max_blocksize: u32,
) -> Result<u32, block::BlockError> {
let (_written, crc) = block::decode_block_into(reader, max_blocksize, output)?;
Ok(crc)
}
#[cfg(test)]
mod tests {
use super::*;
fn max_blocksize(data: &[u8]) -> u32 {
100_000 * (data[3] - b'0') as u32
}
#[test]
fn chunk_hello() {
let data = include_bytes!("../test_data/hello.bz2");
let split = split_chunk(data, 4, max_blocksize(data), true).unwrap();
let total_bits = data.len() as u64 * 8;
let mut output = Vec::new();
for i in 0..split.decode_segments {
let start = split.segment_starts[i].bit_offset;
let end = if i + 1 < split.segment_starts.len() {
split.segment_starts[i + 1].bit_offset
} else {
total_bits
};
output.extend_from_slice(&decode_segment(data, start, end, max_blocksize(data)));
}
assert_eq!(&output, b"Hello, World!\n");
}
/// Decode the whole stream as ONE segment (so every block lands in the same
/// interleave-batch sequence) with the chosen [`DecodeMode`].
fn decode_whole(data: &[u8], mode: DecodeMode) -> Vec<u8> {
let split = split_chunk(data, 1, max_blocksize(data), true).unwrap();
let total_bits = data.len() as u64 * 8;
let mut out_all = Vec::new();
let mut buf = Vec::new();
for i in 0..split.decode_segments {
let start = split.segment_starts[i].bit_offset;
let end = if i + 1 < split.segment_starts.len() {
split.segment_starts[i + 1].bit_offset
} else {
total_bits
};
decode_segment_checked_into(data, start, end, max_blocksize(data), &mut buf, mode)
.unwrap();
out_all.extend_from_slice(&buf);
}
out_all
}
/// The single-thread `DecodeMode::Interleaved` (MLP) path must produce
/// byte-identical output to the serial baseline and the sequential reference.
#[test]
fn interleaved_matches_serial() {
let data = include_bytes!("../test_data/liechtenstein.osm.bz2");
let reference = crate::stream::decompress(data).unwrap();
let serial = decode_whole(data, DecodeMode::Serial);
let interleaved = decode_whole(data, DecodeMode::Interleaved);
assert_eq!(serial, reference, "serial mode diverged from reference");
assert_eq!(
interleaved.len(),
reference.len(),
"interleaved size mismatch"
);
assert_eq!(
interleaved, reference,
"interleaved mode diverged from reference"
);
}
/// A valid `bzip2 -9` stream of 4 MiB of a single byte. RLE1 collapses the
/// 4 MiB to ~49 bytes on disk, and the whole thing is ONE block whose final
/// RLE1-*decode* expands back to 4 MiB — ~3.7× the old `max_blocksize + 25%`
/// (1.125 MiB) segment-buffer bound. Before the growable-append fix this
/// overflowed the fixed slice: an assert-panic on the checked path, a raw
/// heap overrun / SIGSEGV on the lenient planet-convert path. Regression for
/// the "non-deterministic planet bz2 crash" — planet OSM-XML has exactly this
/// shape (long whitespace / repeated-tag runs). See `block::rle2_decode_append`.
const REP_4M_A_BZ2: [u8; 49] = [
0x42, 0x5a, 0x68, 0x39, 0x31, 0x41, 0x59, 0x26, 0x53, 0x59, 0xf7, 0x09, 0x6c, 0x45, 0x00,
0x20, 0x22, 0x0c, 0x00, 0x80, 0x04, 0x20, 0x00, 0x00, 0x08, 0x20, 0x00, 0x30, 0xcc, 0x05,
0x49, 0xea, 0x71, 0x01, 0x80, 0x50, 0x06, 0x01, 0xe2, 0xee, 0x48, 0xa7, 0x0a, 0x12, 0x1e,
0xe1, 0x2d, 0x88, 0xa0,
];
#[test]
fn rle1_expansion_far_exceeds_blocksize() {
let data = &REP_4M_A_BZ2;
let expected = vec![b'A'; 4 * 1024 * 1024];
// Sequential whole-block path.
assert_eq!(crate::stream::decompress(data).unwrap(), expected);
// Parallel block-decode path (the one that crashed on the planet).
assert_eq!(
crate::parallel::decompress_parallel(data).unwrap(),
expected
);
// Segment path directly (checked wrapper — must NOT panic/overflow).
let split = split_chunk(data, 4, max_blocksize(data), true).unwrap();
let total_bits = data.len() as u64 * 8;
let mut out = Vec::new();
for i in 0..split.decode_segments {
let start = split.segment_starts[i].bit_offset;
let end = if i + 1 < split.segment_starts.len() {
split.segment_starts[i + 1].bit_offset
} else {
total_bits
};
let seg = decode_segment_checked(data, start, end, max_blocksize(data)).unwrap();
out.extend_from_slice(&seg.data);
}
assert_eq!(out, expected);
// Single-thread interleaved (MLP) path must survive the same RLE1-run
// expansion (the per-chain capacity-growth invariant, exercised here).
assert_eq!(decode_whole(data, DecodeMode::Interleaved), expected);
}
#[test]
fn chunk_liechtenstein() {
let data = include_bytes!("../test_data/liechtenstein.osm.bz2");
let n = std::thread::available_parallelism()
.map(|n| n.get())
.unwrap_or(4);
let split = split_chunk(data, n, max_blocksize(data), true).unwrap();
let total_bits = data.len() as u64 * 8;
let mut output = Vec::new();
for i in 0..split.decode_segments {
let start = split.segment_starts[i].bit_offset;
let end = if i + 1 < split.segment_starts.len() {
split.segment_starts[i + 1].bit_offset
} else {
total_bits
};
output.extend_from_slice(&decode_segment(data, start, end, max_blocksize(data)));
}
let reference = crate::stream::decompress(data).unwrap();
assert_eq!(output.len(), reference.len());
assert_eq!(output, reference);
}
#[test]
fn chunk_split_simulation() {
let data = include_bytes!("../test_data/liechtenstein.osm.bz2");
let mbs = max_blocksize(data);
let mid = data.len() / 2;
let n = std::thread::available_parallelism()
.map(|n| n.get())
.unwrap_or(4);
let split1 = split_chunk(&data[..mid], n, mbs, false).unwrap();
let consumed1 = split1.consumed;
assert!(consumed1 <= mid);
let total_bits1 = mid as u64 * 8;
let mut out1 = Vec::new();
for i in 0..split1.decode_segments {
let start = split1.segment_starts[i].bit_offset;
let end = if i + 1 < split1.segment_starts.len() {
split1.segment_starts[i + 1].bit_offset
} else {
total_bits1
};
out1.extend_from_slice(&decode_segment(&data[..mid], start, end, mbs));
}
assert!(!out1.is_empty());
let chunk2 = &data[consumed1..];
let split2 = split_chunk(chunk2, n, mbs, true).unwrap();
let total_bits2 = chunk2.len() as u64 * 8;
let mut out2 = Vec::new();
for i in 0..split2.decode_segments {
let start = split2.segment_starts[i].bit_offset;
let end = if i + 1 < split2.segment_starts.len() {
split2.segment_starts[i + 1].bit_offset
} else {
total_bits2
};
out2.extend_from_slice(&decode_segment(chunk2, start, end, mbs));
}
let mut combined = out1;
combined.extend_from_slice(&out2);
let reference = crate::stream::decompress(data).unwrap();
assert_eq!(combined.len(), reference.len());
assert_eq!(combined, reference);
}
}