fix(transformer): fix spans and scopes in TS enum transform

overlookmotel · overlookmotel · commit d512cf4cbfb4 · 2024-06-26T02:35:48.000+01:00
diff --git a/crates/oxc_transformer/src/helpers/bindings.rs b/crates/oxc_transformer/src/helpers/bindings.rs
@@ -37,7 +37,7 @@ impl<'a> BoundIdentifier<'a> {
 
     /// Create `IdentifierReference` referencing this binding which is read from
     /// in current scope
-    pub fn create_read_reference(&self, ctx: &mut TraverseCtx) -> IdentifierReference<'a> {
+    pub fn create_read_reference(&self, ctx: &mut TraverseCtx<'a>) -> IdentifierReference<'a> {
         self.create_spanned_read_reference(SPAN, ctx)
     }
 
@@ -46,14 +46,9 @@ impl<'a> BoundIdentifier<'a> {
     pub fn create_spanned_read_reference(
         &self,
         span: Span,
-        ctx: &mut TraverseCtx,
+        ctx: &mut TraverseCtx<'a>,
     ) -> 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))
+        ctx.create_bound_reference_id(span, self.name.clone(), self.symbol_id, ReferenceFlag::Read)
     }
 
     /// Create `BindingIdentifier` for this binding
diff --git a/crates/oxc_transformer/src/typescript/enum.rs b/crates/oxc_transformer/src/typescript/enum.rs
@@ -1,9 +1,13 @@
+use std::cell::Cell;
+
 use oxc_allocator::Vec;
 use oxc_ast::{ast::*, visit::walk_mut, VisitMut};
-use oxc_span::{Atom, SPAN};
+use oxc_span::{Atom, Span, SPAN};
 use oxc_syntax::{
     number::{NumberBase, ToJsInt32, ToJsString},
     operator::{AssignmentOperator, BinaryOperator, LogicalOperator, UnaryOperator},
+    reference::ReferenceFlag,
+    symbol::SymbolFlags,
 };
 use oxc_traverse::TraverseCtx;
 use rustc_hash::FxHashMap;
@@ -20,14 +24,14 @@ impl<'a> TypeScriptEnum<'a> {
         Self { ctx, enums: FxHashMap::default() }
     }
 
-    pub fn transform_statement(&mut self, stmt: &mut Statement<'a>, ctx: &TraverseCtx<'a>) {
+    pub fn transform_statement(&mut self, stmt: &mut Statement<'a>, ctx: &mut TraverseCtx<'a>) {
         let new_stmt = match stmt {
             Statement::TSEnumDeclaration(ts_enum_decl) => {
-                self.transform_ts_enum(ts_enum_decl, false, ctx)
+                self.transform_ts_enum(ts_enum_decl, None, ctx)
             }
             Statement::ExportNamedDeclaration(decl) => {
                 if let Some(Declaration::TSEnumDeclaration(ts_enum_decl)) = &decl.declaration {
-                    self.transform_ts_enum(ts_enum_decl, true, ctx)
+                    self.transform_ts_enum(ts_enum_decl, Some(decl.span), ctx)
                 } else {
                     None
                 }
@@ -56,16 +60,29 @@ impl<'a> TypeScriptEnum<'a> {
     fn transform_ts_enum(
         &mut self,
         decl: &TSEnumDeclaration<'a>,
-        is_export: bool,
-        ctx: &TraverseCtx<'a>,
+        export_span: Option<Span>,
+        ctx: &mut TraverseCtx<'a>,
     ) -> Option<Statement<'a>> {
         if decl.declare {
             return None;
         }
 
+        let is_export = export_span.is_some();
         let is_not_top_scope = !ctx.scopes().get_flags(ctx.current_scope_id()).is_top();
-        let span = decl.span;
-        let ident = decl.id.clone();
+
+        let enum_name = decl.id.name.clone();
+        let func_scope_id = decl.scope_id.get().unwrap();
+        let param_symbol_id = ctx.symbols_mut().create_symbol(
+            decl.id.span,
+            enum_name.to_compact_str(),
+            SymbolFlags::FunctionScopedVariable,
+            func_scope_id,
+        );
+        let ident = BindingIdentifier {
+            span: decl.id.span,
+            name: decl.id.name.clone(),
+            symbol_id: Cell::new(Some(param_symbol_id)),
+        };
         let kind = self.ctx.ast.binding_pattern_identifier(ident);
         let id = self.ctx.ast.binding_pattern(kind, None, false);
 
@@ -81,25 +98,39 @@ impl<'a> TypeScriptEnum<'a> {
         );
 
         // Foo[Foo["X"] = 0] = "X";
-        let enum_name = decl.id.name.clone();
         let is_already_declared = self.enums.contains_key(&enum_name);
         let statements = self.transform_ts_enum_members(&decl.members, enum_name.clone(), ctx);
         let body = self.ctx.ast.function_body(decl.span, self.ctx.ast.new_vec(), statements);
-        let r#type = FunctionType::FunctionExpression;
-        let callee = self.ctx.ast.plain_function(r#type, SPAN, None, params, Some(body));
-        let callee = Expression::FunctionExpression(callee);
+        let callee = Expression::FunctionExpression(ctx.alloc(Function {
+            r#type: FunctionType::FunctionExpression,
+            span: SPAN,
+            id: None,
+            generator: false,
+            r#async: false,
+            declare: false,
+            this_param: None,
+            params,
+            body: Some(body),
+            type_parameters: None,
+            return_type: None,
+            scope_id: Cell::new(Some(func_scope_id)),
+        }));
 
+        let var_symbol_id = decl.id.symbol_id.get().unwrap();
         let arguments = if (is_export || is_not_top_scope) && !is_already_declared {
             // }({});
             let object_expr = self.ctx.ast.object_expression(SPAN, self.ctx.ast.new_vec(), None);
             self.ctx.ast.new_vec_single(Argument::from(object_expr))
         } else {
             // }(Foo || {});
             let op = LogicalOperator::Or;
-            let left = self
-                .ctx
-                .ast
-                .identifier_reference_expression(IdentifierReference::new(SPAN, enum_name.clone()));
+            let left = ctx.create_bound_reference_id(
+                decl.id.span,
+                enum_name.clone(),
+                var_symbol_id,
+                ReferenceFlag::Read,
+            );
+            let left = ctx.ast.identifier_reference_expression(left);
             let right = self.ctx.ast.object_expression(SPAN, self.ctx.ast.new_vec(), None);
             let expression = self.ctx.ast.logical_expression(SPAN, left, op, right);
             self.ctx.ast.new_vec_single(Argument::from(expression))
@@ -109,12 +140,15 @@ impl<'a> TypeScriptEnum<'a> {
 
         if is_already_declared {
             let op = AssignmentOperator::Assign;
-            let left = self.ctx.ast.simple_assignment_target_identifier(IdentifierReference::new(
-                SPAN,
+            let left = ctx.create_bound_reference_id(
+                decl.id.span,
                 enum_name.clone(),
-            ));
+                var_symbol_id,
+                ReferenceFlag::Write,
+            );
+            let left = ctx.ast.simple_assignment_target_identifier(left);
             let expr = self.ctx.ast.assignment_expression(SPAN, op, left, call_expression);
-            return Some(self.ctx.ast.expression_statement(SPAN, expr));
+            return Some(self.ctx.ast.expression_statement(decl.span, expr));
         }
 
         let kind = if is_export || is_not_top_scope {
@@ -123,25 +157,21 @@ impl<'a> TypeScriptEnum<'a> {
             VariableDeclarationKind::Var
         };
         let decls = {
-            let mut decls = self.ctx.ast.new_vec();
-
-            let binding_identifier = BindingIdentifier::new(SPAN, enum_name.clone());
+            let binding_identifier = decl.id.clone();
             let binding_pattern_kind = self.ctx.ast.binding_pattern_identifier(binding_identifier);
             let binding = self.ctx.ast.binding_pattern(binding_pattern_kind, None, false);
             let decl =
                 self.ctx.ast.variable_declarator(SPAN, kind, binding, Some(call_expression), false);
-
-            decls.push(decl);
-            decls
+            ctx.ast.new_vec_single(decl)
         };
-        let variable_declaration = self.ctx.ast.variable_declaration(span, kind, decls, false);
+        let variable_declaration = self.ctx.ast.variable_declaration(decl.span, kind, decls, false);
         let variable_declaration = Declaration::VariableDeclaration(variable_declaration);
 
-        let stmt = if is_export {
-            let declaration =
-                self.ctx.ast.plain_export_named_declaration_declaration(SPAN, variable_declaration);
-
-            self.ctx.ast.module_declaration(ModuleDeclaration::ExportNamedDeclaration(declaration))
+        let stmt = if let Some(export_span) = export_span {
+            let declaration = ctx
+                .ast
+                .plain_export_named_declaration_declaration(export_span, variable_declaration);
+            Statement::ExportNamedDeclaration(declaration)
         } else {
             Statement::from(variable_declaration)
         };
@@ -154,6 +184,8 @@ impl<'a> TypeScriptEnum<'a> {
         enum_name: Atom<'a>,
         ctx: &TraverseCtx<'a>,
     ) -> Vec<'a, Statement<'a>> {
+        // TODO: Set `span` and `references_id` on all `IdentifierReference`s created here
+
         let mut statements = self.ctx.ast.new_vec();
         let mut prev_constant_value = Some(ConstantValue::Number(-1.0));
         let mut previous_enum_members = self.enums.entry(enum_name.clone()).or_default().clone();
diff --git a/crates/oxc_transformer/src/typescript/mod.rs b/crates/oxc_transformer/src/typescript/mod.rs
@@ -167,7 +167,7 @@ impl<'a> TypeScript<'a> {
         self.annotations.transform_statements_on_exit(stmts, ctx);
     }
 
-    pub fn transform_statement(&mut self, stmt: &mut Statement<'a>, ctx: &TraverseCtx<'a>) {
+    pub fn transform_statement(&mut self, stmt: &mut Statement<'a>, ctx: &mut TraverseCtx<'a>) {
         self.r#enum.transform_statement(stmt, ctx);
     }
 
diff --git a/crates/oxc_traverse/src/context/mod.rs b/crates/oxc_traverse/src/context/mod.rs
@@ -1,10 +1,10 @@
 use oxc_allocator::{Allocator, Box};
 use oxc_ast::{
-    ast::{Expression, Statement},
+    ast::{Expression, IdentifierReference, Statement},
     AstBuilder,
 };
 use oxc_semantic::{ScopeTree, SymbolTable};
-use oxc_span::CompactStr;
+use oxc_span::{Atom, CompactStr, Span};
 use oxc_syntax::{
     reference::{ReferenceFlag, ReferenceId},
     scope::{ScopeFlags, ScopeId},
@@ -358,6 +358,19 @@ impl<'a> TraverseCtx<'a> {
         self.scoping.create_bound_reference(name, symbol_id, flag)
     }
 
+    /// Create an `IdentifierReference` bound to a `SymbolId`.
+    ///
+    /// This is a shortcut for `ctx.scoping.create_bound_reference_id`.
+    pub fn create_bound_reference_id(
+        &mut self,
+        span: Span,
+        name: Atom<'a>,
+        symbol_id: SymbolId,
+        flag: ReferenceFlag,
+    ) -> IdentifierReference<'a> {
+        self.scoping.create_bound_reference_id(span, name, symbol_id, flag)
+    }
+
     /// Create an unbound reference.
     ///
     /// This is a shortcut for `ctx.scoping.create_unbound_reference`.
@@ -369,6 +382,18 @@ impl<'a> TraverseCtx<'a> {
         self.scoping.create_unbound_reference(name, flag)
     }
 
+    /// Create an unbound `IdentifierReference`.
+    ///
+    /// This is a shortcut for `ctx.scoping.create_unbound_reference_id`.
+    pub fn create_unbound_reference_id(
+        &mut self,
+        span: Span,
+        name: Atom<'a>,
+        flag: ReferenceFlag,
+    ) -> IdentifierReference<'a> {
+        self.scoping.create_unbound_reference_id(span, name, flag)
+    }
+
     /// Create a reference optionally bound to a `SymbolId`.
     ///
     /// If you know if there's a `SymbolId` or not, prefer `TraverseCtx::create_bound_reference`
@@ -384,6 +409,22 @@ impl<'a> TraverseCtx<'a> {
         self.scoping.create_reference(name, symbol_id, flag)
     }
 
+    /// Create an `IdentifierReference` optionally bound to a `SymbolId`.
+    ///
+    /// If you know if there's a `SymbolId` or not, prefer `TraverseCtx::create_bound_reference_id`
+    /// or `TraverseCtx::create_unbound_reference_id`.
+    ///
+    /// This is a shortcut for `ctx.scoping.create_reference_id`.
+    pub fn create_reference_id(
+        &mut self,
+        span: Span,
+        name: Atom<'a>,
+        symbol_id: Option<SymbolId>,
+        flag: ReferenceFlag,
+    ) -> IdentifierReference<'a> {
+        self.scoping.create_reference_id(span, name, symbol_id, flag)
+    }
+
     /// Create reference in current scope, looking up binding for `name`,
     ///
     /// This is a shortcut for `ctx.scoping.create_reference_in_current_scope`.
diff --git a/crates/oxc_traverse/src/context/scoping.rs b/crates/oxc_traverse/src/context/scoping.rs
@@ -1,4 +1,4 @@
-use std::str;
+use std::{cell::Cell, str};
 
 use compact_str::{format_compact, CompactString};
 #[allow(clippy::wildcard_imports)]
@@ -7,7 +7,7 @@ use oxc_ast::{
     visit::{walk, Visit},
 };
 use oxc_semantic::{AstNodeId, Reference, ScopeTree, SymbolTable};
-use oxc_span::{CompactStr, SPAN};
+use oxc_span::{Atom, CompactStr, Span, SPAN};
 use oxc_syntax::{
     reference::{ReferenceFlag, ReferenceId},
     scope::{ScopeFlags, ScopeId},
@@ -270,6 +270,23 @@ impl TraverseScoping {
         reference_id
     }
 
+    /// Create an `IdentifierReference` bound to a `SymbolId`
+    pub fn create_bound_reference_id<'a>(
+        &mut self,
+        span: Span,
+        name: Atom<'a>,
+        symbol_id: SymbolId,
+        flag: ReferenceFlag,
+    ) -> IdentifierReference<'a> {
+        let reference_id = self.create_bound_reference(name.to_compact_str(), symbol_id, flag);
+        IdentifierReference {
+            span,
+            name,
+            reference_id: Cell::new(Some(reference_id)),
+            reference_flag: flag,
+        }
+    }
+
     /// Create an unbound reference
     pub fn create_unbound_reference(
         &mut self,
@@ -282,6 +299,22 @@ impl TraverseScoping {
         reference_id
     }
 
+    /// Create an unbound `IdentifierReference`
+    pub fn create_unbound_reference_id<'a>(
+        &mut self,
+        span: Span,
+        name: Atom<'a>,
+        flag: ReferenceFlag,
+    ) -> IdentifierReference<'a> {
+        let reference_id = self.create_unbound_reference(name.to_compact_str(), flag);
+        IdentifierReference {
+            span,
+            name,
+            reference_id: Cell::new(Some(reference_id)),
+            reference_flag: flag,
+        }
+    }
+
     /// Create a reference optionally bound to a `SymbolId`.
     ///
     /// If you know if there's a `SymbolId` or not, prefer `TraverseCtx::create_bound_reference`
@@ -299,6 +332,24 @@ impl TraverseScoping {
         }
     }
 
+    /// Create an `IdentifierReference` optionally bound to a `SymbolId`.
+    ///
+    /// If you know if there's a `SymbolId` or not, prefer `TraverseCtx::create_bound_reference_id`
+    /// or `TraverseCtx::create_unbound_reference_id`.
+    pub fn create_reference_id<'a>(
+        &mut self,
+        span: Span,
+        name: Atom<'a>,
+        symbol_id: Option<SymbolId>,
+        flag: ReferenceFlag,
+    ) -> IdentifierReference<'a> {
+        if let Some(symbol_id) = symbol_id {
+            self.create_bound_reference_id(span, name, symbol_id, flag)
+        } else {
+            self.create_unbound_reference_id(span, name, flag)
+        }
+    }
+
     /// Create reference in current scope, looking up binding for `name`
     pub fn create_reference_in_current_scope(
         &mut self,

Original file line number	Diff line number	Diff line change
`@@ -167,7 +167,7 @@ impl<'a> TypeScript<'a> {`
`167`	`167`	`self.annotations.transform_statements_on_exit(stmts, ctx);`
`168`	`168`	`}`
`169`	`169`
`170`		`- pub fn transform_statement(&mut self, stmt: &mut Statement<'a>, ctx: &TraverseCtx<'a>) {`
	`170`	`+ pub fn transform_statement(&mut self, stmt: &mut Statement<'a>, ctx: &mut TraverseCtx<'a>) {`
`171`	`171`	`self.r#enum.transform_statement(stmt, ctx);`
`172`	`172`	`}`
`173`	`173`