| 1 | //! Builds an index mapping identifier node-IDs to source locations. |
| 2 | //! |
| 3 | //! Walks the function's AST to collect `(node_id, start, SourceLocation, is_jsx)` |
| 4 | //! for every Identifier and JSXIdentifier node. Keyed by node_id for identity |
| 5 | //! lookups; each entry also stores `start` (byte offset) for range-containment |
| 6 | //! checks in `gather_captured_context`. |
| 7 | |
| 8 | use rustc_hash::FxHashMap; |
| 9 | |
| 10 | use react_compiler_ast::expressions::*; |
| 11 | use react_compiler_ast::jsx::JSXIdentifier; |
| 12 | use react_compiler_ast::jsx::JSXOpeningElement; |
| 13 | use react_compiler_ast::scope::ScopeId; |
| 14 | use react_compiler_ast::scope::ScopeInfo; |
| 15 | use react_compiler_ast::statements::FunctionDeclaration; |
| 16 | use react_compiler_ast::visitor::AstWalker; |
| 17 | use react_compiler_ast::visitor::Visitor; |
| 18 | use react_compiler_hir::SourceLocation; |
| 19 | |
| 20 | use crate::FunctionNode; |
| 21 | |
| 22 | /// Source location and whether the identifier is a JSXIdentifier. |
| 23 | pub struct IdentifierLocEntry { |
| 24 | /// The byte offset of the identifier (base.start). Stored here so that |
| 25 | /// callers iterating by node_id can still do position-range containment |
| 26 | /// checks without a separate bridge map. |
| 27 | pub start: u32, |
| 28 | pub loc: SourceLocation, |
| 29 | pub is_jsx: bool, |
| 30 | /// For JSX identifiers that are the root name of a JSXOpeningElement, |
| 31 | /// stores the JSXOpeningElement's loc (which spans the full tag). |
| 32 | pub opening_element_loc: Option<SourceLocation>, |
| 33 | /// True if this identifier is the name of a function/class declaration |
| 34 | /// (not an expression reference). Used by `gather_captured_context` to |
| 35 | /// skip non-expression positions, matching the TS behavior where the |
| 36 | /// Expression visitor doesn't visit declaration names. |
| 37 | pub is_declaration_name: bool, |
| 38 | /// True if this identifier sits inside a type annotation subtree |
| 39 | /// (TypeAnnotation/TSTypeAnnotation/TypeAlias/TSTypeAliasDeclaration). |
| 40 | /// `gather_captured_context` skips these to match the TS |
| 41 | /// gatherCapturedContext traverse, which skips those subtrees; the |
| 42 | /// hoisting analysis and FindContextIdentifiers do NOT skip them in TS. |
| 43 | pub in_type_annotation: bool, |
| 44 | } |
| 45 | |
| 46 | /// Index mapping node_id → IdentifierLocEntry for all Identifier |
| 47 | /// and JSXIdentifier nodes in a function's AST. |
| 48 | pub type IdentifierLocIndex = FxHashMap<u32, IdentifierLocEntry>; |
| 49 | |
| 50 | struct IdentifierLocVisitor { |
| 51 | index: IdentifierLocIndex, |
| 52 | /// Tracks the current JSXOpeningElement's loc while walking its name. |
| 53 | current_opening_element_loc: Option<SourceLocation>, |
| 54 | } |
| 55 | |
| 56 | fn convert_loc(loc: &react_compiler_ast::common::SourceLocation) -> SourceLocation { |
| 57 | SourceLocation { |
| 58 | start: react_compiler_hir::Position { |
| 59 | line: loc.start.line, |
| 60 | column: loc.start.column, |
| 61 | index: loc.start.index, |
| 62 | }, |
| 63 | end: react_compiler_hir::Position { |
| 64 | line: loc.end.line, |
| 65 | column: loc.end.column, |
| 66 | index: loc.end.index, |
| 67 | }, |
| 68 | } |
| 69 | } |
| 70 | |
| 71 | impl IdentifierLocVisitor { |
| 72 | fn insert_identifier(&mut self, node: &Identifier, is_declaration_name: bool) { |
| 73 | if let (Some(nid), Some(start), Some(loc)) = |
| 74 | (node.base.node_id, node.base.start, &node.base.loc) |
| 75 | { |
| 76 | self.index.insert( |
| 77 | nid, |
| 78 | IdentifierLocEntry { |
| 79 | start, |
| 80 | loc: convert_loc(loc), |
| 81 | is_jsx: false, |
| 82 | opening_element_loc: None, |
| 83 | is_declaration_name, |
| 84 | in_type_annotation: false, |
| 85 | }, |
| 86 | ); |
| 87 | } |
| 88 | } |
| 89 | |
| 90 | /// Recursively walk a serde_json::Value tree to find and index all Identifier |
| 91 | /// and JSXIdentifier nodes. Used for class bodies which are stored as untyped |
| 92 | /// JSON and not walked by the typed AstWalker. This matches the TS behavior |
| 93 | /// where gatherCapturedContext's Babel traverse walks into class bodies. |
| 94 | /// |
| 95 | /// `in_annotation` is true once the walk has descended through a type |
| 96 | /// annotation container node; identifiers found there are flagged so |
| 97 | /// `gather_captured_context` can mirror TS's TypeAnnotation subtree skip. |
| 98 | fn walk_json_for_identifiers(&mut self, value: &serde_json::Value, in_annotation: bool) { |
| 99 | match value { |
| 100 | serde_json::Value::Object(obj) => { |
| 101 | let node_in_annotation = in_annotation |
| 102 | || matches!( |
| 103 | obj.get("type").and_then(|t| t.as_str()), |
| 104 | Some( |
| 105 | "TypeAnnotation" |
| 106 | | "TSTypeAnnotation" |
| 107 | | "TypeAlias" |
| 108 | | "TSTypeAliasDeclaration" |
| 109 | ) |
| 110 | ); |
| 111 | if let Some(serde_json::Value::String(ty)) = obj.get("type") { |
| 112 | if ty == "Identifier" || ty == "JSXIdentifier" { |
| 113 | if let (Some(nid), Some(start)) = ( |
| 114 | obj.get("_nodeId").and_then(|s| s.as_u64()), |
| 115 | obj.get("start").and_then(|s| s.as_u64()), |
| 116 | ) { |
| 117 | if let Some(loc) = Self::extract_loc_from_json(obj) { |
| 118 | let is_jsx = ty == "JSXIdentifier"; |
| 119 | self.index.entry(nid as u32).or_insert(IdentifierLocEntry { |
| 120 | start: start as u32, |
| 121 | loc, |
| 122 | is_jsx, |
| 123 | opening_element_loc: None, |
| 124 | is_declaration_name: false, |
| 125 | in_type_annotation: node_in_annotation, |
| 126 | }); |
| 127 | } |
| 128 | } |
| 129 | } |
| 130 | } |
| 131 | for (_, v) in obj { |
| 132 | self.walk_json_for_identifiers(v, node_in_annotation); |
| 133 | } |
| 134 | } |
| 135 | serde_json::Value::Array(arr) => { |
| 136 | for v in arr { |
| 137 | self.walk_json_for_identifiers(v, in_annotation); |
| 138 | } |
| 139 | } |
| 140 | _ => {} |
| 141 | } |
| 142 | } |
| 143 | |
| 144 | fn extract_loc_from_json( |
| 145 | obj: &serde_json::Map<String, serde_json::Value>, |
| 146 | ) -> Option<SourceLocation> { |
| 147 | let loc = obj.get("loc")?.as_object()?; |
| 148 | let start = loc.get("start")?.as_object()?; |
| 149 | let end = loc.get("end")?.as_object()?; |
| 150 | Some(SourceLocation { |
| 151 | start: react_compiler_hir::Position { |
| 152 | line: start.get("line")?.as_u64()? as u32, |
| 153 | column: start.get("column")?.as_u64()? as u32, |
| 154 | index: start |
| 155 | .get("index") |
| 156 | .and_then(|i| i.as_u64()) |
| 157 | .map(|i| i as u32), |
| 158 | }, |
| 159 | end: react_compiler_hir::Position { |
| 160 | line: end.get("line")?.as_u64()? as u32, |
| 161 | column: end.get("column")?.as_u64()? as u32, |
| 162 | index: end.get("index").and_then(|i| i.as_u64()).map(|i| i as u32), |
| 163 | }, |
| 164 | }) |
| 165 | } |
| 166 | } |
| 167 | |
| 168 | impl<'ast> Visitor<'ast> for IdentifierLocVisitor { |
| 169 | fn enter_identifier(&mut self, node: &'ast Identifier, _scope_stack: &[ScopeId]) { |
| 170 | self.insert_identifier(node, false); |
| 171 | } |
| 172 | |
| 173 | fn enter_jsx_identifier(&mut self, node: &'ast JSXIdentifier, _scope_stack: &[ScopeId]) { |
| 174 | if let (Some(nid), Some(start), Some(loc)) = |
| 175 | (node.base.node_id, node.base.start, &node.base.loc) |
| 176 | { |
| 177 | self.index.insert( |
| 178 | nid, |
| 179 | IdentifierLocEntry { |
| 180 | start, |
| 181 | loc: convert_loc(loc), |
| 182 | is_jsx: true, |
| 183 | opening_element_loc: self.current_opening_element_loc.clone(), |
| 184 | is_declaration_name: false, |
| 185 | in_type_annotation: false, |
| 186 | }, |
| 187 | ); |
| 188 | } |
| 189 | } |
| 190 | |
| 191 | fn enter_jsx_opening_element( |
| 192 | &mut self, |
| 193 | node: &'ast JSXOpeningElement, |
| 194 | _scope_stack: &[ScopeId], |
| 195 | ) { |
| 196 | self.current_opening_element_loc = node.base.loc.as_ref().map(|loc| convert_loc(loc)); |
| 197 | } |
| 198 | |
| 199 | fn leave_jsx_opening_element( |
| 200 | &mut self, |
| 201 | _node: &'ast JSXOpeningElement, |
| 202 | _scope_stack: &[ScopeId], |
| 203 | ) { |
| 204 | self.current_opening_element_loc = None; |
| 205 | } |
| 206 | |
| 207 | // Visit function/class declaration and expression name identifiers, |
| 208 | // which are not walked by the generic walker (to avoid affecting |
| 209 | // other Visitor consumers like find_context_identifiers). |
| 210 | fn enter_function_declaration( |
| 211 | &mut self, |
| 212 | node: &'ast FunctionDeclaration, |
| 213 | _scope_stack: &[ScopeId], |
| 214 | ) { |
| 215 | if let Some(id) = &node.id { |
| 216 | self.insert_identifier(id, true); |
| 217 | } |
| 218 | } |
| 219 | |
| 220 | fn enter_function_expression( |
| 221 | &mut self, |
| 222 | node: &'ast FunctionExpression, |
| 223 | _scope_stack: &[ScopeId], |
| 224 | ) { |
| 225 | if let Some(id) = &node.id { |
| 226 | self.insert_identifier(id, true); |
| 227 | } |
| 228 | } |
| 229 | |
| 230 | fn enter_class_declaration( |
| 231 | &mut self, |
| 232 | node: &'ast react_compiler_ast::statements::ClassDeclaration, |
| 233 | _scope_stack: &[ScopeId], |
| 234 | ) { |
| 235 | if let Some(id) = &node.id { |
| 236 | self.insert_identifier(id, true); |
| 237 | } |
| 238 | // Walk class body JSON to index identifiers inside class methods. |
| 239 | // The typed AstWalker skips class bodies (stored as Vec<serde_json::Value>), |
| 240 | // but gatherCapturedContext in TS traverses them via Babel's traverse. |
| 241 | for member in &node.body.body { |
| 242 | self.walk_json_for_identifiers(&member.parse_value(), false); |
| 243 | } |
| 244 | } |
| 245 | |
| 246 | fn enter_class_expression( |
| 247 | &mut self, |
| 248 | node: &'ast react_compiler_ast::expressions::ClassExpression, |
| 249 | _scope_stack: &[ScopeId], |
| 250 | ) { |
| 251 | if let Some(id) = &node.id { |
| 252 | self.insert_identifier(id, true); |
| 253 | } |
| 254 | // Walk class body JSON to index identifiers inside class methods |
| 255 | for member in &node.body.body { |
| 256 | self.walk_json_for_identifiers(&member.parse_value(), false); |
| 257 | } |
| 258 | } |
| 259 | } |
| 260 | |
| 261 | /// Build an index of all Identifier and JSXIdentifier positions in a function's AST. |
| 262 | pub fn build_identifier_loc_index( |
| 263 | func: &FunctionNode<'_>, |
| 264 | scope_info: &ScopeInfo, |
| 265 | ) -> IdentifierLocIndex { |
| 266 | let func_scope = scope_info |
| 267 | .resolve_scope_for_node(func.node_id()) |
| 268 | .unwrap_or(scope_info.program_scope); |
| 269 | |
| 270 | let mut visitor = IdentifierLocVisitor { |
| 271 | index: FxHashMap::default(), |
| 272 | current_opening_element_loc: None, |
| 273 | }; |
| 274 | let mut walker = AstWalker::with_initial_scope(scope_info, func_scope); |
| 275 | |
| 276 | // Visit the top-level function's own name identifier (if any), |
| 277 | // since the walker only walks params + body, not the function node itself. |
| 278 | match func { |
| 279 | FunctionNode::FunctionDeclaration(d) => { |
| 280 | if let Some(id) = &d.id { |
| 281 | visitor.enter_identifier(id, &[]); |
| 282 | } |
| 283 | for param in &d.params { |
| 284 | walker.walk_pattern(&mut visitor, param); |
| 285 | } |
| 286 | walker.walk_block_statement(&mut visitor, &d.body); |
| 287 | } |
| 288 | FunctionNode::FunctionExpression(e) => { |
| 289 | if let Some(id) = &e.id { |
| 290 | visitor.enter_identifier(id, &[]); |
| 291 | } |
| 292 | for param in &e.params { |
| 293 | walker.walk_pattern(&mut visitor, param); |
| 294 | } |
| 295 | walker.walk_block_statement(&mut visitor, &e.body); |
| 296 | } |
| 297 | FunctionNode::ArrowFunctionExpression(a) => { |
| 298 | for param in &a.params { |
| 299 | walker.walk_pattern(&mut visitor, param); |
| 300 | } |
| 301 | match a.body.as_ref() { |
| 302 | ArrowFunctionBody::BlockStatement(block) => { |
| 303 | walker.walk_block_statement(&mut visitor, block); |
| 304 | } |
| 305 | ArrowFunctionBody::Expression(expr) => { |
| 306 | walker.walk_expression(&mut visitor, expr); |
| 307 | } |
| 308 | } |
| 309 | } |
| 310 | } |
| 311 | |
| 312 | // Walk type annotations that the AST walker skips. |
| 313 | // The walker skips TypeAlias, TSTypeAliasDeclaration, and similar statements, |
| 314 | // but Babel's isReferencedIdentifier() returns true for identifiers inside them |
| 315 | // (e.g., typeof x in `type T = ReturnType<typeof x>`). The TS compiler's |
| 316 | // FindContextIdentifiers includes these via its Identifier visitor. We match by |
| 317 | // serializing the function body to JSON and walking the full JSON tree. |
| 318 | // The walk_json_for_identifiers method uses entry().or_insert() so it won't |
| 319 | // overwrite entries already added by the typed walker above. |
| 320 | let body_json: Option<serde_json::Value> = match func { |
| 321 | FunctionNode::FunctionDeclaration(d) => serde_json::to_value(&d.body).ok(), |
| 322 | FunctionNode::FunctionExpression(e) => serde_json::to_value(&e.body).ok(), |
| 323 | FunctionNode::ArrowFunctionExpression(a) => serde_json::to_value(&a.body).ok(), |
| 324 | }; |
| 325 | if let Some(json) = body_json { |
| 326 | visitor.walk_json_for_identifiers(&json, false); |
| 327 | } |
| 328 | |
| 329 | visitor.index |
| 330 | } |