lbzip2 0.5.8

//! CLI parallel bzip2 decompressor — pure Rust, all available cores.
//!
//! Usage: lbunzip2 <input.bz2> [output]
//!   If output is omitted, strips .bz2 extension.
//!
//! Streaming decode on the shared **gatling** ordered engine
//! ([`gatling::gatling::ordered::run_ordered_sink`]). A single dispatch thread is
//! the I/O producer: it reads compressed chunks, prepends the carry (the chunk's
//! undecoded tail), `split_chunk`s each chunk into per-block segments, and yields
//! one segment at a time. N self-dispatching map workers CRC-check-decode segments;
//! the ordered sink writes decoded output in strict stream order. There is no
//! hand-rolled worker pool and no `Arc<Mutex<Receiver>>` — the shared engine is the
//! one parallel primitive in the constellation.

use std::{
    fs::File,
    io::{BufWriter, Read, Write},
    sync::{Arc, Mutex},
    time::Instant,
};

use anyhow::Result;
use gatling::gatling::ordered::{run_ordered_sink, OrderedSink};

use lbzip2::chunk::{self, ChunkSplit, DecodeMode};

/// Compressed bytes read per chunk. This is the streaming unit: a chunk's blocks
/// are split into per-worker segments, decoded, and flushed (in order) to the
/// writer before the next chunk is admitted, so the decoded output resident at any
/// instant is bounded by roughly the few segments in flight (gatling's `cap`) — NOT
/// the whole file. 8 MB compressed ≈ ~60 MB decoded/chunk and still ~60 bzip2
/// blocks per chunk, so every worker stays fed while memory stays bounded.
const CHUNK_SIZE_DEFAULT: usize = 4 * 1024 * 1024;
const BUF_CAP: usize = 4 * 1024 * 1024;

/// Detect physical (non-SMT) core count via Linux sysfs topology.
/// Falls back to available_parallelism() on non-Linux or if sysfs is unavailable.
fn physical_cores() -> Option<usize> {
    #[cfg(target_os = "linux")]
    {
        use std::collections::HashSet;
        let mut ids = HashSet::new();
        for entry in std::fs::read_dir("/sys/devices/system/cpu").ok()? {
            let entry = entry.ok()?;
            let name = entry.file_name();
            let name = name.to_str()?;
            if !name.starts_with("cpu") || !name[3..].chars().next()?.is_ascii_digit() {
                continue;
            }
            let base = entry.path().join("topology");
            let pkg = std::fs::read_to_string(base.join("physical_package_id"))
                .ok()
                .and_then(|s| s.trim().parse::<u32>().ok())
                .unwrap_or(0);
            if let Ok(s) = std::fs::read_to_string(base.join("core_id")) {
                if let Ok(core) = s.trim().parse::<u32>() {
                    ids.insert((pkg, core));
                }
            }
        }
        if !ids.is_empty() {
            return Some(ids.len());
        }
    }
    std::thread::available_parallelism().map(|n| n.get()).ok()
}

fn read_chunk(reader: &mut impl Read, buf: &mut [u8]) -> usize {
    let mut got = 0;
    while got < buf.len() {
        match reader.read(&mut buf[got..]) {
            Ok(0) => break,
            Ok(k) => got += k,
            Err(e) if e.kind() == std::io::ErrorKind::Interrupted => continue,
            Err(e) => {
                eprintln!("read error: {e}");
                break;
            }
        }
    }
    got
}

/// One decode segment handed to a gatling map worker: a shared reference to the
/// chunk's compressed bytes plus the `[start_bit, end_bit)` block range within them.
/// The `Arc` keeps the chunk buffer alive across every segment that references it —
/// the ownership-based replacement for the old raw-pointer-into-a-recycled-slot.
struct SegIn {
    data: Arc<Vec<u8>>,
    start_bit: u64,
    end_bit: u64,
    max_blocksize: u32,
}

/// The chunk the producer is currently yielding segments from.
struct CurChunk {
    data: Arc<Vec<u8>>,
    split: ChunkSplit,
    next_seg: usize,
    total_bits: u64,
}

/// Recycled decode-slot pool. A worker `take()`s a buffer, decodes its segment INTO
/// it (`decode_segment_checked_into`), and hands it to the ordered writer, which
/// `give()`s it back after writing. After warmup (≈ `cap` + `n_workers` buffers) the
/// steady state performs **zero allocation** — the "decompress into a pre-allocated
/// bigger-than-max slot, writer takes the bytes later" design. Each slot keeps its
/// capacity across reuse; a pathological (RLE1-expanded) block grows its slot once and
/// that capacity is then retained.
struct BufPool {
    free: Mutex<Vec<Vec<u8>>>,
    slot_cap: usize,
}

impl BufPool {
    fn new(slot_cap: usize) -> Self {
        Self { free: Mutex::new(Vec::new()), slot_cap }
    }
    fn take(&self) -> Vec<u8> {
        self.free
            .lock()
            .unwrap()
            .pop()
            .unwrap_or_else(|| Vec::with_capacity(self.slot_cap))
    }
    fn give(&self, buf: Vec<u8>) {
        // Capacity retained; `decode_segment_checked_into` clears on next use.
        self.free.lock().unwrap().push(buf);
    }
}

/// Ordered sink: write decoded segments to the output file in strict stream order,
/// then recycle each slot back to the pool. Runs on the calling thread (the gatling
/// collector), which owns the writer — so it is the single-core writer that "takes
/// the offset" (its running position) from each in-order slot.
struct WriteSink<'a, W: Write> {
    w: W,
    total: u64,
    pool: &'a BufPool,
}

impl<W: Write> OrderedSink<Vec<u8>> for WriteSink<'_, W> {
    fn emit(&mut self, _seq: u64, data: Vec<u8>) -> Result<()> {
        if !data.is_empty() {
            self.total += data.len() as u64;
            self.w.write_all(&data)?;
        }
        self.pool.give(data);
        Ok(())
    }
}

fn main() {
    let args: Vec<String> = std::env::args().collect();
    if args.len() < 2 {
        eprintln!("usage: lbunzip2 <input.bz2> [output]");
        std::process::exit(1);
    }

    let input_path = &args[1];
    let output_path = if args.len() > 2 {
        args[2].clone()
    } else if input_path.ends_with(".bz2") {
        input_path[..input_path.len() - 4].to_string()
    } else {
        format!("{input_path}.out")
    };

    let in_file = File::open(input_path).expect("open input");
    let in_size = in_file.metadata().map(|m| m.len()).unwrap_or(0);
    // Read straight from the File — no BufReader. `read_chunk` fills the chunk buffer
    // in ≤ CHUNK_SIZE reads directly, so a BufReader would only add a second userspace
    // copy (kernel page-cache → its buffer → our `data`). Direct = one copy.
    let mut reader = in_file;

    // Read bz2 header (4 bytes: "BZhN")
    let mut header = [0u8; 4];
    reader.read_exact(&mut header).expect("read bz2 header");

    let bz2_level = header[3];
    if &header[..2] != b"BZ" || header[2] != b'h' || !(b'1'..=b'9').contains(&bz2_level) {
        eprintln!("invalid bzip2 header");
        std::process::exit(1);
    }
    let max_blocksize = 100_000 * (bz2_level - b'0') as u32;

    let n_workers: usize = std::env::var("LBZIP2_THREADS")
        .ok()
        .and_then(|v| v.parse().ok())
        .unwrap_or_else(|| physical_cores().unwrap_or(4));

    // Chunk **segmentation granularity** — DISTINCT from worker parallelism.
    //
    // The streaming producer holds back the LAST segment of every non-final
    // chunk as `carry` (its trailing bzip2 block may continue into the next
    // chunk), decoding the first `decode_segments = n - 1`. That requires every
    // chunk to split into ≥2 segments: with a single segment `split_chunk`
    // returns `None` for non-last chunks, so the producer's `None` branch
    // discards those bytes and only the FINAL chunk is ever decoded — a silent
    // truncation to the last ~CHUNK_SIZE (and, in the bench, an impossible
    // "single-core faster than all cores" throughput because the full corpus
    // size is billed against that near-instant partial decode).
    //
    // `split_boundaries_parallel(_, 1, _)` returns no interior splits ⇒ n = 1,
    // so `LBZIP2_THREADS=1` used to trip exactly that path. Floor the segment
    // count at 2 so a single worker still streams the WHOLE file (segments are
    // just decoded one at a time, in order); worker parallelism stays exactly
    // `n_workers`.
    // Oversubscription knob: segments/chunk = n_workers * LBZIP2_SEG_MULT (default 1).
    // >1 gives the ordered-sink pool more units than cores so the per-chunk tail
    // self-balances instead of stranding a core. Floor at 2 (carry needs ≥2).
    let seg_mult: usize = std::env::var("LBZIP2_SEG_MULT")
        .ok()
        .and_then(|v| v.parse().ok())
        .filter(|&m| m >= 1)
        .unwrap_or(3);
    let n_segments = (n_workers * seg_mult).max(2);

    // Compressed read unit per chunk. LBZIP2_CHUNK_MB overrides (in MB).
    let chunk_size: usize = std::env::var("LBZIP2_CHUNK_MB")
        .ok()
        .and_then(|v| v.parse::<usize>().ok())
        .filter(|&m| m >= 1)
        .map(|m| m * 1024 * 1024)
        .unwrap_or(CHUNK_SIZE_DEFAULT);

    eprintln!(
        "lbunzip2: {} ({} MB) → {}  [{} gatling workers, {} segments/chunk]",
        input_path,
        in_size / (1024 * 1024),
        output_path,
        n_workers,
        n_segments,
    );

    let t0 = Instant::now();

    // ── Producer: read → carry → split → yield one segment at a time ──────────
    // Runs on gatling's dispatch thread, so its I/O overlaps decode. Keeps the
    // reader, the carry (undecoded tail prepended into the next chunk), and the
    // in-progress chunk's split state across pulls.
    let mut carry: Vec<u8> = header.to_vec();
    let mut cur: Option<CurChunk> = None;
    let mut finished = false;

    let producer = move || -> Option<((), SegIn)> {
        loop {
            // Yield the next undelivered segment of the current chunk.
            if let Some(c) = cur.as_mut() {
                if c.next_seg < c.split.decode_segments {
                    let i = c.next_seg;
                    c.next_seg += 1;
                    let n_seg = c.split.segment_starts.len();
                    let start_bit = c.split.segment_starts[i].bit_offset;
                    let end_bit = if i + 1 < n_seg {
                        c.split.segment_starts[i + 1].bit_offset
                    } else {
                        c.total_bits
                    };
                    return Some((
                        (),
                        SegIn {
                            data: Arc::clone(&c.data),
                            start_bit,
                            end_bit,
                            max_blocksize,
                        },
                    ));
                }
            }
            cur = None;
            if finished {
                return None;
            }

            // Read the next chunk into a carry-prefixed buffer.
            let base = carry.len();
            let mut data = std::mem::take(&mut carry);
            data.resize(base + chunk_size, 0);
            let got = read_chunk(&mut reader, &mut data[base..]);
            data.truncate(base + got);
            let is_last = got < chunk_size;

            if got == 0 && base <= 4 {
                finished = true;
                return None;
            }

            match chunk::split_chunk(&data, n_segments, max_blocksize, is_last) {
                None => {
                    if is_last {
                        finished = true;
                        return None;
                    }
                    // No block boundary in this chunk: drop the read bytes and keep
                    // the pre-read carry unchanged (the first `base` bytes are it).
                    data.truncate(base);
                    carry = data;
                    continue;
                }
                Some(split) => {
                    let total_bits = data.len() as u64 * 8;
                    carry = data[split.consumed..].to_vec();
                    if is_last {
                        finished = true;
                    }
                    cur = Some(CurChunk {
                        data: Arc::new(data),
                        split,
                        next_seg: 0,
                        total_bits,
                    });
                    // loop back to yield this chunk's first segment
                }
            }
        }
    };

    // Recycled decode-slot pool. Slot capacity is generous (a segment is ~one bzip2
    // block ≈ chunk/segments decoded); each slot grows once to its true size then is
    // reused, so steady-state decode allocates nothing and never re-zeroes fresh pages.
    let pool = BufPool::new(2 * 1024 * 1024);

    // ── Map: CRC-checked segment decode INTO a recycled slot ──────────────────
    // SECURITY: a per-block CRC mismatch is fatal — we exit non-zero rather than
    // emit silently-truncated output. Matches the old worker behaviour exactly.
    // GATE: memory-level-parallel interleaved decode is a single-core win but
    // neutral-to-harmful once the socket is memory-bandwidth-bound, so use it ONLY
    // when decoding with a single worker (`LBZIP2_THREADS=1` / the `_st` bench).
    // Every multi-worker decode takes the byte-identical serial baseline.
    let decode_mode = if n_workers == 1 {
        DecodeMode::Interleaved
    } else {
        DecodeMode::Serial
    };

    let map = |_label: (), s: SegIn| -> Vec<u8> {
        let mut buf = pool.take();
        match chunk::decode_segment_checked_into(
            &s.data, s.start_bit, s.end_bit, s.max_blocksize, &mut buf, decode_mode,
        ) {
            Ok(()) => buf,
            Err(e) => {
                eprintln!("lbunzip2: corrupt input: {e}");
                std::process::exit(1);
            }
        }
    };

    // ── Sink: ordered streaming write on the calling thread ───────────────────
    let out_file = File::create(&output_path).expect("create output");
    let mut sink = WriteSink {
        w: BufWriter::with_capacity(BUF_CAP, out_file),
        total: 0,
        pool: &pool,
    };

    // Ring depth (segments in flight). 0 ⇒ run_ordered_sink default = 2*n_workers.
    let cap: usize = std::env::var("LBZIP2_CAP").ok().and_then(|v| v.parse().ok()).unwrap_or(0);
    run_ordered_sink(producer, n_workers, cap, map, &mut sink).expect("decode pipeline");
    sink.w.flush().expect("flush output");
    let total_out = sink.total;

    let elapsed = t0.elapsed().as_secs_f64();
    let out_mb = total_out / (1024 * 1024);
    let in_mb = in_size / (1024 * 1024);

    eprintln!(
        "done: {:.1}s  {} MB → {} MB  ({:.0} MB/s decompressed)",
        elapsed,
        in_mb,
        out_mb,
        out_mb as f64 / elapsed,
    );

    // Introspection marker: the parallel bzip2 decode pipeline drained cleanly.
    lbzip2::functional_status(
        "lbunzip2",
        "decode",
        true,
        &format!("{out_mb} MB decompressed in {elapsed:.1}s"),
    );
}