fix(react-compiler): Fix fast check (#10538)

Author: kdy1

.changeset/nice-grapes-mix.md

@@ -0,0 +1,6 @@
+---
+swc_core: patch
+swc_ecma_react_compiler: patch
+---
+
+fix(react-compiler): Fix fast check

Cargo.lock

@@ -17,3 +17,7 @@ swc_atoms      = { version = "5.0.0", path = "../swc_atoms" }
 swc_common     = { version = "11.0.0", path = "../swc_common" }
 swc_ecma_ast   = { version = "11.0.0", path = "../swc_ecma_ast" }
 swc_ecma_visit = { version = "11.0.0", path = "../swc_ecma_visit" }
+
+[dev-dependencies]
+swc_ecma_parser = { version = "14.0.0", path = "../swc_ecma_parser" }
+testing         = { version = "12.0.0", path = "../testing" }

crates/swc_ecma_react_compiler/Cargo.toml

@@ -1,4 +1,7 @@
-use swc_ecma_ast::{Callee, Expr, FnDecl, FnExpr, Pat, Program, ReturnStmt, Stmt, VarDeclarator};
+use swc_ecma_ast::{
+    Callee, ExportDefaultDecl, ExportDefaultExpr, Expr, FnDecl, FnExpr, Pat, Program, Stmt,
+    VarDeclarator,
+};
 use swc_ecma_visit::{Visit, VisitWith};
 pub fn is_required(program: &Program) -> bool {
     let mut finder = Finder::default();
@@ -31,18 +34,40 @@ impl Visit for Finder {
         node.visit_children_with(self);
     }
 
+    fn visit_export_default_decl(&mut self, node: &ExportDefaultDecl) {
+        let old = self.is_interested;
+
+        self.is_interested = true;
+
+        node.visit_children_with(self);
+
+        self.is_interested = old;
+    }
+
+    fn visit_export_default_expr(&mut self, node: &ExportDefaultExpr) {
+        let old = self.is_interested;
+
+        self.is_interested = true;
+
+        node.visit_children_with(self);
+
+        self.is_interested = old;
+    }
+
     fn visit_expr(&mut self, node: &Expr) {
         if self.found {
             return;
         }
-        if matches!(
-            node,
-            Expr::JSXMember(..)
-                | Expr::JSXNamespacedName(..)
-                | Expr::JSXEmpty(..)
-                | Expr::JSXElement(..)
-                | Expr::JSXFragment(..)
-        ) {
+        if self.is_interested
+            && matches!(
+                node,
+                Expr::JSXMember(..)
+                    | Expr::JSXNamespacedName(..)
+                    | Expr::JSXEmpty(..)
+                    | Expr::JSXElement(..)
+                    | Expr::JSXFragment(..)
+            )
+        {
             self.found = true;
             return;
         }
@@ -52,6 +77,7 @@ impl Visit for Finder {
 
     fn visit_fn_decl(&mut self, node: &FnDecl) {
         let old = self.is_interested;
+
         self.is_interested = node.ident.sym.starts_with("use")
             || node.ident.sym.starts_with(|c: char| c.is_ascii_uppercase());
 
@@ -72,19 +98,6 @@ impl Visit for Finder {
         self.is_interested = old;
     }
 
-    fn visit_return_stmt(&mut self, node: &ReturnStmt) {
-        if self.is_interested {
-            if let Some(Expr::JSXElement(..) | Expr::JSXEmpty(..) | Expr::JSXFragment(..)) =
-                node.arg.as_deref()
-            {
-                self.found = true;
-                return;
-            }
-        }
-
-        node.visit_children_with(self);
-    }
-
     fn visit_stmt(&mut self, node: &Stmt) {
         if self.found {
             return;
@@ -95,15 +108,170 @@ impl Visit for Finder {
     fn visit_var_declarator(&mut self, node: &VarDeclarator) {
         let old = self.is_interested;
 
-        if let Pat::Ident(ident) = &node.name {
-            self.is_interested = ident.sym.starts_with("use")
-                || ident.sym.starts_with(|c: char| c.is_ascii_uppercase());
-        } else {
-            self.is_interested = false;
+        if matches!(node.init.as_deref(), Some(Expr::Fn(..) | Expr::Arrow(..))) {
+            if let Pat::Ident(ident) = &node.name {
+                self.is_interested = ident.sym.starts_with("use")
+                    || ident.sym.starts_with(|c: char| c.is_ascii_uppercase());
+            } else {
+                self.is_interested = false;
+            }
         }
 
         node.visit_children_with(self);
 
         self.is_interested = old;
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use swc_common::FileName;
+    use swc_ecma_parser::{parse_file_as_program, EsSyntax, Syntax};
+    use testing::run_test2;
+
+    use super::*;
+
+    fn assert_required(code: &str, required: bool) {
+        run_test2(false, |cm, _| {
+            let fm = cm.new_source_file(FileName::Custom("test.tsx".into()).into(), code.into());
+
+            let program = parse_file_as_program(
+                &fm,
+                Syntax::Es(EsSyntax {
+                    jsx: true,
+                    ..Default::default()
+                }),
+                Default::default(),
+                Default::default(),
+                &mut vec![],
+            )
+            .unwrap();
+
+            assert_eq!(is_required(&program), required);
+
+            Ok(())
+        })
+        .unwrap();
+    }
+
+    #[test]
+    fn lazy_return() {
+        assert_required(
+            "
+            function Foo() {
+                const a = <div>Hello</div>;
+
+                return a
+            }
+            ",
+            true,
+        );
+
+        assert_required(
+            "
+            function Foo() {
+            ",
+            false,
+        );
+    }
+
+    #[test]
+    fn return_jsx() {
+        assert_required(
+            "
+            function Foo() {
+                return <div>Hello</div>;
+            }
+            ",
+            true,
+        );
+    }
+
+    #[test]
+    fn use_hooks() {
+        assert_required(
+            "
+            function Foo(props) {
+                const [a, b] = useState(0);
+
+                return props.children;
+            }
+            ",
+            true,
+        );
+    }
+
+    #[test]
+    fn arrow_function() {
+        assert_required(
+            "
+            const Foo = () => <div>Hello</div>;
+            ",
+            true,
+        );
+
+        assert_required(
+            "
+            const Foo = () => {
+                return <div>Hello</div>;
+            };
+            ",
+            true,
+        );
+    }
+
+    #[test]
+    fn export_const_arrow_function() {
+        assert_required(
+            "
+            export const Foo = () => <div>Hello</div>;
+            ",
+            true,
+        );
+
+        assert_required(
+            "
+            export const Foo = () => {
+                return <div>Hello</div>;
+            };
+            ",
+            true,
+        );
+    }
+
+    #[test]
+    fn normal_arrow_function() {
+        assert_required(
+            "
+            const Foo = () => {
+                const a = 1;
+                console.log(a);
+            };
+            ",
+            false,
+        );
+    }
+
+    #[test]
+    fn export_default_arrow_function() {
+        assert_required(
+            "
+            export default () => <div>Hello</div>;
+            ",
+            true,
+        );
+    }
+
+    #[test]
+    fn not_required_arrow_function() {
+        assert_required(
+            "
+            export default () => {
+                const a = 1;
+                console.log(a);
+            };
+            ",
+            false,
+        );
+    }
+}