Skip to content

Commit 267c144

Browse files
kdy1kdy1
committed
Inject!
1 parent 77cb29c commit 267c144

File tree

3 files changed

+314
-5
lines changed

3 files changed

+314
-5
lines changed

‎crates/swc_ecma_transformer/src/common/statement_injector.rs‎

Lines changed: 86 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
//! self.ctx.statement_injector.insert_many_after(address, statements);
1616
//! ```
1717
18+
use rustc_hash::FxHashMap;
1819
use swc_ecma_ast::*;
1920
use swc_ecma_hooks::VisitMutHook;
2021

@@ -36,7 +37,90 @@ struct AdjacentStmt {
3637
}
3738

3839
/// Store for Stmts to be added to the Stmts.
40+
///
41+
/// The key is the address of the statement in the AST, represented as a
42+
/// pointer.
3943
#[derive(Default)]
40-
pub struct StmtInjectorStore {}
44+
pub struct StmtInjectorStore {
45+
/// Map from statement address to adjacent statements to insert
46+
stmts: FxHashMap<usize, Vec<AdjacentStmt>>,
47+
}
48+
49+
impl StmtInjectorStore {
50+
/// Insert a statement before the statement at the given address
51+
pub fn insert_before(&mut self, address: usize, stmt: Stmt) {
52+
self.stmts.entry(address).or_default().push(AdjacentStmt {
53+
stmt,
54+
direction: Direction::Before,
55+
});
56+
}
57+
58+
/// Insert a statement after the statement at the given address
59+
pub fn insert_after(&mut self, address: usize, stmt: Stmt) {
60+
self.stmts.entry(address).or_default().push(AdjacentStmt {
61+
stmt,
62+
direction: Direction::After,
63+
});
64+
}
65+
66+
/// Insert multiple statements after the statement at the given address
67+
pub fn insert_many_after(&mut self, address: usize, stmts: Vec<Stmt>) {
68+
let entry = self.stmts.entry(address).or_default();
69+
for stmt in stmts {
70+
entry.push(AdjacentStmt {
71+
stmt,
72+
direction: Direction::After,
73+
});
74+
}
75+
}
76+
77+
/// Get all statements to be inserted at the given address
78+
fn take_stmts(&mut self, address: usize) -> Option<Vec<AdjacentStmt>> {
79+
self.stmts.remove(&address)
80+
}
81+
}
4182

42-
impl VisitMutHook<TraverseCtx> for StmtInjector {}
83+
impl VisitMutHook<TraverseCtx> for StmtInjector {
84+
fn enter_stmts(&mut self, stmts: &mut Vec<Stmt>, ctx: &mut TraverseCtx) {
85+
let mut i = 0;
86+
while i < stmts.len() {
87+
let stmt = &stmts[i];
88+
let address = stmt as *const Stmt as usize;
89+
90+
// Check if there are any statements to insert at this address
91+
if let Some(adjacent_stmts) = ctx.statement_injector.take_stmts(address) {
92+
let mut before_stmts = Vec::new();
93+
let mut after_stmts = Vec::new();
94+
95+
// Separate statements by direction
96+
for adjacent in adjacent_stmts {
97+
match adjacent.direction {
98+
Direction::Before => before_stmts.push(adjacent.stmt),
99+
Direction::After => after_stmts.push(adjacent.stmt),
100+
}
101+
}
102+
103+
// Insert statements before
104+
let before_count = before_stmts.len();
105+
if before_count > 0 {
106+
// Insert all before statements at position i
107+
for (offset, stmt) in before_stmts.into_iter().enumerate() {
108+
stmts.insert(i + offset, stmt);
109+
}
110+
// Move index forward by the number of inserted statements
111+
i += before_count;
112+
}
113+
114+
// Insert statements after
115+
if !after_stmts.is_empty() {
116+
// Insert all after statements at position i + 1
117+
for (offset, stmt) in after_stmts.into_iter().enumerate() {
118+
stmts.insert(i + 1 + offset, stmt);
119+
}
120+
}
121+
}
122+
123+
i += 1;
124+
}
125+
}
126+
}

‎crates/swc_ecma_transformer/src/es2016/exponentiation_operator.rs‎

Lines changed: 213 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,38 @@
1+
//! ES2016: Exponentiation Operator
2+
//!
3+
//! This plugin transforms the exponentiation operator (`**`) to `Math.pow`.
4+
//!
5+
//! > This plugin is included in `preset-env`, in ES2016
6+
//!
7+
//! ## Example
8+
//!
9+
//! Input:
10+
//! ```js
11+
//! let x = 10 ** 2;
12+
//! x **= 3;
13+
//! obj.prop **= 4;
14+
//! ```
15+
//!
16+
//! Output:
17+
//! ```js
18+
//! let x = Math.pow(10, 2);
19+
//! x = Math.pow(x, 3);
20+
//! obj["prop"] = Math.pow(obj["prop"], 4);
21+
//! ```
22+
//!
23+
//! ## Implementation
24+
//!
25+
//! Implementation based on [@babel/plugin-transform-exponentiation-operator](https://babel.dev/docs/babel-plugin-transform-exponentiation-operator).
26+
//!
27+
//! ## References:
28+
//!
29+
//! * Babel plugin implementation: <https://github.com/babel/babel/blob/v7.26.2/packages/babel-plugin-transform-exponentiation-operator>
30+
//! <https://github.com/babel/babel/tree/v7.26.2/packages/babel-helper-builder-binary-assignment-operator-visitor>
31+
//! * Exponentiation operator TC39 proposal: <https://github.com/tc39/proposal-exponentiation-operator>
32+
//! * Exponentiation operator specification: <https://tc39.es/ecma262/#sec-exp-operator>
33+
34+
use swc_common::{util::take::Take, SyntaxContext, DUMMY_SP};
35+
use swc_ecma_ast::*;
136
use swc_ecma_hooks::VisitMutHook;
237

338
use crate::TraverseCtx;
@@ -9,4 +44,181 @@ pub fn hook() -> impl VisitMutHook<TraverseCtx> {
944
#[derive(Default)]
1045
struct ExponentiationOperatorPass {}
1146

12-
impl VisitMutHook<TraverseCtx> for ExponentiationOperatorPass {}
47+
impl VisitMutHook<TraverseCtx> for ExponentiationOperatorPass {
48+
fn enter_expr(&mut self, expr: &mut Expr, _ctx: &mut TraverseCtx) {
49+
match expr {
50+
// `left ** right` -> `Math.pow(left, right)`
51+
Expr::Bin(bin_expr) if bin_expr.op == BinaryOp::Exp => {
52+
// Do not transform BigInt
53+
if is_bigint_literal(&bin_expr.left) || is_bigint_literal(&bin_expr.right) {
54+
return;
55+
}
56+
57+
let left = bin_expr.left.take();
58+
let right = bin_expr.right.take();
59+
*expr = create_math_pow(left, right);
60+
}
61+
// `left **= right` -> various transformations
62+
Expr::Assign(assign_expr) if assign_expr.op == AssignOp::ExpAssign => {
63+
// Do not transform BigInt
64+
if is_bigint_literal(&assign_expr.right) {
65+
return;
66+
}
67+
68+
self.transform_exp_assign(expr);
69+
}
70+
_ => {}
71+
}
72+
}
73+
}
74+
75+
impl ExponentiationOperatorPass {
76+
/// Transform exponentiation assignment expression
77+
///
78+
/// Transforms based on the type of left side:
79+
/// - Identifier: `x **= 2` -> `x = Math.pow(x, 2)`
80+
/// - Member expression: `obj.prop **= 2` -> `obj.prop = Math.pow(obj.prop,
81+
/// 2)` or with temp var
82+
fn transform_exp_assign(&self, expr: &mut Expr) {
83+
let Expr::Assign(assign_expr) = expr else {
84+
return;
85+
};
86+
87+
match &assign_expr.left {
88+
// Simple identifier: `x **= 2` -> `x = Math.pow(x, 2)`
89+
AssignTarget::Simple(SimpleAssignTarget::Ident(ident)) => {
90+
let left = ident.id.clone();
91+
let right = assign_expr.right.take();
92+
93+
let pow_left = Box::new(Expr::Ident(left.clone()));
94+
let pow_call = create_math_pow(pow_left, right);
95+
96+
*expr = Expr::Assign(AssignExpr {
97+
span: DUMMY_SP,
98+
op: AssignOp::Assign,
99+
left: AssignTarget::Simple(SimpleAssignTarget::Ident(BindingIdent {
100+
id: left,
101+
type_ann: None,
102+
})),
103+
right: Box::new(pow_call),
104+
});
105+
}
106+
// Member expression: needs special handling to avoid side effects
107+
AssignTarget::Simple(SimpleAssignTarget::Member(member_expr)) => {
108+
// Clone the data we need before borrowing mutably
109+
let obj = member_expr.obj.clone();
110+
let prop = member_expr.prop.clone();
111+
let right = assign_expr.right.take();
112+
113+
self.transform_member_exp_assign_impl(expr, obj, prop, right);
114+
}
115+
_ => {}
116+
}
117+
}
118+
119+
/// Transform member expression exponentiation assignment implementation
120+
///
121+
/// `obj.prop **= 2` or `obj[prop] **= 2`
122+
fn transform_member_exp_assign_impl(
123+
&self,
124+
expr: &mut Expr,
125+
obj: Box<Expr>,
126+
prop: MemberProp,
127+
right: Box<Expr>,
128+
) {
129+
// For simplicity, handle the basic case without temp variables
130+
// A more complete implementation would add temp variables to avoid side effects
131+
match prop {
132+
// obj.prop **= right -> obj.prop = Math.pow(obj.prop, right)
133+
MemberProp::Ident(prop_name) => {
134+
// Create Math.pow(obj.prop, right)
135+
let pow_left = Box::new(Expr::Member(MemberExpr {
136+
span: DUMMY_SP,
137+
obj: obj.clone(),
138+
prop: MemberProp::Ident(prop_name.clone()),
139+
}));
140+
let pow_call = create_math_pow(pow_left, right);
141+
142+
// obj.prop = Math.pow(obj.prop, right)
143+
*expr = Expr::Assign(AssignExpr {
144+
span: DUMMY_SP,
145+
op: AssignOp::Assign,
146+
left: AssignTarget::Simple(SimpleAssignTarget::Member(MemberExpr {
147+
span: DUMMY_SP,
148+
obj,
149+
prop: MemberProp::Ident(prop_name),
150+
})),
151+
right: Box::new(pow_call),
152+
});
153+
}
154+
// obj[computed] **= right -> obj[computed] = Math.pow(obj[computed], right)
155+
MemberProp::Computed(computed) => {
156+
let computed_prop = computed.expr;
157+
158+
// Create Math.pow(obj[computed], right)
159+
let pow_left = Box::new(Expr::Member(MemberExpr {
160+
span: DUMMY_SP,
161+
obj: obj.clone(),
162+
prop: MemberProp::Computed(ComputedPropName {
163+
span: DUMMY_SP,
164+
expr: computed_prop.clone(),
165+
}),
166+
}));
167+
let pow_call = create_math_pow(pow_left, right);
168+
169+
// obj[computed] = Math.pow(obj[computed], right)
170+
*expr = Expr::Assign(AssignExpr {
171+
span: DUMMY_SP,
172+
op: AssignOp::Assign,
173+
left: AssignTarget::Simple(SimpleAssignTarget::Member(MemberExpr {
174+
span: DUMMY_SP,
175+
obj,
176+
prop: MemberProp::Computed(ComputedPropName {
177+
span: DUMMY_SP,
178+
expr: computed_prop,
179+
}),
180+
})),
181+
right: Box::new(pow_call),
182+
});
183+
}
184+
_ => {}
185+
}
186+
}
187+
}
188+
189+
/// Create `Math.pow(left, right)` call expression
190+
fn create_math_pow(left: Box<Expr>, right: Box<Expr>) -> Expr {
191+
Expr::Call(CallExpr {
192+
span: DUMMY_SP,
193+
ctxt: SyntaxContext::empty(),
194+
callee: Callee::Expr(Box::new(Expr::Member(MemberExpr {
195+
span: DUMMY_SP,
196+
obj: Box::new(Expr::Ident(Ident {
197+
span: DUMMY_SP,
198+
ctxt: SyntaxContext::empty(),
199+
sym: "Math".into(),
200+
optional: false,
201+
})),
202+
prop: MemberProp::Ident(IdentName {
203+
span: DUMMY_SP,
204+
sym: "pow".into(),
205+
}),
206+
}))),
207+
args: vec![
208+
ExprOrSpread {
209+
spread: None,
210+
expr: left,
211+
},
212+
ExprOrSpread {
213+
spread: None,
214+
expr: right,
215+
},
216+
],
217+
type_args: None,
218+
})
219+
}
220+
221+
/// Check if an expression is a BigInt literal
222+
fn is_bigint_literal(expr: &Expr) -> bool {
223+
matches!(expr, Expr::Lit(Lit::BigInt(_)))
224+
}

‎crates/swc_ecma_transformer/src/lib.rs‎

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,17 @@ mod regexp;
2626
mod typescript;
2727
mod utils;
2828

29-
pub struct TraverseCtx {}
29+
pub struct TraverseCtx {
30+
pub statement_injector: common::StmtInjectorStore,
31+
}
32+
33+
impl Default for TraverseCtx {
34+
fn default() -> Self {
35+
Self {
36+
statement_injector: Default::default(),
37+
}
38+
}
39+
}
3040

3141
pub fn transform_hook(options: Options) -> impl VisitMutHook<TraverseCtx> {
3242
let hook = HookBuilder::new(NoopHook);
@@ -47,11 +57,14 @@ pub fn transform_hook(options: Options) -> impl VisitMutHook<TraverseCtx> {
4757
let hook = hook.chain_optional(crate::regexp::hook(options.env.regexp));
4858
let hook = hook.chain(crate::bugfix::hook(options.env.bugfix));
4959

60+
// Statement injector must be last to process all injected statements
61+
let hook = hook.chain(common::StmtInjector::default());
62+
5063
hook.build()
5164
}
5265

5366
pub fn hook_pass<H: VisitMutHook<TraverseCtx>>(hook: H) -> impl Pass {
54-
let ctx = TraverseCtx {};
67+
let ctx = TraverseCtx::default();
5568

5669
visit_mut_pass(VisitMutWithHook { hook, context: ctx })
5770
}

0 commit comments

Comments
 (0)