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 =
142                cm.new_source_file(FileName::Custom("test.tsx".into()).into(), code.to_string());
143
144            let program = parse_file_as_program(
145                &fm,
146                swc_core::ecma::parser::Syntax::Es(EsSyntax {
147                    jsx: true,
148                    ..Default::default()
149                }),
150                Default::default(),
151                Default::default(),
152                &mut vec![],
153            )
154            .unwrap();
155
156            assert_eq!(is_required(&program), required);
157
158            Ok(())
159        })
160        .unwrap();
161    }
162
163    #[test]
164    fn lazy_return() {
165        assert_required(
166            "
167            function Foo() {
168                const a = <div>Hello</div>;
169
170                return a
171            }
172            ",
173            true,
174        );
175
176        assert_required(
177            "
178            function Foo() {
179            ",
180            false,
181        );
182    }
183
184    #[test]
185    fn return_jsx() {
186        assert_required(
187            "
188            function Foo() {
189                return <div>Hello</div>;
190            }
191            ",
192            true,
193        );
194    }
195
196    #[test]
197    fn use_hooks() {
198        assert_required(
199            "
200            function Foo(props) {
201                const [a, b] = useState(0);
202
203                return props.children;
204            }
205            ",
206            true,
207        );
208    }
209
210    #[test]
211    fn arrow_function() {
212        assert_required(
213            "
214            const Foo = () => <div>Hello</div>;
215            ",
216            true,
217        );
218
219        assert_required(
220            "
221            const Foo = () => {
222                return <div>Hello</div>;
223            };
224            ",
225            true,
226        );
227    }
228
229    #[test]
230    fn export_const_arrow_function() {
231        assert_required(
232            "
233            export const Foo = () => <div>Hello</div>;
234            ",
235            true,
236        );
237
238        assert_required(
239            "
240            export const Foo = () => {
241                return <div>Hello</div>;
242            };
243            ",
244            true,
245        );
246    }
247
248    #[test]
249    fn normal_arrow_function() {
250        assert_required(
251            "
252            const Foo = () => {
253                const a = 1;
254                console.log(a);
255            };
256            ",
257            false,
258        );
259    }
260
261    #[test]
262    fn export_default_arrow_function() {
263        assert_required(
264            "
265            export default () => <div>Hello</div>;
266            ",
267            true,
268        );
269    }
270
271    #[test]
272    fn not_required_arrow_function() {
273        assert_required(
274            "
275            export default () => {
276                const a = 1;
277                console.log(a);
278            };
279            ",
280            false,
281        );
282    }
283}
next_custom_transforms/react_compiler.rs

next_custom_transforms/
react_compiler.rs
