From 3eaff2a51cf76e1af2c0bc8b660262d11e0ad5fc Mon Sep 17 00:00:00 2001
From: Dunqing <29533304+Dunqing@users.noreply.github.com>
Date: Tue, 31 Dec 2024 01:29:53 +0000
Subject: [PATCH] fix(transformer): ensure last expression statement in arrow
 function expression is wrapped in return (#8192)

The class properties plugin will insert new statements after the expression statement, which causes the expression statement to no longer be at the end. This PR fixed it by reverse looping the statements to find the expression statement and transforming it to wrapped in `return`.

This is not a good way to do it, a potential problem is if there is a plugin that inserts a new expression statement after the original expression statement, which will cause a wrong expression wrapped in a `return`
---
 crates/oxc_transformer/src/lib.rs | 34 +++++++++++++++++++++----------
 1 file changed, 23 insertions(+), 11 deletions(-)

diff --git a/crates/oxc_transformer/src/lib.rs b/crates/oxc_transformer/src/lib.rs
index b0a200c16..5aa997ffc 100644
--- a/crates/oxc_transformer/src/lib.rs
+++ b/crates/oxc_transformer/src/lib.rs
@@ -482,18 +482,30 @@ impl<'a, 'ctx> Traverse<'a> for TransformerImpl<'a, 'ctx> {
         // which can cause issues with the `() => x;` case, as it only allows a single statement.
         // To address this, we wrap the last statement in a return statement and set the expression to false.
         // This transforms the arrow function into the form `() => { return x; };`.
-        if arrow.expression && arrow.body.statements.len() > 1 {
-            let Statement::ExpressionStatement(statement) = arrow.body.statements.pop().unwrap()
-            else {
-                unreachable!(
-                    "The last statement in an ArrowFunctionExpression should always be an ExpressionStatement."
-                )
-            };
-            arrow
-                .body
-                .statements
-                .push(ctx.ast.statement_return(SPAN, Some(statement.unbox().expression)));
+        let statements = &mut arrow.body.statements;
+        if arrow.expression && statements.len() > 1 {
             arrow.expression = false;
+
+            // Reverse looping to find the expression statement, because other plugins could
+            // insert new statements after the expression statement.
+            // `() => x;`
+            // ->
+            // ```
+            // () => {
+            //    var new_insert_variable;
+            //    return x;
+            //    function new_insert_function() {}
+            // };
+            // ```
+            for stmt in statements.iter_mut().rev() {
+                let Statement::ExpressionStatement(expr_stmt) = stmt else {
+                    continue;
+                };
+                let expression = Some(ctx.ast.move_expression(&mut expr_stmt.expression));
+                *stmt = ctx.ast.statement_return(SPAN, expression);
+                return;
+            }
+            unreachable!("At least one statement should be expression statement")
         }
     }