fix(transformer): arrow function transform use UIDs for _this vars

crates/oxc_transformer/src/es2015/arrow_functions.rs

@@ -2,12 +2,12 @@ use std::cell::Cell;
 
 use oxc_allocator::Vec;
 use oxc_ast::ast::*;
-use oxc_span::{Atom, SPAN};
-use oxc_syntax::scope::ScopeFlags;
+use oxc_span::SPAN;
+use oxc_syntax::{scope::ScopeFlags, symbol::SymbolFlags};
 use oxc_traverse::TraverseCtx;
 use serde::Deserialize;
 
-use crate::context::Ctx;
+use crate::{context::Ctx, helpers::bindings::BoundIdentifier};
 
 #[derive(Debug, Default, Clone, Deserialize)]
 pub struct ArrowFunctionsOptions {
@@ -31,11 +31,39 @@ pub struct ArrowFunctionsOptions {
 /// * <https://github.com/babel/babel/tree/main/packages/babel-plugin-transform-arrow-functions>
 //
 // TODO: The `spec` option is not currently supported. Add support for it.
+//
+// TODO: We create `var _this = this;` in parent block, whereas we should create it in
+// parent vars block like Babel:
+// ```js
+// // Input
+// function foo() {
+//   { let f = () => this; }
+//   { let f2 = () => this; }
+// }
+//
+// // Babel output
+// function foo() {
+//   var _this = this;
+//   { let f = function () { return _this; } }
+//   { let f2 = function () { return _this; } }
+// }
+//
+// // Our output
+// function foo() {
+//   {
+//     var _this = this;
+//     let f = function () { return _this; }
+//   }
+//   {
+//     var _this2 = this;
+//     let f2 = function () { return _this2; }
+//   }
+// }
+// ```
 pub struct ArrowFunctions<'a> {
     ctx: Ctx<'a>,
     _options: ArrowFunctionsOptions,
-    uid: usize,
-    has_this: bool,
+    this_var: 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;
@@ -44,30 +72,23 @@ pub struct ArrowFunctions<'a> {
 
 impl<'a> ArrowFunctions<'a> {
     pub fn new(options: ArrowFunctionsOptions, ctx: Ctx<'a>) -> Self {
-        Self {
-            ctx,
-            _options: options,
-            uid: 0,
-            has_this: false,
-            stacks: vec![],
-            this_statements: vec![],
-        }
+        Self { ctx, _options: options, this_var: None, stacks: vec![], this_statements: vec![] }
     }
 
     fn is_inside_arrow_function(&self) -> bool {
         self.stacks.last().copied().unwrap_or(false)
     }
 
-    fn get_this_name(&self) -> Atom<'a> {
-        let uid = if self.uid == 1 { String::new() } else { self.uid.to_string() };
-        self.ctx.ast.new_atom(&format!("_this{uid}"))
-    }
-
-    fn mark_this_as_found(&mut self) {
-        if !self.has_this {
-            self.has_this = true;
-            self.uid += 1;
+    fn get_this_name(&mut self, ctx: &mut TraverseCtx<'a>) -> BoundIdentifier<'a> {
+        if self.this_var.is_none() {
+            self.this_var = Some(BoundIdentifier::new_uid(
+                "this",
+                ctx.current_scope_id(),
+                SymbolFlags::FunctionScopedVariable,
+                ctx,
+            ));
         }
+        self.this_var.as_ref().unwrap().clone()
     }
 
     pub fn transform_statements(&mut self, _stmts: &mut Vec<'a, Statement<'a>>) {
@@ -91,11 +112,9 @@ impl<'a> ArrowFunctions<'a> {
             stmts.insert(0, stmt);
         }
 
-        if self.has_this {
+        if let Some(id) = &self.this_var {
             let binding_pattern = self.ctx.ast.binding_pattern(
-                self.ctx
-                    .ast
-                    .binding_pattern_identifier(BindingIdentifier::new(SPAN, self.get_this_name())),
+                self.ctx.ast.binding_pattern_identifier(id.create_binding_identifier()),
                 None,
                 false,
             );
@@ -118,12 +137,24 @@ impl<'a> ArrowFunctions<'a> {
             let stmt = Statement::VariableDeclaration(stmt);
             // store it, insert it in last statements
             self.this_statements.last_mut().unwrap().replace(stmt);
-            self.has_this = false;
+
+            // 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
-    pub fn transform_jsx_element_name(&mut self, name: &mut JSXElementName<'a>) {
+    pub fn transform_jsx_element_name(
+        &mut self,
+        name: &mut JSXElementName<'a>,
+        ctx: &mut TraverseCtx<'a>,
+    ) {
         if !self.is_inside_arrow_function() {
             return;
         }
@@ -136,8 +167,14 @@ impl<'a> ArrowFunctions<'a> {
             JSXElementName::NamespacedName(_) => return,
         };
         if ident.name == "this" {
-            self.mark_this_as_found();
-            ident.name = self.get_this_name();
+            // We can't produce a proper identifier with a `ReferenceId` because `JSXIdentifier`
+            // lacks that field. https://github.com/oxc-project/oxc/issues/3528
+            // So generate a reference and just use its name.
+            // If JSX transform is enabled, that transform runs before this and will have converted
+            // this to a proper `ThisExpression`, and this visitor won't run.
+            // So only a problem if JSX transform is disabled.
+            let new_ident = self.get_this_name(ctx).create_read_reference(ctx);
+            ident.name = new_ident.name;
         }
     }
 
@@ -215,11 +252,9 @@ impl<'a> ArrowFunctions<'a> {
                     return;
                 }
 
-                self.mark_this_as_found();
-                *expr = self.ctx.ast.identifier_reference_expression(IdentifierReference::new(
-                    this_expr.span,
-                    self.get_this_name(),
-                ));
+                let ident =
+                    self.get_this_name(ctx).create_spanned_read_reference(this_expr.span, ctx);
+                *expr = self.ctx.ast.identifier_reference_expression(ident);
             }
             Expression::ArrowFunctionExpression(arrow_function_expr) => {
                 *expr = self.transform_arrow_function_expression(arrow_function_expr, ctx);

crates/oxc_transformer/src/es2015/mod.rs

@@ -44,9 +44,13 @@ impl<'a> ES2015<'a> {
         }
     }
 
-    pub fn transform_jsx_element_name(&mut self, elem: &mut JSXElementName<'a>) {
+    pub fn transform_jsx_element_name(
+        &mut self,
+        elem: &mut JSXElementName<'a>,
+        ctx: &mut TraverseCtx<'a>,
+    ) {
         if self.options.arrow_function.is_some() {
-            self.arrow_functions.transform_jsx_element_name(elem);
+            self.arrow_functions.transform_jsx_element_name(elem, ctx);
         }
     }
 

crates/oxc_transformer/src/helpers/bindings.rs

@@ -1,9 +1,10 @@
 use std::cell::Cell;
 
 use oxc_ast::ast::{BindingIdentifier, IdentifierReference};
-use oxc_span::{Atom, SPAN};
+use oxc_span::{Atom, Span, SPAN};
 use oxc_syntax::{
     reference::ReferenceFlag,
+    scope::ScopeId,
     symbol::{SymbolFlags, SymbolId},
 };
 use oxc_traverse::TraverseCtx;
@@ -16,22 +17,43 @@ pub struct BoundIdentifier<'a> {
 }
 
 impl<'a> BoundIdentifier<'a> {
-    /// Create `BoundIdentifier` for new binding in root scope
-    pub fn new_root_uid(name: &str, flags: SymbolFlags, ctx: &mut TraverseCtx<'a>) -> Self {
-        let symbol_id = ctx.generate_uid_in_root_scope(name, flags);
+    /// Create `BoundIdentifier` for new binding
+    pub fn new_uid(
+        name: &str,
+        scope_id: ScopeId,
+        flags: SymbolFlags,
+        ctx: &mut TraverseCtx<'a>,
+    ) -> Self {
+        let symbol_id = ctx.generate_uid(name, scope_id, flags);
         let name = ctx.ast.new_atom(&ctx.symbols().names[symbol_id]);
         Self { name, symbol_id }
     }
 
+    /// Create `BoundIdentifier` for new binding in root scope
+    pub fn new_root_uid(name: &str, flags: SymbolFlags, ctx: &mut TraverseCtx<'a>) -> Self {
+        let scope_id = ctx.scopes().root_scope_id();
+        Self::new_uid(name, scope_id, flags, ctx)
+    }
+
     /// Create `IdentifierReference` referencing this binding which is read from
     /// in current scope
     pub fn create_read_reference(&self, ctx: &mut TraverseCtx) -> IdentifierReference<'a> {
+        self.create_spanned_read_reference(SPAN, ctx)
+    }
+
+    /// Create `IdentifierReference` referencing this binding which is read from
+    /// in current scope
+    pub fn create_spanned_read_reference(
+        &self,
+        span: Span,
+        ctx: &mut TraverseCtx,
+    ) -> IdentifierReference<'a> {
         let reference_id = ctx.create_bound_reference(
             self.name.to_compact_str(),
             self.symbol_id,
             ReferenceFlag::Read,
         );
-        IdentifierReference::new_read(SPAN, self.name.clone(), Some(reference_id))
+        IdentifierReference::new_read(span, self.name.clone(), Some(reference_id))
     }
 
     /// Create `BindingIdentifier` for this binding

crates/oxc_transformer/src/lib.rs

@@ -203,12 +203,8 @@ impl<'a> Traverse<'a> for Transformer<'a> {
         self.x1_react.transform_jsx_opening_element(elem, ctx);
     }
 
-    fn enter_jsx_element_name(
-        &mut self,
-        elem: &mut JSXElementName<'a>,
-        _ctx: &mut TraverseCtx<'a>,
-    ) {
-        self.x3_es2015.transform_jsx_element_name(elem);
+    fn enter_jsx_element_name(&mut self, elem: &mut JSXElementName<'a>, ctx: &mut TraverseCtx<'a>) {
+        self.x3_es2015.transform_jsx_element_name(elem, ctx);
     }
 
     fn enter_method_definition(