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 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 **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 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}