fix(es/transform_react): should force refresh if code contains @refresh reset

crates/swc_ecma_transforms_react/src/refresh/mod.rs

@@ -249,6 +249,12 @@ where
                 });
             }
 
+            comments.with_leading(n.lo, |comments| {
+                if comments.iter().any(|c| c.text.contains("@refresh reset")) {
+                    should_refresh = true
+                }
+            });
+
             comments.with_trailing(n.lo, |comments| {
                 if comments.iter().any(|c| c.text.contains("@refresh reset")) {
                     should_refresh = true
@@ -265,14 +271,18 @@ impl<C: Comments> VisitMut for Refresh<C> {
     // Does anyone write react without esmodule?
     // fn visit_mut_script(&mut self, _: &mut Script) {}
 
-    fn visit_mut_module_items(&mut self, module_items: &mut Vec<ModuleItem>) {
+    fn visit_mut_module(&mut self, n: &mut Module) {
         if !self.enable {
             return;
         }
 
         // to collect comments
-        self.visit_module_items(module_items);
+        self.visit_module(n);
+
+        self.visit_mut_module_items(&mut n.body);
+    }
 
+    fn visit_mut_module_items(&mut self, module_items: &mut Vec<ModuleItem>) {
         let used_in_jsx = collect_ident_in_jsx(module_items);
 
         let mut items = Vec::with_capacity(module_items.len());

crates/swc_ecma_transforms_react/src/refresh/tests.rs

@@ -1535,3 +1535,35 @@ const a = (a)=>{
 }
 "#
 );
+
+test!(
+    ::swc_ecma_parser::Syntax::Es(::swc_ecma_parser::EsConfig {
+        jsx: true,
+        ..Default::default()
+    }),
+    tr,
+    issue_6022,
+    r#"/* @refresh reset */
+    import { useState } from 'react';
+
+    function Counter() {
+      const [count, setCount] = useState(0);
+
+      return (
+        <button type="button" onClick={() => setCount(c => c + 1)}>{count}</button>
+      );
+    }
+"#,
+    r#"var _s = $RefreshSig$();
+import { useState } from 'react';
+function Counter() {
+    _s();
+    const [count, setCount] = useState(0);
+    return <button type="button" onClick={()=>setCount((c)=>c + 1)}>{count}</button>;
+}
+_s(Counter, "useState{[count, setCount](0)}", true);
+_c = Counter;
+var _c;
+$RefreshReg$(_c, "Counter");
+"#
+);