1use 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 fn create(&self) -> Self;
17
18 fn merge(&mut self, other: Self);
20
21 fn after_stmts(&mut self, _stmts: &mut Vec<Stmt>) {}
26
27 fn after_module_items(&mut self, _stmts: &mut Vec<ModuleItem>) {}
32}
33
34pub trait ParallelExt: Parallel {
35 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 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 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#[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