react_compiler/
fixture_utils.rs

1use react_compiler_ast::File;
2use react_compiler_ast::declarations::{Declaration, ExportDefaultDecl};
3use react_compiler_ast::expressions::Expression;
4use react_compiler_ast::statements::Statement;
5use react_compiler_lowering::FunctionNode;
6
7/// Count the number of top-level functions in an AST file.
8///
9/// "Top-level" means:
10/// - FunctionDeclaration at program body level
11/// - FunctionExpression/ArrowFunctionExpression in a VariableDeclarator at program body level
12/// - FunctionDeclaration inside ExportNamedDeclaration
13/// - FunctionDeclaration/FunctionExpression/ArrowFunctionExpression inside ExportDefaultDeclaration
14/// - VariableDeclaration with function expressions inside ExportNamedDeclaration
15///
16/// This matches the TS test binary's traversal behavior.
17pub fn count_top_level_functions(ast: &File) -> usize {
18    let mut count = 0;
19    for stmt in &ast.program.body {
20        count += count_functions_in_statement(stmt);
21    }
22    count
23}
24
25fn count_functions_in_statement(stmt: &Statement) -> usize {
26    match stmt {
27        Statement::FunctionDeclaration(_) => 1,
28        Statement::VariableDeclaration(var_decl) => {
29            let mut count = 0;
30            for declarator in &var_decl.declarations {
31                if let Some(init) = &declarator.init {
32                    if is_function_expression(init) {
33                        count += 1;
34                    }
35                }
36            }
37            count
38        }
39        Statement::ExportNamedDeclaration(export) => {
40            if let Some(decl) = &export.declaration {
41                match decl.as_ref() {
42                    Declaration::FunctionDeclaration(_) => 1,
43                    Declaration::VariableDeclaration(var_decl) => {
44                        let mut count = 0;
45                        for declarator in &var_decl.declarations {
46                            if let Some(init) = &declarator.init {
47                                if is_function_expression(init) {
48                                    count += 1;
49                                }
50                            }
51                        }
52                        count
53                    }
54                    _ => 0,
55                }
56            } else {
57                0
58            }
59        }
60        Statement::ExportDefaultDeclaration(export) => match export.declaration.as_ref() {
61            ExportDefaultDecl::FunctionDeclaration(_) => 1,
62            ExportDefaultDecl::Expression(expr) => {
63                if is_function_expression(expr) {
64                    1
65                } else {
66                    0
67                }
68            }
69            _ => 0,
70        },
71        // Expression statements with function expressions (uncommon but possible)
72        Statement::ExpressionStatement(expr_stmt) => {
73            if is_function_expression(&expr_stmt.expression) {
74                1
75            } else {
76                0
77            }
78        }
79        _ => 0,
80    }
81}
82
83fn is_function_expression(expr: &Expression) -> bool {
84    matches!(
85        expr,
86        Expression::FunctionExpression(_) | Expression::ArrowFunctionExpression(_)
87    )
88}
89
90/// Extract the nth top-level function from an AST file as a `FunctionNode`.
91/// Also returns the inferred name (e.g. from a variable declarator).
92/// Returns None if function_index is out of bounds.
93pub fn extract_function(
94    ast: &File,
95    function_index: usize,
96) -> Option<(FunctionNode<'_>, Option<&str>)> {
97    let mut index = 0usize;
98
99    for stmt in &ast.program.body {
100        match stmt {
101            Statement::FunctionDeclaration(func_decl) => {
102                if index == function_index {
103                    let name = func_decl.id.as_ref().map(|id| id.name.as_str());
104                    return Some((FunctionNode::FunctionDeclaration(func_decl), name));
105                }
106                index += 1;
107            }
108            Statement::VariableDeclaration(var_decl) => {
109                for declarator in &var_decl.declarations {
110                    if let Some(init) = &declarator.init {
111                        match init.as_ref() {
112                            Expression::FunctionExpression(func) => {
113                                if index == function_index {
114                                    let name = match &declarator.id {
115                                        react_compiler_ast::patterns::PatternLike::Identifier(
116                                            ident,
117                                        ) => Some(ident.name.as_str()),
118                                        _ => func.id.as_ref().map(|id| id.name.as_str()),
119                                    };
120                                    return Some((FunctionNode::FunctionExpression(func), name));
121                                }
122                                index += 1;
123                            }
124                            Expression::ArrowFunctionExpression(arrow) => {
125                                if index == function_index {
126                                    let name = match &declarator.id {
127                                        react_compiler_ast::patterns::PatternLike::Identifier(
128                                            ident,
129                                        ) => Some(ident.name.as_str()),
130                                        _ => None,
131                                    };
132                                    return Some((
133                                        FunctionNode::ArrowFunctionExpression(arrow),
134                                        name,
135                                    ));
136                                }
137                                index += 1;
138                            }
139                            _ => {}
140                        }
141                    }
142                }
143            }
144            Statement::ExportNamedDeclaration(export) => {
145                if let Some(decl) = &export.declaration {
146                    match decl.as_ref() {
147                        Declaration::FunctionDeclaration(func_decl) => {
148                            if index == function_index {
149                                let name = func_decl.id.as_ref().map(|id| id.name.as_str());
150                                return Some((FunctionNode::FunctionDeclaration(func_decl), name));
151                            }
152                            index += 1;
153                        }
154                        Declaration::VariableDeclaration(var_decl) => {
155                            for declarator in &var_decl.declarations {
156                                if let Some(init) = &declarator.init {
157                                    match init.as_ref() {
158                                        Expression::FunctionExpression(func) => {
159                                            if index == function_index {
160                                                let name = match &declarator.id {
161                                                    react_compiler_ast::patterns::PatternLike::Identifier(ident) => Some(ident.name.as_str()),
162                                                    _ => func.id.as_ref().map(|id| id.name.as_str()),
163                                                };
164                                                return Some((
165                                                    FunctionNode::FunctionExpression(func),
166                                                    name,
167                                                ));
168                                            }
169                                            index += 1;
170                                        }
171                                        Expression::ArrowFunctionExpression(arrow) => {
172                                            if index == function_index {
173                                                let name = match &declarator.id {
174                                                    react_compiler_ast::patterns::PatternLike::Identifier(ident) => Some(ident.name.as_str()),
175                                                    _ => None,
176                                                };
177                                                return Some((
178                                                    FunctionNode::ArrowFunctionExpression(arrow),
179                                                    name,
180                                                ));
181                                            }
182                                            index += 1;
183                                        }
184                                        _ => {}
185                                    }
186                                }
187                            }
188                        }
189                        _ => {}
190                    }
191                }
192            }
193            Statement::ExportDefaultDeclaration(export) => match export.declaration.as_ref() {
194                ExportDefaultDecl::FunctionDeclaration(func_decl) => {
195                    if index == function_index {
196                        let name = func_decl.id.as_ref().map(|id| id.name.as_str());
197                        return Some((FunctionNode::FunctionDeclaration(func_decl), name));
198                    }
199                    index += 1;
200                }
201                ExportDefaultDecl::Expression(expr) => match expr.as_ref() {
202                    Expression::FunctionExpression(func) => {
203                        if index == function_index {
204                            let name = func.id.as_ref().map(|id| id.name.as_str());
205                            return Some((FunctionNode::FunctionExpression(func), name));
206                        }
207                        index += 1;
208                    }
209                    Expression::ArrowFunctionExpression(arrow) => {
210                        if index == function_index {
211                            return Some((FunctionNode::ArrowFunctionExpression(arrow), None));
212                        }
213                        index += 1;
214                    }
215                    _ => {}
216                },
217                _ => {}
218            },
219            Statement::ExpressionStatement(expr_stmt) => match expr_stmt.expression.as_ref() {
220                Expression::FunctionExpression(func) => {
221                    if index == function_index {
222                        let name = func.id.as_ref().map(|id| id.name.as_str());
223                        return Some((FunctionNode::FunctionExpression(func), name));
224                    }
225                    index += 1;
226                }
227                Expression::ArrowFunctionExpression(arrow) => {
228                    if index == function_index {
229                        return Some((FunctionNode::ArrowFunctionExpression(arrow), None));
230                    }
231                    index += 1;
232                }
233                _ => {}
234            },
235            _ => {}
236        }
237    }
238    None
239}
react_compiler/fixture_utils.rs

react_compiler/
fixture_utils.rs
