1use std::mem::take;
2
3use rustc_hash::FxHashMap;
4use swc_common::{SourceMap, SyntaxContext, DUMMY_SP};
5use swc_ecma_ast::*;
6use swc_ecma_visit::{Fold, FoldWith, Visit, VisitMut, VisitMutWith, VisitWith};
7
8use crate::ModuleId;
9
10mod sort;
11
12#[derive(Debug, Clone)]
13pub struct Modules {
14 pub(crate) injected_ctxt: SyntaxContext,
18
19 modules: Vec<(ModuleId, Module)>,
21 prepended_stmts: FxHashMap<ModuleId, Vec<ModuleItem>>,
22 appended_stmts: FxHashMap<ModuleId, Vec<ModuleItem>>,
23}
24
25impl Modules {
26 pub fn empty(injected_ctxt: SyntaxContext) -> Self {
27 Self {
28 injected_ctxt,
29 modules: Default::default(),
30 prepended_stmts: Default::default(),
31 appended_stmts: Default::default(),
32 }
33 }
34
35 pub fn from(id: ModuleId, module: Module, injected_ctxt: SyntaxContext) -> Self {
36 let mut ret = Self::empty(injected_ctxt);
37 ret.modules.push((id, module));
38 ret
39 }
40
41 fn into_items(self) -> Vec<ModuleItem> {
42 debug_assert!(
43 self.prepended_stmts.is_empty(),
44 "sort should be called before calling into_items"
45 );
46 debug_assert!(
47 self.appended_stmts.is_empty(),
48 "sort should be called before calling into_items"
49 );
50 self.modules.into_iter().flat_map(|v| v.1.body).collect()
51 }
52
53 pub fn add_dep(&mut self, dep: Modules) {
54 self.push_all(dep)
55 }
56
57 pub fn push_all(&mut self, other: Modules) {
58 for (id, stmts) in other.prepended_stmts {
59 self.prepended_stmts.entry(id).or_default().extend(stmts);
60 }
61
62 for (id, stmts) in other.appended_stmts {
63 self.appended_stmts.entry(id).or_default().extend(stmts);
64 }
65
66 for (id, module) in other.modules {
67 if let Some(prev) = self.modules.iter_mut().find(|prev| prev.0 == id) {
68 prev.1.body.extend(module.body);
69 } else {
70 self.modules.push((id, module));
71 }
72 }
73 }
74
75 pub fn iter(&self) -> impl Iterator<Item = (ModuleId, &ModuleItem)> {
76 self.prepended_stmts
77 .iter()
78 .flat_map(|(id, stmts)| stmts.iter().map(move |stmt| (*id, stmt)))
79 .chain(
80 self.modules
81 .iter()
82 .flat_map(|(id, m)| m.body.iter().map(move |v| (*id, v))),
83 )
84 .chain(
85 self.appended_stmts
86 .iter()
87 .flat_map(|(id, stmts)| stmts.iter().map(move |stmt| (*id, stmt))),
88 )
89 }
90
91 pub fn iter_mut(&mut self) -> impl Iterator<Item = (ModuleId, &mut ModuleItem)> {
92 self.prepended_stmts
93 .iter_mut()
94 .flat_map(|(id, stmts)| stmts.iter_mut().map(move |stmt| (*id, stmt)))
95 .chain(
96 self.modules
97 .iter_mut()
98 .flat_map(|(id, m)| m.body.iter_mut().map(move |v| (*id, v))),
99 )
100 .chain(
101 self.appended_stmts
102 .iter_mut()
103 .flat_map(|(id, stmts)| stmts.iter_mut().map(move |stmt| (*id, stmt))),
104 )
105 }
106
107 pub fn map_any_items<F>(&mut self, mut op: F)
108 where
109 F: FnMut(ModuleId, Vec<ModuleItem>) -> Vec<ModuleItem>,
110 {
111 let p = take(&mut self.prepended_stmts);
112 self.prepended_stmts = p
113 .into_iter()
114 .map(|(id, items)| (id, op(id, items)))
115 .collect();
116
117 self.modules = take(&mut self.modules)
118 .into_iter()
119 .map(|mut m| {
120 let body = op(m.0, take(&mut m.1.body));
121
122 (m.0, Module { body, ..m.1 })
123 })
124 .collect();
125
126 let a = take(&mut self.appended_stmts);
127 self.appended_stmts = a
128 .into_iter()
129 .map(|(id, items)| (id, op(id, items)))
130 .collect();
131 }
132
133 pub fn map_items_mut<F>(&mut self, mut op: F)
134 where
135 F: FnMut(ModuleId, &mut ModuleItem),
136 {
137 self.iter_mut().for_each(|(id, item)| op(id, item))
138 }
139
140 pub fn append_all(&mut self, items: impl IntoIterator<Item = (ModuleId, ModuleItem)>) {
141 for v in items {
142 self.append(v.0, v.1);
143 }
144 }
145
146 pub fn append(&mut self, module_id: ModuleId, item: ModuleItem) {
147 self.appended_stmts.entry(module_id).or_default().push(item);
148 }
149
150 pub fn prepend(&mut self, module_id: ModuleId, item: ModuleItem) {
151 self.prepended_stmts
152 .entry(module_id)
153 .or_default()
154 .push(item);
155 }
156
157 #[cfg(not(feature = "concurrent"))]
158 pub(crate) fn par_visit_mut_with<V>(&mut self, v: &mut V)
159 where
160 V: VisitMut,
161 {
162 self.visit_mut_with(v)
163 }
164
165 #[cfg(feature = "concurrent")]
166 pub(crate) fn par_visit_mut_with<V>(&mut self, v: &mut V)
167 where
168 V: Clone + VisitMut + Send + Sync,
169 {
170 use rayon::prelude::*;
171
172 let pre = &mut self.prepended_stmts;
173 let modules = &mut self.modules;
174 let app = &mut self.appended_stmts;
175
176 rayon::scope(|s| {
177 s.spawn(|_| {
178 pre.par_iter_mut()
179 .for_each(|(_, stmts)| stmts.visit_mut_with(&mut v.clone()));
180 });
181
182 s.spawn(|_| {
183 modules
184 .par_iter_mut()
185 .for_each(|(_, stmts)| stmts.visit_mut_with(&mut v.clone()));
186 });
187
188 s.spawn(|_| {
189 app.par_iter_mut()
190 .for_each(|(_, stmts)| stmts.visit_mut_with(&mut v.clone()));
191 });
192 });
193 }
194
195 pub fn visit_mut_with<V>(&mut self, v: &mut V)
196 where
197 V: VisitMut,
198 {
199 self.iter_mut().for_each(|item| item.1.visit_mut_with(v));
200 }
201
202 pub fn fold_with<V>(mut self, v: &mut V) -> Self
203 where
204 V: Fold,
205 {
206 self.prepended_stmts = self
207 .prepended_stmts
208 .into_iter()
209 .map(|(id, items)| (id, items.fold_with(&mut *v)))
210 .collect();
211 self.modules = self
212 .modules
213 .into_iter()
214 .map(|m| (m.0, m.1.fold_with(&mut *v)))
215 .collect();
216
217 self.appended_stmts = self
218 .appended_stmts
219 .into_iter()
220 .map(|(id, items)| (id, items.fold_with(&mut *v)))
221 .collect();
222
223 self
224 }
225
226 pub fn visit_with<V>(&self, v: &mut V)
227 where
228 V: Visit,
229 {
230 self.iter().for_each(|item| item.1.visit_with(v));
231 }
232
233 pub fn retain_mut<F>(&mut self, mut op: F)
234 where
235 F: FnMut(ModuleId, &mut ModuleItem) -> bool,
236 {
237 self.prepended_stmts
238 .iter_mut()
239 .for_each(|(id, v)| v.retain_mut(|item| op(*id, item)));
240
241 for module in &mut self.modules {
242 let id = module.0;
243 module.1.body.retain_mut(|item| op(id, item));
244 }
245
246 self.appended_stmts
247 .iter_mut()
248 .for_each(|(id, v)| v.retain_mut(|item| op(*id, item)));
249 }
250
251 #[allow(unused)]
252 #[cfg(debug_assertions)]
253 pub(crate) fn print(
254 &self,
255 cm: &swc_common::sync::Lrc<SourceMap>,
256 event: impl std::fmt::Display,
257 ) {
258 let files = self
259 .modules
260 .iter()
261 .map(|(_, m)| m.span)
262 .filter_map(|module_span| {
263 if module_span.is_dummy() {
264 return None;
265 }
266 Some(format!("{}\n", cm.lookup_source_file(module_span.lo).name))
267 })
268 .collect::<String>();
269 let mut cloned = self.clone();
270 let mut stmts = Vec::new();
271
272 for (id, mut module) in cloned.modules {
273 swc_ecma_utils::prepend_stmts(
274 &mut module.body,
275 cloned
276 .prepended_stmts
277 .get(&id)
278 .cloned()
279 .unwrap_or_default()
280 .into_iter(),
281 );
282
283 module
284 .body
285 .extend(cloned.appended_stmts.get(&id).cloned().unwrap_or_default());
286
287 stmts.extend(module.body);
288 }
289
290 crate::debug::print_hygiene(
291 &format!("{event}\n{files}"),
292 cm,
293 &Module {
294 span: DUMMY_SP,
295 body: stmts,
296 shebang: None,
297 },
298 );
299 }
300}
301
302impl From<Modules> for Module {
303 fn from(modules: Modules) -> Self {
304 Self {
306 span: DUMMY_SP,
307 body: modules.into_items(),
308 shebang: None,
309 }
310 }
311}