feat(transformer/arrow-functions): the output that uses this inside blocks doesn't match Babel (#5188)

Fixes 666282a13b/crates/oxc_transformer/src/es2015/arrow_functions.rs (L35-L62)

crates/oxc_transformer/src/es2015/arrow_functions.rs

@@ -83,26 +83,33 @@ pub struct ArrowFunctionsOptions {
 pub struct ArrowFunctions<'a> {
     ctx: Ctx<'a>,
     _options: ArrowFunctionsOptions,
-    this_var: Option<BoundIdentifier<'a>>,
+    this_vars: std::vec::Vec<Option<BoundIdentifier<'a>>>,
     /// Stack to keep track of whether we are inside an arrow function or not.
     stacks: std::vec::Vec<bool>,
-    // var _this = this;
-    this_statements: std::vec::Vec<Option<Statement<'a>>>,
 }
 
 impl<'a> ArrowFunctions<'a> {
     pub fn new(options: ArrowFunctionsOptions, ctx: Ctx<'a>) -> Self {
-        Self { ctx, _options: options, this_var: None, stacks: vec![], this_statements: vec![] }
+        Self {
+            ctx,
+            _options: options,
+            // Reserve for the global scope
+            this_vars: vec![None],
+            stacks: vec![],
+        }
     }
 }
 
 impl<'a> Traverse<'a> for ArrowFunctions<'a> {
-    fn enter_statements(
+    /// Insert `var _this = this;` for the global scope.
-        &mut self,
+    fn exit_program(&mut self, program: &mut Program<'a>, _ctx: &mut TraverseCtx<'a>) {
-        _stmts: &mut Vec<'a, Statement<'a>>,
+        self.insert_this_var_statement_at_the_top_of_statements(&mut program.body);
-        _ctx: &mut TraverseCtx<'a>,
+    }
-    ) {
+
-        self.this_statements.push(None);
+    fn enter_function(&mut self, func: &mut Function<'a>, _ctx: &mut TraverseCtx<'a>) {
+        if func.body.is_some() {
+            self.this_vars.push(None);
+        }
     }
 
     /// ```ts
@@ -116,49 +123,12 @@ impl<'a> Traverse<'a> for ArrowFunctions<'a> {
     /// }
     /// ```
     /// Insert the var _this = this; statement outside the arrow function
-    fn exit_statements(&mut self, stmts: &mut Vec<'a, Statement<'a>>, _ctx: &mut TraverseCtx<'a>) {
+    fn exit_function(&mut self, func: &mut Function<'a>, _ctx: &mut TraverseCtx<'a>) {
-        // Insert the var _this = this;
+        let Some(body) = func.body.as_mut() else {
-        if let Some(Some(stmt)) = self.this_statements.pop() {
+            return;
-            stmts.insert(0, stmt);
+        };
-        }
 
-        if let Some(id) = &self.this_var {
+        self.insert_this_var_statement_at_the_top_of_statements(&mut body.statements);
-            let binding_pattern = self.ctx.ast.binding_pattern(
-                self.ctx
-                    .ast
-                    .binding_pattern_kind_from_binding_identifier(id.create_binding_identifier()),
-                Option::<TSTypeAnnotation>::None,
-                false,
-            );
-
-            let variable_declarator = self.ctx.ast.variable_declarator(
-                SPAN,
-                VariableDeclarationKind::Var,
-                binding_pattern,
-                Some(self.ctx.ast.expression_this(SPAN)),
-                false,
-            );
-
-            let stmt = self.ctx.ast.alloc_variable_declaration(
-                SPAN,
-                VariableDeclarationKind::Var,
-                self.ctx.ast.vec1(variable_declarator),
-                false,
-            );
-
-            let stmt = Statement::VariableDeclaration(stmt);
-            // store it, insert it in last statements
-            self.this_statements.last_mut().unwrap().replace(stmt);
-
-            // TODO: This isn't quite right. In this case, output is invalid:
-            // ```js
-            // function foo() {
-            //   let f = () => this;
-            //   let f2 = () => this;
-            // }
-            // ```
-            self.this_var = None;
-        }
     }
 
     /// Change <this></this> to <_this></_this>, and mark it as found
@@ -245,14 +215,15 @@ impl<'a> ArrowFunctions<'a> {
     }
 
     fn get_this_name(&mut self, ctx: &mut TraverseCtx<'a>) -> BoundIdentifier<'a> {
-        if self.this_var.is_none() {
+        let this_var = self.this_vars.last_mut().unwrap();
-            self.this_var = Some(BoundIdentifier::new_uid_in_current_scope(
+        if this_var.is_none() {
+            this_var.replace(BoundIdentifier::new_uid_in_current_scope(
                 "this",
                 SymbolFlags::FunctionScopedVariable,
                 ctx,
             ));
         }
-        self.this_var.as_ref().unwrap().clone()
+        this_var.as_ref().unwrap().clone()
     }
 
     fn transform_arrow_function_expression(
@@ -315,4 +286,39 @@ impl<'a> ArrowFunctions<'a> {
         // `() => {};` => `(function () {});`
         self.ctx.ast.expression_parenthesized(SPAN, expr)
     }
+
+    /// Insert `var _this = this;` at the top of the statements.
+    fn insert_this_var_statement_at_the_top_of_statements(
+        &mut self,
+        statements: &mut Vec<'a, Statement<'a>>,
+    ) {
+        if let Some(id) = &self.this_vars.pop().unwrap() {
+            let binding_pattern = self.ctx.ast.binding_pattern(
+                self.ctx
+                    .ast
+                    .binding_pattern_kind_from_binding_identifier(id.create_binding_identifier()),
+                Option::<TSTypeAnnotation>::None,
+                false,
+            );
+
+            let variable_declarator = self.ctx.ast.variable_declarator(
+                SPAN,
+                VariableDeclarationKind::Var,
+                binding_pattern,
+                Some(self.ctx.ast.expression_this(SPAN)),
+                false,
+            );
+
+            let stmt = self.ctx.ast.alloc_variable_declaration(
+                SPAN,
+                VariableDeclarationKind::Var,
+                self.ctx.ast.vec1(variable_declarator),
+                false,
+            );
+
+            let stmt = Statement::VariableDeclaration(stmt);
+
+            statements.insert(0, stmt);
+        }
+    }
 }

crates/oxc_transformer/src/es2015/mod.rs

@@ -5,7 +5,6 @@ use std::rc::Rc;
 
 pub use arrow_functions::{ArrowFunctions, ArrowFunctionsOptions};
 pub use options::ES2015Options;
-use oxc_allocator::Vec;
 use oxc_ast::ast::*;
 use oxc_traverse::{Traverse, TraverseCtx};
 
@@ -34,15 +33,21 @@ impl<'a> ES2015<'a> {
 }
 
 impl<'a> Traverse<'a> for ES2015<'a> {
-    fn enter_statements(&mut self, stmts: &mut Vec<'a, Statement<'a>>, ctx: &mut TraverseCtx<'a>) {
+    fn exit_program(&mut self, program: &mut Program<'a>, ctx: &mut TraverseCtx<'a>) {
         if self.options.arrow_function.is_some() {
-            self.arrow_functions.enter_statements(stmts, ctx);
+            self.arrow_functions.exit_program(program, ctx);
         }
     }
 
-    fn exit_statements(&mut self, stmts: &mut Vec<'a, Statement<'a>>, ctx: &mut TraverseCtx<'a>) {
+    fn enter_function(&mut self, func: &mut Function<'a>, ctx: &mut TraverseCtx<'a>) {
         if self.options.arrow_function.is_some() {
-            self.arrow_functions.exit_statements(stmts, ctx);
+            self.arrow_functions.enter_function(func, ctx);
+        }
+    }
+
+    fn exit_function(&mut self, func: &mut Function<'a>, ctx: &mut TraverseCtx<'a>) {
+        if self.options.arrow_function.is_some() {
+            self.arrow_functions.exit_function(func, ctx);
         }
     }
 

crates/oxc_transformer/src/lib.rs

@@ -123,6 +123,7 @@ impl<'a> Traverse<'a> for Transformer<'a> {
     fn exit_program(&mut self, program: &mut Program<'a>, ctx: &mut TraverseCtx<'a>) {
         self.x1_react.transform_program_on_exit(program, ctx);
         self.x0_typescript.transform_program_on_exit(program, ctx);
+        self.x3_es2015.exit_program(program, ctx);
     }
 
     // ALPHASORT
@@ -202,8 +203,14 @@ impl<'a> Traverse<'a> for Transformer<'a> {
         self.x0_typescript.transform_formal_parameter(param);
     }
 
-    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>) {
         self.x0_typescript.transform_function(func);
+        self.x3_es2015.enter_function(func, ctx);
+    }
+
+    fn exit_function(&mut self, func: &mut Function<'a>, ctx: &mut TraverseCtx<'a>) {
+        self.x0_typescript.transform_function(func);
+        self.x3_es2015.exit_function(func, ctx);
     }
 
     fn enter_jsx_element(&mut self, node: &mut JSXElement<'a>, _ctx: &mut TraverseCtx<'a>) {
@@ -261,7 +268,6 @@ impl<'a> Traverse<'a> for Transformer<'a> {
         self.x2_es2021.enter_statements(stmts, ctx);
         self.x2_es2020.enter_statements(stmts, ctx);
         self.x2_es2016.enter_statements(stmts, ctx);
-        self.x3_es2015.enter_statements(stmts, ctx);
     }
 
     fn exit_statements(&mut self, stmts: &mut Vec<'a, Statement<'a>>, ctx: &mut TraverseCtx<'a>) {
@@ -270,7 +276,6 @@ impl<'a> Traverse<'a> for Transformer<'a> {
         self.x2_es2021.exit_statements(stmts, ctx);
         self.x2_es2020.exit_statements(stmts, ctx);
         self.x2_es2016.exit_statements(stmts, ctx);
-        self.x3_es2015.exit_statements(stmts, ctx);
     }
 
     fn enter_tagged_template_expression(

tasks/transform_conformance/oxc.snap.md

@@ -1,11 +1,37 @@
 commit: 12619ffe
 
-Passed: 9/35
+Passed: 9/36
 
 # All Passed:
 * babel-plugin-transform-optional-catch-binding
 
 
+# babel-plugin-transform-arrow-functions (0/1)
+* use-this-inside-blocks/input.js
+  x Bindings mismatch:
+  | after transform: ScopeId(1): []
+  | rebuilt        : ScopeId(1): ["_this"]
+
+  x Bindings mismatch:
+  | after transform: ScopeId(3): ["_this"]
+  | rebuilt        : ScopeId(3): []
+
+  x Symbol scope ID mismatch:
+  | after transform: SymbolId(3): ScopeId(3)
+  | rebuilt        : SymbolId(1): ScopeId(1)
+
+  x Symbol flags mismatch:
+  | after transform: SymbolId(1): SymbolFlags(BlockScopedVariable |
+  | ArrowFunction)
+  | rebuilt        : SymbolId(2): SymbolFlags(BlockScopedVariable)
+
+  x Symbol flags mismatch:
+  | after transform: SymbolId(2): SymbolFlags(BlockScopedVariable |
+  | ArrowFunction)
+  | rebuilt        : SymbolId(3): SymbolFlags(BlockScopedVariable)
+
+
+
 # babel-plugin-transform-typescript (2/7)
 * computed-constant-value/input.ts
   x Semantic Collector failed after transform

tasks/transform_conformance/tests/babel-plugin-transform-arrow-functions/test/fixtures/options.json

@@ -0,0 +1,3 @@
+{
+  "plugins": [["transform-arrow-functions"]]
+}

tasks/transform_conformance/tests/babel-plugin-transform-arrow-functions/test/fixtures/use-this-inside-blocks/input.js

@@ -0,0 +1,4 @@
+function foo() {
+  { let f = () => this; }
+  { let f2 = () => this; }
+}

tasks/transform_conformance/tests/babel-plugin-transform-arrow-functions/test/fixtures/use-this-inside-blocks/output.js

@@ -0,0 +1,13 @@
+function foo() {
+	var _this = this;
+	{
+					let f = function() {
+									return _this;
+					};
+	}
+	{
+					let f2 = function() {
+									return _this;
+					};
+	}
+}