next_custom_transforms/transforms/
cjs_optimizer.rs

1use rustc_hash::{FxHashMap, FxHashSet};
2use serde::Deserialize;
3use swc_core::{
4    atoms::atom,
5    common::{util::take::Take, SyntaxContext, DUMMY_SP},
6    ecma::{
7        ast::{
8            CallExpr, Callee, Decl, Expr, Id, Ident, IdentName, Lit, MemberExpr, MemberProp,
9            Module, ModuleItem, Pat, Script, Stmt, VarDecl, VarDeclKind, VarDeclarator,
10        },
11        atoms::Atom,
12        utils::{prepend_stmts, private_ident, ExprFactory, IdentRenamer},
13        visit::{noop_visit_mut_type, noop_visit_type, Visit, VisitMut, VisitMutWith, VisitWith},
14    },
15};
16
17pub fn cjs_optimizer(config: Config, unresolved_ctxt: SyntaxContext) -> CjsOptimizer {
18    CjsOptimizer {
19        data: State::default(),
20        packages: config.packages,
21        unresolved_ctxt,
22    }
23}
24
25#[derive(Clone, Debug, Deserialize)]
26pub struct Config {
27    pub packages: FxHashMap<Atom, PackageConfig>,
28}
29
30#[derive(Clone, Debug, Deserialize)]
31#[serde(rename_all = "camelCase")]
32pub struct PackageConfig {
33    pub transforms: FxHashMap<Atom, Atom>,
34}
35
36pub struct CjsOptimizer {
37    data: State,
38    packages: FxHashMap<Atom, PackageConfig>,
39    unresolved_ctxt: SyntaxContext,
40}
41
42#[derive(Debug, Default)]
43struct State {
44    /// List of `require` calls **which should be replaced**.
45    ///
46    ///  `(identifier): (module_record)`
47    imports: FxHashMap<Id, ImportRecord>,
48
49    /// `(module_specifier, property): (identifier)`
50    replaced: FxHashMap<(Atom, Atom), Id>,
51
52    extra_stmts: Vec<Stmt>,
53
54    rename_map: FxHashMap<Id, Id>,
55
56    /// Ignored identifiers for `obj` of [MemberExpr].
57    ignored: FxHashSet<Id>,
58
59    is_prepass: bool,
60}
61
62#[derive(Debug)]
63struct ImportRecord {
64    module_specifier: Atom,
65}
66
67impl CjsOptimizer {
68    fn should_rewrite(&self, module_specifier: &Atom) -> Option<&FxHashMap<Atom, Atom>> {
69        self.packages.get(module_specifier).map(|v| &v.transforms)
70    }
71}
72
73impl VisitMut for CjsOptimizer {
74    noop_visit_mut_type!();
75
76    fn visit_mut_module_items(&mut self, stmts: &mut Vec<ModuleItem>) {
77        self.data.is_prepass = true;
78        stmts.visit_mut_children_with(self);
79        self.data.is_prepass = false;
80        stmts.visit_mut_children_with(self);
81    }
82
83    fn visit_mut_expr(&mut self, e: &mut Expr) {
84        e.visit_mut_children_with(self);
85
86        if let Expr::Member(n) = e {
87            if let MemberProp::Ident(prop) = &n.prop {
88                if let Expr::Ident(obj) = &*n.obj {
89                    let key = obj.to_id();
90                    if self.data.ignored.contains(&key) {
91                        return;
92                    }
93
94                    if let Some(record) = self.data.imports.get(&key) {
95                        let mut replaced = false;
96
97                        let new_id = self
98                            .data
99                            .replaced
100                            .entry((record.module_specifier.clone(), prop.sym.clone()))
101                            .or_insert_with(|| private_ident!(prop.sym.clone()).to_id())
102                            .clone();
103
104                        if let Some(map) = self.should_rewrite(&record.module_specifier) {
105                            if let Some(renamed) = map.get(&prop.sym) {
106                                replaced = true;
107                                if !self.data.is_prepass {
108                                    // Transform as `require('foo').bar`
109                                    let var = VarDeclarator {
110                                        span: DUMMY_SP,
111                                        name: Pat::Ident(new_id.clone().into()),
112                                        init: Some(Box::new(Expr::Member(MemberExpr {
113                                            span: DUMMY_SP,
114                                            obj: Box::new(Expr::Call(CallExpr {
115                                                span: DUMMY_SP,
116                                                callee: Ident::new(
117                                                    atom!("require"),
118                                                    DUMMY_SP,
119                                                    self.unresolved_ctxt,
120                                                )
121                                                .as_callee(),
122                                                args: vec![Expr::Lit(Lit::Str(
123                                                    renamed.clone().into(),
124                                                ))
125                                                .as_arg()],
126                                                ..Default::default()
127                                            })),
128                                            prop: MemberProp::Ident(IdentName::new(
129                                                prop.sym.clone(),
130                                                DUMMY_SP,
131                                            )),
132                                        }))),
133                                        definite: false,
134                                    };
135
136                                    if !self.data.extra_stmts.iter().any(|s| {
137                                        if let Stmt::Decl(Decl::Var(v)) = &s {
138                                            v.decls.iter().any(|d| d.name == var.name)
139                                        } else {
140                                            false
141                                        }
142                                    }) {
143                                        self.data.extra_stmts.push(Stmt::Decl(Decl::Var(
144                                            Box::new(VarDecl {
145                                                span: DUMMY_SP,
146                                                kind: VarDeclKind::Const,
147                                                decls: vec![var],
148                                                ..Default::default()
149                                            }),
150                                        )));
151                                    }
152
153                                    *e = Expr::Ident(new_id.into());
154                                }
155                            }
156                        }
157
158                        if !replaced {
159                            self.data.ignored.insert(key);
160                        }
161                    }
162                }
163            }
164        }
165    }
166
167    fn visit_mut_module(&mut self, n: &mut Module) {
168        n.visit_children_with(&mut Analyzer {
169            data: &mut self.data,
170            in_member_or_var: false,
171        });
172
173        n.visit_mut_children_with(self);
174
175        prepend_stmts(
176            &mut n.body,
177            self.data.extra_stmts.drain(..).map(ModuleItem::Stmt),
178        );
179
180        n.visit_mut_children_with(&mut IdentRenamer::new(&self.data.rename_map));
181    }
182
183    fn visit_mut_script(&mut self, n: &mut Script) {
184        n.visit_children_with(&mut Analyzer {
185            data: &mut self.data,
186            in_member_or_var: false,
187        });
188
189        n.visit_mut_children_with(self);
190
191        prepend_stmts(&mut n.body, self.data.extra_stmts.drain(..));
192
193        n.visit_mut_children_with(&mut IdentRenamer::new(&self.data.rename_map));
194    }
195
196    fn visit_mut_stmt(&mut self, n: &mut Stmt) {
197        n.visit_mut_children_with(self);
198
199        if let Stmt::Decl(Decl::Var(v)) = n {
200            if v.decls.is_empty() {
201                n.take();
202            }
203        }
204    }
205
206    fn visit_mut_var_declarator(&mut self, n: &mut VarDeclarator) {
207        n.visit_mut_children_with(self);
208
209        // Find `require('foo')`
210        if let Some(Expr::Call(CallExpr {
211            callee: Callee::Expr(callee),
212            args,
213            ..
214        })) = n.init.as_deref()
215        {
216            if let Expr::Ident(ident) = &**callee {
217                if ident.ctxt == self.unresolved_ctxt && ident.sym == *"require" {
218                    if let Some(arg) = args.first() {
219                        if let Expr::Lit(Lit::Str(v)) = &*arg.expr {
220                            // TODO: Config
221
222                            if let Pat::Ident(name) = &n.name {
223                                if self.should_rewrite(&v.value).is_some() {
224                                    let key = name.to_id();
225
226                                    if !self.data.is_prepass {
227                                        if !self.data.ignored.contains(&key) {
228                                            // Drop variable declarator.
229                                            n.name.take();
230                                        }
231                                    } else {
232                                        self.data.imports.insert(
233                                            key,
234                                            ImportRecord {
235                                                module_specifier: v.value.clone(),
236                                            },
237                                        );
238                                    }
239                                }
240                            }
241                        }
242                    }
243                }
244            }
245        }
246    }
247
248    fn visit_mut_var_declarators(&mut self, n: &mut Vec<VarDeclarator>) {
249        n.visit_mut_children_with(self);
250
251        // We make `name` invalid if we should drop it.
252        n.retain(|v| !v.name.is_invalid());
253    }
254}
255
256struct Analyzer<'a> {
257    in_member_or_var: bool,
258    data: &'a mut State,
259}
260
261impl Visit for Analyzer<'_> {
262    noop_visit_type!();
263
264    fn visit_var_declarator(&mut self, n: &VarDeclarator) {
265        let mut safe_to_ignore = false;
266
267        // Ignore the require itself (foo = require('foo'))
268        if let Some(Expr::Call(CallExpr {
269            callee: Callee::Expr(callee),
270            ..
271        })) = n.init.as_deref()
272        {
273            if let Expr::Ident(ident) = &**callee {
274                if ident.sym == *"require" {
275                    safe_to_ignore = true;
276                }
277            }
278        }
279
280        if safe_to_ignore {
281            self.in_member_or_var = true;
282            n.visit_children_with(self);
283            self.in_member_or_var = false;
284        } else {
285            n.visit_children_with(self);
286        }
287    }
288
289    fn visit_member_expr(&mut self, e: &MemberExpr) {
290        self.in_member_or_var = true;
291        e.visit_children_with(self);
292        self.in_member_or_var = false;
293
294        if let (Expr::Ident(obj), MemberProp::Computed(..)) = (&*e.obj, &e.prop) {
295            self.data.ignored.insert(obj.to_id());
296        }
297    }
298
299    fn visit_ident(&mut self, i: &Ident) {
300        i.visit_children_with(self);
301        if !self.in_member_or_var {
302            self.data.ignored.insert(i.to_id());
303        }
304    }
305}
next_custom_transforms/transforms/cjs_optimizer.rs

next_custom_transforms/transforms/
cjs_optimizer.rs
