next_custom_transforms/
react_compiler.rs

1use swc_core::ecma::{
2    ast::{
3        Callee, ExportDefaultDecl, ExportDefaultExpr, Expr, FnDecl, FnExpr, Pat, Program, Stmt,
4        VarDeclarator,
5    },
6    visit::{Visit, VisitWith},
7};
8
9pub fn is_required(program: &Program) -> bool {
10    let mut finder = Finder::default();
11    finder.visit_program(program);
12    finder.found
13}
14
15#[derive(Default)]
16struct Finder {
17    found: bool,
18
19    /// We are in a function that starts with a capital letter or it's a function that starts with
20    /// `use`
21    is_interested: bool,
22}
23
24impl Visit for Finder {
25    fn visit_callee(&mut self, node: &Callee) {
26        if self.is_interested {
27            if let Callee::Expr(e) = node {
28                if let Expr::Ident(c) = &**e {
29                    if c.sym.starts_with("use") {
30                        self.found = true;
31                        return;
32                    }
33                }
34            }
35        }
36
37        node.visit_children_with(self);
38    }
39
40    fn visit_export_default_decl(&mut self, node: &ExportDefaultDecl) {
41        let old = self.is_interested;
42
43        self.is_interested = true;
44
45        node.visit_children_with(self);
46
47        self.is_interested = old;
48    }
49
50    fn visit_export_default_expr(&mut self, node: &ExportDefaultExpr) {
51        let old = self.is_interested;
52
53        self.is_interested = true;
54
55        node.visit_children_with(self);
56
57        self.is_interested = old;
58    }
59
60    fn visit_expr(&mut self, node: &Expr) {
61        if self.found {
62            return;
63        }
64        if self.is_interested
65            && matches!(
66                node,
67                Expr::JSXMember(..)
68                    | Expr::JSXNamespacedName(..)
69                    | Expr::JSXEmpty(..)
70                    | Expr::JSXElement(..)
71                    | Expr::JSXFragment(..)
72            )
73        {
74            self.found = true;
75            return;
76        }
77
78        node.visit_children_with(self);
79    }
80
81    fn visit_fn_decl(&mut self, node: &FnDecl) {
82        let old = self.is_interested;
83
84        self.is_interested = node.ident.sym.starts_with("use")
85            || node.ident.sym.starts_with(|c: char| c.is_ascii_uppercase());
86
87        node.visit_children_with(self);
88
89        self.is_interested = old;
90    }
91
92    fn visit_fn_expr(&mut self, node: &FnExpr) {
93        let old = self.is_interested;
94
95        self.is_interested |= node.ident.as_ref().is_some_and(|ident| {
96            ident.sym.starts_with("use") || ident.sym.starts_with(|c: char| c.is_ascii_uppercase())
97        });
98
99        node.visit_children_with(self);
100
101        self.is_interested = old;
102    }
103
104    fn visit_stmt(&mut self, node: &Stmt) {
105        if self.found {
106            return;
107        }
108        node.visit_children_with(self);
109    }
110
111    fn visit_var_declarator(&mut self, node: &VarDeclarator) {
112        let old = self.is_interested;
113
114        if matches!(node.init.as_deref(), Some(Expr::Fn(..) | Expr::Arrow(..))) {
115            if let Pat::Ident(ident) = &node.name {
116                self.is_interested = ident.sym.starts_with("use")
117                    || ident.sym.starts_with(|c: char| c.is_ascii_uppercase());
118            } else {
119                self.is_interested = false;
120            }
121        }
122
123        node.visit_children_with(self);
124
125        self.is_interested = old;
126    }
127}
128
129#[cfg(test)]
130mod tests {
131    use swc_core::{
132        common::FileName,
133        ecma::parser::{parse_file_as_program, EsSyntax},
134    };
135    use testing::run_test2;
136
137    use super::*;
138
139    fn assert_required(code: &str, required: bool) {
140        run_test2(false, |cm, _| {
141            let fm = cm.new_source_file(FileName::Custom("test.tsx".into()).into(), code.into());
142
143            let program = parse_file_as_program(
144                &fm,
145                swc_core::ecma::parser::Syntax::Es(EsSyntax {
146                    jsx: true,
147                    ..Default::default()
148                }),
149                Default::default(),
150                Default::default(),
151                &mut vec![],
152            )
153            .unwrap();
154
155            assert_eq!(is_required(&program), required);
156
157            Ok(())
158        })
159        .unwrap();
160    }
161
162    #[test]
163    fn lazy_return() {
164        assert_required(
165            "
166            function Foo() {
167                const a = <div>Hello</div>;
168
169                return a
170            }
171            ",
172            true,
173        );
174
175        assert_required(
176            "
177            function Foo() {
178            ",
179            false,
180        );
181    }
182
183    #[test]
184    fn return_jsx() {
185        assert_required(
186            "
187            function Foo() {
188                return <div>Hello</div>;
189            }
190            ",
191            true,
192        );
193    }
194
195    #[test]
196    fn use_hooks() {
197        assert_required(
198            "
199            function Foo(props) {
200                const [a, b] = useState(0);
201
202                return props.children;
203            }
204            ",
205            true,
206        );
207    }
208
209    #[test]
210    fn arrow_function() {
211        assert_required(
212            "
213            const Foo = () => <div>Hello</div>;
214            ",
215            true,
216        );
217
218        assert_required(
219            "
220            const Foo = () => {
221                return <div>Hello</div>;
222            };
223            ",
224            true,
225        );
226    }
227
228    #[test]
229    fn export_const_arrow_function() {
230        assert_required(
231            "
232            export const Foo = () => <div>Hello</div>;
233            ",
234            true,
235        );
236
237        assert_required(
238            "
239            export const Foo = () => {
240                return <div>Hello</div>;
241            };
242            ",
243            true,
244        );
245    }
246
247    #[test]
248    fn normal_arrow_function() {
249        assert_required(
250            "
251            const Foo = () => {
252                const a = 1;
253                console.log(a);
254            };
255            ",
256            false,
257        );
258    }
259
260    #[test]
261    fn export_default_arrow_function() {
262        assert_required(
263            "
264            export default () => <div>Hello</div>;
265            ",
266            true,
267        );
268    }
269
270    #[test]
271    fn not_required_arrow_function() {
272        assert_required(
273            "
274            export default () => {
275                const a = 1;
276                console.log(a);
277            };
278            ",
279            false,
280        );
281    }
282}