lbzip2 0.5.2

//! Scoped-thread parallel map — rayon-free work-stealing helper.
//!
//! Replaces lbzip2's old dedicated `rayon::ThreadPool`. Spawns a fixed
//! pool of scoped threads (sized to physical cores, overridable via
//! `LBZIP2_THREADS`) that dynamically claim work items from `0..n`, so a
//! few cheap items and a few expensive ones still balance across cores.

use std::cell::UnsafeCell;
use std::mem::MaybeUninit;
use std::sync::atomic::{AtomicUsize, Ordering};

/// Worker budget: `LBZIP2_THREADS` if set, else available parallelism.
///
/// Separate from any caller pool so bzip2 decode doesn't oversubscribe.
pub fn worker_count() -> usize {
    std::env::var("LBZIP2_THREADS")
        .ok()
        .and_then(|v| v.parse::<usize>().ok())
        .unwrap_or_else(|| {
            std::thread::available_parallelism()
                .map(|n| n.get())
                .unwrap_or(4)
        })
        .max(1)
}

/// Apply `f` to every index in `0..n` in parallel, returning the results
/// in index order.
pub fn par_map<T, F>(n: usize, f: F) -> Vec<T>
where
    T: Send,
    F: Fn(usize) -> T + Sync,
{
    if n == 0 {
        return Vec::new();
    }
    let threads = worker_count().min(n);
    if threads <= 1 {
        return (0..n).map(f).collect();
    }

    // Perf: a disjoint-index scatter needs no per-item synchronisation. The
    // atomic `next` hands each index to exactly one worker, so each output slot
    // is written by exactly one thread — there is no contention to guard. We
    // drop the old `Vec<Mutex<Option<T>>>` (a mutex + heap alloc + Option per
    // item) for a plain `Vec<UnsafeCell<MaybeUninit<T>>>` written via raw,
    // unsynchronised pointer stores.
    //
    // `Slots` is a thin Sync wrapper so the cell vec can be shared across the
    // scope; safety is upheld by the disjoint-index invariant below.
    struct Slots<T>(Vec<UnsafeCell<MaybeUninit<T>>>);
    // Safety: workers only ever write the unique index they claimed from
    // `next`; no two threads touch the same cell, and all writes complete
    // before the scope joins and we read the cells.
    unsafe impl<T: Send> Sync for Slots<T> {}

    let next = AtomicUsize::new(0);
    let slots = Slots((0..n).map(|_| UnsafeCell::new(MaybeUninit::uninit())).collect());
    let next_ref = &next;
    let slots_ref = &slots;
    let f_ref = &f;

    std::thread::scope(|s| {
        for _ in 0..threads {
            s.spawn(move || loop {
                let i = next_ref.fetch_add(1, Ordering::Relaxed);
                if i >= n {
                    break;
                }
                let v = f_ref(i);
                // Safety: `i` is unique to this worker (atomic fetch_add), so
                // this cell is written exactly once and never aliased.
                unsafe { (*slots_ref.0[i].get()).write(v); }
            });
        }
    });

    // Every index in 0..n was claimed and written exactly once (the workers
    // exhaust `next` before the scope joins), so all cells are initialised.
    slots
        .0
        .into_iter()
        .map(|cell| unsafe { cell.into_inner().assume_init() })
        .collect()
}