swc_ecma_utils/
parallel.rs

1//! Module for parallel processing
2
3use once_cell::sync::Lazy;
4use swc_common::GLOBALS;
5use swc_ecma_ast::*;
6use swc_parallel::{
7    items::{IntoItems, Items},
8    join,
9};
10
11static CPU_COUNT: Lazy<usize> = Lazy::new(num_cpus::get);
12
13pub fn cpu_count() -> usize {
14    *CPU_COUNT
15}
16
17pub trait Parallel: swc_common::sync::Send + swc_common::sync::Sync {
18    /// Used to create visitor.
19    fn create(&self) -> Self;
20
21    /// This can be called in anytime.
22    fn merge(&mut self, other: Self);
23
24    /// Invoked after visiting all [Stmt]s, possibly in parallel.
25    ///
26    ///
27    /// Note: `visit_*_par` never calls this.
28    fn after_stmts(&mut self, _stmts: &mut Vec<Stmt>) {}
29
30    /// Invoked after visiting all [ModuleItem]s, possibly in parallel.
31    ///
32    ///
33    /// Note: `visit_*_par` never calls this.
34    fn after_module_items(&mut self, _stmts: &mut Vec<ModuleItem>) {}
35}
36
37pub trait ParallelExt: Parallel {
38    /// Invoke `op` in parallel, if `swc_ecma_utils` is compiled with
39    /// concurrent feature enabled and `nodes.len()` is bigger than threshold.
40    ///
41    ///
42    /// This configures [GLOBALS], while not configuring [HANDLER] nor [HELPERS]
43    fn maybe_par<I, F>(&mut self, threshold: usize, nodes: I, op: F)
44    where
45        I: IntoItems,
46        F: Send + Sync + Fn(&mut Self, I::Elem),
47    {
48        self.maybe_par_idx(threshold, nodes, |v, _, n| op(v, n))
49    }
50
51    /// Invoke `op` in parallel, if `swc_ecma_utils` is compiled with
52    /// concurrent feature enabled and `nodes.len()` is bigger than threshold.
53    ///
54    ///
55    /// This configures [GLOBALS], while not configuring [HANDLER] nor [HELPERS]
56    fn maybe_par_idx<I, F>(&mut self, threshold: usize, nodes: I, op: F)
57    where
58        I: IntoItems,
59        F: Send + Sync + Fn(&mut Self, usize, I::Elem),
60    {
61        self.maybe_par_idx_raw(threshold, nodes.into_items(), &op)
62    }
63
64    /// If you don't have a special reason, use [`ParallelExt::maybe_par`] or
65    /// [`ParallelExt::maybe_par_idx`] instead.
66    fn maybe_par_idx_raw<I, F>(&mut self, threshold: usize, nodes: I, op: &F)
67    where
68        I: Items,
69        F: Send + Sync + Fn(&mut Self, usize, I::Elem);
70}
71
72#[cfg(feature = "concurrent")]
73impl<T> ParallelExt for T
74where
75    T: Parallel,
76{
77    fn maybe_par_idx_raw<I, F>(&mut self, threshold: usize, nodes: I, op: &F)
78    where
79        I: Items,
80        F: Send + Sync + Fn(&mut Self, usize, I::Elem),
81    {
82        if nodes.len() >= threshold {
83            GLOBALS.with(|globals| {
84                let len = nodes.len();
85                if len == 0 {
86                    return;
87                }
88
89                if len == 1 {
90                    op(self, 0, nodes.into_iter().next().unwrap());
91                    return;
92                }
93
94                let (na, nb) = nodes.split_at(len / 2);
95                let mut vb = Parallel::create(&*self);
96
97                let (_, vb) = join(
98                    || {
99                        GLOBALS.set(globals, || {
100                            self.maybe_par_idx_raw(threshold, na, op);
101                        })
102                    },
103                    || {
104                        GLOBALS.set(globals, || {
105                            vb.maybe_par_idx_raw(threshold, nb, op);
106
107                            vb
108                        })
109                    },
110                );
111
112                Parallel::merge(self, vb);
113            });
114
115            return;
116        }
117
118        for (idx, n) in nodes.into_iter().enumerate() {
119            op(self, idx, n);
120        }
121    }
122}
123
124#[cfg(not(feature = "concurrent"))]
125impl<T> ParallelExt for T
126where
127    T: Parallel,
128{
129    fn maybe_par_idx_raw<I, F>(&mut self, _threshold: usize, nodes: I, op: &F)
130    where
131        I: Items,
132        F: Send + Sync + Fn(&mut Self, usize, I::Elem),
133    {
134        for (idx, n) in nodes.into_iter().enumerate() {
135            op(self, idx, n);
136        }
137    }
138}