swc_ecma_transforms_base/
perf.rs

1use swc_common::util::move_map::MoveMap;
2#[cfg(feature = "concurrent")]
3use swc_common::{errors::HANDLER, GLOBALS};
4use swc_ecma_ast::*;
5pub use swc_ecma_utils::parallel::*;
6use swc_ecma_visit::{Fold, FoldWith, Visit, VisitMut, VisitMutWith, VisitWith};
7
8use crate::helpers::Helpers;
9#[cfg(feature = "concurrent")]
10use crate::helpers::HELPERS;
11
12pub trait Check: Visit + Default {
13    fn should_handle(&self) -> bool;
14}
15
16pub fn should_work<C, T>(n: &T) -> bool
17where
18    C: Check,
19    T: VisitWith<C>,
20{
21    let mut checker = C::default();
22    n.visit_with(&mut checker);
23    checker.should_handle()
24}
25
26pub trait ParExplode: Parallel {
27    /// Invoked after visiting each statements.
28    ///
29    /// Implementor should not delete/prepend to `stmts`.
30    fn after_one_stmt(&mut self, stmts: &mut Vec<Stmt>);
31
32    /// Invoked after visiting each statements.
33    ///
34    /// Implementor should not delete/prepend to `stmts`.
35    fn after_one_module_item(&mut self, stmts: &mut Vec<ModuleItem>);
36}
37
38pub trait ParVisit: Visit + Parallel {
39    fn visit_par<N>(&mut self, threshold: usize, nodes: &[N])
40    where
41        N: Send + Sync + VisitWith<Self>;
42}
43
44#[cfg(feature = "concurrent")]
45impl<T> ParVisit for T
46where
47    T: Visit + Parallel,
48{
49    fn visit_par<N>(&mut self, threshold: usize, nodes: &[N])
50    where
51        N: Send + Sync + VisitWith<Self>,
52    {
53        if nodes.len() >= threshold {
54            HELPERS.with(|helpers| {
55                let helpers = helpers.data();
56
57                HANDLER.with(|handler| {
58                    self.maybe_par(threshold, nodes, |visitor, node| {
59                        let helpers = Helpers::from_data(helpers);
60                        HELPERS.set(&helpers, || {
61                            HANDLER.set(handler, || {
62                                node.visit_with(visitor);
63                            });
64                        });
65                    });
66                })
67            });
68            return;
69        }
70
71        for n in nodes {
72            n.visit_with(self);
73        }
74    }
75}
76
77pub trait ParVisitMut: VisitMut + Parallel {
78    fn visit_mut_par<N>(&mut self, threshold: usize, nodes: &mut [N])
79    where
80        N: Send + Sync + VisitMutWith<Self>;
81}
82
83#[cfg(feature = "concurrent")]
84impl<T> ParVisitMut for T
85where
86    T: VisitMut + Parallel,
87{
88    fn visit_mut_par<N>(&mut self, threshold: usize, nodes: &mut [N])
89    where
90        N: Send + Sync + VisitMutWith<Self>,
91    {
92        if nodes.len() >= threshold {
93            HELPERS.with(|helpers| {
94                let helpers = helpers.data();
95
96                HANDLER.with(|handler| {
97                    self.maybe_par(threshold, nodes, |visitor, node| {
98                        let helpers = Helpers::from_data(helpers);
99                        HELPERS.set(&helpers, || {
100                            HANDLER.set(handler, || {
101                                node.visit_mut_with(visitor);
102                            });
103                        });
104                    });
105                })
106            });
107
108            return;
109        }
110
111        for n in nodes {
112            n.visit_mut_with(self);
113        }
114    }
115}
116
117pub trait ParFold: Fold + Parallel {
118    fn fold_par<N>(&mut self, threshold: usize, nodes: Vec<N>) -> Vec<N>
119    where
120        N: Send + Sync + FoldWith<Self>;
121}
122
123#[cfg(feature = "concurrent")]
124impl<T> ParFold for T
125where
126    T: Fold + Parallel,
127{
128    fn fold_par<N>(&mut self, threshold: usize, nodes: Vec<N>) -> Vec<N>
129    where
130        N: Send + Sync + FoldWith<Self>,
131    {
132        if nodes.len() >= threshold {
133            use rayon::prelude::*;
134
135            let (visitor, nodes) = GLOBALS.with(|globals| {
136                HELPERS.with(|helpers| {
137                    let helpers = helpers.data();
138                    HANDLER.with(|handler| {
139                        nodes
140                            .into_par_iter()
141                            .map(|node| {
142                                let helpers = Helpers::from_data(helpers);
143                                GLOBALS.set(globals, || {
144                                    HELPERS.set(&helpers, || {
145                                        HANDLER.set(handler, || {
146                                            let mut visitor = Parallel::create(&*self);
147                                            let node = node.fold_with(&mut visitor);
148
149                                            (visitor, node)
150                                        })
151                                    })
152                                })
153                            })
154                            .fold(
155                                || (Parallel::create(&*self), Vec::new()),
156                                |mut a, b| {
157                                    Parallel::merge(&mut a.0, b.0);
158
159                                    a.1.push(b.1);
160
161                                    a
162                                },
163                            )
164                            .reduce(
165                                || (Parallel::create(&*self), Vec::new()),
166                                |mut a, b| {
167                                    Parallel::merge(&mut a.0, b.0);
168
169                                    a.1.extend(b.1);
170
171                                    a
172                                },
173                            )
174                    })
175                })
176            });
177
178            Parallel::merge(self, visitor);
179
180            return nodes;
181        }
182
183        nodes.move_map(|n| n.fold_with(self))
184    }
185}
186
187#[cfg(not(feature = "concurrent"))]
188impl<T> ParVisit for T
189where
190    T: Visit + Parallel,
191{
192    fn visit_par<N>(&mut self, _: usize, nodes: &[N])
193    where
194        N: Send + Sync + VisitWith<Self>,
195    {
196        for n in nodes {
197            n.visit_with(self);
198        }
199    }
200}
201
202#[cfg(not(feature = "concurrent"))]
203impl<T> ParVisitMut for T
204where
205    T: VisitMut + Parallel,
206{
207    fn visit_mut_par<N>(&mut self, _: usize, nodes: &mut [N])
208    where
209        N: Send + Sync + VisitMutWith<Self>,
210    {
211        for n in nodes {
212            n.visit_mut_with(self);
213        }
214    }
215}
216
217#[cfg(not(feature = "concurrent"))]
218impl<T> ParFold for T
219where
220    T: Fold + Parallel,
221{
222    fn fold_par<N>(&mut self, _: usize, nodes: Vec<N>) -> Vec<N>
223    where
224        N: Send + Sync + FoldWith<Self>,
225    {
226        nodes.move_map(|n| n.fold_with(self))
227    }
228}