loam_time/
par.rs

1use std::sync::{Mutex, OnceLock};
2
3use thiserror::Error;
4
5pub trait ChunkSource {
6    fn run_next(&self) -> bool;
7}
8
9pub trait Executor: Send + Sync {
10    fn parallelism(&self) -> usize;
11
12    /// `workers` is a request; each worker drains `chunks` until `run_next` is false.
13    fn for_each_chunk(&self, workers: usize, chunks: &(dyn ChunkSource + Sync));
14}
15
16pub struct Sequential;
17
18impl Executor for Sequential {
19    fn parallelism(&self) -> usize {
20        1
21    }
22
23    fn for_each_chunk(&self, _workers: usize, chunks: &(dyn ChunkSource + Sync)) {
24        while chunks.run_next() {}
25    }
26}
27
28#[derive(Debug, Error)]
29#[error("a parallel executor is already installed")]
30pub struct AlreadyInstalled;
31
32static INSTALLED: OnceLock<Box<dyn Executor>> = OnceLock::new();
33
34static SEQUENTIAL: Sequential = Sequential;
35
36/// Once per process; before it, and after a refusal, [`executor`] is [`Sequential`].
37pub fn install(executor: Box<dyn Executor>) -> Result<(), AlreadyInstalled> {
38    INSTALLED.set(executor).map_err(|_| AlreadyInstalled)
39}
40
41pub fn executor() -> &'static dyn Executor {
42    match INSTALLED.get() {
43        Some(installed) => installed.as_ref(),
44        None => &SEQUENTIAL,
45    }
46}
47
48struct Chunks<'a, T, F> {
49    rest: Mutex<&'a mut [T]>,
50    chunk: usize,
51    task: F,
52}
53
54impl<T: Send, F: Fn(&mut [T]) + Sync> ChunkSource for Chunks<'_, T, F> {
55    fn run_next(&self) -> bool {
56        let part = {
57            let Ok(mut rest) = self.rest.lock() else {
58                return false;
59            };
60            if rest.is_empty() {
61                return false;
62            }
63            let take = rest.len().min(self.chunk);
64            let (head, tail) = std::mem::take(&mut *rest).split_at_mut(take);
65            *rest = tail;
66            head
67        };
68        (self.task)(part);
69        true
70    }
71}
72
73/// Runs `task` on disjoint chunks of `data` in no fixed order.
74pub fn for_each_chunk<T: Send>(data: &mut [T], chunk: usize, task: impl Fn(&mut [T]) + Sync) {
75    let chunk = chunk.max(1);
76    let executor = executor();
77    let workers = executor
78        .parallelism()
79        .min(data.len().div_ceil(chunk))
80        .max(1);
81    if workers == 1 {
82        for part in data.chunks_mut(chunk) {
83            task(part);
84        }
85        return;
86    }
87    let chunks = Chunks {
88        rest: Mutex::new(data),
89        chunk,
90        task,
91    };
92    executor.for_each_chunk(workers, &chunks);
93}
loam_time/par.rs

loam_time/
par.rs
