From 03f0cf2b2802f06247017083ffa2b1f55a2f98ad Mon Sep 17 00:00:00 2001
From: Dunqing <dengqing0821@gmail.com>
Date: Fri, 18 Oct 2024 15:20:30 +0800
Subject: [PATCH] refactor(transformer/typescript): insert assignments after
 super by StatementInjector

---
 .../src/common/statement_injector.rs          |  1 -
 .../src/typescript/annotations.rs             | 49 +++++++++----------
 crates/oxc_transformer/src/typescript/mod.rs  |  4 ++
 3 files changed, 28 insertions(+), 26 deletions(-)

diff --git a/crates/oxc_transformer/src/common/statement_injector.rs b/crates/oxc_transformer/src/common/statement_injector.rs
index 9c90329a1284fc..ed8e96375069f5 100644
--- a/crates/oxc_transformer/src/common/statement_injector.rs
+++ b/crates/oxc_transformer/src/common/statement_injector.rs
@@ -86,7 +86,6 @@ impl<'a> StatementInjectorStore<'a> {
     }
 
     /// Add multiple statements to be inserted immediately after the target statement.
-    #[expect(dead_code)]
     pub fn insert_many_after(&self, target: Address, stmts: Vec<Statement<'a>>) {
         let mut insertions = self.insertions.borrow_mut();
         let adjacent_stmts = insertions.entry(target).or_default();
diff --git a/crates/oxc_transformer/src/typescript/annotations.rs b/crates/oxc_transformer/src/typescript/annotations.rs
index aa5b529dcf8db5..1a2e15d037f3c4 100644
--- a/crates/oxc_transformer/src/typescript/annotations.rs
+++ b/crates/oxc_transformer/src/typescript/annotations.rs
@@ -4,7 +4,7 @@ use std::cell::Cell;
 
 use rustc_hash::FxHashSet;
 
-use oxc_allocator::Vec as ArenaVec;
+use oxc_allocator::{GetAddress, Vec as ArenaVec};
 use oxc_ast::ast::*;
 use oxc_diagnostics::OxcDiagnostic;
 use oxc_semantic::SymbolFlags;
@@ -403,10 +403,32 @@ impl<'a, 'ctx> Traverse<'a> for TypeScriptAnnotations<'a, 'ctx> {
         );
     }
 
+    fn exit_statement(&mut self, stmt: &mut Statement<'a>, ctx: &mut TraverseCtx<'a>) {
+        // Add assignments after super calls
+        if self.assignments.is_empty() {
+            return;
+        }
+
+        let has_super_call = matches!(stmt, Statement::ExpressionStatement(stmt) if stmt.expression.is_super_call_expression());
+        if !has_super_call {
+            return;
+        }
+
+        // Add assignments after super calls
+        self.ctx.statement_injector.insert_many_after(
+            stmt.address(),
+            self.assignments
+                .iter()
+                .map(|assignment| assignment.create_this_property_assignment(ctx))
+                .collect::<Vec<_>>(),
+        );
+        self.has_super_call = true;
+    }
+
     fn exit_statements(
         &mut self,
         stmts: &mut ArenaVec<'a, Statement<'a>>,
-        ctx: &mut TraverseCtx<'a>,
+        _ctx: &mut TraverseCtx<'a>,
     ) {
         // Remove TS specific statements
         stmts.retain(|stmt| match stmt {
@@ -417,29 +439,6 @@ impl<'a, 'ctx> Traverse<'a> for TypeScriptAnnotations<'a, 'ctx> {
             // Ignore ModuleDeclaration as it's handled in the program
             _ => true,
         });
-
-        // Add assignments after super calls
-        if !self.assignments.is_empty() {
-            let has_super_call = stmts.iter().any(|stmt| {
-                matches!(stmt, Statement::ExpressionStatement(stmt) if stmt.expression.is_super_call_expression())
-            });
-            if has_super_call {
-                let mut new_stmts = ctx.ast.vec();
-                for stmt in stmts.drain(..) {
-                    let is_super_call = matches!(stmt, Statement::ExpressionStatement(ref stmt) if stmt.expression.is_super_call_expression());
-                    new_stmts.push(stmt);
-                    if is_super_call {
-                        new_stmts.extend(
-                            self.assignments
-                                .iter()
-                                .map(|assignment| assignment.create_this_property_assignment(ctx)),
-                        );
-                    }
-                }
-                self.has_super_call = true;
-                *stmts = new_stmts;
-            }
-        }
     }
 
     /// Transform if statement's consequent and alternate to block statements if they are super calls
diff --git a/crates/oxc_transformer/src/typescript/mod.rs b/crates/oxc_transformer/src/typescript/mod.rs
index 5bf1711485896c..4b2bf1a562e801 100644
--- a/crates/oxc_transformer/src/typescript/mod.rs
+++ b/crates/oxc_transformer/src/typescript/mod.rs
@@ -205,6 +205,10 @@ impl<'a, 'ctx> Traverse<'a> for TypeScript<'a, 'ctx> {
         self.annotations.exit_statements(stmts, ctx);
     }
 
+    fn exit_statement(&mut self, stmt: &mut Statement<'a>, ctx: &mut TraverseCtx<'a>) {
+        self.annotations.exit_statement(stmt, ctx);
+    }
+
     fn enter_statement(&mut self, stmt: &mut Statement<'a>, ctx: &mut TraverseCtx<'a>) {
         self.r#enum.enter_statement(stmt, ctx);
     }