diff --git a/crates/oxc_transformer/src/lib.rs b/crates/oxc_transformer/src/lib.rs index e2fbd1ed01f11..d6d20c76694d2 100644 --- a/crates/oxc_transformer/src/lib.rs +++ b/crates/oxc_transformer/src/lib.rs @@ -125,11 +125,11 @@ impl<'a> Transformer<'a> { impl<'a> Traverse<'a> for Transformer<'a> { fn enter_program(&mut self, program: &mut Program<'a>, ctx: &mut TraverseCtx<'a>) { self.x0_typescript.enter_program(program, ctx); - self.x1_react.transform_program(program, ctx); + self.x1_react.enter_program(program, ctx); } fn exit_program(&mut self, program: &mut Program<'a>, ctx: &mut TraverseCtx<'a>) { - self.x1_react.transform_program_on_exit(program, ctx); + self.x1_react.exit_program(program, ctx); self.x0_typescript.exit_program(program, ctx); self.x3_es2015.exit_program(program, ctx); } @@ -150,7 +150,7 @@ impl<'a> Traverse<'a> for Transformer<'a> { fn enter_call_expression(&mut self, expr: &mut CallExpression<'a>, ctx: &mut TraverseCtx<'a>) { self.x0_typescript.enter_call_expression(expr, ctx); - self.x1_react.transform_call_expression(expr, ctx); + self.x1_react.enter_call_expression(expr, ctx); } fn enter_class(&mut self, class: &mut Class<'a>, ctx: &mut TraverseCtx<'a>) { @@ -175,7 +175,7 @@ impl<'a> Traverse<'a> for Transformer<'a> { fn enter_expression(&mut self, expr: &mut Expression<'a>, ctx: &mut TraverseCtx<'a>) { self.x0_typescript.enter_expression(expr, ctx); - self.x1_react.transform_expression(expr, ctx); + self.x1_react.enter_expression(expr, ctx); self.x2_es2021.enter_expression(expr, ctx); self.x2_es2020.enter_expression(expr, ctx); self.x2_es2018.enter_expression(expr, ctx); @@ -185,7 +185,7 @@ impl<'a> Traverse<'a> for Transformer<'a> { } fn exit_expression(&mut self, expr: &mut Expression<'a>, ctx: &mut TraverseCtx<'a>) { - self.x1_react.transform_expression_on_exit(expr, ctx); + self.x1_react.exit_expression(expr, ctx); self.x3_es2015.exit_expression(expr, ctx); } @@ -219,7 +219,7 @@ impl<'a> Traverse<'a> for Transformer<'a> { fn exit_function(&mut self, func: &mut Function<'a>, ctx: &mut TraverseCtx<'a>) { self.x0_typescript.exit_function(func, ctx); - self.x1_react.transform_function_on_exit(func, ctx); + self.x1_react.exit_function(func, ctx); self.x3_es2015.exit_function(func, ctx); } @@ -237,7 +237,7 @@ impl<'a> Traverse<'a> for Transformer<'a> { ctx: &mut TraverseCtx<'a>, ) { self.x0_typescript.enter_jsx_opening_element(elem, ctx); - self.x1_react.transform_jsx_opening_element(elem, ctx); + self.x1_react.enter_jsx_opening_element(elem, ctx); } fn enter_method_definition( @@ -278,7 +278,7 @@ impl<'a> Traverse<'a> for Transformer<'a> { fn enter_statements(&mut self, stmts: &mut Vec<'a, Statement<'a>>, ctx: &mut TraverseCtx<'a>) { self.x0_typescript.enter_statements(stmts, ctx); - self.x1_react.transform_statements(stmts, ctx); + self.x1_react.enter_statements(stmts, ctx); self.x2_es2021.enter_statements(stmts, ctx); self.x2_es2020.enter_statements(stmts, ctx); self.x2_es2016.enter_statements(stmts, ctx); @@ -308,7 +308,7 @@ impl<'a> Traverse<'a> for Transformer<'a> { fn exit_statements(&mut self, stmts: &mut Vec<'a, Statement<'a>>, ctx: &mut TraverseCtx<'a>) { self.x0_typescript.exit_statements(stmts, ctx); - self.x1_react.transform_statements_on_exit(stmts, ctx); + self.x1_react.exit_statements(stmts, ctx); self.x2_es2021.exit_statements(stmts, ctx); self.x2_es2020.exit_statements(stmts, ctx); self.x2_es2016.exit_statements(stmts, ctx); diff --git a/crates/oxc_transformer/src/react/display_name.rs b/crates/oxc_transformer/src/react/display_name.rs index 068f33249d866..9cc733cc78940 100644 --- a/crates/oxc_transformer/src/react/display_name.rs +++ b/crates/oxc_transformer/src/react/display_name.rs @@ -1,7 +1,7 @@ use oxc_allocator::Box; use oxc_ast::ast::*; use oxc_span::{Atom, SPAN}; -use oxc_traverse::{Ancestor, TraverseCtx}; +use oxc_traverse::{Ancestor, Traverse, TraverseCtx}; use crate::context::Ctx; @@ -23,12 +23,11 @@ impl<'a> ReactDisplayName<'a> { } } -// Transforms -impl<'a> ReactDisplayName<'a> { - pub fn transform_call_expression( - &self, +impl<'a> Traverse<'a> for ReactDisplayName<'a> { + fn enter_call_expression( + &mut self, call_expr: &mut CallExpression<'a>, - ctx: &TraverseCtx<'a>, + ctx: &mut TraverseCtx<'a>, ) { let Some(obj_expr) = Self::get_object_from_create_class(call_expr) else { return; diff --git a/crates/oxc_transformer/src/react/jsx.rs b/crates/oxc_transformer/src/react/jsx.rs index 12c549bc04d33..286cd9903d236 100644 --- a/crates/oxc_transformer/src/react/jsx.rs +++ b/crates/oxc_transformer/src/react/jsx.rs @@ -9,7 +9,7 @@ use oxc_syntax::{ symbol::SymbolFlags, xml_entities::XML_ENTITIES, }; -use oxc_traverse::TraverseCtx; +use oxc_traverse::{Traverse, TraverseCtx}; use super::{diagnostics, utils::get_line_column}; pub use super::{ @@ -291,7 +291,6 @@ impl<'a> Pragma<'a> { } } -// Transforms impl<'a> ReactJsx<'a> { pub fn new(options: ReactOptions, ctx: Ctx<'a>) -> Self { let bindings = match options.runtime { @@ -364,31 +363,31 @@ impl<'a> ReactJsx<'a> { bindings, } } +} - pub fn transform_program_on_exit( - &mut self, - program: &mut Program<'a>, - ctx: &mut TraverseCtx<'a>, - ) { +impl<'a> Traverse<'a> for ReactJsx<'a> { + fn exit_program(&mut self, program: &mut Program<'a>, ctx: &mut TraverseCtx<'a>) { self.add_runtime_imports(program, ctx); } - pub fn transform_jsx_element( - &mut self, - e: &JSXElement<'a>, - ctx: &mut TraverseCtx<'a>, - ) -> Expression<'a> { - self.transform_jsx(&JSXElementOrFragment::Element(e), ctx) - } + fn enter_expression(&mut self, expr: &mut Expression<'a>, ctx: &mut TraverseCtx<'a>) { + let new_expr = match expr { + Expression::JSXElement(e) => { + Some(self.transform_jsx(&JSXElementOrFragment::Element(e), ctx)) + } + Expression::JSXFragment(e) => { + Some(self.transform_jsx(&JSXElementOrFragment::Fragment(e), ctx)) + } + _ => None, + }; - pub fn transform_jsx_fragment( - &mut self, - e: &JSXFragment<'a>, - ctx: &mut TraverseCtx<'a>, - ) -> Expression<'a> { - self.transform_jsx(&JSXElementOrFragment::Fragment(e), ctx) + if let Some(new_expr) = new_expr { + *expr = new_expr; + } } +} +impl<'a> ReactJsx<'a> { fn is_script(&self) -> bool { self.ctx.source_type.is_script() } @@ -396,11 +395,8 @@ impl<'a> ReactJsx<'a> { fn ast(&self) -> AstBuilder<'a> { self.ctx.ast } -} -// Add imports -impl<'a> ReactJsx<'a> { - pub fn add_runtime_imports(&mut self, program: &mut Program<'a>, ctx: &mut TraverseCtx<'a>) { + fn add_runtime_imports(&mut self, program: &mut Program<'a>, ctx: &mut TraverseCtx<'a>) { if self.bindings.is_classic() { if let Some(stmt) = self.jsx_source.get_var_file_name_statement() { program.body.insert(0, stmt); @@ -426,50 +422,7 @@ impl<'a> ReactJsx<'a> { program.body.splice(index..index, imports); } -} - -enum JSXElementOrFragment<'a, 'b> { - Element(&'b JSXElement<'a>), - Fragment(&'b JSXFragment<'a>), -} - -impl<'a, 'b> JSXElementOrFragment<'a, 'b> { - fn span(&self) -> Span { - match self { - Self::Element(e) => e.span, - Self::Fragment(e) => e.span, - } - } - - fn children(&self) -> &'b Vec<'a, JSXChild<'a>> { - match self { - Self::Element(e) => &e.children, - Self::Fragment(e) => &e.children, - } - } - - fn is_fragment(&self) -> bool { - matches!(self, Self::Fragment(_)) - } - /// The react jsx/jsxs transform falls back to `createElement` when an explicit `key` argument comes after a spread - /// - fn has_key_after_props_spread(&self) -> bool { - let Self::Element(e) = self else { return false }; - let mut spread = false; - for attr in &e.opening_element.attributes { - if matches!(attr, JSXAttributeItem::SpreadAttribute(_)) { - spread = true; - } else if spread && matches!(attr, JSXAttributeItem::Attribute(a) if a.is_key()) { - return true; - } - } - false - } -} - -// Transform jsx -impl<'a> ReactJsx<'a> { /// ## Automatic /// ### Element /// Builds JSX into: @@ -998,6 +951,46 @@ impl<'a> ReactJsx<'a> { } } +enum JSXElementOrFragment<'a, 'b> { + Element(&'b JSXElement<'a>), + Fragment(&'b JSXFragment<'a>), +} + +impl<'a, 'b> JSXElementOrFragment<'a, 'b> { + fn span(&self) -> Span { + match self { + Self::Element(e) => e.span, + Self::Fragment(e) => e.span, + } + } + + fn children(&self) -> &'b Vec<'a, JSXChild<'a>> { + match self { + Self::Element(e) => &e.children, + Self::Fragment(e) => &e.children, + } + } + + fn is_fragment(&self) -> bool { + matches!(self, Self::Fragment(_)) + } + + /// The react jsx/jsxs transform falls back to `createElement` when an explicit `key` argument comes after a spread + /// + fn has_key_after_props_spread(&self) -> bool { + let Self::Element(e) = self else { return false }; + let mut spread = false; + for attr in &e.opening_element.attributes { + if matches!(attr, JSXAttributeItem::SpreadAttribute(_)) { + spread = true; + } else if spread && matches!(attr, JSXAttributeItem::Attribute(a) if a.is_key()) { + return true; + } + } + false + } +} + /// Create `IdentifierReference` for var name in current scope which is read from fn get_read_identifier_reference<'a>( span: Span, diff --git a/crates/oxc_transformer/src/react/jsx_self.rs b/crates/oxc_transformer/src/react/jsx_self.rs index 8616bfd5a4860..60979bd7f5b10 100644 --- a/crates/oxc_transformer/src/react/jsx_self.rs +++ b/crates/oxc_transformer/src/react/jsx_self.rs @@ -1,7 +1,7 @@ use oxc_ast::ast::*; use oxc_diagnostics::OxcDiagnostic; use oxc_span::{Span, SPAN}; -use oxc_traverse::{Ancestor, TraverseCtx}; +use oxc_traverse::{Ancestor, Traverse, TraverseCtx}; use crate::context::Ctx; @@ -23,20 +23,19 @@ impl<'a> ReactJsxSelf<'a> { pub fn new(ctx: Ctx<'a>) -> Self { Self { ctx } } +} - pub fn transform_jsx_opening_element(&self, elem: &mut JSXOpeningElement<'a>) { +impl<'a> Traverse<'a> for ReactJsxSelf<'a> { + fn enter_jsx_opening_element( + &mut self, + elem: &mut JSXOpeningElement<'a>, + _ctx: &mut TraverseCtx<'a>, + ) { self.add_self_this_attribute(elem); } +} - pub fn get_object_property_kind_for_jsx_plugin(&self) -> ObjectPropertyKind<'a> { - let kind = PropertyKind::Init; - let key = self.ctx.ast.property_key_identifier_name(SPAN, SELF); - let value = self.ctx.ast.expression_this(SPAN); - self.ctx - .ast - .object_property_kind_object_property(SPAN, kind, key, value, None, false, false, false) - } - +impl<'a> ReactJsxSelf<'a> { pub fn report_error(&self, span: Span) { let error = OxcDiagnostic::warn("Duplicate __self prop found.").with_label(span); self.ctx.error(error); @@ -63,12 +62,19 @@ impl<'a> ReactJsxSelf<'a> { true } + pub fn get_object_property_kind_for_jsx_plugin(&self) -> ObjectPropertyKind<'a> { + let kind = PropertyKind::Init; + let key = self.ctx.ast.property_key_identifier_name(SPAN, SELF); + let value = self.ctx.ast.expression_this(SPAN); + self.ctx + .ast + .object_property_kind_object_property(SPAN, kind, key, value, None, false, false, false) + } + pub fn can_add_self_attribute(&self, ctx: &TraverseCtx<'a>) -> bool { !self.is_inside_constructor(ctx) || Self::has_no_super_class(ctx) } -} -impl<'a> ReactJsxSelf<'a> { /// `
` /// ^^^^^^^^^^^^^ fn add_self_this_attribute(&self, elem: &mut JSXOpeningElement<'a>) { diff --git a/crates/oxc_transformer/src/react/jsx_source.rs b/crates/oxc_transformer/src/react/jsx_source.rs index 9a6d23b4f054e..d5b0924a0bec8 100644 --- a/crates/oxc_transformer/src/react/jsx_source.rs +++ b/crates/oxc_transformer/src/react/jsx_source.rs @@ -2,7 +2,7 @@ use oxc_ast::ast::*; use oxc_diagnostics::OxcDiagnostic; use oxc_span::{Span, SPAN}; use oxc_syntax::{number::NumberBase, symbol::SymbolFlags}; -use oxc_traverse::TraverseCtx; +use oxc_traverse::{Traverse, TraverseCtx}; use super::utils::get_line_column; use crate::{context::Ctx, helpers::bindings::BoundIdentifier}; @@ -27,15 +27,19 @@ impl<'a> ReactJsxSource<'a> { pub fn new(ctx: Ctx<'a>) -> Self { Self { ctx, filename_var: None } } +} - pub fn transform_jsx_opening_element( +impl<'a> Traverse<'a> for ReactJsxSource<'a> { + fn enter_jsx_opening_element( &mut self, elem: &mut JSXOpeningElement<'a>, ctx: &mut TraverseCtx<'a>, ) { self.add_source_attribute(elem, ctx); } +} +impl<'a> ReactJsxSource<'a> { pub fn get_object_property_kind_for_jsx_plugin( &mut self, line: usize, @@ -54,9 +58,7 @@ impl<'a> ReactJsxSource<'a> { let error = OxcDiagnostic::warn("Duplicate __source prop found.").with_label(span); self.ctx.error(error); } -} -impl<'a> ReactJsxSource<'a> { /// `` /// ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ fn add_source_attribute( diff --git a/crates/oxc_transformer/src/react/mod.rs b/crates/oxc_transformer/src/react/mod.rs index 3790e41e8b61d..46926f2354ec8 100644 --- a/crates/oxc_transformer/src/react/mod.rs +++ b/crates/oxc_transformer/src/react/mod.rs @@ -11,7 +11,7 @@ use std::rc::Rc; use oxc_allocator::Vec; use oxc_ast::ast::*; -use oxc_traverse::TraverseCtx; +use oxc_traverse::{Traverse, TraverseCtx}; use refresh::ReactRefresh; pub use self::{ @@ -68,105 +68,76 @@ impl<'a> React<'a> { } } -// Transforms -impl<'a> React<'a> { - pub fn transform_program(&mut self, program: &mut Program<'a>, ctx: &mut TraverseCtx<'a>) { +impl<'a> Traverse<'a> for React<'a> { + fn enter_program(&mut self, program: &mut Program<'a>, ctx: &mut TraverseCtx<'a>) { if self.refresh_plugin { - self.refresh.transform_program(program, ctx); + self.refresh.enter_program(program, ctx); } } - pub fn transform_program_on_exit( - &mut self, - program: &mut Program<'a>, - ctx: &mut TraverseCtx<'a>, - ) { + fn exit_program(&mut self, program: &mut Program<'a>, ctx: &mut TraverseCtx<'a>) { if self.refresh_plugin { - self.refresh.transform_program_on_exit(program, ctx); + self.refresh.exit_program(program, ctx); } if self.jsx_plugin { - self.jsx.transform_program_on_exit(program, ctx); + self.jsx.exit_program(program, ctx); } } - pub fn transform_statements( - &mut self, - stmts: &mut Vec<'a, Statement<'a>>, - ctx: &mut TraverseCtx<'a>, - ) { + fn enter_statements(&mut self, stmts: &mut Vec<'a, Statement<'a>>, ctx: &mut TraverseCtx<'a>) { if self.refresh_plugin { - self.refresh.transform_statements(stmts, ctx); + self.refresh.enter_statements(stmts, ctx); } } - pub fn transform_statements_on_exit( - &mut self, - stmts: &mut Vec<'a, Statement<'a>>, - ctx: &mut TraverseCtx<'a>, - ) { + fn exit_statements(&mut self, stmts: &mut Vec<'a, Statement<'a>>, ctx: &mut TraverseCtx<'a>) { if self.refresh_plugin { - self.refresh.transform_statements_on_exit(stmts, ctx); + self.refresh.exit_statements(stmts, ctx); } } - pub fn transform_expression(&mut self, expr: &mut Expression<'a>, ctx: &mut TraverseCtx<'a>) { + fn enter_expression(&mut self, expr: &mut Expression<'a>, ctx: &mut TraverseCtx<'a>) { if self.jsx_plugin { - match expr { - Expression::JSXElement(e) => { - *expr = self.jsx.transform_jsx_element(e, ctx); - } - Expression::JSXFragment(e) => { - *expr = self.jsx.transform_jsx_fragment(e, ctx); - } - _ => {} - } + self.jsx.enter_expression(expr, ctx); } } - pub fn transform_call_expression( + fn enter_call_expression( &mut self, call_expr: &mut CallExpression<'a>, ctx: &mut TraverseCtx<'a>, ) { if self.display_name_plugin { - self.display_name.transform_call_expression(call_expr, ctx); + self.display_name.enter_call_expression(call_expr, ctx); } if self.refresh_plugin { - self.refresh.transform_call_expression(call_expr, ctx); + self.refresh.enter_call_expression(call_expr, ctx); } } - pub fn transform_jsx_opening_element( + fn enter_jsx_opening_element( &mut self, elem: &mut JSXOpeningElement<'a>, ctx: &mut TraverseCtx<'a>, ) { if self.jsx_self_plugin && self.jsx.jsx_self.can_add_self_attribute(ctx) { - self.jsx.jsx_self.transform_jsx_opening_element(elem); + self.jsx.jsx_self.enter_jsx_opening_element(elem, ctx); } if self.jsx_source_plugin { - self.jsx.jsx_source.transform_jsx_opening_element(elem, ctx); + self.jsx.jsx_source.enter_jsx_opening_element(elem, ctx); } } - pub fn transform_expression_on_exit( - &mut self, - expr: &mut Expression<'a>, - ctx: &mut TraverseCtx<'a>, - ) { + fn exit_expression(&mut self, expr: &mut Expression<'a>, ctx: &mut TraverseCtx<'a>) { if self.refresh_plugin { - self.refresh.transform_expression_on_exit(expr, ctx); + self.refresh.exit_expression(expr, ctx); } } - pub fn transform_function_on_exit( - &mut self, - func: &mut Function<'a>, - ctx: &mut TraverseCtx<'a>, - ) { + fn exit_function(&mut self, func: &mut Function<'a>, ctx: &mut TraverseCtx<'a>) { if self.refresh_plugin { - self.refresh.transform_function_on_exit(func, ctx); + self.refresh.exit_function(func, ctx); } } } diff --git a/crates/oxc_transformer/src/react/refresh.rs b/crates/oxc_transformer/src/react/refresh.rs index 82f372aad9df3..22328dc318211 100644 --- a/crates/oxc_transformer/src/react/refresh.rs +++ b/crates/oxc_transformer/src/react/refresh.rs @@ -5,11 +5,10 @@ use oxc_ast::{ast::*, match_expression, match_member_expression}; use oxc_semantic::{ReferenceFlags, ScopeId, SymbolFlags, SymbolId}; use oxc_span::{Atom, GetSpan, SPAN}; use oxc_syntax::operator::AssignmentOperator; -use oxc_traverse::{Ancestor, TraverseCtx}; +use oxc_traverse::{Ancestor, Traverse, TraverseCtx}; use rustc_hash::FxHashMap; use super::options::ReactRefreshOptions; - use crate::context::Ctx; /// React Fast Refresh @@ -52,262 +51,261 @@ impl<'a> ReactRefresh<'a> { non_builtin_hooks_callee: FxHashMap::default(), } } +} - fn create_registration( - &mut self, - persistent_id: Atom<'a>, - ctx: &mut TraverseCtx<'a>, - ) -> AssignmentTarget<'a> { - let symbol_id = ctx.generate_uid_in_root_scope("c", SymbolFlags::FunctionScopedVariable); - self.registrations.push((symbol_id, persistent_id)); - let name = ctx.ast.atom(ctx.symbols().get_name(symbol_id)); - let ident = ctx.create_reference_id(SPAN, name, Some(symbol_id), ReferenceFlags::Write); - let ident = ctx.ast.simple_assignment_target_from_identifier_reference(ident); - ctx.ast.assignment_target_simple(ident) - } - - /// Similar to the `findInnerComponents` function in `react-refresh/babel`. - fn replace_inner_components( - &mut self, - inferred_name: &str, - expr: &mut Expression<'a>, - is_variable_declarator: bool, - ctx: &mut TraverseCtx<'a>, - ) -> bool { - match expr { - Expression::Identifier(ref ident) => { - // For case like: - // export const Something = hoc(Foo) - // we don't want to wrap Foo inside the call. - // Instead we assume it's registered at definition. - return is_componentish_name(&ident.name); - } - Expression::FunctionExpression(_) => {} - Expression::ArrowFunctionExpression(arrow) => { - // Don't transform `() => () => {}` - if arrow - .get_expression() - .is_some_and(|expr| matches!(expr, Expression::ArrowFunctionExpression(_))) - { - return false; - } +impl<'a> Traverse<'a> for ReactRefresh<'a> { + fn enter_program(&mut self, program: &mut Program<'a>, ctx: &mut TraverseCtx<'a>) { + let mut new_statements = ctx.ast.vec_with_capacity(program.body.len()); + for mut statement in program.body.drain(..) { + let next_statement = self.process_statement(&mut statement, ctx); + new_statements.push(statement); + if let Some(assignment_expression) = next_statement { + new_statements.push(assignment_expression); } - Expression::CallExpression(ref mut call_expr) => { - let allowed_callee = matches!( - call_expr.callee, - Expression::Identifier(_) - | Expression::ComputedMemberExpression(_) - | Expression::StaticMemberExpression(_) - ); - - if allowed_callee { - let callee_span = call_expr.callee.span(); - - let Some(argument_expr) = - call_expr.arguments.first_mut().and_then(|e| e.as_expression_mut()) - else { - return false; - }; - - let found_inside = self.replace_inner_components( - format!( - "{}${}", - inferred_name, - callee_span.source_text(self.ctx.source_text) - ) - .as_str(), - argument_expr, - /* is_variable_declarator */ false, - ctx, - ); - - if !found_inside { - return false; - } + } + program.body = new_statements; + } - // const Foo = hoc1(hoc2(() => {})) - // export default memo(React.forwardRef(function() {})) - if is_variable_declarator { - return true; - } - } else { - return false; - } - } - _ => { - return false; - } + fn exit_program(&mut self, program: &mut Program<'a>, ctx: &mut TraverseCtx<'a>) { + if self.registrations.is_empty() { + return; } - *expr = ctx.ast.expression_assignment( - SPAN, - AssignmentOperator::Assign, - self.create_registration(ctx.ast.atom(inferred_name), ctx), - ctx.ast.move_expression(expr), - ); + let mut variable_declarator_items = ctx.ast.vec_with_capacity(self.registrations.len()); + let mut new_statements = ctx.ast.vec_with_capacity(self.registrations.len() + 1); + for (symbol_id, persistent_id) in self.registrations.drain(..) { + let name = ctx.ast.atom(ctx.symbols().get_name(symbol_id)); + let binding_identifier = BindingIdentifier { + name: name.clone(), + symbol_id: Cell::new(Some(symbol_id)), + span: SPAN, + }; - true - } + variable_declarator_items.push( + ctx.ast.variable_declarator( + SPAN, + VariableDeclarationKind::Var, + ctx.ast.binding_pattern( + ctx.ast.binding_pattern_kind_from_binding_identifier( + binding_identifier.clone(), + ), + None::>, + false, + ), + None, + false, + ), + ); - /// Create an identifier reference from a binding identifier. - fn create_identifier_reference_from_binding_identifier( - id: &BindingIdentifier<'a>, - ctx: &mut TraverseCtx<'a>, - ) -> Expression<'a> { - ctx.ast.expression_from_identifier_reference(ctx.create_reference_id( + let refresh_reg_ident = ctx.create_reference_id( + SPAN, + self.refresh_reg.clone(), + Some(symbol_id), + ReferenceFlags::Read, + ); + let callee = ctx.ast.expression_from_identifier_reference(refresh_reg_ident); + let mut arguments = ctx.ast.vec_with_capacity(2); + arguments.push(ctx.ast.argument_expression( + Self::create_identifier_reference_from_binding_identifier(&binding_identifier, ctx), + )); + arguments.push(ctx.ast.argument_expression( + ctx.ast.expression_string_literal(SPAN, self.ctx.ast.atom(&persistent_id)), + )); + new_statements.push(ctx.ast.statement_expression( + SPAN, + ctx.ast.expression_call( + SPAN, + callee, + Option::::None, + arguments, + false, + ), + )); + } + program.body.push(Statement::from(ctx.ast.declaration_variable( SPAN, - id.name.clone(), - id.symbol_id.get(), - ReferenceFlags::Read, - )) + VariableDeclarationKind::Var, + variable_declarator_items, + false, + ))); + program.body.extend(new_statements); } - /// _c = id.name; - fn create_assignment_expression( + fn enter_statements( &mut self, - id: &BindingIdentifier<'a>, - ctx: &mut TraverseCtx<'a>, - ) -> Statement<'a> { - let left = self.create_registration(id.name.clone(), ctx); - let right = ctx.create_bound_reference_id( - SPAN, - id.name.clone(), - id.symbol_id.get().unwrap(), - ReferenceFlags::Read, - ); - let right = ctx.ast.expression_from_identifier_reference(right); - let expr = ctx.ast.expression_assignment(SPAN, AssignmentOperator::Assign, left, right); - ctx.ast.statement_expression(SPAN, expr) + _stmts: &mut oxc_allocator::Vec<'a, Statement<'a>>, + _ctx: &mut TraverseCtx<'a>, + ) { + self.signature_declarator_items.push(self.ctx.ast.vec()); } - fn create_signature_call_expression( + fn exit_statements( &mut self, - scope_id: ScopeId, - body: &mut FunctionBody<'a>, + stmts: &mut oxc_allocator::Vec<'a, Statement<'a>>, ctx: &mut TraverseCtx<'a>, - ) -> Option<(BindingIdentifier<'a>, oxc_allocator::Vec<'a, Argument<'a>>)> { - let fn_hook_calls = self.hook_calls.remove(&scope_id)?; + ) { + // TODO: check is there any function declaration - let key = fn_hook_calls - .into_iter() - .map(|(hook_name, hook_key)| format!("{hook_name}{{{hook_key}}}")) - .collect::>() - .join("\\n"); + let mut new_stmts = ctx.ast.vec_with_capacity(stmts.len() + 1); - let callee_list = self.non_builtin_hooks_callee.remove(&scope_id).unwrap_or_default(); - let callee_len = callee_list.len(); - let custom_hooks_in_scope = ctx.ast.vec_from_iter( - callee_list + let declarations = self.signature_declarator_items.pop().unwrap(); + if !declarations.is_empty() { + new_stmts.push(Statement::from(ctx.ast.declaration_variable( + SPAN, + VariableDeclarationKind::Var, + declarations, + false, + ))); + } + new_stmts.extend(stmts.drain(..).flat_map(move |stmt| { + let symbol_ids = get_symbol_id_from_function_and_declarator(&stmt); + let extra_stmts = symbol_ids .into_iter() - .filter_map(|e| e.map(|e| ctx.ast.array_expression_element_expression(e))), - ); - - let force_reset = custom_hooks_in_scope.len() != callee_len; + .filter_map(|symbol_id| self.extra_statements.remove(&symbol_id)) + .flatten() + .collect::>(); + once(stmt).chain(extra_stmts) + })); - let mut arguments = ctx.ast.vec(); - arguments.push( - ctx.ast - .argument_expression(ctx.ast.expression_string_literal(SPAN, ctx.ast.atom(&key))), - ); + *stmts = new_stmts; + } - if force_reset || !custom_hooks_in_scope.is_empty() { - arguments.push( - self.ctx.ast.argument_expression( - self.ctx.ast.expression_boolean_literal(SPAN, force_reset), - ), - ); - } + fn exit_expression(&mut self, expr: &mut Expression<'a>, ctx: &mut TraverseCtx<'a>) { + let signature = match expr { + Expression::FunctionExpression(func) => self.create_signature_call_expression( + func.scope_id.get().unwrap(), + func.body.as_mut().unwrap(), + ctx, + ), + Expression::ArrowFunctionExpression(arrow) => { + let call_fn = self.create_signature_call_expression( + arrow.scope_id.get().unwrap(), + &mut arrow.body, + ctx, + ); - if !custom_hooks_in_scope.is_empty() { - // function () { return custom_hooks_in_scope } - let formal_parameters = self.ctx.ast.formal_parameters( - SPAN, - FormalParameterKind::FormalParameter, - self.ctx.ast.vec(), - Option::::None, - ); - let function_body = self.ctx.ast.function_body( - SPAN, - self.ctx.ast.vec(), - self.ctx.ast.vec1(self.ctx.ast.statement_return( + // If the signature is found, we will push a new statement to the arrow function body. So it's not an expression anymore. + if call_fn.is_some() { + Self::transform_arrow_function_to_block(arrow, ctx); + } + call_fn + } + // hoc(_c = function() { }) + Expression::AssignmentExpression(_) => return, + // hoc1(hoc2(...)) + Expression::CallExpression(_) => self.last_signature.take(), + _ => None, + }; + + let Some((binding_identifier, mut arguments)) = signature else { + return; + }; + + if !matches!(expr, Expression::CallExpression(_)) { + if let Ancestor::VariableDeclaratorInit(declarator) = ctx.parent() { + // Special case when a function would get an inferred name: + // let Foo = () => {} + // let Foo = function() {} + // We'll add signature it on next line so that + // we don't mess up the inferred 'Foo' function name. + + // Result: let Foo = () => {}; __signature(Foo, ...); + let id = declarator.id().get_binding_identifier().unwrap(); + let symbol_id = id.symbol_id.get().unwrap(); + let first_argument = Argument::from(ctx.ast.expression_from_identifier_reference( + ctx.create_reference_id( + SPAN, + id.name.clone(), + Some(symbol_id), + ReferenceFlags::Read, + ), + )); + arguments.insert(0, first_argument); + + let statement = ctx.ast.statement_expression( SPAN, - Some(self.ctx.ast.expression_array(SPAN, custom_hooks_in_scope, None)), - )), - ); - let fn_expr = self.ctx.ast.expression_function( - FunctionType::FunctionExpression, - SPAN, - None, - false, - false, - false, - Option::::None, - Option::::None, - formal_parameters, - Option::::None, - Some(function_body), - ); - arguments.push(self.ctx.ast.argument_expression(fn_expr)); + ctx.ast.expression_call( + SPAN, + Self::create_identifier_reference_from_binding_identifier( + &binding_identifier, + ctx, + ), + Option::::None, + arguments, + false, + ), + ); + self.extra_statements.entry(symbol_id).or_insert(ctx.ast.vec()).push(statement); + return; + } } - let symbol_id = - ctx.generate_uid("s", ctx.current_scope_id(), SymbolFlags::FunctionScopedVariable); + let mut found_call_expression = false; + for ancestor in ctx.ancestors() { + if ancestor.is_assignment_expression() { + continue; + } + if ancestor.is_call_expression() { + found_call_expression = true; + } + break; + } - let symbol_name = ctx.ast.atom(ctx.symbols().get_name(symbol_id)); + if found_call_expression { + self.last_signature = + Some((binding_identifier.clone(), arguments.clone_in(ctx.ast.allocator))); + } - let binding_identifier = BindingIdentifier { - span: SPAN, - name: symbol_name.clone(), - symbol_id: Cell::new(Some(symbol_id)), + arguments.insert(0, Argument::from(ctx.ast.move_expression(expr))); + *expr = self.ctx.ast.expression_call( + SPAN, + Self::create_identifier_reference_from_binding_identifier(&binding_identifier, ctx), + Option::::None, + arguments, + false, + ); + } + + fn exit_function(&mut self, func: &mut Function<'a>, ctx: &mut TraverseCtx<'a>) { + if !func.is_function_declaration() { + return; + } + + let Some((binding_identifier, mut arguments)) = self.create_signature_call_expression( + func.scope_id.get().unwrap(), + func.body.as_mut().unwrap(), + ctx, + ) else { + return; }; - let sig_identifier_reference = ctx.create_reference_id( - SPAN, - self.refresh_sig.clone(), - Some(symbol_id), - ReferenceFlags::Read, + let Some(id) = func.id.as_ref() else { + return; + }; + + arguments.insert( + 0, + Argument::from(Self::create_identifier_reference_from_binding_identifier(id, ctx)), ); - // _s(); - let call_expression = ctx.ast.statement_expression( - SPAN, - ctx.ast.expression_call( + self.extra_statements.entry(id.symbol_id.get().unwrap()).or_insert(ctx.ast.vec()).push( + ctx.ast.statement_expression( SPAN, - Self::create_identifier_reference_from_binding_identifier(&binding_identifier, ctx), - Option::::None, - ctx.ast.vec(), - false, + ctx.ast.expression_call( + SPAN, + Self::create_identifier_reference_from_binding_identifier( + &binding_identifier, + ctx, + ), + Option::::None, + arguments, + false, + ), ), ); - - body.statements.insert(0, call_expression); - - // _s = refresh_sig(); - self.signature_declarator_items.last_mut().unwrap().push(ctx.ast.variable_declarator( - SPAN, - VariableDeclarationKind::Var, - ctx.ast.binding_pattern( - ctx.ast.binding_pattern_kind_from_binding_identifier(binding_identifier.clone()), - Option::::None, - false, - ), - Some(ctx.ast.expression_call( - SPAN, - ctx.ast.expression_from_identifier_reference(sig_identifier_reference.clone()), - Option::::None, - ctx.ast.vec(), - false, - )), - false, - )); - - // Following is the signature call expression, will be generated in call site. - // _s(App, signature_key, false, function() { return [] }); - // ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ custom hooks only - Some((binding_identifier, arguments)) } - pub fn transform_call_expression( + fn enter_call_expression( &mut self, call_expr: &mut CallExpression<'a>, ctx: &mut TraverseCtx<'a>, @@ -405,61 +403,302 @@ impl<'a> ReactRefresh<'a> { } } -// Internal Methods for transforming +// Internal Methods impl<'a> ReactRefresh<'a> { - /// Process statement and return a statement(if any) to insert it after current statement. - /// - /// ```js - /// const Foo = styled("div")`color: hotpink`; - /// function Bar() {} - /// ``` - /// to - /// ```js - /// const Foo = styled("div")`color: hotpink`; - /// _c = Foo; - /// function Bar() { } - /// _c1 = Bar; - /// ``` - fn process_statement( + fn create_registration( &mut self, - statement: &mut Statement<'a>, + persistent_id: Atom<'a>, ctx: &mut TraverseCtx<'a>, - ) -> Option> { - match statement { - Statement::VariableDeclaration(variable) => { - self.handle_variable_declaration(variable, ctx) + ) -> AssignmentTarget<'a> { + let symbol_id = ctx.generate_uid_in_root_scope("c", SymbolFlags::FunctionScopedVariable); + self.registrations.push((symbol_id, persistent_id)); + let name = ctx.ast.atom(ctx.symbols().get_name(symbol_id)); + let ident = ctx.create_reference_id(SPAN, name, Some(symbol_id), ReferenceFlags::Write); + let ident = ctx.ast.simple_assignment_target_from_identifier_reference(ident); + ctx.ast.assignment_target_simple(ident) + } + + /// Similar to the `findInnerComponents` function in `react-refresh/babel`. + fn replace_inner_components( + &mut self, + inferred_name: &str, + expr: &mut Expression<'a>, + is_variable_declarator: bool, + ctx: &mut TraverseCtx<'a>, + ) -> bool { + match expr { + Expression::Identifier(ref ident) => { + // For case like: + // export const Something = hoc(Foo) + // we don't want to wrap Foo inside the call. + // Instead we assume it's registered at definition. + return is_componentish_name(&ident.name); } - Statement::FunctionDeclaration(func) => self.handle_function_declaration(func, ctx), - Statement::ExportNamedDeclaration(export_decl) => { - if let Some(declaration) = &mut export_decl.declaration { - match declaration { - Declaration::FunctionDeclaration(func) => { - self.handle_function_declaration(func, ctx) - } - Declaration::VariableDeclaration(variable) => { - self.handle_variable_declaration(variable, ctx) - } - _ => None, - } - } else { - None + Expression::FunctionExpression(_) => {} + Expression::ArrowFunctionExpression(arrow) => { + // Don't transform `() => () => {}` + if arrow + .get_expression() + .is_some_and(|expr| matches!(expr, Expression::ArrowFunctionExpression(_))) + { + return false; } } - Statement::ExportDefaultDeclaration(ref mut stmt_decl) => { - match &mut stmt_decl.declaration { - declaration @ match_expression!(ExportDefaultDeclarationKind) => { - let expression = declaration.to_expression_mut(); - if !matches!(expression, Expression::CallExpression(_)) { - // For now, we only support possible HOC calls here. - // Named function declarations are handled in FunctionDeclaration. - // Anonymous direct exports like export default function() {} - // are currently ignored. - return None; - } + Expression::CallExpression(ref mut call_expr) => { + let allowed_callee = matches!( + call_expr.callee, + Expression::Identifier(_) + | Expression::ComputedMemberExpression(_) + | Expression::StaticMemberExpression(_) + ); - // This code path handles nested cases like: - // export default memo(() => {}) - // In those cases it is more plausible people will omit names + if allowed_callee { + let callee_span = call_expr.callee.span(); + + let Some(argument_expr) = + call_expr.arguments.first_mut().and_then(|e| e.as_expression_mut()) + else { + return false; + }; + + let found_inside = self.replace_inner_components( + format!( + "{}${}", + inferred_name, + callee_span.source_text(self.ctx.source_text) + ) + .as_str(), + argument_expr, + /* is_variable_declarator */ false, + ctx, + ); + + if !found_inside { + return false; + } + + // const Foo = hoc1(hoc2(() => {})) + // export default memo(React.forwardRef(function() {})) + if is_variable_declarator { + return true; + } + } else { + return false; + } + } + _ => { + return false; + } + } + + *expr = ctx.ast.expression_assignment( + SPAN, + AssignmentOperator::Assign, + self.create_registration(ctx.ast.atom(inferred_name), ctx), + ctx.ast.move_expression(expr), + ); + + true + } + + /// Create an identifier reference from a binding identifier. + fn create_identifier_reference_from_binding_identifier( + id: &BindingIdentifier<'a>, + ctx: &mut TraverseCtx<'a>, + ) -> Expression<'a> { + ctx.ast.expression_from_identifier_reference(ctx.create_reference_id( + SPAN, + id.name.clone(), + id.symbol_id.get(), + ReferenceFlags::Read, + )) + } + + /// _c = id.name; + fn create_assignment_expression( + &mut self, + id: &BindingIdentifier<'a>, + ctx: &mut TraverseCtx<'a>, + ) -> Statement<'a> { + let left = self.create_registration(id.name.clone(), ctx); + let right = ctx.create_bound_reference_id( + SPAN, + id.name.clone(), + id.symbol_id.get().unwrap(), + ReferenceFlags::Read, + ); + let right = ctx.ast.expression_from_identifier_reference(right); + let expr = ctx.ast.expression_assignment(SPAN, AssignmentOperator::Assign, left, right); + ctx.ast.statement_expression(SPAN, expr) + } + + fn create_signature_call_expression( + &mut self, + scope_id: ScopeId, + body: &mut FunctionBody<'a>, + ctx: &mut TraverseCtx<'a>, + ) -> Option<(BindingIdentifier<'a>, oxc_allocator::Vec<'a, Argument<'a>>)> { + let fn_hook_calls = self.hook_calls.remove(&scope_id)?; + + let key = fn_hook_calls + .into_iter() + .map(|(hook_name, hook_key)| format!("{hook_name}{{{hook_key}}}")) + .collect::>() + .join("\\n"); + + let callee_list = self.non_builtin_hooks_callee.remove(&scope_id).unwrap_or_default(); + let callee_len = callee_list.len(); + let custom_hooks_in_scope = ctx.ast.vec_from_iter( + callee_list + .into_iter() + .filter_map(|e| e.map(|e| ctx.ast.array_expression_element_expression(e))), + ); + + let force_reset = custom_hooks_in_scope.len() != callee_len; + + let mut arguments = ctx.ast.vec(); + arguments.push( + ctx.ast + .argument_expression(ctx.ast.expression_string_literal(SPAN, ctx.ast.atom(&key))), + ); + + if force_reset || !custom_hooks_in_scope.is_empty() { + arguments.push( + self.ctx.ast.argument_expression( + self.ctx.ast.expression_boolean_literal(SPAN, force_reset), + ), + ); + } + + if !custom_hooks_in_scope.is_empty() { + // function () { return custom_hooks_in_scope } + let formal_parameters = self.ctx.ast.formal_parameters( + SPAN, + FormalParameterKind::FormalParameter, + self.ctx.ast.vec(), + Option::::None, + ); + let function_body = self.ctx.ast.function_body( + SPAN, + self.ctx.ast.vec(), + self.ctx.ast.vec1(self.ctx.ast.statement_return( + SPAN, + Some(self.ctx.ast.expression_array(SPAN, custom_hooks_in_scope, None)), + )), + ); + let fn_expr = self.ctx.ast.expression_function( + FunctionType::FunctionExpression, + SPAN, + None, + false, + false, + false, + Option::::None, + Option::::None, + formal_parameters, + Option::::None, + Some(function_body), + ); + arguments.push(self.ctx.ast.argument_expression(fn_expr)); + } + + let symbol_id = + ctx.generate_uid("s", ctx.current_scope_id(), SymbolFlags::FunctionScopedVariable); + + let symbol_name = ctx.ast.atom(ctx.symbols().get_name(symbol_id)); + + let binding_identifier = BindingIdentifier { + span: SPAN, + name: symbol_name.clone(), + symbol_id: Cell::new(Some(symbol_id)), + }; + + let sig_identifier_reference = ctx.create_reference_id( + SPAN, + self.refresh_sig.clone(), + Some(symbol_id), + ReferenceFlags::Read, + ); + + // _s(); + let call_expression = ctx.ast.statement_expression( + SPAN, + ctx.ast.expression_call( + SPAN, + Self::create_identifier_reference_from_binding_identifier(&binding_identifier, ctx), + Option::::None, + ctx.ast.vec(), + false, + ), + ); + + body.statements.insert(0, call_expression); + + // _s = refresh_sig(); + self.signature_declarator_items.last_mut().unwrap().push(ctx.ast.variable_declarator( + SPAN, + VariableDeclarationKind::Var, + ctx.ast.binding_pattern( + ctx.ast.binding_pattern_kind_from_binding_identifier(binding_identifier.clone()), + Option::::None, + false, + ), + Some(ctx.ast.expression_call( + SPAN, + ctx.ast.expression_from_identifier_reference(sig_identifier_reference.clone()), + Option::::None, + ctx.ast.vec(), + false, + )), + false, + )); + + // Following is the signature call expression, will be generated in call site. + // _s(App, signature_key, false, function() { return [] }); + // ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ custom hooks only + Some((binding_identifier, arguments)) + } + + fn process_statement( + &mut self, + statement: &mut Statement<'a>, + ctx: &mut TraverseCtx<'a>, + ) -> Option> { + match statement { + Statement::VariableDeclaration(variable) => { + self.handle_variable_declaration(variable, ctx) + } + Statement::FunctionDeclaration(func) => self.handle_function_declaration(func, ctx), + Statement::ExportNamedDeclaration(export_decl) => { + if let Some(declaration) = &mut export_decl.declaration { + match declaration { + Declaration::FunctionDeclaration(func) => { + self.handle_function_declaration(func, ctx) + } + Declaration::VariableDeclaration(variable) => { + self.handle_variable_declaration(variable, ctx) + } + _ => None, + } + } else { + None + } + } + Statement::ExportDefaultDeclaration(ref mut stmt_decl) => { + match &mut stmt_decl.declaration { + declaration @ match_expression!(ExportDefaultDeclarationKind) => { + let expression = declaration.to_expression_mut(); + if !matches!(expression, Expression::CallExpression(_)) { + // For now, we only support possible HOC calls here. + // Named function declarations are handled in FunctionDeclaration. + // Anonymous direct exports like export default function() {} + // are currently ignored. + return None; + } + + // This code path handles nested cases like: + // export default memo(() => {}) + // In those cases it is more plausible people will omit names // so they're worth handling despite possible false positives. // More importantly, it handles the named case: // export default memo(function Named() {}) @@ -573,8 +812,6 @@ impl<'a> ReactRefresh<'a> { Some(self.create_assignment_expression(id, ctx)) } - // --------------------------- refresh sig --------------------------- - /// Convert arrow function expression to normal arrow function /// /// ```js @@ -605,302 +842,6 @@ impl<'a> ReactRefresh<'a> { } } -// Transform -impl<'a> ReactRefresh<'a> { - /// Mutate statements and insert new assignment statements; - pub fn transform_program(&mut self, program: &mut Program<'a>, ctx: &mut TraverseCtx<'a>) { - let mut new_statements = ctx.ast.vec_with_capacity(program.body.len()); - for mut statement in program.body.drain(..) { - let next_statement = self.process_statement(&mut statement, ctx); - new_statements.push(statement); - if let Some(assignment_expression) = next_statement { - new_statements.push(assignment_expression); - } - } - program.body = new_statements; - } - - /// Insert all registrations at the end of the program. - /// - /// ``` - /// _c1 = refresh_reg(Foo, ...); - /// _c2 = refresh_reg(Foo, ...); - /// ``` - pub fn transform_program_on_exit( - &mut self, - program: &mut Program<'a>, - ctx: &mut TraverseCtx<'a>, - ) { - if self.registrations.is_empty() { - return; - } - - let mut variable_declarator_items = ctx.ast.vec_with_capacity(self.registrations.len()); - let mut new_statements = ctx.ast.vec_with_capacity(self.registrations.len() + 1); - for (symbol_id, persistent_id) in self.registrations.drain(..) { - let name = ctx.ast.atom(ctx.symbols().get_name(symbol_id)); - let binding_identifier = BindingIdentifier { - name: name.clone(), - symbol_id: Cell::new(Some(symbol_id)), - span: SPAN, - }; - - variable_declarator_items.push( - ctx.ast.variable_declarator( - SPAN, - VariableDeclarationKind::Var, - ctx.ast.binding_pattern( - ctx.ast.binding_pattern_kind_from_binding_identifier( - binding_identifier.clone(), - ), - None::>, - false, - ), - None, - false, - ), - ); - - let refresh_reg_ident = ctx.create_reference_id( - SPAN, - self.refresh_reg.clone(), - Some(symbol_id), - ReferenceFlags::Read, - ); - let callee = ctx.ast.expression_from_identifier_reference(refresh_reg_ident); - let mut arguments = ctx.ast.vec_with_capacity(2); - arguments.push(ctx.ast.argument_expression( - Self::create_identifier_reference_from_binding_identifier(&binding_identifier, ctx), - )); - arguments.push(ctx.ast.argument_expression( - ctx.ast.expression_string_literal(SPAN, self.ctx.ast.atom(&persistent_id)), - )); - new_statements.push(ctx.ast.statement_expression( - SPAN, - ctx.ast.expression_call( - SPAN, - callee, - Option::::None, - arguments, - false, - ), - )); - } - program.body.push(Statement::from(ctx.ast.declaration_variable( - SPAN, - VariableDeclarationKind::Var, - variable_declarator_items, - false, - ))); - program.body.extend(new_statements); - } - - pub fn transform_statements( - &mut self, - _stmts: &mut oxc_allocator::Vec<'a, Statement<'a>>, - ctx: &mut TraverseCtx<'a>, - ) { - self.signature_declarator_items.push(ctx.ast.vec()); - } - - pub fn transform_statements_on_exit( - &mut self, - stmts: &mut oxc_allocator::Vec<'a, Statement<'a>>, - ctx: &mut TraverseCtx<'a>, - ) { - // TODO: check is there any function declaration - - let mut new_stmts = ctx.ast.vec_with_capacity(stmts.len() + 1); - - let declarations = self.signature_declarator_items.pop().unwrap(); - if !declarations.is_empty() { - new_stmts.push(Statement::from(ctx.ast.declaration_variable( - SPAN, - VariableDeclarationKind::Var, - declarations, - false, - ))); - } - new_stmts.extend(stmts.drain(..).flat_map(move |stmt| { - let symbol_ids = get_symbol_id_from_function_and_declarator(&stmt); - let extra_stmts = symbol_ids - .into_iter() - .filter_map(|symbol_id| self.extra_statements.remove(&symbol_id)) - .flatten() - .collect::>(); - once(stmt).chain(extra_stmts) - })); - - *stmts = new_stmts; - } - - /// Transform an expression to insert a signature call, - /// and wrap it with a signature call - /// - /// ```js - /// Foo(() => {}) - /// let Foo = React.forwardRef(React.memo(() => {})); - /// ``` - /// to - /// ```js - /// Foo(_s1(() => {})) - /// let Foo = s1(React.forwardRef(_s1(React.memo(_s1(() => {_s1()}, ...), ...))), ...); - /// ``` - pub fn transform_expression_on_exit( - &mut self, - expr: &mut Expression<'a>, - ctx: &mut TraverseCtx<'a>, - ) { - let signature = match expr { - Expression::FunctionExpression(func) => self.create_signature_call_expression( - func.scope_id.get().unwrap(), - func.body.as_mut().unwrap(), - ctx, - ), - Expression::ArrowFunctionExpression(arrow) => { - let call_fn = self.create_signature_call_expression( - arrow.scope_id.get().unwrap(), - &mut arrow.body, - ctx, - ); - - // If the signature is found, we will push a new statement to the arrow function body. So it's not an expression anymore. - if call_fn.is_some() { - Self::transform_arrow_function_to_block(arrow, ctx); - } - call_fn - } - // hoc(_c = function() { }) - Expression::AssignmentExpression(_) => return, - // hoc1(hoc2(...)) - // Result: let Foo = __signature(hoc(__signature(() => {}, ...)), ...) - Expression::CallExpression(_) => self.last_signature.take(), - _ => None, - }; - - let Some((binding_identifier, mut arguments)) = signature else { - return; - }; - - if !matches!(expr, Expression::CallExpression(_)) { - if let Ancestor::VariableDeclaratorInit(declarator) = ctx.parent() { - // Special case when a function would get an inferred name: - // let Foo = () => {} - // let Foo = function() {} - // We'll add signature it on next line so that - // we don't mess up the inferred 'Foo' function name. - - // Result: let Foo = () => {}; __signature(Foo, ...); - let id = declarator.id().get_binding_identifier().unwrap(); - let symbol_id = id.symbol_id.get().unwrap(); - let first_argument = Argument::from(ctx.ast.expression_from_identifier_reference( - ctx.create_reference_id( - SPAN, - id.name.clone(), - Some(symbol_id), - ReferenceFlags::Read, - ), - )); - arguments.insert(0, first_argument); - - let statement = ctx.ast.statement_expression( - SPAN, - ctx.ast.expression_call( - SPAN, - Self::create_identifier_reference_from_binding_identifier( - &binding_identifier, - ctx, - ), - Option::::None, - arguments, - false, - ), - ); - self.extra_statements.entry(symbol_id).or_insert(ctx.ast.vec()).push(statement); - return; - } - } - - let mut found_call_expression = false; - for ancestor in ctx.ancestors() { - if ancestor.is_assignment_expression() { - continue; - } - if ancestor.is_call_expression() { - found_call_expression = true; - } - break; - } - - if found_call_expression { - self.last_signature = - Some((binding_identifier.clone(), arguments.clone_in(ctx.ast.allocator))); - } - - arguments.insert(0, Argument::from(ctx.ast.move_expression(expr))); - *expr = self.ctx.ast.expression_call( - SPAN, - Self::create_identifier_reference_from_binding_identifier(&binding_identifier, ctx), - Option::::None, - arguments, - false, - ); - } - - /// Modify a function to insert a signature call, - /// and return a statement to insert it after current statement - /// - /// ```js - /// function Foo() {}; - /// ``` - /// to - /// ```js - /// function Foo() { _s() }; _s(Foo, ...); - /// ``` - pub fn transform_function_on_exit( - &mut self, - func: &mut Function<'a>, - ctx: &mut TraverseCtx<'a>, - ) { - if !func.is_function_declaration() { - return; - } - - let Some((binding_identifier, mut arguments)) = self.create_signature_call_expression( - func.scope_id.get().unwrap(), - func.body.as_mut().unwrap(), - ctx, - ) else { - return; - }; - - let Some(id) = func.id.as_ref() else { - return; - }; - - arguments.insert( - 0, - Argument::from(Self::create_identifier_reference_from_binding_identifier(id, ctx)), - ); - - self.extra_statements.entry(id.symbol_id.get().unwrap()).or_insert(ctx.ast.vec()).push( - ctx.ast.statement_expression( - SPAN, - ctx.ast.expression_call( - SPAN, - Self::create_identifier_reference_from_binding_identifier( - &binding_identifier, - ctx, - ), - Option::::None, - arguments, - false, - ), - ), - ); - } -} - fn is_componentish_name(name: &str) -> bool { name.chars().next().unwrap().is_ascii_uppercase() }