Skip to main content

grit_lib/parallel/
map.rs

1//! Ordered parallel map helpers built on the scoped thread pool.
2
3use std::convert::Infallible;
4use std::num::NonZeroUsize;
5
6use super::policy::ParallelOptions;
7use super::pool::{run_ordered_pool, run_ordered_pool_with_local, DEFAULT_CHUNK};
8use super::ParallelHashError;
9
10#[must_use]
11pub(crate) fn pool_worthwhile(len: usize, total_bytes: usize, options: ParallelOptions) -> bool {
12    if options.threads.get() <= 1 || len < 2 {
13        return false;
14    }
15    if let Some(threshold) = options.signed_entry_threshold {
16        if (len as i64) < threshold {
17            return false;
18        }
19    }
20    if options.min_items == 0 {
21        return true;
22    }
23    if len < options.min_items {
24        return false;
25    }
26    if let Some(min_bytes) = options.min_total_bytes {
27        if total_bytes < min_bytes {
28            return false;
29        }
30    }
31    true
32}
33
34fn worker_count(options: ParallelOptions, len: usize) -> usize {
35    options.threads.get().min(len).max(1)
36}
37
38/// Run `f` over `items` with up to `options.threads` workers, preserving order.
39///
40/// Falls back to a serial loop when [`ParallelOptions::min_items`] / `min_total_bytes`
41/// thresholds are not met.
42#[must_use]
43pub fn par_map<T, R, F>(items: &[T], options: ParallelOptions, total_bytes: usize, f: F) -> Vec<R>
44where
45    T: Sync,
46    R: Send,
47    F: Fn(&T) -> R + Sync,
48{
49    try_par_map::<T, R, Infallible, _>(items, options, total_bytes, |item| Ok(f(item)))
50        .unwrap_or_else(|e| match e {
51            ParallelHashError::Task(infallible) => match infallible {},
52        })
53}
54
55/// Fallible ordered parallel map; first error cancels the pool.
56///
57/// # Errors
58///
59/// Returns [`ParallelHashError::Task`] with the first closure error.
60pub fn try_par_map<T, R, E, F>(
61    items: &[T],
62    options: ParallelOptions,
63    total_bytes: usize,
64    f: F,
65) -> Result<Vec<R>, ParallelHashError<E>>
66where
67    T: Sync,
68    R: Send,
69    E: Send + 'static,
70    F: Fn(&T) -> Result<R, E> + Sync,
71{
72    let len = items.len();
73    if len == 0 {
74        return Ok(Vec::new());
75    }
76    let n_workers = worker_count(options, len);
77    if !pool_worthwhile(len, total_bytes, options) {
78        return items
79            .iter()
80            .map(&f)
81            .collect::<Result<Vec<R>, E>>()
82            .map_err(ParallelHashError::Task);
83    }
84    run_ordered_pool(items, n_workers, DEFAULT_CHUNK, f).map_err(ParallelHashError::Task)
85}
86
87/// Like [`try_par_map`], but ignores batch thresholds (still requires `threads` > 1).
88///
89/// # Errors
90///
91/// Returns [`ParallelHashError::Task`] with the first closure error.
92pub fn try_par_map_force<T, R, E, F>(
93    items: &[T],
94    options: ParallelOptions,
95    f: F,
96) -> Result<Vec<R>, ParallelHashError<E>>
97where
98    T: Sync,
99    R: Send,
100    E: Send + 'static,
101    F: Fn(&T) -> Result<R, E> + Sync,
102{
103    let len = items.len();
104    if len == 0 {
105        return Ok(Vec::new());
106    }
107    let n_workers = worker_count(options, len);
108    if n_workers <= 1 {
109        return items
110            .iter()
111            .map(&f)
112            .collect::<Result<Vec<R>, E>>()
113            .map_err(ParallelHashError::Task);
114    }
115    run_ordered_pool(items, n_workers, DEFAULT_CHUNK, f).map_err(ParallelHashError::Task)
116}
117
118/// Parallel map with local state from `init` for each item.
119///
120/// # Errors
121///
122/// Returns [`ParallelHashError::Task`] with the first closure error.
123pub fn try_par_map_with_local<T, R, E, S, Init, F>(
124    items: &[T],
125    options: ParallelOptions,
126    total_bytes: usize,
127    init: Init,
128    f: F,
129) -> Result<Vec<R>, ParallelHashError<E>>
130where
131    T: Sync,
132    R: Send,
133    E: Send + 'static,
134    S: Send,
135    Init: Fn() -> S + Sync,
136    F: Fn(&mut S, &T) -> Result<R, E> + Sync,
137{
138    let len = items.len();
139    if len == 0 {
140        return Ok(Vec::new());
141    }
142    let n_workers = worker_count(options, len);
143    if !pool_worthwhile(len, total_bytes, options) {
144        let mut state = init();
145        return items
146            .iter()
147            .map(|item| f(&mut state, item))
148            .collect::<Result<Vec<R>, E>>()
149            .map_err(ParallelHashError::Task);
150    }
151    run_ordered_pool_with_local(items, n_workers, DEFAULT_CHUNK, init, f)
152        .map_err(ParallelHashError::Task)
153}
154
155/// Legacy helper: returns true when parallel hashing is expected to beat a serial loop.
156#[must_use]
157pub fn parallel_worthwhile(
158    item_count: usize,
159    total_bytes: usize,
160    threads: NonZeroUsize,
161    min_items: usize,
162    min_total_bytes: usize,
163) -> bool {
164    pool_worthwhile(
165        item_count,
166        total_bytes,
167        ParallelOptions {
168            threads,
169            min_items,
170            min_total_bytes: Some(min_total_bytes),
171            signed_entry_threshold: None,
172        },
173    )
174}