lzip-parallel 0.2.2

//! lzip — parallel ZIP extractor.
//!
//! Pipeline (mirrors lbzip2-rs/src/bin/lbunzip2.rs):
//!
//!   Reader ──→ Main(carry+split) ──→ N Workers ──→ Collector ──→ Writer
//!     ↑                                                  │
//!     └────────────── ring-slot recycle ─────────────────┘
//!
//! - Reader thread: pulls free 232 MB ring slots, fills with compressed bytes
//!   from the ZIP file (sequential disk I/O over entries sorted by offset).
//! - Main thread: prepends carry, calls split_chunk (rayon for the parallel
//!   forward-scan), posts N WorkItems (raw pointer into slot) to a shared
//!   mpsc channel + an InFlightSlot to the collector.
//! - N dedicated decoder threads (NOT rayon par_iter): each loops on
//!   `Arc<Mutex<Receiver<WorkItem>>>::recv()` → decode → send SegmentResult.
//!   A worker that finishes a small segment immediately pulls the next item —
//!   no fork-join barrier, no idle time while work is pending.
//! - Collector: tracks `Vec<InFlightSlot>` keyed by chunk_id, applies results,
//!   flushes completed chunks in order to the writer, recycles slot back to
//!   the reader's free pool.
//! - Writer: single thread, mpsc channel, no locking.
//!
//! Slot lifetime safety invariant (load-bearing): a slot is recycled ONLY
//! when `done_count == decode_segments` for that chunk_id. Workers hold
//! `*const u8` into the slot — valid because the slot lives inside
//! InFlightSlot owned by the collector until that condition is met.

use std::collections::HashSet;
use std::fs::{self, File};
use std::io::{Read, Seek, SeekFrom, Write};
use std::path::PathBuf;
use std::sync::{mpsc, Arc, Mutex};
use std::thread;

use lzip_parallel::central_dir::{read_central_directory_from, EntryLocation};
use lzip_parallel::chunk::{decode_segment_into, split_chunk};
use lzip_parallel::entry::ZipError;

// ── Constants (DESIGN.md Rule 2.1) ───────────────────────────────────────────

const CHUNK_SIZE: usize       = 200 * 1024 * 1024;
const CARRY_HEADROOM: usize   =  32 * 1024 * 1024;
const SLOT_SIZE: usize        = CARRY_HEADROOM + CHUNK_SIZE;
/// 6 ring slots = up to 6 chunks in flight simultaneously.
/// Entries larger than 6 × CHUNK_SIZE (1.2 GB) still work: the reader blocks
/// on `free_rx.recv()` until a slot is recycled, then fills the next chunk.
/// Pipeline depth is bounded at 6 but throughput is the same.
const RING_SLOTS: usize       = 6;
const LARGE_THRESHOLD: usize  = 4 * 1024 * 1024;
const LOCAL_HEADER_SIG: u32   = 0x04034b50;

// ── Resolved entry (compressed-data offset already looked up) ────────────────

#[derive(Debug, Clone)]
struct ResolvedEntry {
    name: String,
    data_offset: u64,
    compressed_size: u64,
    uncompressed_size: u64,
    compression_method: u16,
}

// ── Pipeline messages ────────────────────────────────────────────────────────

/// Slot filled by the reader, plus per-chunk metadata.
struct FilledSlot {
    slot: Vec<u8>,
    read_len: usize,
    is_ring_slot: bool,         // false for ad-hoc small-entry slots
    entry_name: String,
    is_first_of_entry: bool,
    is_last_of_entry: bool,
    // For STORE entries we skip decode entirely and just copy the bytes.
    is_store: bool,
}

/// Work item dispatched to a decoder thread.
///
/// `data_ptr` points into a slot owned by an `InFlightSlot` in the collector.
/// Safe as long as the collector holds the slot until all segments are done.
struct WorkItem {
    chunk_id: u64,
    segment_id: usize,
    data_ptr: *const u8,
    data_len: usize,
    start_byte: usize,
    end_byte: usize,
    is_store: bool,
    /// Hint for pre-allocating output buffer. For single-segment entries this
    /// equals uncompressed_size (exact). For multi-segment, an estimate.
    output_size_hint: usize,
}
unsafe impl Send for WorkItem {}

struct SegmentResult {
    chunk_id: u64,
    segment_id: usize,
    output: Result<Vec<u8>, ZipError>,
}

enum CollectorMsg {
    NewSlot(InFlightSlot),
    Result(SegmentResult),
}

/// State held by the collector for one in-flight chunk.
struct InFlightSlot {
    chunk_id: u64,
    slot: Vec<u8>,
    is_ring_slot: bool,
    decode_segments: usize,
    results: Vec<Option<Result<Vec<u8>, ZipError>>>,
    done_count: usize,
    entry_name: String,
    is_first_of_entry: bool,
    is_last_of_entry: bool,
}

/// Writer-thread message stream — one EntryStart per entry, one or more Chunks,
/// then EntryEnd. Allows streaming writes without buffering the whole entry.
enum WriteMsg {
    EntryStart(String),
    Chunk(Vec<u8>),
    EntryEnd,
}

// ── Physical core detection (Linux sysfs) ────────────────────────────────────

fn physical_cores() -> usize {
    let n = (|| -> Option<usize> {
        let mut seen = HashSet::new();
        for e in fs::read_dir("/sys/devices/system/cpu").ok()? {
            let e = e.ok()?;
            let fname = e.file_name();
            let s = fname.to_str()?;
            if !s.starts_with("cpu") || s[3..].is_empty()
                || !s[3..].bytes().all(|b| b.is_ascii_digit())
            { continue; }
            let pkg  = fs::read_to_string(e.path().join("topology/physical_package_id")).ok()?;
            let core = fs::read_to_string(e.path().join("topology/core_id")).ok()?;
            seen.insert((pkg.trim().to_string(), core.trim().to_string()));
        }
        Some(seen.len()).filter(|&n| n > 0)
    })();
    n.unwrap_or_else(|| thread::available_parallelism().map(|n| n.get()).unwrap_or(4))
}

// ── Local file header → compressed data offset ───────────────────────────────

fn local_data_offset(file: &mut File, loc: &EntryLocation) -> Result<u64, ZipError> {
    file.seek(SeekFrom::Start(loc.local_header_offset))
        .map_err(|_| ZipError("seek to local header failed"))?;
    let mut hdr = [0u8; 30];
    file.read_exact(&mut hdr)
        .map_err(|_| ZipError("read local header failed"))?;
    if u32::from_le_bytes(hdr[0..4].try_into().unwrap()) != LOCAL_HEADER_SIG {
        return Err(ZipError("invalid local file header signature"));
    }
    let name_len  = u16::from_le_bytes(hdr[26..28].try_into().unwrap()) as u64;
    let extra_len = u16::from_le_bytes(hdr[28..30].try_into().unwrap()) as u64;
    Ok(loc.local_header_offset + 30 + name_len + extra_len)
}

// ── Pipeline: reader thread ──────────────────────────────────────────────────

/// Reader thread: iterates entries in file-offset order. For large entries
/// (≥ LARGE_THRESHOLD) it pulls ring slots from `free_rx` and fills them in
/// CHUNK_SIZE pieces. For small entries it allocates an ad-hoc Vec sized to
/// the compressed payload.
///
/// Sends `FilledSlot` messages tagged with entry name and first/last flags.
fn reader_thread(
    mut file: File,
    entries: Vec<ResolvedEntry>,
    free_rx: mpsc::Receiver<Vec<u8>>,
    filled_tx: mpsc::SyncSender<FilledSlot>,
) {
    for entry in entries {
        if let Err(e) = file.seek(SeekFrom::Start(entry.data_offset)) {
            eprintln!("seek to {}: {}", entry.name, e);
            continue;
        }

        let is_store = entry.compression_method == 0;

        if (entry.compressed_size as usize) < LARGE_THRESHOLD {
            // Small entry: one ad-hoc slot, no ring.
            let mut buf = vec![0u8; entry.compressed_size as usize];
            if let Err(e) = file.read_exact(&mut buf) {
                eprintln!("read {}: {}", entry.name, e);
                continue;
            }
            let read_len = buf.len();
            let msg = FilledSlot {
                slot: buf,
                read_len,
                is_ring_slot: false,
                entry_name: entry.name.clone(),
                is_first_of_entry: true,
                is_last_of_entry: true,
                is_store,
            };
            if filled_tx.send(msg).is_err() { return; }
            continue;
        }

        // Large entry: pull ring slots, read CHUNK_SIZE bytes each.
        let mut remaining = entry.compressed_size as usize;
        let mut is_first = true;
        while remaining > 0 {
            let mut slot = match free_rx.recv() {
                Ok(s) => s,
                Err(_) => return,
            };
            let to_read = remaining.min(CHUNK_SIZE);
            if let Err(e) = file.read_exact(&mut slot[CARRY_HEADROOM..CARRY_HEADROOM + to_read]) {
                eprintln!("read chunk for {}: {}", entry.name, e);
                return;
            }
            remaining -= to_read;
            let is_last = remaining == 0;
            let msg = FilledSlot {
                slot,
                read_len: to_read,
                is_ring_slot: true,
                entry_name: entry.name.clone(),
                is_first_of_entry: is_first,
                is_last_of_entry: is_last,
                is_store,
            };
            is_first = false;
            if filled_tx.send(msg).is_err() { return; }
        }
    }
}

// ── Pipeline: collector thread helpers ──────────────────────────────────────

fn apply_result(in_flight: &mut [InFlightSlot], result: SegmentResult) {
    for s in in_flight.iter_mut() {
        if s.chunk_id == result.chunk_id {
            s.results[result.segment_id] = Some(result.output);
            s.done_count += 1;
            return;
        }
    }
}

fn flush_completed(
    in_flight: &mut Vec<InFlightSlot>,
    next_write_id: &mut u64,
    write_tx: &mpsc::SyncSender<WriteMsg>,
    slot_return: &mpsc::SyncSender<Vec<u8>>,
) {
    loop {
        let idx = match in_flight.iter().position(|s| s.chunk_id == *next_write_id) {
            Some(i) => i,
            None => break,
        };
        if in_flight[idx].done_count < in_flight[idx].decode_segments { break; }

        let mut completed = in_flight.remove(idx);

        if completed.is_first_of_entry {
            if write_tx.send(WriteMsg::EntryStart(completed.entry_name.clone())).is_err() {
                return;
            }
        }
        for seg in completed.results.drain(..) {
            match seg {
                Some(Ok(data)) => {
                    if !data.is_empty() {
                        if write_tx.send(WriteMsg::Chunk(data)).is_err() { return; }
                    }
                }
                Some(Err(e)) => {
                    eprintln!("decode error in {}: {}", completed.entry_name, e);
                }
                None => {
                    eprintln!("missing segment result in {}", completed.entry_name);
                }
            }
        }
        if completed.is_last_of_entry {
            if write_tx.send(WriteMsg::EntryEnd).is_err() { return; }
        }
        if completed.is_ring_slot {
            slot_return.send(completed.slot).ok();
        }
        *next_write_id += 1;
    }
}

// ── main ─────────────────────────────────────────────────────────────────────

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

    let zip_path = &args[1];
    let out_dir: Option<PathBuf> = args.get(2).map(PathBuf::from);
    let n_workers = physical_cores();

    let mut file = File::open(zip_path).unwrap_or_else(|e| {
        eprintln!("error opening {zip_path}: {e}");
        std::process::exit(1);
    });

    let entries = read_central_directory_from(&mut file).unwrap_or_else(|e| {
        eprintln!("error: {e}");
        std::process::exit(1);
    });

    // ── List mode: no decompression, just print CD info ────────────────────
    if out_dir.is_none() {
        let files: Vec<_> = entries.iter().filter(|e| !e.is_directory).collect();
        for e in &files {
            println!("{:>10}  {}", e.uncompressed_size, e.name);
        }
        println!("{} entries", files.len());
        return;
    }
    let out_dir = out_dir.unwrap();

    // ── Resolve entries (compute data offsets via local headers) ──────────
    let mut resolved: Vec<ResolvedEntry> = Vec::new();
    let mut sorted: Vec<&EntryLocation> = entries.iter().filter(|e| !e.is_directory).collect();
    sorted.sort_by_key(|e| e.local_header_offset);
    for loc in &sorted {
        match local_data_offset(&mut file, loc) {
            Ok(off) => resolved.push(ResolvedEntry {
                name: loc.name.clone(),
                data_offset: off,
                compressed_size: loc.compressed_size,
                uncompressed_size: loc.uncompressed_size,
                compression_method: loc.compression_method,
            }),
            Err(e) => eprintln!("skipping {}: {}", loc.name, e),
        }
    }

    eprintln!(
        "lzip: {} entries, {} workers, {} ring slots ({} MB each)",
        resolved.len(), n_workers, RING_SLOTS, SLOT_SIZE / (1024 * 1024),
    );

    // ── Channels ──────────────────────────────────────────────────────────
    let (free_tx, free_rx) = mpsc::sync_channel::<Vec<u8>>(RING_SLOTS);
    let (filled_tx, filled_rx) = mpsc::sync_channel::<FilledSlot>(RING_SLOTS);
    let (work_tx, work_rx) = mpsc::sync_channel::<WorkItem>(n_workers * 2);
    let (collector_tx, collector_rx) = mpsc::sync_channel::<CollectorMsg>(n_workers * 4);
    let (write_tx, write_rx) = mpsc::sync_channel::<WriteMsg>(RING_SLOTS);

    // Pre-allocate RING_SLOTS ring slots.
    for _ in 0..RING_SLOTS {
        free_tx.send(vec![0u8; SLOT_SIZE]).unwrap();
    }

    // ── Reader thread ─────────────────────────────────────────────────────
    let reader_handle = thread::Builder::new()
        .name("lzip-reader".into())
        .spawn(move || reader_thread(file, resolved, free_rx, filled_tx))
        .expect("spawn reader");

    // ── Writer thread ─────────────────────────────────────────────────────
    let out_dir_w = out_dir.clone();
    let writer_handle = thread::Builder::new()
        .name("lzip-writer".into())
        .spawn(move || {
            let mut current: Option<(String, File)> = None;
            for msg in write_rx {
                match msg {
                    WriteMsg::EntryStart(name) => {
                        let dest = out_dir_w.join(&name);
                        if let Some(parent) = dest.parent() {
                            if let Err(e) = fs::create_dir_all(parent) {
                                eprintln!("mkdir {}: {e}", parent.display());
                                current = None;
                                continue;
                            }
                        }
                        match File::create(&dest) {
                            Ok(f) => current = Some((name, f)),
                            Err(e) => {
                                eprintln!("create {}: {e}", dest.display());
                                current = None;
                            }
                        }
                    }
                    WriteMsg::Chunk(data) => {
                        if let Some((_, ref mut f)) = current {
                            if let Err(e) = f.write_all(&data) {
                                eprintln!("write error: {e}");
                            }
                        }
                    }
                    WriteMsg::EntryEnd => {
                        if let Some((name, _)) = current.take() {
                            println!("{name}");
                        }
                    }
                }
            }
        })
        .expect("spawn writer");

    // ── N decoder worker threads ──────────────────────────────────────────
    let work_rx = Arc::new(Mutex::new(work_rx));
    let mut worker_handles = Vec::with_capacity(n_workers);
    for w in 0..n_workers {
        let work_rx = work_rx.clone();
        let collector_tx = collector_tx.clone();
        worker_handles.push(
            thread::Builder::new()
                .name(format!("lzip-w{w}"))
                .spawn(move || {
                    loop {
                        let item = {
                            let rx = work_rx.lock().unwrap();
                            match rx.recv() {
                                Ok(i) => i,
                                Err(_) => break,
                            }
                        };
                        // Safety: data_ptr is valid for data_len bytes until the
                        // collector recycles the slot (after done_count ==
                        // decode_segments). The collector controls slot lifetime.
                        let data = unsafe {
                            std::slice::from_raw_parts(item.data_ptr, item.data_len)
                        };
                        let segment = &data[item.start_byte..item.end_byte];

                        let output = if item.is_store {
                            Ok(segment.to_vec())
                        } else {
                            let mut out = Vec::new();
                            decode_segment_into(segment, &mut out).map(|_| out)
                        };

                        let _ = collector_tx.send(CollectorMsg::Result(SegmentResult {
                            chunk_id: item.chunk_id,
                            segment_id: item.segment_id,
                            output,
                        }));
                    }
                })
                .expect("spawn worker"),
        );
    }
    drop(work_rx);

    // ── Collector thread ──────────────────────────────────────────────────
    let inflight_tx = collector_tx.clone();
    drop(collector_tx);

    let slot_return_for_collector = free_tx.clone();
    let collector_handle = thread::Builder::new()
        .name("lzip-collect".into())
        .spawn(move || {
            let mut in_flight: Vec<InFlightSlot> = Vec::new();
            let mut next_write_id: u64 = 0;
            while let Ok(msg) = collector_rx.recv() {
                match msg {
                    CollectorMsg::NewSlot(s) => in_flight.push(s),
                    CollectorMsg::Result(r)  => apply_result(&mut in_flight, r),
                }
                flush_completed(&mut in_flight, &mut next_write_id,
                                &write_tx, &slot_return_for_collector);
            }
            flush_completed(&mut in_flight, &mut next_write_id,
                            &write_tx, &slot_return_for_collector);
        })
        .expect("spawn collector");

    // ── Main: carry + split + post work ───────────────────────────────────
    let mut carry: Vec<u8> = Vec::new();
    let mut chunk_id: u64 = 0;
    // Fallback buffer: accumulates compressed data for multi-chunk entries
    // that lack Z_FULL_FLUSH boundaries. Flushed as a single segment on
    // the last chunk of the entry.
    let mut fallback_buf: Vec<u8> = Vec::new();

    for filled in filled_rx {
        let FilledSlot {
            slot, read_len, is_ring_slot, entry_name,
            is_first_of_entry, is_last_of_entry, is_store,
        } = filled;

        if is_first_of_entry {
            assert!(carry.is_empty());
            fallback_buf.clear();
        }

        // Compute data region: ring slots have carry headroom, ad-hoc slots don't.
        let (data_start, data_end) = if is_ring_slot {
            let carry_len = carry.len();
            assert!(carry_len <= CARRY_HEADROOM);
            let start = CARRY_HEADROOM - carry_len;
            (start, CARRY_HEADROOM + read_len)
        } else {
            (0, read_len)
        };

        // Copy carry into headroom (must happen before we read data).
        let mut work_slot = slot;
        if is_ring_slot {
            let carry_len = carry.len();
            let start = CARRY_HEADROOM - carry_len;
            work_slot[start..CARRY_HEADROOM].copy_from_slice(&carry);
        }

        // ── Segmentation strategy ────────────────────────────────────────
        //
        // STORE entries and small single-chunk DEFLATE entries are trivial.
        // Large entries try split_chunk for parallel decode. If no flush
        // boundaries exist in a multi-chunk entry, accumulate into
        // fallback_buf and decompress as one segment when the entry ends.

        // Check if we're already accumulating (prior chunk had no flush).
        if !fallback_buf.is_empty() {
            fallback_buf.extend_from_slice(&work_slot[data_start..data_end]);
            if is_ring_slot { free_tx.send(work_slot).ok(); }
            carry.clear();
            if !is_last_of_entry { continue; }
            // Last chunk: flush the accumulated buffer as a single segment.
            let fb = std::mem::take(&mut fallback_buf);
            let fb_len = fb.len();
            let inflight = InFlightSlot {
                chunk_id,
                slot: fb,
                is_ring_slot: false,
                decode_segments: 1,
                results: vec![None],
                done_count: 0,
                entry_name: entry_name.clone(),
                is_first_of_entry: false,
                is_last_of_entry: true,
            };
            let data_ptr = inflight.slot.as_ptr();
            if inflight_tx.send(CollectorMsg::NewSlot(inflight)).is_err() { break; }
            if work_tx.send(WorkItem {
                chunk_id, segment_id: 0, data_ptr,
                data_len: fb_len, start_byte: 0, end_byte: fb_len, is_store,
                output_size_hint: 0,
            }).is_err() { break; }
            chunk_id += 1;
            continue;
        }

        // Determine segments. `final_slot` is what the InFlightSlot will own.
        // `final_start` / `final_len` describe the data region within it.
        let (final_slot, final_is_ring, final_start, final_len,
             segment_starts, decode_segments, consumed);

        if is_store {
            final_slot = work_slot;
            final_is_ring = is_ring_slot;
            final_start = data_start;
            final_len = data_end - data_start;
            segment_starts = vec![0usize];
            decode_segments = 1;
            consumed = final_len;
        } else if is_first_of_entry && is_last_of_entry {
            // Single-chunk entry: try parallel, fall back to one segment.
            let data = &work_slot[data_start..data_end];
            match split_chunk(data, n_workers, true) {
                Some(s) => {
                    segment_starts = s.segment_starts;
                    decode_segments = s.decode_segments;
                    consumed = s.consumed;
                }
                None => {
                    segment_starts = vec![0usize];
                    decode_segments = 1;
                    consumed = data.len();
                }
            }
            final_slot = work_slot;
            final_is_ring = is_ring_slot;
            final_start = data_start;
            final_len = data_end - data_start;
        } else {
            // Multi-chunk entry: try to find flush boundaries.
            let data = &work_slot[data_start..data_end];
            match split_chunk(data, n_workers, is_last_of_entry) {
                Some(s) => {
                    segment_starts = s.segment_starts;
                    decode_segments = s.decode_segments;
                    consumed = s.consumed;
                    final_slot = work_slot;
                    final_is_ring = is_ring_slot;
                    final_start = data_start;
                    final_len = data_end - data_start;
                }
                None => {
                    // No flush boundaries — enter fallback mode.
                    fallback_buf.extend_from_slice(data);
                    if is_ring_slot { free_tx.send(work_slot).ok(); }
                    carry.clear();
                    if !is_last_of_entry { continue; }
                    // Last chunk and no boundaries anywhere — single segment.
                    let fb = std::mem::take(&mut fallback_buf);
                    let fb_len = fb.len();
                    final_slot = fb;
                    final_is_ring = false;
                    final_start = 0;
                    final_len = fb_len;
                    segment_starts = vec![0usize];
                    decode_segments = 1;
                    consumed = fb_len;
                }
            }
        };

        // Carry: unconsumed tail flows into next slot's headroom.
        carry.clear();
        if !is_last_of_entry {
            carry.extend_from_slice(&final_slot[final_start + consumed..final_start + final_len]);
        }

        // Register InFlightSlot with collector, then post work items.
        let inflight = InFlightSlot {
            chunk_id,
            slot: final_slot,
            is_ring_slot: final_is_ring,
            decode_segments,
            results: (0..decode_segments).map(|_| None).collect(),
            done_count: 0,
            entry_name: entry_name.clone(),
            is_first_of_entry,
            is_last_of_entry,
        };
        let data_ptr = unsafe { inflight.slot.as_ptr().add(final_start) };
        let data_len = final_len;

        if inflight_tx.send(CollectorMsg::NewSlot(inflight)).is_err() { break; }

        for i in 0..decode_segments {
            let start_byte = segment_starts[i];
            let end_byte = if i + 1 < segment_starts.len() {
                segment_starts[i + 1]
            } else {
                data_len
            };
            let item = WorkItem {
                chunk_id,
                segment_id: i,
                data_ptr,
                data_len,
                start_byte,
                end_byte,
                is_store,
                output_size_hint: 0,
            };
            if work_tx.send(item).is_err() { break; }
        }

        chunk_id += 1;
    }

    // Shut down the pipeline in order.
    drop(work_tx);       // workers exit
    drop(inflight_tx);   // collector exits when all workers' senders also drop
    drop(free_tx);       // reader exits if it was still holding the recv side
    for h in worker_handles {
        h.join().expect("worker panicked");
    }
    collector_handle.join().expect("collector panicked");
    reader_handle.join().expect("reader panicked");
    writer_handle.join().expect("writer panicked");
}