refactor(transformer/react): move all entry points to implementation of Traverse trait (#5473)

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);

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> Traverse<'a> for ReactDisplayName<'a> {
-impl<'a> ReactDisplayName<'a> {
+    fn enter_call_expression(
-    pub fn transform_call_expression(
+        &mut self,
-        &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;

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(
+impl<'a> Traverse<'a> for ReactJsx<'a> {
-        &mut self,
+    fn exit_program(&mut self, program: &mut Program<'a>, ctx: &mut TraverseCtx<'a>) {
-        program: &mut Program<'a>,
-        ctx: &mut TraverseCtx<'a>,
-    ) {
         self.add_runtime_imports(program, ctx);
     }
 
-    pub fn transform_jsx_element(
+    fn enter_expression(&mut self, expr: &mut Expression<'a>, ctx: &mut TraverseCtx<'a>) {
-        &mut self,
+        let new_expr = match expr {
-        e: &JSXElement<'a>,
+            Expression::JSXElement(e) => {
-        ctx: &mut TraverseCtx<'a>,
+                Some(self.transform_jsx(&JSXElementOrFragment::Element(e), ctx))
-    ) -> Expression<'a> {
-        self.transform_jsx(&JSXElementOrFragment::Element(e), ctx)
             }
-
+            Expression::JSXFragment(e) => {
-    pub fn transform_jsx_fragment(
+                Some(self.transform_jsx(&JSXElementOrFragment::Fragment(e), ctx))
-        &mut self,
-        e: &JSXFragment<'a>,
-        ctx: &mut TraverseCtx<'a>,
-    ) -> Expression<'a> {
-        self.transform_jsx(&JSXElementOrFragment::Fragment(e), ctx)
             }
+            _ => None,
+        };
 
+        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
+    fn add_runtime_imports(&mut self, program: &mut Program<'a>, ctx: &mut TraverseCtx<'a>) {
-impl<'a> ReactJsx<'a> {
-    pub 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
-    /// <https://github.com/microsoft/TypeScript/blob/6134091642f57c32f50e7b5604635e4d37dd19e8/src/compiler/transformers/jsx.ts#L264-L278>
-    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
+    /// <https://github.com/microsoft/TypeScript/blob/6134091642f57c32f50e7b5604635e4d37dd19e8/src/compiler/transformers/jsx.ts#L264-L278>
+    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,

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> {
+impl<'a> ReactJsxSelf<'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 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> {
     /// `<div __self={this} />`
     ///       ^^^^^^^^^^^^^
     fn add_self_this_attribute(&self, elem: &mut JSXOpeningElement<'a>) {

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> {
     /// `<sometag __source={ { fileName: 'this/file.js', lineNumber: 10, columnNumber: 1 } } />`
     ///           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
     fn add_source_attribute(

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> Traverse<'a> for React<'a> {
-impl<'a> React<'a> {
+    fn enter_program(&mut self, program: &mut Program<'a>, ctx: &mut TraverseCtx<'a>) {
-    pub fn transform_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(
+    fn exit_program(&mut self, program: &mut Program<'a>, ctx: &mut TraverseCtx<'a>) {
-        &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(
+    fn enter_statements(&mut self, stmts: &mut Vec<'a, Statement<'a>>, ctx: &mut TraverseCtx<'a>) {
-        &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(
+    fn exit_statements(&mut self, stmts: &mut Vec<'a, Statement<'a>>, ctx: &mut TraverseCtx<'a>) {
-        &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 {
+            self.jsx.enter_expression(expr, ctx);
-                Expression::JSXElement(e) => {
-                    *expr = self.jsx.transform_jsx_element(e, ctx);
-                }
-                Expression::JSXFragment(e) => {
-                    *expr = self.jsx.transform_jsx_fragment(e, 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(
+    fn exit_expression(&mut self, expr: &mut Expression<'a>, ctx: &mut TraverseCtx<'a>) {
-        &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(
+    fn exit_function(&mut self, func: &mut Function<'a>, ctx: &mut TraverseCtx<'a>) {
-        &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);
         }
     }
 }

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,7 +51,360 @@ impl<'a> ReactRefresh<'a> {
             non_builtin_hooks_callee: FxHashMap::default(),
         }
     }
+}
 
+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);
+            }
+        }
+        program.body = new_statements;
+    }
+
+    fn exit_program(&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::<TSTypeAnnotation<'a>>,
+                        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::<TSTypeParameterInstantiation>::None,
+                    arguments,
+                    false,
+                ),
+            ));
+        }
+        program.body.push(Statement::from(ctx.ast.declaration_variable(
+            SPAN,
+            VariableDeclarationKind::Var,
+            variable_declarator_items,
+            false,
+        )));
+        program.body.extend(new_statements);
+    }
+
+    fn enter_statements(
+        &mut self,
+        _stmts: &mut oxc_allocator::Vec<'a, Statement<'a>>,
+        _ctx: &mut TraverseCtx<'a>,
+    ) {
+        self.signature_declarator_items.push(self.ctx.ast.vec());
+    }
+
+    fn exit_statements(
+        &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::<Vec<_>>();
+            once(stmt).chain(extra_stmts)
+        }));
+
+        *stmts = new_stmts;
+    }
+
+    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 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,
+                    ctx.ast.expression_call(
+                        SPAN,
+                        Self::create_identifier_reference_from_binding_identifier(
+                            &binding_identifier,
+                            ctx,
+                        ),
+                        Option::<TSTypeParameterInstantiation>::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::<TSTypeParameterInstantiation>::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 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::<TSTypeParameterInstantiation>::None,
+                    arguments,
+                    false,
+                ),
+            ),
+        );
+    }
+
+    fn enter_call_expression(
+        &mut self,
+        call_expr: &mut CallExpression<'a>,
+        ctx: &mut TraverseCtx<'a>,
+    ) {
+        let current_scope_id = ctx.current_scope_id();
+        if !ctx.scopes().get_flags(current_scope_id).is_function() {
+            return;
+        }
+
+        let name = match &call_expr.callee {
+            Expression::Identifier(ident) => Some(ident.name.clone()),
+            Expression::StaticMemberExpression(ref member) => Some(member.property.name.clone()),
+            _ => None,
+        };
+
+        let Some(hook_name) = name else {
+            return;
+        };
+
+        if !is_use_hook_name(&hook_name) {
+            return;
+        }
+
+        if !is_builtin_hook(&hook_name) {
+            let (binding_name, hook_name) = match &call_expr.callee {
+                Expression::Identifier(ident) => (ident.name.clone(), None),
+                callee @ match_member_expression!(Expression) => {
+                    let member_expr = callee.to_member_expression();
+                    match member_expr.object() {
+                        Expression::Identifier(ident) => {
+                            (ident.name.clone(), Some(hook_name.clone()))
+                        }
+                        _ => unreachable!(),
+                    }
+                }
+                _ => unreachable!(),
+            };
+
+            let callees = self.non_builtin_hooks_callee.entry(current_scope_id).or_default();
+
+            callees.push(
+                ctx.scopes()
+                    .find_binding(
+                        ctx.scopes().get_parent_id(ctx.current_scope_id()).unwrap(),
+                        binding_name.as_str(),
+                    )
+                    .map(|symbol_id| {
+                        let ident = ctx.create_reference_id(
+                            SPAN,
+                            binding_name.clone(),
+                            Some(symbol_id),
+                            ReferenceFlags::Read,
+                        );
+
+                        let mut expr = self.ctx.ast.expression_from_identifier_reference(ident);
+
+                        if let Some(hook_name) = hook_name {
+                            // binding_name.hook_name
+                            expr = Expression::from(self.ctx.ast.member_expression_static(
+                                SPAN,
+                                expr,
+                                self.ctx.ast.identifier_name(SPAN, hook_name),
+                                false,
+                            ));
+                        }
+                        expr
+                    }),
+            );
+        }
+
+        let key = if let Ancestor::VariableDeclaratorInit(declarator) = ctx.parent() {
+            // TODO: if there is no LHS, consider some other heuristic.
+            declarator.id().span().source_text(self.ctx.source_text)
+        } else {
+            ""
+        };
+
+        let args = &call_expr.arguments;
+        let args_key = if hook_name == "useState" && args.len() > 0 {
+            args[0].span().source_text(self.ctx.source_text)
+        } else if hook_name == "useReducer" && args.len() > 1 {
+            args[1].span().source_text(self.ctx.source_text)
+        } else {
+            ""
+        };
+
+        let key = format!(
+            "{}{}{args_key}{}",
+            key,
+            if args_key.is_empty() { "" } else { "(" },
+            if args_key.is_empty() { "" } else { ")" }
+        );
+
+        self.hook_calls.entry(current_scope_id).or_default().push((hook_name, ctx.ast.atom(&key)));
+    }
+}
+
+// Internal Methods
+impl<'a> ReactRefresh<'a> {
     fn create_registration(
         &mut self,
         persistent_id: Atom<'a>,
@@ -307,119 +659,6 @@ impl<'a> ReactRefresh<'a> {
         Some((binding_identifier, arguments))
     }
 
-    pub fn transform_call_expression(
-        &mut self,
-        call_expr: &mut CallExpression<'a>,
-        ctx: &mut TraverseCtx<'a>,
-    ) {
-        let current_scope_id = ctx.current_scope_id();
-        if !ctx.scopes().get_flags(current_scope_id).is_function() {
-            return;
-        }
-
-        let name = match &call_expr.callee {
-            Expression::Identifier(ident) => Some(ident.name.clone()),
-            Expression::StaticMemberExpression(ref member) => Some(member.property.name.clone()),
-            _ => None,
-        };
-
-        let Some(hook_name) = name else {
-            return;
-        };
-
-        if !is_use_hook_name(&hook_name) {
-            return;
-        }
-
-        if !is_builtin_hook(&hook_name) {
-            let (binding_name, hook_name) = match &call_expr.callee {
-                Expression::Identifier(ident) => (ident.name.clone(), None),
-                callee @ match_member_expression!(Expression) => {
-                    let member_expr = callee.to_member_expression();
-                    match member_expr.object() {
-                        Expression::Identifier(ident) => {
-                            (ident.name.clone(), Some(hook_name.clone()))
-                        }
-                        _ => unreachable!(),
-                    }
-                }
-                _ => unreachable!(),
-            };
-
-            let callees = self.non_builtin_hooks_callee.entry(current_scope_id).or_default();
-
-            callees.push(
-                ctx.scopes()
-                    .find_binding(
-                        ctx.scopes().get_parent_id(ctx.current_scope_id()).unwrap(),
-                        binding_name.as_str(),
-                    )
-                    .map(|symbol_id| {
-                        let ident = ctx.create_reference_id(
-                            SPAN,
-                            binding_name.clone(),
-                            Some(symbol_id),
-                            ReferenceFlags::Read,
-                        );
-
-                        let mut expr = self.ctx.ast.expression_from_identifier_reference(ident);
-
-                        if let Some(hook_name) = hook_name {
-                            // binding_name.hook_name
-                            expr = Expression::from(self.ctx.ast.member_expression_static(
-                                SPAN,
-                                expr,
-                                self.ctx.ast.identifier_name(SPAN, hook_name),
-                                false,
-                            ));
-                        }
-                        expr
-                    }),
-            );
-        }
-
-        let key = if let Ancestor::VariableDeclaratorInit(declarator) = ctx.parent() {
-            // TODO: if there is no LHS, consider some other heuristic.
-            declarator.id().span().source_text(self.ctx.source_text)
-        } else {
-            ""
-        };
-
-        let args = &call_expr.arguments;
-        let args_key = if hook_name == "useState" && args.len() > 0 {
-            args[0].span().source_text(self.ctx.source_text)
-        } else if hook_name == "useReducer" && args.len() > 1 {
-            args[1].span().source_text(self.ctx.source_text)
-        } else {
-            ""
-        };
-
-        let key = format!(
-            "{}{}{args_key}{}",
-            key,
-            if args_key.is_empty() { "" } else { "(" },
-            if args_key.is_empty() { "" } else { ")" }
-        );
-
-        self.hook_calls.entry(current_scope_id).or_default().push((hook_name, ctx.ast.atom(&key)));
-    }
-}
-
-// Internal Methods for transforming
-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(
         &mut self,
         statement: &mut Statement<'a>,
@@ -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::<TSTypeAnnotation<'a>>,
-                        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::<TSTypeParameterInstantiation>::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::<Vec<_>>();
-            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::<TSTypeParameterInstantiation>::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::<TSTypeParameterInstantiation>::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::<TSTypeParameterInstantiation>::None,
-                    arguments,
-                    false,
-                ),
-            ),
-        );
-    }
-}
-
 fn is_componentish_name(name: &str) -> bool {
     name.chars().next().unwrap().is_ascii_uppercase()
 }