From 36b26d3496fd9ad7e3ccb41c164600657476a8a3 Mon Sep 17 00:00:00 2001
From: HeYunfei <i.heyunfei@gmail.com>
Date: Wed, 4 Jan 2023 22:17:01 +0800
Subject: [PATCH] fix(es/react): Force refresh if code contains `@refresh
 reset` (#6749)

---
 .../src/refresh/mod.rs                        | 14 ++++++--
 .../src/refresh/tests.rs                      | 32 +++++++++++++++++++
 2 files changed, 44 insertions(+), 2 deletions(-)

diff --git a/crates/swc_ecma_transforms_react/src/refresh/mod.rs b/crates/swc_ecma_transforms_react/src/refresh/mod.rs
index 5194bb274b4a..58b423c2bdec 100644
--- a/crates/swc_ecma_transforms_react/src/refresh/mod.rs
+++ b/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());
diff --git a/crates/swc_ecma_transforms_react/src/refresh/tests.rs b/crates/swc_ecma_transforms_react/src/refresh/tests.rs
index b537aeb0046a..128037f21c4f 100644
--- a/crates/swc_ecma_transforms_react/src/refresh/tests.rs
+++ b/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");
+"#
+);