swc_bundler/modules/
mod.rs

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    /// Indicates that a statement is injected.
15    ///
16    /// Note: This context should be shared for a bundle.
17    pub(crate) injected_ctxt: SyntaxContext,
18
19    // We will change this into `Vec<Module>`.
20    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        // TODO
305        Self {
306            span: DUMMY_SP,
307            body: modules.into_items(),
308            shebang: None,
309        }
310    }
311}