react_compiler_lowering/
identifier_loc_index.rs

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
8use std::collections::HashMap;
9
10use react_compiler_ast::expressions::*;
11use react_compiler_ast::jsx::JSXIdentifier;
12use react_compiler_ast::jsx::JSXOpeningElement;
13use react_compiler_ast::scope::ScopeId;
14use react_compiler_ast::scope::ScopeInfo;
15use react_compiler_ast::statements::FunctionDeclaration;
16use react_compiler_ast::visitor::AstWalker;
17use react_compiler_ast::visitor::Visitor;
18use react_compiler_hir::SourceLocation;
19
20use crate::FunctionNode;
21
22/// Source location and whether the identifier is a JSXIdentifier.
23pub 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.
48pub type IdentifierLocIndex = HashMap<u32, IdentifierLocEntry>;
49
50struct IdentifierLocVisitor {
51    index: IdentifierLocIndex,
52    /// Tracks the current JSXOpeningElement's loc while walking its name.
53    current_opening_element_loc: Option<SourceLocation>,
54}
55
56fn 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
71impl 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
168impl<'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, 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, false);
257        }
258    }
259}
260
261/// Build an index of all Identifier and JSXIdentifier positions in a function's AST.
262pub 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: HashMap::new(),
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}
react_compiler_lowering/identifier_loc_index.rs

react_compiler_lowering/
identifier_loc_index.rs
