winter_utils/
iterators.rs

1// Copyright (c) Facebook, Inc. and its affiliates.
2//
3// This source code is licensed under the MIT license found in the
4// LICENSE file in the root directory of this source tree.
5
6//! Components needed for parallel iterators.
7//!
8//! When `concurrent` feature is enabled, this module re-exports `rayon::prelude`. Otherwise,
9//! this is an empty module.
10
11#[cfg(feature = "concurrent")]
12pub use rayon::{current_num_threads as rayon_num_threads, prelude::*};
13
14/// Returns either a regular or a parallel iterator depending on whether `concurrent` feature
15/// is enabled.
16///
17/// When `concurrent` feature is enabled, creates a parallel iterator; otherwise, creates a
18/// regular iterator. Optionally, `min_length` can be used to specify the minimum length of
19/// iterator to be processed in each thread.
20///
21/// Adapted from: <https://github.com/arkworks-rs/utils/blob/master/src/lib.rs>
22#[macro_export]
23macro_rules! iter {
24    ($e:expr) => {{
25        #[cfg(feature = "concurrent")]
26        let result = $e.par_iter();
27
28        #[cfg(not(feature = "concurrent"))]
29        let result = $e.iter();
30
31        result
32    }};
33    ($e:expr, $min_len:expr) => {{
34        #[cfg(feature = "concurrent")]
35        let result = $e.par_iter().with_min_len($min_len);
36
37        #[cfg(not(feature = "concurrent"))]
38        let result = $e.iter();
39
40        result
41    }};
42}
43
44/// Returns either a regular or a parallel mutable iterator depending on whether `concurrent`
45/// feature is enabled.
46///
47/// When `concurrent` feature is enabled, creates a mutable parallel iterator; otherwise,
48/// creates a regular mutable iterator. Optionally, `min_length` can be used to specify the
49/// minimum length of iterator to be processed in each thread.
50///
51/// Adapted from: <https://github.com/arkworks-rs/utils/blob/master/src/lib.rs>
52#[macro_export]
53macro_rules! iter_mut {
54    ($e:expr) => {{
55        #[cfg(feature = "concurrent")]
56        let result = $e.par_iter_mut();
57
58        #[cfg(not(feature = "concurrent"))]
59        let result = $e.iter_mut();
60
61        result
62    }};
63    ($e:expr, $min_len:expr) => {{
64        #[cfg(feature = "concurrent")]
65        let result = $e.par_iter_mut().with_min_len($min_len);
66
67        #[cfg(not(feature = "concurrent"))]
68        let result = $e.iter_mut();
69
70        result
71    }};
72}
73
74/// Applies a procedure to the provided slice either in a single thread or multiple threads
75/// based on whether `concurrent` feature is enabled.
76///
77/// When `concurrent` feature is enabled, breaks the slice into batches and processes each
78/// batch in a separate thread; otherwise, the entire slice is processed as a single batch
79/// in one thread. Optionally, `min_batch_size` can be used to specify the minimum size of
80/// the resulting batches.
81#[macro_export]
82macro_rules! batch_iter_mut {
83    ($e: expr, $c: expr) => {
84        #[cfg(feature = "concurrent")]
85        {
86            let batch_size = $e.len() / rayon_num_threads().next_power_of_two();
87            if batch_size < 1 {
88                $c($e, 0);
89            }
90            else {
91                $e.par_chunks_mut(batch_size).enumerate().for_each(|(i, batch)| {
92                    $c(batch, i * batch_size);
93                });
94            }
95        }
96
97        #[cfg(not(feature = "concurrent"))]
98        $c($e, 0);
99    };
100    ($e: expr, $min_batch_size: expr, $c: expr) => {
101        #[cfg(feature = "concurrent")]
102        {
103            let batch_size = $e.len() / rayon_num_threads().next_power_of_two();
104            if batch_size < $min_batch_size {
105                $c($e, 0);
106            }
107            else {
108                $e.par_chunks_mut(batch_size).enumerate().for_each(|(i, batch)| {
109                    $c(batch, i * batch_size);
110                });
111            }
112        }
113
114        #[cfg(not(feature = "concurrent"))]
115        $c($e, 0);
116    };
117}
winter_utils/iterators.rs

winter_utils/
iterators.rs
