Skip to main content

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