halo2_axiom/
multicore.rs

1#[cfg(all(
2    feature = "multicore",
3    target_arch = "wasm32",
4    not(target_feature = "atomics")
5))]
6compile_error!(
7    "The multicore feature flag is not supported on wasm32 architectures without atomics"
8);
9
10pub use maybe_rayon::{
11    iter::{IntoParallelIterator, IntoParallelRefMutIterator, ParallelIterator},
12    join, scope, Scope,
13};
14
15#[cfg(feature = "multicore")]
16pub use maybe_rayon::{
17    current_num_threads,
18    iter::{IndexedParallelIterator, IntoParallelRefIterator},
19    slice::{ParallelSlice, ParallelSliceMut},
20};
21
22#[cfg(not(feature = "multicore"))]
23pub fn current_num_threads() -> usize {
24    1
25}
26
27pub trait TryFoldAndReduce<T, E> {
28    /// Implements `iter.try_fold().try_reduce()` for `rayon::iter::ParallelIterator`,
29    /// falling back on `Iterator::try_fold` when the `multicore` feature flag is
30    /// disabled.
31    /// The `try_fold_and_reduce` function can only be called by a iter with
32    /// `Result<T, E>` item type because the `fold_op` must meet the trait
33    /// bounds of both `try_fold` and `try_reduce` from rayon.
34    fn try_fold_and_reduce(
35        self,
36        identity: impl Fn() -> T + Send + Sync,
37        fold_op: impl Fn(T, Result<T, E>) -> Result<T, E> + Send + Sync,
38    ) -> Result<T, E>;
39}
40
41#[cfg(feature = "multicore")]
42impl<T, E, I> TryFoldAndReduce<T, E> for I
43where
44    T: Send + Sync,
45    E: Send + Sync,
46    I: maybe_rayon::iter::ParallelIterator<Item = Result<T, E>>,
47{
48    fn try_fold_and_reduce(
49        self,
50        identity: impl Fn() -> T + Send + Sync,
51        fold_op: impl Fn(T, Result<T, E>) -> Result<T, E> + Send + Sync,
52    ) -> Result<T, E> {
53        self.try_fold(&identity, &fold_op)
54            .try_reduce(&identity, |a, b| fold_op(a, Ok(b)))
55    }
56}
57
58#[cfg(not(feature = "multicore"))]
59impl<T, E, I> TryFoldAndReduce<T, E> for I
60where
61    I: std::iter::Iterator<Item = Result<T, E>>,
62{
63    fn try_fold_and_reduce(
64        mut self,
65        identity: impl Fn() -> T + Send + Sync,
66        fold_op: impl Fn(T, Result<T, E>) -> Result<T, E> + Send + Sync,
67    ) -> Result<T, E> {
68        self.try_fold(identity(), fold_op)
69    }
70}