48 lines
1.8 KiB
Rust
48 lines
1.8 KiB
Rust
//! Tiny data-parallel primitive shared across the crate — no work-stealing dependency, just scoped
|
|
//! threads pulling work by an atomic index. Lives in the library (not the binary) so lib modules
|
|
//! (`xref`, …) and the CI pipeline can parallelise directly, not only the CLI front-end.
|
|
|
|
use std::sync::Mutex;
|
|
use std::sync::atomic::{AtomicUsize, Ordering};
|
|
|
|
/// Run `f` over `items` across `nthreads` scoped threads, pulling work by atomic index
|
|
/// (dynamic load-balancing without a work-stealing dep). Results returned in input order, so callers
|
|
/// that merge them stay deterministic regardless of which thread finished which item.
|
|
pub fn parallel_map<T, R, F>(items: &[T], nthreads: usize, f: F) -> Vec<R>
|
|
where
|
|
T: Sync,
|
|
R: Send,
|
|
F: Fn(&T) -> R + Sync,
|
|
{
|
|
let len = items.len();
|
|
let nthreads = nthreads.clamp(1, len.max(1));
|
|
let next = AtomicUsize::new(0);
|
|
let out: Mutex<Vec<(usize, R)>> = Mutex::new(Vec::with_capacity(len));
|
|
std::thread::scope(|s| {
|
|
for _ in 0..nthreads {
|
|
s.spawn(|| {
|
|
let mut local = Vec::new();
|
|
loop {
|
|
let i = next.fetch_add(1, Ordering::Relaxed);
|
|
if i >= len {
|
|
break;
|
|
}
|
|
local.push((i, f(&items[i])));
|
|
}
|
|
out.lock().unwrap().extend(local);
|
|
});
|
|
}
|
|
});
|
|
let mut v = out.into_inner().unwrap();
|
|
v.sort_by_key(|(i, _)| *i);
|
|
v.into_iter().map(|(_, r)| r).collect()
|
|
}
|
|
|
|
/// Requested thread count, or the machine's available parallelism.
|
|
pub fn default_threads(threads: Option<usize>) -> usize {
|
|
threads.unwrap_or_else(|| {
|
|
std::thread::available_parallelism()
|
|
.map(|n| n.get())
|
|
.unwrap_or(4)
|
|
})
|
|
}
|