swc_ecma_transforms_optimization/
const_modules.rs

1use std::{
2    collections::{HashMap, HashSet},
3    sync::Arc,
4};
5
6use bytes_str::BytesStr;
7use dashmap::DashMap;
8use once_cell::sync::Lazy;
9use rustc_hash::{FxBuildHasher, FxHashMap};
10use swc_atoms::{atom, Atom, Wtf8Atom};
11use swc_common::{
12    errors::HANDLER,
13    sync::Lrc,
14    util::{move_map::MoveMap, take::Take},
15    FileName, SourceMap,
16};
17use swc_ecma_ast::*;
18use swc_ecma_parser::parse_file_as_expr;
19use swc_ecma_utils::drop_span;
20use swc_ecma_visit::{noop_visit_mut_type, visit_mut_pass, VisitMut, VisitMutWith};
21
22pub fn const_modules(
23    cm: Lrc<SourceMap>,
24    globals: FxHashMap<Atom, FxHashMap<Atom, BytesStr>>,
25) -> impl Pass {
26    visit_mut_pass(ConstModules {
27        globals: globals
28            .into_iter()
29            .map(|(src, map)| {
30                let map = map
31                    .into_iter()
32                    .map(|(key, value)| {
33                        let value = parse_option(&cm, &key, value);
34
35                        (key.into(), value)
36                    })
37                    .collect();
38
39                (src.into(), map)
40            })
41            .collect(),
42        scope: Default::default(),
43    })
44}
45
46fn parse_option(cm: &SourceMap, name: &str, src: BytesStr) -> Arc<Expr> {
47    static CACHE: Lazy<DashMap<BytesStr, Arc<Expr>, FxBuildHasher>> = Lazy::new(DashMap::default);
48
49    let fm = cm.new_source_file(
50        FileName::Internal(format!("<const-module-{name}.js>")).into(),
51        src,
52    );
53    if let Some(expr) = CACHE.get(&fm.src) {
54        return expr.clone();
55    }
56
57    let expr = parse_file_as_expr(
58        &fm,
59        Default::default(),
60        Default::default(),
61        None,
62        &mut Vec::new(),
63    )
64    .map_err(|e| {
65        if HANDLER.is_set() {
66            HANDLER.with(|h| e.into_diagnostic(h).emit())
67        }
68    })
69    .map(drop_span)
70    .unwrap_or_else(|()| {
71        panic!(
72            "failed to parse jsx option {}: '{}' is not an expression",
73            name, fm.src,
74        )
75    });
76
77    let expr = Arc::new(*expr);
78
79    CACHE.insert(fm.src.clone(), expr.clone());
80
81    expr
82}
83
84struct ConstModules {
85    globals: HashMap<Wtf8Atom, HashMap<Wtf8Atom, Arc<Expr>>>,
86    scope: Scope,
87}
88
89#[derive(Default)]
90struct Scope {
91    namespace: HashSet<Id>,
92    imported: HashMap<Wtf8Atom, Arc<Expr>>,
93}
94
95impl VisitMut for ConstModules {
96    noop_visit_mut_type!(fail);
97
98    fn visit_mut_module_items(&mut self, n: &mut Vec<ModuleItem>) {
99        *n = n.take().move_flat_map(|item| match item {
100            ModuleItem::ModuleDecl(ModuleDecl::Import(import)) => {
101                let entry = self.globals.get(&import.src.value);
102                if let Some(entry) = entry {
103                    for s in &import.specifiers {
104                        match *s {
105                            ImportSpecifier::Named(ref s) => {
106                                let imported = s
107                                    .imported
108                                    .as_ref()
109                                    .map(|m| match m {
110                                        ModuleExportName::Ident(id) => id.sym.clone().into(),
111                                        ModuleExportName::Str(s) => s.value.clone(),
112                                        #[cfg(swc_ast_unknown)]
113                                        _ => panic!("unable to access unknown nodes"),
114                                    })
115                                    .unwrap_or_else(|| s.local.sym.clone().into());
116                                let value = entry.get(&imported).cloned().unwrap_or_else(|| {
117                                    panic!(
118                                        "The requested const_module `{:?}` does not provide an \
119                                         export named `{:?}`",
120                                        import.src.value, imported
121                                    )
122                                });
123                                self.scope.imported.insert(imported.clone(), value);
124                            }
125                            ImportSpecifier::Namespace(ref s) => {
126                                self.scope.namespace.insert(s.local.to_id());
127                            }
128                            ImportSpecifier::Default(ref s) => {
129                                let imported: Wtf8Atom = s.local.sym.clone().into();
130                                let default_import_key: Wtf8Atom = atom!("default").into();
131                                let value =
132                                    entry.get(&default_import_key).cloned().unwrap_or_else(|| {
133                                        panic!(
134                                            "The requested const_module `{:?}` does not provide \
135                                             default export",
136                                            import.src.value
137                                        )
138                                    });
139                                self.scope.imported.insert(imported, value);
140                            }
141                            #[cfg(swc_ast_unknown)]
142                            _ => panic!("unable to access unknown nodes"),
143                        };
144                    }
145
146                    None
147                } else {
148                    Some(import.into())
149                }
150            }
151            _ => Some(item),
152        });
153
154        if self.scope.imported.is_empty() && self.scope.namespace.is_empty() {
155            return;
156        }
157
158        n.visit_mut_children_with(self);
159    }
160
161    fn visit_mut_expr(&mut self, n: &mut Expr) {
162        match n {
163            Expr::Ident(ref id @ Ident { ref sym, .. }) => {
164                let sym_wtf8: Wtf8Atom = sym.clone().into();
165                if let Some(value) = self.scope.imported.get(&sym_wtf8) {
166                    *n = (**value).clone();
167                    return;
168                }
169
170                if self.scope.namespace.contains(&id.to_id()) {
171                    panic!(
172                        "The const_module namespace `{sym}` cannot be used without member accessor"
173                    )
174                }
175            }
176            Expr::Member(MemberExpr { obj, prop, .. }) if obj.is_ident() => {
177                if let Some(module_name) = obj
178                    .as_ident()
179                    .filter(|member_obj| self.scope.namespace.contains(&member_obj.to_id()))
180                    .map(|member_obj| &member_obj.sym)
181                {
182                    let imported_name: Wtf8Atom = match prop {
183                        MemberProp::Ident(ref id) => id.sym.clone().into(),
184                        MemberProp::Computed(ref p) => match &*p.expr {
185                            Expr::Lit(Lit::Str(s)) => s.value.clone(),
186                            _ => return,
187                        },
188                        MemberProp::PrivateName(..) => return,
189                        #[cfg(swc_ast_unknown)]
190                        _ => panic!("unable to access unknown nodes"),
191                    };
192
193                    let module_name_wtf8: Wtf8Atom = module_name.clone().into();
194                    let value = self
195                        .globals
196                        .get(&module_name_wtf8)
197                        .and_then(|entry| entry.get(&imported_name))
198                        .unwrap_or_else(|| {
199                            panic!(
200                                "The requested const_module `{module_name}` does not provide an \
201                                 export named `{imported_name:?}`"
202                            )
203                        });
204
205                    *n = (**value).clone();
206                } else {
207                    n.visit_mut_children_with(self);
208                }
209            }
210            _ => {
211                n.visit_mut_children_with(self);
212            }
213        };
214    }
215
216    fn visit_mut_prop(&mut self, n: &mut Prop) {
217        match n {
218            Prop::Shorthand(id) => {
219                let sym_wtf8: Wtf8Atom = id.sym.clone().into();
220                if let Some(value) = self.scope.imported.get(&sym_wtf8) {
221                    *n = Prop::KeyValue(KeyValueProp {
222                        key: id.take().into(),
223                        value: Box::new((**value).clone()),
224                    });
225                    return;
226                }
227
228                if self.scope.namespace.contains(&id.to_id()) {
229                    panic!(
230                        "The const_module namespace `{}` cannot be used without member accessor",
231                        id.sym
232                    )
233                }
234            }
235            _ => n.visit_mut_children_with(self),
236        }
237    }
238}