next_custom_transforms/transforms/
cjs_optimizer.rs

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

next_custom_transforms/transforms/
cjs_optimizer.rs
