lbzip2/src/parallel.rs
//! Parallel bzip2 decompression — decode N blocks concurrently.
//!
//! Uses `block_scan::find_all_blocks()` to locate every block boundary,
//! then decodes all blocks in parallel (scoped threads). Results are concatenated
//! in order.
use crate::block::BlockError;
use crate::block_scan;
use crate::chunk;
use crate::limits::{self, OutputBudget};
/// Decompress a complete bzip2 stream using parallel block decode, into at most
/// [`limits::DEFAULT_MAX_OUTPUT`] bytes.
///
/// **The cap is on by default.** This function was the worst of the two
/// whole-stream entry points: it holds every decoded segment resident *and then*
/// copies them all into an assembled output, so it asked for the claimed size
/// twice. A 50 KB file of concatenated all-zero bzip2 streams (1 GiB of zeros is
/// 785 compressed bytes — measured, see [`limits`]) asked for 128 GiB.
///
/// The ceiling is a single [`OutputBudget`] shared by every worker and charged
/// once per decoded block; segments are uneven, so no per-segment share could
/// bound the total without rejecting legitimate streams.
///
/// [`decompress_parallel_uncapped`] is the pre-cap behaviour, by name;
/// [`decompress_parallel_capped`] takes any other ceiling; [`crate::reader`]
/// streams an output larger than memory and is not capped.
pub fn decompress_parallel(data: &[u8]) -> Result<Vec<u8>, BlockError> {
decompress_parallel_capped(data, limits::DEFAULT_MAX_OUTPUT)
}
/// [`decompress_parallel`] with no output ceiling at all — the pre-cap
/// behaviour, now reachable only by asking for it by name.
pub fn decompress_parallel_uncapped(data: &[u8]) -> Result<Vec<u8>, BlockError> {
decompress_parallel_budgeted(data, &OutputBudget::unlimited())
}
/// [`decompress_parallel`] against a caller-chosen ceiling.
///
/// Fails with [`limits::OUTPUT_TOO_LARGE`] as soon as the workers have between
/// them produced more than `max_out` bytes. Peak stays near `max_out` plus one
/// block per worker, rather than the size the stream claimed.
pub fn decompress_parallel_capped(data: &[u8], max_out: usize) -> Result<Vec<u8>, BlockError> {
decompress_parallel_budgeted(data, &OutputBudget::new(max_out))
}
fn decompress_parallel_budgeted(data: &[u8], budget: &OutputBudget) -> Result<Vec<u8>, BlockError> {
// ── Parse stream header ─────────────────────────────────────────────
if data.len() < 4 {
return Err(BlockError("input too short for bzip2 header"));
}
if &data[..2] != b"BZ" {
return Err(BlockError("bad bzip2 signature"));
}
if data[2] != b'h' {
return Err(BlockError("only huffman bzip2 supported"));
}
let level = data[3];
if !(b'1'..=b'9').contains(&level) {
return Err(BlockError("invalid bzip2 block size level"));
}
let max_blocksize = 100_000 * (level - b'0') as u32;
// ── Scan for all block boundaries ───────────────────────────────────
// Parallel bit-scan for BZh block magics. The serial `find_all_blocks`
// walks the whole stream on one thread (a serial prefix that dominates on
// large inputs); the parallel scan splits the buffer across cores and
// merges, producing the identical boundary set.
let n_threads = std::thread::available_parallelism()
.map(|n| n.get())
.unwrap_or(4);
let boundaries = block_scan::find_all_blocks_parallel(data, n_threads);
if boundaries.is_empty() {
// Might be an empty file or just EOS marker
return Ok(Vec::new());
}
// ── Group blocks into one segment per worker ────────────────────────
// Perf: rather than decode every block into its own `Vec<u8>` and then
// concatenate ~N_blocks buffers (≈1 GB of extra copying for a large stream,
// plus a heap allocation per block), we split the blocks into `n_threads`
// contiguous segments. Each worker runs the `chunk::decode_segment` path,
// which appends every block in its range directly into a single growing
// buffer via `decode_block_into` — no per-block Vec. The final assembly
// then concatenates only `n_threads` segment buffers instead of N_blocks.
let total_bits = data.len() as u64 * 8;
let n_blocks = boundaries.len();
// Oversubscribe segments relative to thread count so the balanced pool's
// heaviest-first (LPT) claim can absorb uneven block decode costs (a single
// fat segment must not strand a whole worker on the tail). ~8 segments/thread
// keeps each segment to a few blocks — fine enough for the scheduler to keep
// all cores busy to the end — while still amortising the per-block-Vec + concat
// savings (only `n_segments` buffers get concatenated, not `n_blocks`).
let n_segments = (n_threads * 8).min(n_blocks).max(1);
// Segment i owns blocks [seg_start(i) .. seg_start(i+1)). We record each
// segment's first-block bit offset and the end bit (= next segment's first
// block, or end-of-stream for the last segment).
let mut segments: Vec<(u64, u64)> = Vec::with_capacity(n_segments);
for i in 0..n_segments {
let first_block = i * n_blocks / n_segments;
let next_first_block = (i + 1) * n_blocks / n_segments;
let start_bit = boundaries[first_block].bit_offset;
let end_bit = if next_first_block < n_blocks {
boundaries[next_first_block].bit_offset
} else {
total_bits
};
segments.push((start_bit, end_bit));
}
// ── Parallel decode: one segment buffer per worker ──────────────────
// Heaviest-first (LPT) balanced dispatch: weight each segment by its
// compressed bit-span (`end_bit - start_bit`), a proxy for decode cost, so
// the fat segments are claimed first and the cheap ones fill the tail. With
// plain `gatling_for_each` (in-index-order claim) a heavy segment claimed late
// stranded one core while the rest idled — the multicore tail that cost us
// ~0.3 cores-busy vs lbzip2. Results still come back in original segment order.
let results: Vec<Result<chunk::SegmentDecoded, BlockError>> =
gatling::gatling_forkjoin::gatling_for_each_balanced(
segments.len(),
0, // one worker per core
1, // batch=1 — segments are already coarse; claim one at a time for the tightest tail
|si| segments[si].1 - segments[si].0,
|si| {
let (start_bit, end_bit) = segments[si];
chunk::decode_segment_checked_budgeted(
data,
start_bit,
end_bit,
max_blocksize,
budget,
)
},
);
// ── Integrity: propagate any per-block CRC failure ──────────────────
// A tampered/truncated block whose stored CRC-32 does not match is now a
// hard error instead of silently-truncated output (the security finding).
let mut decoded: Vec<chunk::SegmentDecoded> = Vec::with_capacity(results.len());
for r in results {
decoded.push(r?);
}
// Fold every block's verified CRC into the combined stream CRC, in stream
// order (segments are ordered by start_bit; blocks within a segment are in
// order), matching the sequential decoder's rotate-left-1-then-XOR scheme.
let mut combined_crc: u32 = 0;
for seg in &decoded {
for &c in &seg.block_crcs {
combined_crc = (combined_crc << 1) | (combined_crc >> 31);
combined_crc ^= c;
}
}
// Verify the whole-stream combined CRC for a single bzip2 stream (exactly
// one FINAL_MAGIC + stored CRC). Concatenated (pbzip2) inputs carry one CRC
// per sub-stream; the per-block CRCs above already validate their integrity,
// so we skip the whole-stream fold there rather than risk a false rejection.
let stream_crcs = block_scan::find_stream_crcs_parallel(data, n_threads);
if stream_crcs.len() == 1 && stream_crcs[0] != combined_crc {
return Err(BlockError("stream CRC mismatch"));
}
// ── Assemble output in order (only n_segments copies) ───────────────
let total_size: usize = decoded.iter().map(|s| s.data.len()).sum();
let mut output = Vec::with_capacity(total_size);
for seg in decoded {
output.extend_from_slice(&seg.data);
}
Ok(output)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn parallel_hello() {
let compressed = include_bytes!("../test_data/hello.bz2");
let output = decompress_parallel(compressed).unwrap();
assert_eq!(&output, b"Hello, World!\n");
}
#[test]
fn parallel_liechtenstein() {
let compressed = include_bytes!("../test_data/liechtenstein.osm.bz2");
let output = decompress_parallel(compressed).unwrap();
// Compare with sequential decode
let sequential = crate::stream::decompress(compressed).unwrap();
assert_eq!(output.len(), sequential.len(), "size mismatch");
assert_eq!(output, sequential, "content mismatch");
}
#[test]
fn parallel_matches_sequential() {
let compressed = include_bytes!("../test_data/liechtenstein.osm.bz2");
let par = decompress_parallel(compressed).unwrap();
let seq = crate::stream::decompress(compressed).unwrap();
assert_eq!(par, seq);
}
}