swc_ecma_utils/
parallel.rs

1//! Module for parallel processing
2
3use once_cell::sync::Lazy;
4use par_core::join;
5use swc_common::GLOBALS;
6use swc_ecma_ast::*;
7
8static CPU_COUNT: Lazy<usize> = Lazy::new(num_cpus::get);
9
10pub fn cpu_count() -> usize {
11    *CPU_COUNT
12}
13
14pub trait Parallel: swc_common::sync::Send + swc_common::sync::Sync {
15    /// Used to create visitor.
16    fn create(&self) -> Self;
17
18    /// This can be called in anytime.
19    fn merge(&mut self, other: Self);
20
21    /// Invoked after visiting all [Stmt]s, possibly in parallel.
22    ///
23    ///
24    /// Note: `visit_*_par` never calls this.
25    fn after_stmts(&mut self, _stmts: &mut Vec<Stmt>) {}
26
27    /// Invoked after visiting all [ModuleItem]s, possibly in parallel.
28    ///
29    ///
30    /// Note: `visit_*_par` never calls this.
31    fn after_module_items(&mut self, _stmts: &mut Vec<ModuleItem>) {}
32}
33
34pub trait ParallelExt: Parallel {
35    /// Invoke `op` in parallel, if `swc_ecma_utils` is compiled with
36    /// concurrent feature enabled and `nodes.len()` is bigger than threshold.
37    ///
38    ///
39    /// This configures [GLOBALS], while not configuring [HANDLER] nor [HELPERS]
40    fn maybe_par<I, F>(&mut self, threshold: usize, nodes: I, op: F)
41    where
42        I: IntoItems,
43        F: Send + Sync + Fn(&mut Self, I::Elem),
44    {
45        self.maybe_par_idx(threshold, nodes, |v, _, n| op(v, n))
46    }
47
48    /// Invoke `op` in parallel, if `swc_ecma_utils` is compiled with
49    /// concurrent feature enabled and `nodes.len()` is bigger than threshold.
50    ///
51    ///
52    /// This configures [GLOBALS], while not configuring [HANDLER] nor [HELPERS]
53    fn maybe_par_idx<I, F>(&mut self, threshold: usize, nodes: I, op: F)
54    where
55        I: IntoItems,
56        F: Send + Sync + Fn(&mut Self, usize, I::Elem),
57    {
58        self.maybe_par_idx_raw(threshold, nodes.into_items(), &op)
59    }
60
61    /// If you don't have a special reason, use [`ParallelExt::maybe_par`] or
62    /// [`ParallelExt::maybe_par_idx`] instead.
63    fn maybe_par_idx_raw<I, F>(&mut self, threshold: usize, nodes: I, op: &F)
64    where
65        I: Items,
66        F: Send + Sync + Fn(&mut Self, usize, I::Elem);
67}
68
69#[cfg(feature = "concurrent")]
70impl<T> ParallelExt for T
71where
72    T: Parallel,
73{
74    fn maybe_par_idx_raw<I, F>(&mut self, threshold: usize, nodes: I, op: &F)
75    where
76        I: Items,
77        F: Send + Sync + Fn(&mut Self, usize, I::Elem),
78    {
79        if nodes.len() >= threshold {
80            GLOBALS.with(|globals| {
81                let len = nodes.len();
82                if len == 0 {
83                    return;
84                }
85
86                if len == 1 {
87                    op(self, 0, nodes.into_iter().next().unwrap());
88                    return;
89                }
90
91                let (na, nb) = nodes.split_at(len / 2);
92                let mut vb = Parallel::create(&*self);
93
94                let (_, vb) = join(
95                    || {
96                        GLOBALS.set(globals, || {
97                            self.maybe_par_idx_raw(threshold, na, op);
98                        })
99                    },
100                    || {
101                        GLOBALS.set(globals, || {
102                            vb.maybe_par_idx_raw(threshold, nb, op);
103
104                            vb
105                        })
106                    },
107                );
108
109                Parallel::merge(self, vb);
110            });
111
112            return;
113        }
114
115        for (idx, n) in nodes.into_iter().enumerate() {
116            op(self, idx, n);
117        }
118    }
119}
120
121#[cfg(not(feature = "concurrent"))]
122impl<T> ParallelExt for T
123where
124    T: Parallel,
125{
126    fn maybe_par_idx_raw<I, F>(&mut self, _threshold: usize, nodes: I, op: &F)
127    where
128        I: Items,
129        F: Send + Sync + Fn(&mut Self, usize, I::Elem),
130    {
131        for (idx, n) in nodes.into_iter().enumerate() {
132            op(self, idx, n);
133        }
134    }
135}
136
137use self::private::Sealed;
138
139mod private {
140    pub trait Sealed {}
141
142    impl<T> Sealed for Vec<T> {}
143    impl<T> Sealed for &mut Vec<T> {}
144    impl<T> Sealed for &mut [T] {}
145    impl<T> Sealed for &[T] {}
146}
147pub trait IntoItems: Sealed {
148    type Elem;
149    type Items: Items<Elem = Self::Elem>;
150
151    fn into_items(self) -> Self::Items;
152}
153
154impl<T, I> IntoItems for I
155where
156    I: Items<Elem = T>,
157{
158    type Elem = T;
159    type Items = I;
160
161    fn into_items(self) -> Self::Items {
162        self
163    }
164}
165
166impl<'a, T> IntoItems for &'a mut Vec<T>
167where
168    T: Send + Sync,
169{
170    type Elem = &'a mut T;
171    type Items = &'a mut [T];
172
173    fn into_items(self) -> Self::Items {
174        self
175    }
176}
177
178/// This is considered as a private type and it's NOT A PUBLIC API.
179#[allow(clippy::len_without_is_empty)]
180pub trait Items: Sized + IntoIterator<Item = Self::Elem> + Send + Sealed {
181    type Elem: Send + Sync;
182
183    fn len(&self) -> usize;
184
185    fn split_at(self, idx: usize) -> (Self, Self);
186}
187
188impl<T> Items for Vec<T>
189where
190    T: Send + Sync,
191{
192    type Elem = T;
193
194    fn len(&self) -> usize {
195        Vec::len(self)
196    }
197
198    fn split_at(mut self, at: usize) -> (Self, Self) {
199        let b = self.split_off(at);
200
201        (self, b)
202    }
203}
204
205impl<'a, T> Items for &'a mut [T]
206where
207    T: Send + Sync,
208{
209    type Elem = &'a mut T;
210
211    fn len(&self) -> usize {
212        <[T]>::len(self)
213    }
214
215    fn split_at(self, at: usize) -> (Self, Self) {
216        let (a, b) = self.split_at_mut(at);
217
218        (a, b)
219    }
220}
221
222impl<'a, T> Items for &'a [T]
223where
224    T: Send + Sync,
225{
226    type Elem = &'a T;
227
228    fn len(&self) -> usize {
229        <[T]>::len(self)
230    }
231
232    fn split_at(self, at: usize) -> (Self, Self) {
233        let (a, b) = self.split_at(at);
234
235        (a, b)
236    }
237}
238//