inserted by double-visit

crates/oxc_transformer/src/common/arrow_function_converter.rs

@@ -89,20 +89,20 @@
 
 use compact_str::CompactString;
 use indexmap::IndexMap;
-use rustc_hash::{FxBuildHasher, FxHashMap, FxHashSet};
+use rustc_hash::{FxBuildHasher, FxHashSet};
 
-use oxc_allocator::{Address, Box as ArenaBox, GetAddress, Vec as ArenaVec};
-use oxc_ast::{ast::*, NONE};
-use oxc_data_structures::stack::{NonEmptyStack, SparseStack};
+use oxc_allocator::{Box as ArenaBox, Vec as ArenaVec};
+use oxc_ast::{ast::*, visit::walk_mut::walk_expression, VisitMut, NONE};
+use oxc_data_structures::stack::{NonEmptyStack, SparseStack, Stack};
 use oxc_semantic::{ReferenceFlags, SymbolId};
-use oxc_span::{CompactStr, SPAN};
+use oxc_span::{CompactStr, GetSpan, SPAN};
 use oxc_syntax::{
     scope::{ScopeFlags, ScopeId},
     symbol::SymbolFlags,
 };
 use oxc_traverse::{Ancestor, BoundIdentifier, Traverse, TraverseCtx};
 
-use crate::{context::TransformCtx, EnvOptions};
+use crate::EnvOptions;
 
 type FxIndexMap<K, V> = IndexMap<K, V, FxBuildHasher>;
 
@@ -135,22 +135,20 @@ struct SuperMethodInfo<'a> {
     is_computed: bool,
 }
 
-pub struct ArrowFunctionConverter<'a, 'ctx> {
-    ctx: &'ctx TransformCtx<'a>,
+pub struct ArrowFunctionConverter<'a> {
     mode: ArrowFunctionConverterMode,
     this_var_stack: SparseStack<BoundIdentifier<'a>>,
-    /// Stores the address of statement of containing `super()` expression
-    super_call_addresses: FxHashMap<ScopeId, Address>,
     arguments_var_stack: SparseStack<BoundIdentifier<'a>>,
+    constructor_super_stack: Stack<bool>,
     arguments_needs_transform_stack: NonEmptyStack<bool>,
     renamed_arguments_symbol_ids: FxHashSet<SymbolId>,
     // TODO(improve-on-babel): `FxHashMap` would suffice here. Iteration order is not important.
     // Only using `FxIndexMap` for predictable iteration order to match Babel's output.
     super_methods: Option<FxIndexMap<SuperMethodKey<'a>, SuperMethodInfo<'a>>>,
 }
 
-impl<'a, 'ctx> ArrowFunctionConverter<'a, 'ctx> {
-    pub fn new(env: &EnvOptions, ctx: &'ctx TransformCtx<'a>) -> Self {
+impl ArrowFunctionConverter<'_> {
+    pub fn new(env: &EnvOptions) -> Self {
         let mode = if env.es2015.arrow_function.is_some() {
             ArrowFunctionConverterMode::Enabled
         } else if env.es2017.async_to_generator || env.es2018.async_generator_functions {
@@ -160,19 +158,18 @@ impl<'a, 'ctx> ArrowFunctionConverter<'a, 'ctx> {
         };
         // `SparseStack`s are created with 1 empty entry, for `Program`
         Self {
-            ctx,
             mode,
             this_var_stack: SparseStack::new(),
-            super_call_addresses: FxHashMap::default(),
             arguments_var_stack: SparseStack::new(),
+            constructor_super_stack: Stack::new(),
             arguments_needs_transform_stack: NonEmptyStack::new(false),
             renamed_arguments_symbol_ids: FxHashSet::default(),
             super_methods: None,
         }
     }
 }
 
-impl<'a, 'ctx> Traverse<'a> for ArrowFunctionConverter<'a, 'ctx> {
+impl<'a> Traverse<'a> for ArrowFunctionConverter<'a> {
     // Note: No visitors for `TSModuleBlock` because `this` is not legal in TS module blocks.
     // <https://www.typescriptlang.org/play/?#code/HYQwtgpgzgDiDGEAEAxA9mpBvAsAKCSXjWCgBckANJAXiQAoBKWgPiTIAsBLKAbnwC++fGDQATAK4AbZACEQAJ2z5CxUhWp0mrdtz6D8QA>
 
@@ -205,6 +202,7 @@ impl<'a, 'ctx> Traverse<'a> for ArrowFunctionConverter<'a, 'ctx> {
 
         self.this_var_stack.push(None);
         self.arguments_var_stack.push(None);
+        self.constructor_super_stack.push(false);
         if self.is_async_only() && func.r#async && Self::is_class_method_like_ancestor(ctx.parent())
         {
             self.super_methods = Some(FxIndexMap::default());
@@ -240,6 +238,7 @@ impl<'a, 'ctx> Traverse<'a> for ArrowFunctionConverter<'a, 'ctx> {
             arguments_var,
             ctx,
         );
+        self.constructor_super_stack.pop();
     }
 
     fn enter_arrow_function_expression(
@@ -336,6 +335,12 @@ impl<'a, 'ctx> Traverse<'a> for ArrowFunctionConverter<'a, 'ctx> {
             Expression::ThisExpression(this) => {
                 self.get_this_identifier(this.span, ctx).map(Expression::Identifier)
             }
+            Expression::Super(_) => {
+                if let Some(v) = self.constructor_super_stack.last_mut() {
+                    *v = true;
+                }
+                return;
+            }
             Expression::CallExpression(call) => self.transform_call_expression_for_super(call, ctx),
             Expression::AssignmentExpression(assignment) => {
                 self.transform_assignment_expression_for_super(assignment, ctx)
@@ -402,7 +407,7 @@ impl<'a, 'ctx> Traverse<'a> for ArrowFunctionConverter<'a, 'ctx> {
     }
 }
 
-impl<'a, 'ctx> ArrowFunctionConverter<'a, 'ctx> {
+impl<'a> ArrowFunctionConverter<'a> {
     /// Check if arrow function conversion is disabled
     fn is_disabled(&self) -> bool {
         self.mode == ArrowFunctionConverterMode::Disabled
@@ -441,26 +446,6 @@ impl<'a, 'ctx> ArrowFunctionConverter<'a, 'ctx> {
                 .unwrap();
             ctx.generate_uid("this", target_scope_id, SymbolFlags::FunctionScopedVariable)
         });
-
-        if !self.super_call_addresses.is_empty() {
-            let address = ctx
-                .scopes()
-                .get_parent_id(arrow_scope_id)
-                .and_then(|scope_id| self.super_call_addresses.remove(&scope_id));
-            if let Some(address) = address {
-                // Insert a dummy address to indicate that should inserting `var _this;`
-                // without `init` at the top of the statements.
-                self.super_call_addresses.insert(arrow_scope_id, Address::DUMMY);
-                let assignment = ctx.ast.expression_assignment(
-                    SPAN,
-                    AssignmentOperator::Assign,
-                    this_var.create_write_target(ctx),
-                    ctx.ast.expression_this(SPAN),
-                );
-                let statement = ctx.ast.statement_expression(SPAN, assignment);
-                self.ctx.statement_injector.insert_after(&address, statement);
-            }
-        }
         Some(ctx.ast.alloc(this_var.create_spanned_read_reference(span, ctx)))
     }
 
@@ -728,25 +713,6 @@ impl<'a, 'ctx> ArrowFunctionConverter<'a, 'ctx> {
         ctx: &mut TraverseCtx<'a>,
     ) -> Option<Expression<'a>> {
         if self.super_methods.is_none() || !call.callee.is_member_expression() {
-            // `super()`
-            // Store the address in case we need to insert `var _this;` after it.
-            if call.callee.is_super() {
-                let scope_id = ctx.current_scope_id();
-                self.super_call_addresses.entry(scope_id).or_insert_with(|| {
-                    ctx.ancestors()
-                        .find(|ancestor| {
-                            matches!(
-                                ancestor,
-                                // const A = super():
-                                Ancestor::VariableDeclarationDeclarations(_)
-                                // super();
-                                | Ancestor::ExpressionStatementExpression(_)
-                            )
-                        })
-                        .unwrap()
-                        .address()
-                });
-            }
             return None;
         }
 
@@ -1077,12 +1043,9 @@ impl<'a, 'ctx> ArrowFunctionConverter<'a, 'ctx> {
         let arguments = self.create_arguments_var_declarator(target_scope_id, arguments_var, ctx);
 
         let is_class_method_like = Self::is_class_method_like_ancestor(ctx.parent());
+        let super_method_count = self.super_methods.as_ref().map_or(0, FxIndexMap::len);
         let declarations_count = usize::from(arguments.is_some())
-            + if is_class_method_like {
-                self.super_methods.as_ref().map_or(0, FxIndexMap::len)
-            } else {
-                0
-            }
+            + if is_class_method_like { super_method_count } else { 0 }
             + usize::from(this_var.is_some());
 
         // Exit if no declarations to be inserted
@@ -1114,12 +1077,18 @@ impl<'a, 'ctx> ArrowFunctionConverter<'a, 'ctx> {
 
         // `_this = this;`
         if let Some(this_var) = this_var {
-            let init = if self.super_call_addresses.is_empty() {
-                Some(ctx.ast.expression_this(SPAN))
-            } else {
-                // Clear the dummy address.
-                self.super_call_addresses.clear();
+            let is_constructor = ctx.scopes().get_flags(target_scope_id).is_constructor();
+            let init = if is_constructor
+                && (super_method_count != 0
+                    || self.constructor_super_stack.last().copied().unwrap_or(false))
+            {
+                // `super()` is called in the constructor body, so we need to insert `_this = this;`
+                // after `super()` call. Because `this` is not available before `super()` call.
+                ConstructorBodyThisAfterSuperInserter::new(&this_var, ctx)
+                    .visit_statements(statements);
                 None
+            } else {
+                Some(Expression::ThisExpression(ctx.ast.alloc_this_expression(SPAN)))
             };
             Self::adjust_binding_scope(target_scope_id, &this_var, ctx);
             let variable_declarator = ctx.ast.variable_declarator(
@@ -1146,3 +1115,74 @@ impl<'a, 'ctx> ArrowFunctionConverter<'a, 'ctx> {
         statements.insert(0, stmt);
     }
 }
+
+/// Visitor for inserting `this` after `super` in constructor body.
+struct ConstructorBodyThisAfterSuperInserter<'a, 'arrow> {
+    this_var_binding: &'arrow BoundIdentifier<'a>,
+    ctx: &'arrow mut TraverseCtx<'a>,
+}
+
+impl<'a, 'b> ConstructorBodyThisAfterSuperInserter<'a, 'b> {
+    fn new(this_var_binding: &'b BoundIdentifier<'a>, ctx: &'b mut TraverseCtx<'a>) -> Self {
+        Self { this_var_binding, ctx }
+    }
+
+    #[inline]
+    fn transform_super_call_expression(&mut self, expr: &Expression<'a>) -> Option<Expression<'a>> {
+        if expr.is_super_call_expression() {
+            let assignment = self.ctx.ast.expression_assignment(
+                SPAN,
+                AssignmentOperator::Assign,
+                self.this_var_binding.create_write_target(self.ctx),
+                self.ctx.ast.expression_this(SPAN),
+            );
+            Some(assignment)
+        } else {
+            None
+        }
+    }
+}
+
+impl<'a> VisitMut<'a> for ConstructorBodyThisAfterSuperInserter<'a, '_> {
+    #[inline]
+    fn visit_class(&mut self, _class: &mut Class<'a>) {
+        // Do not need to insert in nested classes
+    }
+
+    /// `super()` -> `super(); _this = this;`
+    #[inline]
+    fn visit_statements(&mut self, statements: &mut ArenaVec<'a, Statement<'a>>) {
+        let mut new_stmts = vec![];
+
+        for (index, stmt) in statements.iter_mut().enumerate() {
+            let Statement::ExpressionStatement(expr_stmt) = stmt else {
+                self.visit_statement(stmt);
+                continue;
+            };
+            if let Some(assignment) = self.transform_super_call_expression(&expr_stmt.expression) {
+                let new_stmt = self.ctx.ast.statement_expression(SPAN, assignment);
+                new_stmts.push((index, new_stmt));
+            } else {
+                self.visit_statement(stmt);
+            }
+        }
+
+        for (index, new_stmt) in new_stmts.into_iter().rev() {
+            // insert the new statement after the super call
+            statements.insert(index + 1, new_stmt);
+        }
+    }
+
+    /// `const A = super()` -> `const A = (super(), _this = this);`
+    #[inline]
+    fn visit_expression(&mut self, expr: &mut Expression<'a>) {
+        if let Some(assignment) = self.transform_super_call_expression(expr) {
+            let span = expr.span();
+            let exprs =
+                self.ctx.ast.vec_from_array([self.ctx.ast.move_expression(expr), assignment]);
+            *expr = self.ctx.ast.expression_sequence(span, exprs);
+        } else {
+            walk_expression(self, expr);
+        }
+    }
+}

crates/oxc_transformer/src/common/mod.rs

@@ -25,7 +25,7 @@ pub struct Common<'a, 'ctx> {
     var_declarations: VarDeclarations<'a, 'ctx>,
     statement_injector: StatementInjector<'a, 'ctx>,
     top_level_statements: TopLevelStatements<'a, 'ctx>,
-    arrow_function_converter: ArrowFunctionConverter<'a, 'ctx>,
+    arrow_function_converter: ArrowFunctionConverter<'a>,
 }
 
 impl<'a, 'ctx> Common<'a, 'ctx> {
@@ -35,7 +35,7 @@ impl<'a, 'ctx> Common<'a, 'ctx> {
             var_declarations: VarDeclarations::new(ctx),
             statement_injector: StatementInjector::new(ctx),
             top_level_statements: TopLevelStatements::new(ctx),
-            arrow_function_converter: ArrowFunctionConverter::new(options, ctx),
+            arrow_function_converter: ArrowFunctionConverter::new(options),
         }
     }
 }

tasks/transform_conformance/snapshots/oxc.snap.md

@@ -1,6 +1,6 @@
 commit: 54a8389f
 
-Passed: 121/140
+Passed: 122/141
 
 # All Passed:
 * babel-plugin-transform-class-static-block
@@ -45,7 +45,7 @@ after transform: SymbolId(0): [ReferenceId(0), ReferenceId(2), ReferenceId(6), R
 rebuilt        : SymbolId(0): [ReferenceId(0), ReferenceId(2), ReferenceId(6), ReferenceId(10)]
 
 
-# babel-plugin-transform-async-to-generator (16/17)
+# babel-plugin-transform-async-to-generator (17/18)
 * super/nested/input.js
 x Output mismatch
 

tasks/transform_conformance/tests/babel-plugin-transform-async-to-generator/test/fixtures/class/this-after-super-nested/output.js

@@ -10,8 +10,7 @@ class Outer {
         var _this2;
 
         if (condition) {
-          const _super = super();
-          _this2 = this;
+          const _super = (super(), _this2 = this);
           this.fn = babelHelpers.asyncToGenerator(function* () {
             return [_this2, 1];
           });

tasks/transform_conformance/tests/babel-plugin-transform-async-to-generator/test/fixtures/class/this-after-super-with-async-arrow-function/input.js

@@ -0,0 +1,6 @@
+class S {}
+class C extends S {
+  constructor() {
+    super(async () => {this})
+  }
+}

tasks/transform_conformance/tests/babel-plugin-transform-async-to-generator/test/fixtures/class/this-after-super-with-async-arrow-function/output.js

@@ -0,0 +1,10 @@
+class S {}
+class C extends S {
+  constructor() {
+    var _this;
+    super(babelHelpers.asyncToGenerator(function* () {
+      _this;
+    }));
+    _this = this;
+  }
+}

tasks/transform_conformance/tests/babel-plugin-transform-async-to-generator/test/fixtures/class/this-after-super/input.js

@@ -1,6 +1,7 @@
+class S {}
 class C extends S {
   constructor() {
-    if (condition) {
+    if (true) {
       const _super = super()
       this.fn = async () => { return [this, 1]; };
     }

tasks/transform_conformance/tests/babel-plugin-transform-async-to-generator/test/fixtures/class/this-after-super/output.js

@@ -1,14 +1,14 @@
+class S {}
 class C extends S {
   constructor() {
     var _this;
-    if (condition) {
-      const _super = super();
-      _this = this;
+    if (true) {
+      const _super = (super(), _this = this);
       this.fn = babelHelpers.asyncToGenerator(function* () {
         return [_this, 1];
       });
     }
-    
+
     super();
     _this = this;
     babelHelpers.asyncToGenerator(function* () {