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

…3634)

Change arrow function transform to use UIDs for `_this` vars, which can't clash with an existing binding.

I believe there are some problems with our current implementation, but have struggled to fix them due to weird interactions with TS transforms. Have added "TODO" comments, so we can return to this and fix them later on, once scopes in TS transforms are correct.

My failed attempt at fixing is on this branch: https://github.com/oxc-project/oxc/tree/transform-arrow-functions

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(