next_custom_transforms/transforms/
cjs_optimizer.rs

1use rustc_hash::{FxHashMap, FxHashSet};
2use serde::Deserialize;
3use swc_core::{
4    atoms::{Wtf8Atom, atom},
5    common::{DUMMY_SP, SyntaxContext, util::take::Take},
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::{ExprFactory, IdentRenamer, prepend_stmts, private_ident},
13        visit::{Visit, VisitMut, VisitMutWith, VisitWith, noop_visit_mut_type, noop_visit_type},
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<(Wtf8Atom, 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: Wtf8Atom,
65}
66
67impl CjsOptimizer {
68    fn should_rewrite(&self, module_specifier: &Wtf8Atom) -> 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            && let MemberProp::Ident(prop) = &n.prop
88            && let Expr::Ident(obj) = &*n.obj
89        {
90            let key = obj.to_id();
91            if self.data.ignored.contains(&key) {
92                return;
93            }
94
95            if let Some(record) = self.data.imports.get(&key) {
96                let mut replaced = false;
97
98                let new_id = self
99                    .data
100                    .replaced
101                    .entry((record.module_specifier.clone(), prop.sym.clone()))
102                    .or_insert_with(|| private_ident!(prop.sym.clone()).to_id())
103                    .clone();
104
105                if let Some(map) = self.should_rewrite(&record.module_specifier)
106                    && let Some(renamed) = map.get(&prop.sym)
107                {
108                    replaced = true;
109                    if !self.data.is_prepass {
110                        // Transform as `require('foo').bar`
111                        let var = VarDeclarator {
112                            span: DUMMY_SP,
113                            name: Pat::Ident(new_id.clone().into()),
114                            init: Some(Box::new(Expr::Member(MemberExpr {
115                                span: DUMMY_SP,
116                                obj: Box::new(Expr::Call(CallExpr {
117                                    span: DUMMY_SP,
118                                    callee: Ident::new(
119                                        atom!("require"),
120                                        DUMMY_SP,
121                                        self.unresolved_ctxt,
122                                    )
123                                    .as_callee(),
124                                    args: vec![
125                                        Expr::Lit(Lit::Str(renamed.clone().into())).as_arg(),
126                                    ],
127                                    ..Default::default()
128                                })),
129                                prop: MemberProp::Ident(IdentName::new(prop.sym.clone(), DUMMY_SP)),
130                            }))),
131                            definite: false,
132                        };
133
134                        if !self.data.extra_stmts.iter().any(|s| {
135                            if let Stmt::Decl(Decl::Var(v)) = &s {
136                                v.decls.iter().any(|d| d.name == var.name)
137                            } else {
138                                false
139                            }
140                        }) {
141                            self.data
142                                .extra_stmts
143                                .push(Stmt::Decl(Decl::Var(Box::new(VarDecl {
144                                    span: DUMMY_SP,
145                                    kind: VarDeclKind::Const,
146                                    decls: vec![var],
147                                    ..Default::default()
148                                }))));
149                        }
150
151                        *e = Expr::Ident(new_id.into());
152                    }
153                }
154
155                if !replaced {
156                    self.data.ignored.insert(key);
157                }
158            }
159        }
160    }
161
162    fn visit_mut_module(&mut self, n: &mut Module) {
163        n.visit_children_with(&mut Analyzer {
164            data: &mut self.data,
165            in_member_or_var: false,
166        });
167
168        n.visit_mut_children_with(self);
169
170        prepend_stmts(
171            &mut n.body,
172            self.data.extra_stmts.drain(..).map(ModuleItem::Stmt),
173        );
174
175        n.visit_mut_children_with(&mut IdentRenamer::new(&self.data.rename_map));
176    }
177
178    fn visit_mut_script(&mut self, n: &mut Script) {
179        n.visit_children_with(&mut Analyzer {
180            data: &mut self.data,
181            in_member_or_var: false,
182        });
183
184        n.visit_mut_children_with(self);
185
186        prepend_stmts(&mut n.body, self.data.extra_stmts.drain(..));
187
188        n.visit_mut_children_with(&mut IdentRenamer::new(&self.data.rename_map));
189    }
190
191    fn visit_mut_stmt(&mut self, n: &mut Stmt) {
192        n.visit_mut_children_with(self);
193
194        if let Stmt::Decl(Decl::Var(v)) = n
195            && v.decls.is_empty()
196        {
197            n.take();
198        }
199    }
200
201    fn visit_mut_var_declarator(&mut self, n: &mut VarDeclarator) {
202        n.visit_mut_children_with(self);
203
204        // Find `require('foo')`
205        if let Some(Expr::Call(CallExpr {
206            callee: Callee::Expr(callee),
207            args,
208            ..
209        })) = n.init.as_deref()
210            && let Expr::Ident(ident) = &**callee
211            && ident.ctxt == self.unresolved_ctxt
212            && ident.sym == *"require"
213            && let Some(arg) = args.first()
214            && let Expr::Lit(Lit::Str(v)) = &*arg.expr
215        {
216            // TODO: Config
217
218            if let Pat::Ident(name) = &n.name
219                && self.should_rewrite(&v.value).is_some()
220            {
221                let key = name.to_id();
222
223                if !self.data.is_prepass {
224                    if !self.data.ignored.contains(&key) {
225                        // Drop variable declarator.
226                        n.name.take();
227                    }
228                } else {
229                    self.data.imports.insert(
230                        key,
231                        ImportRecord {
232                            module_specifier: v.value.clone(),
233                        },
234                    );
235                }
236            }
237        }
238    }
239
240    fn visit_mut_var_declarators(&mut self, n: &mut Vec<VarDeclarator>) {
241        n.visit_mut_children_with(self);
242
243        // We make `name` invalid if we should drop it.
244        n.retain(|v| !v.name.is_invalid());
245    }
246}
247
248struct Analyzer<'a> {
249    in_member_or_var: bool,
250    data: &'a mut State,
251}
252
253impl Visit for Analyzer<'_> {
254    noop_visit_type!();
255
256    fn visit_var_declarator(&mut self, n: &VarDeclarator) {
257        let mut safe_to_ignore = false;
258
259        // Ignore the require itself (foo = require('foo'))
260        if let Some(Expr::Call(CallExpr {
261            callee: Callee::Expr(callee),
262            ..
263        })) = n.init.as_deref()
264            && let Expr::Ident(ident) = &**callee
265            && ident.sym == *"require"
266        {
267            safe_to_ignore = true;
268        }
269
270        if safe_to_ignore {
271            self.in_member_or_var = true;
272            n.visit_children_with(self);
273            self.in_member_or_var = false;
274        } else {
275            n.visit_children_with(self);
276        }
277    }
278
279    fn visit_member_expr(&mut self, e: &MemberExpr) {
280        self.in_member_or_var = true;
281        e.visit_children_with(self);
282        self.in_member_or_var = false;
283
284        if let (Expr::Ident(obj), MemberProp::Computed(..)) = (&*e.obj, &e.prop) {
285            self.data.ignored.insert(obj.to_id());
286        }
287    }
288
289    fn visit_ident(&mut self, i: &Ident) {
290        i.visit_children_with(self);
291        if !self.in_member_or_var {
292            self.data.ignored.insert(i.to_id());
293        }
294    }
295}
next_custom_transforms/transforms/cjs_optimizer.rs

next_custom_transforms/transforms/
cjs_optimizer.rs
