weevy_swc_core/
lib.rs

1use std::collections::BTreeMap;
2use std::mem::take;
3
4use swc_atoms::Atom;
5use swc_common::Span;
6use swc_common::Spanned;
7use swc_common::SyntaxContext;
8use swc_common::sync::Lrc;
9use swc_ecma_ast::BlockStmt;
10use swc_ecma_ast::CallExpr;
11use swc_ecma_ast::Callee;
12use swc_ecma_ast::ComputedPropName;
13use swc_ecma_ast::Decl;
14use swc_ecma_ast::Expr;
15use swc_ecma_ast::ExprOrSpread;
16use swc_ecma_ast::FnExpr;
17use swc_ecma_ast::Function;
18use swc_ecma_ast::Id;
19use swc_ecma_ast::IdentName;
20use swc_ecma_ast::Lit;
21use swc_ecma_ast::MemberExpr;
22use swc_ecma_ast::MemberProp;
23use swc_ecma_ast::MethodProp;
24use swc_ecma_ast::ModuleDecl;
25use swc_ecma_ast::ModuleItem;
26use swc_ecma_ast::ObjectLit;
27use swc_ecma_ast::Pat;
28use swc_ecma_ast::Program;
29use swc_ecma_ast::PropOrSpread;
30use swc_ecma_ast::ReturnStmt;
31use swc_ecma_ast::Script;
32use swc_ecma_ast::Stmt;
33use swc_ecma_ast::Str;
34use swc_ecma_ast::VarDecl;
35use swc_ecma_ast::VarDeclarator;
36use swc_ecma_ast::{Ident, ImportDecl, Module};
37use swc_ecma_visit::VisitMut;
38use swc_ecma_visit::VisitMutWith;
39
40pub fn wevy(span: Span, ctx: SyntaxContext) -> Expr {
41    Expr::Ident(Ident::new(Atom::new("__WeevyMain"), span, ctx))
42}
43pub fn default_source_mapper(m: &(dyn Spanned + '_), root: SyntaxContext) -> Expr {
44    Expr::Member(MemberExpr {
45        span: m.span(),
46        obj: Box::new(wevy(m.span(), root)),
47        prop: swc_ecma_ast::MemberProp::Ident(IdentName {
48            span: m.span(),
49            sym: Atom::new("newSrcDecompressor"),
50        }),
51    })
52}
53pub fn distanced_source_mapper(m: &(dyn Spanned + '_), name: Atom, root: SyntaxContext) -> Expr {
54    Expr::Member(MemberExpr {
55        span: m.span(),
56        obj: Box::new(Expr::Ident(Ident::new(
57            Atom::new("globalThis"),
58            m.span(),
59            root,
60        ))),
61        prop: swc_ecma_ast::MemberProp::Ident(IdentName {
62            span: m.span(),
63            sym: name,
64        }),
65    })
66}
67pub struct SourceMapper {
68    pub source_mapper: Expr,
69    pub sm: Lrc<swc_common::SourceMap>,
70    pub data: Vec<u8>,
71    pub id: Ident,
72    pub root: SyntaxContext,
73}
74impl VisitMut for SourceMapper {
75    fn visit_mut_module(&mut self, m: &mut Module) {
76        m.visit_mut_children_with(self);
77        let s = take(&mut self.data);
78        let mut t = vec![];
79        brotli::BrotliCompress(&mut &*s, &mut t, &Default::default()).unwrap();
80        let s = base2048::encode(&t);
81        let s = Stmt::Decl(Decl::Var(Box::new(VarDecl {
82            span: m.span(),
83            ctxt: Default::default(),
84            kind: swc_ecma_ast::VarDeclKind::Const,
85            declare: true,
86            decls: vec![VarDeclarator {
87                span: m.span(),
88                name: Pat::Ident(self.id.clone().into()),
89                definite: false,
90                init: Some(Box::new(Expr::Call(CallExpr {
91                    span: m.span(),
92                    ctxt: Default::default(),
93                    callee: swc_ecma_ast::Callee::Expr(Box::new(self.source_mapper.clone())),
94                    args: vec![ExprOrSpread {
95                        spread: None,
96                        expr: Box::new(Expr::Lit(Lit::Str(Str {
97                            span: m.span(),
98                            value: Atom::new(s),
99                            raw: None,
100                        }))),
101                    }],
102                    type_args: None,
103                }))),
104            }],
105        })));
106        m.body = [ModuleItem::Stmt(s)]
107            .into_iter()
108            .chain(m.body.drain(..))
109            .collect();
110    }
111    fn visit_mut_script(&mut self, m: &mut Script) {
112        m.visit_mut_children_with(self);
113        let s = take(&mut self.data);
114        let mut t = vec![];
115        brotli::BrotliCompress(&mut &*s, &mut t, &Default::default()).unwrap();
116        let s = base2048::encode(&t);
117        let s = Stmt::Decl(Decl::Var(Box::new(VarDecl {
118            span: m.span(),
119            ctxt: Default::default(),
120            kind: swc_ecma_ast::VarDeclKind::Const,
121            declare: true,
122            decls: vec![VarDeclarator {
123                span: m.span(),
124                name: Pat::Ident(self.id.clone().into()),
125                definite: false,
126                init: Some(Box::new(Expr::Call(CallExpr {
127                    span: m.span(),
128                    ctxt: Default::default(),
129                    callee: swc_ecma_ast::Callee::Expr(Box::new(self.source_mapper.clone())),
130                    args: vec![ExprOrSpread {
131                        spread: None,
132                        expr: Box::new(Expr::Lit(Lit::Str(Str {
133                            span: m.span(),
134                            value: Atom::new(s),
135                            raw: None,
136                        }))),
137                    }],
138                    type_args: None,
139                }))),
140            }],
141        })));
142        m.body = [s].into_iter().chain(m.body.drain(..)).collect();
143    }
144    fn visit_mut_decl(&mut self, node: &mut Decl) {
145        if let Decl::Fn(f) = node.clone() {
146            *node = Decl::Var(Box::new(VarDecl {
147                span: f.span(),
148                ctxt: f.ident.ctxt.clone(),
149                kind: swc_ecma_ast::VarDeclKind::Const,
150                declare: true,
151                decls: vec![VarDeclarator {
152                    span: f.span(),
153                    name: f.ident.clone().into(),
154                    init: Some(Box::new(Expr::Fn(FnExpr {
155                        ident: Some(f.ident),
156                        function: f.function,
157                    }))),
158                    definite: false,
159                }],
160            }))
161        }
162        node.visit_mut_children_with(self);
163    }
164    fn visit_mut_expr(&mut self, node: &mut Expr) {
165        node.visit_mut_children_with(self);
166        if let Expr::Fn(f) = node.clone() {
167            let l = self.data.len();
168            self.data.extend(self.sm.span_to_string(f.span()).bytes());
169            let m = self.data.len();
170            *node = Expr::Call(CallExpr {
171                span: f.span(),
172                ctxt: SyntaxContext::empty(),
173                callee: swc_ecma_ast::Callee::Expr(Box::new(Expr::Ident(self.id.clone()))),
174                args: vec![
175                    ExprOrSpread {
176                        spread: None,
177                        expr: Box::new(Expr::Lit(Lit::Str(Str {
178                            span: f.span(),
179                            raw: None,
180                            value: Atom::new(format!("{l};{m}")),
181                        }))),
182                    },
183                    ExprOrSpread {
184                        spread: None,
185                        expr: Box::new(take(node)),
186                    },
187                ],
188                type_args: None,
189            })
190        }
191        if let Expr::Arrow(f) = node.clone() {
192            let l = self.data.len();
193            self.data.extend(self.sm.span_to_string(f.span()).bytes());
194            let m = self.data.len();
195            *node = Expr::Call(CallExpr {
196                span: f.span(),
197                ctxt: SyntaxContext::empty(),
198                callee: swc_ecma_ast::Callee::Expr(Box::new(Expr::Ident(self.id.clone()))),
199                args: vec![
200                    ExprOrSpread {
201                        spread: None,
202                        expr: Box::new(Expr::Lit(Lit::Str(Str {
203                            span: f.span(),
204                            raw: None,
205                            value: Atom::new(format!("{l};{m}")),
206                        }))),
207                    },
208                    ExprOrSpread {
209                        spread: None,
210                        expr: Box::new(take(node)),
211                    },
212                ],
213                type_args: None,
214            })
215        }
216    }
217}
218pub struct Wimple {
219    pub root: SyntaxContext,
220    pub guest_id: Atom,
221}
222impl VisitMut for Wimple {
223    fn visit_mut_with_stmt(&mut self, node: &mut swc_ecma_ast::WithStmt) {
224        node.visit_mut_children_with(self);
225        node.obj = Box::new(Expr::Call(CallExpr {
226            span: node.span,
227            ctxt: self.root,
228            callee: swc_ecma_ast::Callee::Expr(Box::new(Expr::Member(MemberExpr {
229                span: node.span,
230                obj: Box::new(wevy(node.span, self.root)),
231                prop: swc_ecma_ast::MemberProp::Ident(IdentName {
232                    span: node.span,
233                    sym: Atom::new("withProxy"),
234                }),
235            }))),
236            args: vec![ExprOrSpread {
237                spread: None,
238                expr: take(&mut node.obj),
239            }],
240            type_args: None,
241        }))
242    }
243    fn visit_mut_expr(&mut self, e: &mut Expr) {
244        'a: {
245            if let Expr::Call(c) = e {
246                if let Callee::Expr(e) = &mut c.callee {
247                    if let Expr::Ident(i) = &**e {
248                        if i.sym.as_str() == "eval" {
249                            c.args.visit_mut_children_with(self);
250                            if let Some(e) = c.args.get_mut(0) {
251                                let e = &mut *e.expr;
252                                let w = wevy(e.span(), self.root);
253                                let w = [
254                                    Atom::new("guests"),
255                                    self.guest_id.clone(),
256                                    Atom::new("rewrite"),
257                                ]
258                                .into_iter()
259                                .fold(w, |e, a| {
260                                    let s = e.span();
261                                    Expr::Member(MemberExpr {
262                                        span: e.span(),
263                                        obj: Box::new(e),
264                                        prop: swc_ecma_ast::MemberProp::Ident(IdentName {
265                                            span: s,
266                                            sym: a,
267                                        }),
268                                    })
269                                });
270                                let e2 = take(e);
271                                *e = Expr::Call(CallExpr {
272                                    span: e2.span(),
273                                    ctxt: self.root,
274                                    callee: Callee::Expr(Box::new(w)),
275                                    args: vec![ExprOrSpread {
276                                        spread: None,
277                                        expr: Box::new(e2),
278                                    }],
279                                    type_args: None,
280                                });
281                            }
282                            break 'a;
283                        }
284                    }
285                    if let Expr::Member(m) = &mut **e {
286                        if let MemberProp::Computed(_) | MemberProp::Ident(_) = &m.prop {
287                            c.args.visit_mut_children_with(self);
288                            m.visit_mut_children_with(self);
289                            if *m.obj.as_ref() == wevy(m.obj.span(), self.root) {
290                                return;
291                            }
292                            // let e = &mut **e;
293                            let w = wevy(m.span(), self.root);
294                            let w = [
295                                Atom::new("guests"),
296                                self.guest_id.clone(),
297                                Atom::new("ofThis"),
298                            ]
299                            .into_iter()
300                            .fold(w, |e, a| {
301                                let s = e.span();
302                                Expr::Member(MemberExpr {
303                                    span: e.span(),
304                                    obj: Box::new(e),
305                                    prop: swc_ecma_ast::MemberProp::Ident(IdentName {
306                                        span: s,
307                                        sym: a,
308                                    }),
309                                })
310                            });
311                            let m = take(m);
312                            // let e2 = take(&mut **e);
313                            **e = Expr::Call(CallExpr {
314                                span: m.span(),
315                                ctxt: self.root,
316                                callee: Callee::Expr(Box::new(w)),
317                                args: vec![
318                                    ExprOrSpread {
319                                        spread: None,
320                                        expr: m.obj,
321                                    },
322                                    ExprOrSpread {
323                                        spread: None,
324                                        expr: match m.prop {
325                                            MemberProp::Ident(ident_name) => {
326                                                Box::new(Expr::Lit(Lit::Str(Str {
327                                                    span: ident_name.span,
328                                                    value: ident_name.sym,
329                                                    raw: None,
330                                                })))
331                                            }
332                                            MemberProp::PrivateName(private_name) => todo!(),
333                                            MemberProp::Computed(computed_prop_name) => {
334                                                computed_prop_name.expr
335                                            }
336                                        },
337                                    },
338                                ],
339                                type_args: None,
340                            });
341                            break 'a;
342                        }
343                    }
344                }
345            }
346            e.visit_mut_children_with(self);
347        }
348        // };
349        if let Expr::Lit(_) | Expr::Cond(_) | Expr::Unary(_) | Expr::Bin(_) = e {
350            return;
351        }
352        let w = wevy(e.span(), self.root);
353        let w = [Atom::new("guests"), self.guest_id.clone(), Atom::new("of")]
354            .into_iter()
355            .fold(w, |e, a| {
356                let s = e.span();
357                Expr::Member(MemberExpr {
358                    span: e.span(),
359                    obj: Box::new(e),
360                    prop: swc_ecma_ast::MemberProp::Ident(IdentName { span: s, sym: a }),
361                })
362            });
363        let e2 = take(e);
364        *e = Expr::Call(CallExpr {
365            span: e2.span(),
366            ctxt: self.root,
367            callee: Callee::Expr(Box::new(w)),
368            args: vec![ExprOrSpread {
369                spread: None,
370                expr: Box::new(e2),
371            }],
372            type_args: None,
373        });
374    }
375}