Skip to content

Commit

Permalink
fix(transformer/async-to-generator): incorrect transform when super e…
Browse files Browse the repository at this point in the history
…xpression is inside async method
  • Loading branch information
Dunqing committed Nov 6, 2024
1 parent c307e1b commit 58f56f8
Show file tree
Hide file tree
Showing 3 changed files with 231 additions and 68 deletions.
283 changes: 226 additions & 57 deletions crates/oxc_transformer/src/common/arrow_function_converter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -87,15 +87,16 @@
//! The Implementation based on
//! <https://github.com/babel/babel/blob/d20b314c14533ab86351ecf6ca6b7296b66a57b3/packages/babel-traverse/src/path/conversion.ts#L170-L247>
use oxc_allocator::{Box as ArenaBox, Vec as ArenaVec};
use oxc_ast::ast::*;
use oxc_allocator::{Box as ArenaBox, String as ArenaString, Vec as ArenaVec};
use oxc_ast::{ast::*, NONE};
use oxc_data_structures::stack::SparseStack;
use oxc_span::SPAN;
use oxc_syntax::{
scope::{ScopeFlags, ScopeId},
symbol::SymbolFlags,
};
use oxc_traverse::{Ancestor, BoundIdentifier, Traverse, TraverseCtx};
use rustc_hash::FxHashMap;

use crate::TransformOptions;

Expand All @@ -112,9 +113,15 @@ pub enum ArrowFunctionConverterMode {
AsyncOnly,
}

struct SuperProp<'a> {
binding: BoundIdentifier<'a>,
init: Expression<'a>,
}

pub struct ArrowFunctionConverter<'a> {
mode: ArrowFunctionConverterMode,
this_var_stack: SparseStack<BoundIdentifier<'a>>,
super_props: Option<FxHashMap<Atom<'a>, SuperProp<'a>>>,
}

impl<'a> ArrowFunctionConverter<'a> {
Expand All @@ -129,7 +136,7 @@ impl<'a> ArrowFunctionConverter<'a> {
ArrowFunctionConverterMode::Disabled
};
// `SparseStack` is created with 1 empty entry, for `Program`
Self { mode, this_var_stack: SparseStack::new() }
Self { mode, this_var_stack: SparseStack::new(), super_props: None }
}
}

Expand All @@ -143,25 +150,26 @@ impl<'a> Traverse<'a> for ArrowFunctionConverter<'a> {
return;
}

if let Some(this_var) = self.this_var_stack.take_last() {
let target_scope_id = program.scope_id();
Self::insert_this_var_statement_at_the_top_of_statements(
&mut program.body,
target_scope_id,
&this_var,
ctx,
);
}
let this_var = self.this_var_stack.take_last();
self.insert_variable_statement_at_the_top_of_statements(
program.scope_id(),
&mut program.body,
this_var,
ctx,
);
debug_assert!(self.this_var_stack.len() == 1);
debug_assert!(self.this_var_stack.last().is_none());
}

fn enter_function(&mut self, _func: &mut Function<'a>, _ctx: &mut TraverseCtx<'a>) {
fn enter_function(&mut self, _func: &mut Function<'a>, ctx: &mut TraverseCtx<'a>) {
if self.is_disabled() {
return;
}

self.this_var_stack.push(None);
if self.is_async_only() && matches!(ctx.parent(), Ancestor::MethodDefinitionValue(_)) {
self.super_props = Some(FxHashMap::default());
}
}

/// ```ts
Expand All @@ -180,17 +188,17 @@ impl<'a> Traverse<'a> for ArrowFunctionConverter<'a> {
return;
}

if let Some(this_var) = self.this_var_stack.pop() {
let target_scope_id = func.scope_id();
let Some(body) = &mut func.body else { unreachable!() };

Self::insert_this_var_statement_at_the_top_of_statements(
&mut body.statements,
target_scope_id,
&this_var,
ctx,
);
}
let scope_id = func.scope_id();
let Some(body) = &mut func.body else {
return;
};
let this_var = self.this_var_stack.pop();
self.insert_variable_statement_at_the_top_of_statements(
scope_id,
&mut body.statements,
this_var,
ctx,
);
}

fn enter_static_block(&mut self, _block: &mut StaticBlock<'a>, _ctx: &mut TraverseCtx<'a>) {
Expand All @@ -206,15 +214,13 @@ impl<'a> Traverse<'a> for ArrowFunctionConverter<'a> {
return;
}

if let Some(this_var) = self.this_var_stack.pop() {
let target_scope_id = block.scope_id();
Self::insert_this_var_statement_at_the_top_of_statements(
&mut block.body,
target_scope_id,
&this_var,
ctx,
);
}
let this_var = self.this_var_stack.pop();
self.insert_variable_statement_at_the_top_of_statements(
block.scope_id(),
&mut block.body,
this_var,
ctx,
);
}

fn enter_jsx_element_name(
Expand Down Expand Up @@ -254,10 +260,19 @@ impl<'a> Traverse<'a> for ArrowFunctionConverter<'a> {
return;
}

if let Expression::ThisExpression(this) = expr {
if let Some(ident) = self.get_this_identifier(this.span, ctx) {
*expr = Expression::Identifier(ident);
let new_expr = match expr {
Expression::ThisExpression(this) => {
self.get_this_identifier(this.span, ctx).map(Expression::Identifier)
}
Expression::CallExpression(_) => self.transform_call_expression_for_super(expr, ctx),
match_member_expression!(Expression) => {
self.transform_member_expression_for_super(expr, false, ctx)
}
_ => return,
};

if let Some(new_expr) = new_expr {
*expr = new_expr;
}
}

Expand Down Expand Up @@ -383,14 +398,16 @@ impl<'a> ArrowFunctionConverter<'a> {
// Arrow function
Ancestor::ArrowFunctionExpressionParams(func) => {
return if self.is_async_only() && !*func.r#async() {
None
// Continue checking the parent to see if it's inside an async function.
continue;
} else {
Some(func.scope_id().get().unwrap())
};
}
Ancestor::ArrowFunctionExpressionBody(func) => {
return if self.is_async_only() && !*func.r#async() {
None
// Continue checking the parent to see if it's inside an async function.
continue;
} else {
Some(func.scope_id().get().unwrap())
};
Expand Down Expand Up @@ -449,39 +466,191 @@ impl<'a> ArrowFunctionConverter<'a> {
))
}

/// Insert `var _this = this;` at the top of the statements.
fn insert_this_var_statement_at_the_top_of_statements(
statements: &mut ArenaVec<'a, Statement<'a>>,
fn get_nested_member_expression<'b>(
expr: &'b MemberExpression<'a>,
) -> &'b MemberExpression<'a> {
let mut member = expr;
loop {
match member.object() {
Expression::ComputedMemberExpression(computed_expr) => {
if let Some(m) = computed_expr.object.as_member_expression() {
member = m;
continue;
}
}
Expression::StaticMemberExpression(static_expr) => {
if let Some(m) = static_expr.object.as_member_expression() {
member = m;
continue;
}
}
Expression::ChainExpression(chain) => {
if let Some(m) = &chain.expression.as_member_expression() {
member = m;
continue;
}
}
_ => (),
}
return member;
}
}

fn transform_member_expression_for_super(
&mut self,
expr: &mut Expression<'a>,
is_assignment: bool,
ctx: &mut TraverseCtx<'a>,
) -> Option<Expression<'a>> {
debug_assert!(expr.as_member_expression().is_some());

let super_props = self.super_props.as_mut()?;
let member = Self::get_nested_member_expression(expr.to_member_expression());

if !matches!(member.object(), Expression::Super(_)) {
return None;
}

let mut binding_name = ArenaString::new_in(ctx.ast.allocator);
binding_name.push_str("superprop_");
if is_assignment {
binding_name.push_str("set");
} else {
binding_name.push_str("get");
}

if let MemberExpression::StaticMemberExpression(static_member) = member {
let property = static_member.property.name.as_str();
if let Some(first_byte) = property.as_bytes().first() {
binding_name.push(first_byte.to_ascii_uppercase() as char);
}
binding_name.push_str(&property[1..]);
};
let binding_name = Atom::from(binding_name.into_bump_str());
let super_prop = super_props.entry(binding_name.clone()).or_insert_with(|| {
let binding = ctx
.generate_uid_in_current_scope(&binding_name, SymbolFlags::FunctionScopedVariable);
SuperProp { binding: binding.clone(), init: ctx.ast.move_expression(expr) }
});

Some(super_prop.binding.create_read_expression(ctx))
}

// TODO: add why?
#[inline]
fn transform_call_expression_for_super(
&mut self,
expr: &mut Expression<'a>,
ctx: &mut TraverseCtx<'a>,
) -> Option<Expression<'a>> {
let Expression::CallExpression(call) = expr else { unreachable!() };
if !call.callee.is_member_expression() {
return None;
}
let callee = self.transform_member_expression_for_super(&mut call.callee, false, ctx)?;
call.callee = callee;
// wrap it with .call(this)
let this_arg = Argument::from(ctx.ast.expression_this(SPAN));
let arguments = ctx.ast.vec1(this_arg);
let object = ctx.ast.move_expression(expr);
let property = ctx.ast.identifier_name(SPAN, "call");
let callee = ctx.ast.member_expression_static(SPAN, object, property, false);
let callee = Expression::from(callee);
Some(ctx.ast.expression_call(SPAN, callee, NONE, arguments, false))
}

/// Insert variable statement at the top of the statements.
fn insert_variable_statement_at_the_top_of_statements(
&mut self,
target_scope_id: ScopeId,
this_var: &BoundIdentifier<'a>,
statements: &mut ArenaVec<'a, Statement<'a>>,
this_var: Option<BoundIdentifier<'a>>,
ctx: &mut TraverseCtx<'a>,
) {
let symbol_id = this_var.symbol_id;
let scope_id = ctx.symbols().get_scope_id(symbol_id);
// Because scope can be moved or deleted, we need to check the scope id again,
// if it's different, we need to move the binding to the target scope.
if target_scope_id != scope_id {
ctx.scopes_mut().move_binding(scope_id, target_scope_id, &this_var.name);
ctx.symbols_mut().set_scope_id(symbol_id, target_scope_id);
// `_superprop_getSomething = () => super.getSomething;`
let is_method_definition_parent =
matches!(ctx.parent(), Ancestor::MethodDefinitionValue(_));
let mut declarations = if is_method_definition_parent {
if let Some(super_props) = self.super_props.as_mut() {
let mut declarations = ctx.ast.vec_with_capacity(super_props.len() + 1);
declarations.extend(super_props.drain().map(|(_, super_prop)| {
Self::adjust_binding_scope(target_scope_id, &super_prop.binding, ctx);
let items = ctx.ast.vec();
let params = ctx.ast.formal_parameters(
SPAN,
FormalParameterKind::ArrowFormalParameters,
items,
NONE,
);
let scope_id = ctx.create_child_scope(
target_scope_id,
ScopeFlags::Arrow | ScopeFlags::Function,
);
let statements =
ctx.ast.vec1(ctx.ast.statement_expression(SPAN, super_prop.init));
let body = ctx.ast.function_body(SPAN, ctx.ast.vec(), statements);
let init = ctx.ast.alloc_arrow_function_expression_with_scope_id(
SPAN, true, false, NONE, params, NONE, body, scope_id,
);
ctx.ast.variable_declarator(
SPAN,
VariableDeclarationKind::Var,
super_prop.binding.create_binding_pattern(ctx),
Some(Expression::ArrowFunctionExpression(init)),
false,
)
}));
declarations
} else {
ctx.ast.vec_with_capacity(1)
}
} else {
ctx.ast.vec_with_capacity(1)
};

// `_this = this;`
if let Some(this_var) = this_var {
Self::adjust_binding_scope(target_scope_id, &this_var, ctx);

let variable_declarator = ctx.ast.variable_declarator(
SPAN,
VariableDeclarationKind::Var,
this_var.create_binding_pattern(ctx),
Some(ctx.ast.expression_this(SPAN)),
false,
);
declarations.push(variable_declarator);
}

let variable_declarator = ctx.ast.variable_declarator(
SPAN,
VariableDeclarationKind::Var,
this_var.create_binding_pattern(ctx),
Some(ctx.ast.expression_this(SPAN)),
false,
);
if declarations.is_empty() {
return;
}

let stmt = ctx.ast.alloc_variable_declaration(
SPAN,
VariableDeclarationKind::Var,
ctx.ast.vec1(variable_declarator),
declarations,
false,
);

let stmt = Statement::VariableDeclaration(stmt);

statements.insert(0, stmt);
}

/// Adjust the scope of the binding.
///
/// Since scope can be moved or deleted, we need to ensure the scope of the binding
/// same as the target scope, if it's mismatch, we need to move the binding to the target scope.
fn adjust_binding_scope(
target_scope_id: ScopeId,
binding: &BoundIdentifier<'a>,
ctx: &mut TraverseCtx<'a>,
) {
let original_scope_id = ctx.symbols().get_scope_id(binding.symbol_id);
if target_scope_id != original_scope_id {
ctx.symbols_mut().set_scope_id(binding.symbol_id, target_scope_id);
ctx.scopes_mut().move_binding(original_scope_id, target_scope_id, &binding.name);
}
}
}
Loading

0 comments on commit 58f56f8

Please sign in to comment.