Feat: Insert field initializers after super()

crates/oxc_transformer/src/es2022/class_properties.rs

@@ -78,14 +78,14 @@ use indexmap::IndexMap;
 use rustc_hash::FxHasher;
 use serde::Deserialize;
 
-use oxc_allocator::{Address, Box as ArenaBox, GetAddress};
-use oxc_ast::{ast::*, NONE};
+use oxc_allocator::{Address, Box as ArenaBox, GetAddress, Vec as ArenaVec};
+use oxc_ast::{ast::*, visit::walk_mut::walk_call_expression, AstBuilder, VisitMut, NONE};
 use oxc_data_structures::stack::{NonEmptyStack, SparseStack};
 use oxc_span::SPAN;
 use oxc_syntax::{
     node::NodeId,
     reference::{ReferenceFlags, ReferenceId},
-    scope::ScopeFlags,
+    scope::{ScopeFlags, ScopeId},
     symbol::{SymbolFlags, SymbolId},
 };
 use oxc_traverse::{
@@ -490,6 +490,13 @@ impl<'a, 'ctx> ClassProperties<'a, 'ctx> {
         // TODO(improve-on-babel): If outer scope is sloppy mode, all code which is moved to outside
         // the class should be wrapped in an IIFE with `'use strict'` directive. Babel doesn't do this.
 
+        struct Constructor {
+            element_index: usize,
+            // TODO: Remove this if don't use it
+            #[expect(dead_code)]
+            scope_id: ScopeId,
+        }
+
         // Check if class has any properties and get index and `ScopeId` of constructor (if class has one)
         let mut instance_prop_count = 0;
         let mut has_static_prop_or_static_block = false;
@@ -538,7 +545,10 @@ impl<'a, 'ctx> ClassProperties<'a, 'ctx> {
                     if method.kind == MethodDefinitionKind::Constructor
                         && method.value.body.is_some()
                     {
-                        constructor = Some((index_not_including_removed, method.value.scope_id()));
+                        constructor = Some(Constructor {
+                            element_index: index_not_including_removed,
+                            scope_id: method.value.scope_id(),
+                        });
                     }
                 }
                 ClassElement::AccessorProperty(_) | ClassElement::TSIndexSignature(_) => {
@@ -601,14 +611,9 @@ impl<'a, 'ctx> ClassProperties<'a, 'ctx> {
         // Insert instance initializers into constructor
         if !instance_inits.is_empty() {
             // TODO: Re-parent any scopes within initializers.
-            if let Some((constructor_index, _)) = constructor {
+            if let Some(Constructor { element_index, .. }) = constructor {
                 // Existing constructor - amend it
-                Self::insert_inits_into_constructor(
-                    &mut class.body,
-                    instance_inits,
-                    constructor_index,
-                    ctx,
-                );
+                Self::insert_inits_into_constructor(class, instance_inits, element_index, ctx);
             } else {
                 // No constructor - create one
                 Self::insert_constructor(class, instance_inits, ctx);
@@ -969,20 +974,33 @@ impl<'a, 'ctx> ClassProperties<'a, 'ctx> {
     }
 
     fn insert_inits_into_constructor(
-        body: &mut ClassBody<'a>,
+        class: &mut Class<'a>,
         inits: Vec<Expression<'a>>,
         constructor_index: usize,
         ctx: &mut TraverseCtx<'a>,
     ) {
-        // TODO: Insert after `super()` if class has super-class.
-        // TODO: Insert as expression sequence if `super()` is used in an expression.
         // TODO: Handle where vars used in property init clash with vars in top scope of constructor.
-        let element = body.body.get_mut(constructor_index).unwrap();
-        let ClassElement::MethodDefinition(method) = element else { unreachable!() };
-        let func_body = method.value.body.as_mut().unwrap();
-        func_body
-            .statements
-            .splice(0..0, inits.into_iter().map(|expr| ctx.ast.statement_expression(SPAN, expr)));
+        // (or maybe do that earlier?)
+        // TODO: Handle private props in constructor params `class C { #x; constructor(x = this.#x) {} }`.
+        let constructor = match class.body.body.get_mut(constructor_index).unwrap() {
+            ClassElement::MethodDefinition(constructor) => constructor.as_mut(),
+            _ => unreachable!(),
+        };
+        debug_assert!(constructor.kind == MethodDefinitionKind::Constructor);
+
+        let constructor_scope_id = constructor.value.scope_id();
+        let func_body = constructor.value.body.as_mut().unwrap();
+
+        if class.super_class.is_some() {
+            // Class has super class. Insert after `super()`.
+            ConstructorInitsInserter::insert(func_body, inits, constructor_scope_id, ctx);
+        } else {
+            // No super class. Insert at top of constructor.
+            func_body.statements.splice(
+                0..0,
+                inits.into_iter().map(|expr| ctx.ast.statement_expression(SPAN, expr)),
+            );
+        }
     }
 
     fn insert_constructor(
@@ -2078,3 +2096,178 @@ impl<'a, 'ctx> ClassProperties<'a, 'ctx> {
         ctx.ast.identifier_name(SPAN, Atom::from("_"))
     }
 }
+
+/// TODO: Doc comment
+struct ConstructorInitsInserter<'a, 'c> {
+    constructor_scope_id: ScopeId,
+    super_binding: Option<BoundIdentifier<'a>>,
+    ctx: &'c mut TraverseCtx<'a>,
+}
+
+impl<'a, 'c> ConstructorInitsInserter<'a, 'c> {
+    fn insert(
+        body: &mut FunctionBody<'a>,
+        inits: Vec<Expression<'a>>,
+        constructor_scope_id: ScopeId,
+        ctx: &'c mut TraverseCtx<'a>,
+    ) {
+        let mut inserter = Self::new(constructor_scope_id, ctx);
+        inserter.insert_inits(body, inits);
+    }
+
+    fn new(constructor_scope_id: ScopeId, ctx: &'c mut TraverseCtx<'a>) -> Self {
+        Self { constructor_scope_id, super_binding: None, ctx }
+    }
+
+    fn insert_inits(&mut self, body: &mut FunctionBody<'a>, inits: Vec<Expression<'a>>) {
+        // TODO: Re-parent child scopes of `init`s
+        let stmts = &mut body.statements;
+        let mut stmts_iter = stmts.iter_mut();
+        let mut index = 0;
+        loop {
+            // Constructor must contain at least one `super()`, so loop must exit before get to end of `stmts`
+            let stmt = stmts_iter.next().unwrap();
+            if let Statement::ExpressionStatement(expr_stmt) = &*stmt {
+                if let Expression::CallExpression(call_expr) = &expr_stmt.expression {
+                    if let Expression::Super(_) = &call_expr.callee {
+                        // `super()` as top level statement
+                        stmts.splice(index..index, Self::exprs_into_stmts(inits, self.ctx.ast));
+                        return;
+                    }
+                }
+            }
+
+            self.visit_statement(stmt);
+            if self.super_binding.is_some() {
+                break;
+            }
+
+            index += 1;
+        }
+
+        // `super()` found not in top level position. Convert all other `super()`s to `_super()`.
+        for stmt in stmts_iter {
+            self.visit_statement(stmt);
+        }
+
+        // Insert `_super` function at top of constructor
+        self.insert_super_func(stmts, inits);
+    }
+
+    /// Insert `_super` function at top of constructor.
+    /// ```js
+    /// var _super = (..._args) => {
+    ///   super(..._args);
+    ///   <inits>
+    ///   return this;
+    /// };
+    /// ```
+    fn insert_super_func(
+        &mut self,
+        stmts: &mut ArenaVec<'a, Statement<'a>>,
+        inits: Vec<Expression<'a>>,
+    ) {
+        let ctx = &mut self.ctx;
+
+        let super_func_scope = ctx.scopes_mut().add_scope(
+            Some(self.constructor_scope_id),
+            NodeId::DUMMY,
+            ScopeFlags::Function | ScopeFlags::Arrow | ScopeFlags::StrictMode,
+        );
+        let args_binding =
+            ctx.generate_uid("args", super_func_scope, SymbolFlags::FunctionScopedVariable);
+        // `super(..._args)`
+        let super_call = ctx.ast.statement_expression(
+            SPAN,
+            ctx.ast.expression_call(
+                SPAN,
+                ctx.ast.expression_super(SPAN),
+                NONE,
+                ctx.ast.vec1(
+                    ctx.ast.argument_spread_element(SPAN, args_binding.create_read_expression(ctx)),
+                ),
+                false,
+            ),
+        );
+        // `return this;`
+        let return_stmt = ctx.ast.statement_return(SPAN, Some(ctx.ast.expression_this(SPAN)));
+        // `super(..._args); <inits>; return this;`
+        let body_stmts = ctx.ast.vec_from_iter(
+            [super_call]
+                .into_iter()
+                .chain(Self::exprs_into_stmts(inits, ctx.ast))
+                .chain([return_stmt]),
+        );
+        // `(...args) => { super(..._args); <inits>; return this; }`
+        let super_func = Expression::ArrowFunctionExpression(
+            ctx.ast.alloc_arrow_function_expression_with_scope_id(
+                SPAN,
+                false,
+                false,
+                NONE,
+                ctx.ast.alloc_formal_parameters(
+                    SPAN,
+                    FormalParameterKind::ArrowFormalParameters,
+                    ctx.ast.vec(),
+                    Some(ctx.ast.alloc_binding_rest_element(
+                        SPAN,
+                        args_binding.create_binding_pattern(ctx),
+                    )),
+                ),
+                NONE,
+                ctx.ast.alloc_function_body(SPAN, ctx.ast.vec(), body_stmts),
+                super_func_scope,
+            ),
+        );
+        // `var _super = (...args) => { ... }`
+        let super_binding = self.super_binding.as_ref().unwrap().clone();
+        let super_func_decl = Statement::from(ctx.ast.declaration_variable(
+            SPAN,
+            VariableDeclarationKind::Var,
+            ctx.ast.vec1(ctx.ast.variable_declarator(
+                SPAN,
+                VariableDeclarationKind::Var,
+                super_binding.create_binding_pattern(ctx),
+                Some(super_func),
+                false,
+            )),
+            false,
+        ));
+
+        stmts.insert(0, super_func_decl);
+    }
+
+    fn exprs_into_stmts(
+        exprs: Vec<Expression<'a>>,
+        ast: AstBuilder<'a>,
+    ) -> impl Iterator<Item = Statement<'a>> {
+        exprs.into_iter().map(move |expr| ast.statement_expression(SPAN, expr))
+    }
+}
+
+impl<'a, 'c> VisitMut<'a> for ConstructorInitsInserter<'a, 'c> {
+    /// Replace `super()` with `_super()`.
+    // `#[inline]` to make hot path for all other function calls as cheap as possible.
+    #[inline]
+    fn visit_call_expression(&mut self, call_expr: &mut CallExpression<'a>) {
+        if let Expression::Super(super_) = &call_expr.callee {
+            let span = super_.span;
+            self.replace_super(call_expr, span);
+        }
+
+        walk_call_expression(self, call_expr);
+    }
+}
+
+impl<'a, 'c> ConstructorInitsInserter<'a, 'c> {
+    fn replace_super(&mut self, call_expr: &mut CallExpression<'a>, span: Span) {
+        let super_binding = self.super_binding.get_or_insert_with(|| {
+            self.ctx.generate_uid(
+                "super",
+                self.constructor_scope_id,
+                SymbolFlags::FunctionScopedVariable,
+            )
+        });
+        call_expr.callee = super_binding.create_spanned_read_expression(span, &mut self.ctx);
+    }
+}