lbzip2 0.5.4

//! 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;

/// Decompress a complete bzip2 stream using parallel block decode.
///
/// `data` must be a complete bzip2 stream (header + blocks + EOS).
/// Blocks are decoded concurrently on scoped worker threads.
pub fn decompress_parallel(data: &[u8]) -> 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 `par_map`'s dynamic
    // work-stealing can balance uneven block sizes (a single fat block must not
    // strand a whole worker). ~4 segments/thread keeps each segment to a few
    // blocks while still amortising the per-block-Vec + concat savings.
    let n_segments = (n_threads * 4).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 ──────────────────
    let results: Vec<Vec<u8>> = crate::par::par_map(segments.len(), |si| {
        let (start_bit, end_bit) = segments[si];
        chunk::decode_segment(data, start_bit, end_bit, max_blocksize)
    });

    // ── Assemble output in order (only n_segments copies) ───────────────
    let total_size: usize = results.iter().map(|v| v.len()).sum();
    let mut output = Vec::with_capacity(total_size);
    for seg in results {
        output.extend_from_slice(&seg);
    }

    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);
    }
}