1//! Rust equivalent of the TypeScript `FindContextIdentifiers` pass.2//!3//! Determines which bindings need StoreContext/LoadContext semantics by4//! walking the AST with scope tracking to find variables that cross5//! function boundaries.67use rustc_hash::{FxHashMap, FxHashSet};89use react_compiler_ast::expressions::*;10use react_compiler_ast::patterns::*;11use react_compiler_ast::scope::*;12use react_compiler_ast::statements::FunctionDeclaration;13use react_compiler_ast::visitor::AstWalker;14use react_compiler_ast::visitor::Visitor;15use react_compiler_diagnostics::CompilerError;16use react_compiler_diagnostics::CompilerErrorDetail;17use react_compiler_diagnostics::ErrorCategory;18use react_compiler_diagnostics::Position;19use react_compiler_diagnostics::SourceLocation;20use react_compiler_hir::environment::Environment;2122use crate::FunctionNode;2324#[derive(Default)]25struct BindingInfo {26    reassigned: bool,27    reassigned_by_inner_fn: bool,28    referenced_by_inner_fn: bool,29}3031struct ContextIdentifierVisitor<'a> {32    scope_info: &'a ScopeInfo,33    env: &'a mut Environment,34    /// Stack of inner function scopes encountered during traversal.35    /// Empty when at the top level of the function being compiled.36    function_stack: Vec<ScopeId>,37    binding_info: FxHashMap<BindingId, BindingInfo>,38    error: Option<CompilerError>,39}4041impl<'a> ContextIdentifierVisitor<'a> {42    fn push_function_scope(&mut self, _start: Option<u32>, node_id: Option<u32>) {43        let scope = self.scope_info.resolve_scope_for_node(node_id);44        if let Some(scope) = scope {45            self.function_stack.push(scope);46        }47    }4849    fn pop_function_scope(&mut self, _start: Option<u32>, node_id: Option<u32>) {50        let has_scope = self.scope_info.resolve_scope_for_node(node_id);51        if has_scope.is_some() {52            self.function_stack.pop();53        }54    }5556    fn check_captured_reference(&mut self, _start: Option<u32>, node_id: Option<u32>) {57        let binding_id = match self.scope_info.resolve_reference_id_for_node(node_id) {58            Some(id) => id,59            None => return,60        };61        let &fn_scope = match self.function_stack.last() {62            Some(s) => s,63            None => return,64        };65        let binding = &self.scope_info.bindings[binding_id.0 as usize];66        if is_captured_by_function(self.scope_info, binding.scope, fn_scope) {67            let info = self.binding_info.entry(binding_id).or_default();68            info.referenced_by_inner_fn = true;69        }70    }7172    fn handle_reassignment_identifier(&mut self, name: &str, current_scope: ScopeId) {73        if let Some(binding_id) = self.scope_info.get_binding(current_scope, name) {74            let info = self.binding_info.entry(binding_id).or_default();75            info.reassigned = true;76            if let Some(&fn_scope) = self.function_stack.last() {77                let binding = &self.scope_info.bindings[binding_id.0 as usize];78                if is_captured_by_function(self.scope_info, binding.scope, fn_scope) {79                    info.reassigned_by_inner_fn = true;80                }81            }82        }83    }84}8586impl<'ast> Visitor<'ast> for ContextIdentifierVisitor<'_> {87    fn enter_function_declaration(&mut self, node: &'ast FunctionDeclaration, _: &[ScopeId]) {88        self.push_function_scope(node.base.start, node.base.node_id);89    }90    fn leave_function_declaration(&mut self, node: &'ast FunctionDeclaration, _: &[ScopeId]) {91        self.pop_function_scope(node.base.start, node.base.node_id);92    }93    fn enter_function_expression(&mut self, node: &'ast FunctionExpression, _: &[ScopeId]) {94        self.push_function_scope(node.base.start, node.base.node_id);95    }96    fn leave_function_expression(&mut self, node: &'ast FunctionExpression, _: &[ScopeId]) {97        self.pop_function_scope(node.base.start, node.base.node_id);98    }99    fn enter_arrow_function_expression(100        &mut self,101        node: &'ast ArrowFunctionExpression,102        _: &[ScopeId],103    ) {104        self.push_function_scope(node.base.start, node.base.node_id);105    }106    fn leave_arrow_function_expression(107        &mut self,108        node: &'ast ArrowFunctionExpression,109        _: &[ScopeId],110    ) {111        self.pop_function_scope(node.base.start, node.base.node_id);112    }113    fn enter_object_method(&mut self, node: &'ast ObjectMethod, _: &[ScopeId]) {114        self.push_function_scope(node.base.start, node.base.node_id);115    }116    fn leave_object_method(&mut self, node: &'ast ObjectMethod, _: &[ScopeId]) {117        self.pop_function_scope(node.base.start, node.base.node_id);118    }119120    fn enter_identifier(&mut self, node: &'ast Identifier, _scope_stack: &[ScopeId]) {121        self.check_captured_reference(node.base.start, node.base.node_id);122    }123124    fn enter_jsx_identifier(125        &mut self,126        node: &'ast react_compiler_ast::jsx::JSXIdentifier,127        _scope_stack: &[ScopeId],128    ) {129        self.check_captured_reference(node.base.start, node.base.node_id);130    }131132    fn enter_assignment_expression(133        &mut self,134        node: &'ast AssignmentExpression,135        scope_stack: &[ScopeId],136    ) {137        let current_scope = scope_stack138            .last()139            .copied()140            .unwrap_or(self.scope_info.program_scope);141        if self.error.is_none() {142            if let Err(error) = walk_lval_for_reassignment(self, &node.left, current_scope) {143                self.error = Some(error);144            }145        }146    }147148    fn enter_update_expression(&mut self, node: &'ast UpdateExpression, scope_stack: &[ScopeId]) {149        if let Expression::Identifier(ident) = node.argument.as_ref() {150            let current_scope = scope_stack151                .last()152                .copied()153                .unwrap_or(self.scope_info.program_scope);154            self.handle_reassignment_identifier(&ident.name, current_scope);155        }156    }157}158159/// Recursively walk an LVal pattern to find all reassignment target identifiers.160fn walk_lval_for_reassignment(161    visitor: &mut ContextIdentifierVisitor<'_>,162    pattern: &PatternLike,163    current_scope: ScopeId,164) -> Result<(), CompilerError> {165    match pattern {166        PatternLike::Identifier(ident) => {167            visitor.handle_reassignment_identifier(&ident.name, current_scope);168        }169        PatternLike::ArrayPattern(pat) => {170            for element in &pat.elements {171                if let Some(el) = element {172                    walk_lval_for_reassignment(visitor, el, current_scope)?;173                }174            }175        }176        PatternLike::ObjectPattern(pat) => {177            for prop in &pat.properties {178                match prop {179                    ObjectPatternProperty::ObjectProperty(p) => {180                        walk_lval_for_reassignment(visitor, &p.value, current_scope)?;181                    }182                    ObjectPatternProperty::RestElement(p) => {183                        walk_lval_for_reassignment(visitor, &p.argument, current_scope)?;184                    }185                }186            }187        }188        PatternLike::AssignmentPattern(pat) => {189            walk_lval_for_reassignment(visitor, &pat.left, current_scope)?;190        }191        PatternLike::RestElement(pat) => {192            walk_lval_for_reassignment(visitor, &pat.argument, current_scope)?;193        }194        PatternLike::MemberExpression(_) => {195            // Interior mutability - not a variable reassignment196        }197        PatternLike::TSAsExpression(node) => {198            record_unsupported_lval(199                visitor.env,200                "TSAsExpression",201                convert_opt_loc(&node.base.loc),202            )?;203        }204        PatternLike::TSSatisfiesExpression(node) => {205            record_unsupported_lval(206                visitor.env,207                "TSSatisfiesExpression",208                convert_opt_loc(&node.base.loc),209            )?;210        }211        PatternLike::TSNonNullExpression(node) => {212            record_unsupported_lval(213                visitor.env,214                "TSNonNullExpression",215                convert_opt_loc(&node.base.loc),216            )?;217        }218        PatternLike::TSTypeAssertion(node) => {219            record_unsupported_lval(220                visitor.env,221                "TSTypeAssertion",222                convert_opt_loc(&node.base.loc),223            )?;224        }225        PatternLike::TypeCastExpression(node) => {226            record_unsupported_lval(227                visitor.env,228                "TypeCastExpression",229                convert_opt_loc(&node.base.loc),230            )?;231        }232    }233    Ok(())234}235236fn convert_loc(loc: &react_compiler_ast::common::SourceLocation) -> SourceLocation {237    SourceLocation {238        start: Position {239            line: loc.start.line,240            column: loc.start.column,241            index: loc.start.index,242        },243        end: Position {244            line: loc.end.line,245            column: loc.end.column,246            index: loc.end.index,247        },248    }249}250251fn convert_opt_loc(252    loc: &Option<react_compiler_ast::common::SourceLocation>,253) -> Option<SourceLocation> {254    loc.as_ref().map(convert_loc)255}256257/// Record the TS-faithful Todo for an unsupported assignment-target wrapper258/// node, mirroring the TypeScript `FindContextIdentifiers` pass. TS throws259/// immediately (CompilerError.throwTodo in handleAssignment's default case),260/// aborting before BuildHIR ever runs or logs, so this must return Err rather261/// than record-and-continue: otherwise Rust emits HIR debug entries for a262/// function TS never lowered.263fn record_unsupported_lval(264    env: &mut Environment,265    type_name: &str,266    loc: Option<SourceLocation>,267) -> Result<(), CompilerError> {268    let _ = env;269    let mut err = CompilerError::new();270    err.push_error_detail(CompilerErrorDetail {271        category: ErrorCategory::Todo,272        reason: format!(273            "[FindContextIdentifiers] Cannot handle Object destructuring assignment target {type_name}"274        ),275        description: None,276        loc,277        suggestions: None,278    });279    Err(err)280}281282/// Check if a binding declared at `binding_scope` is captured by a function at `function_scope`.283/// Returns true if the binding is declared above the function (in the parent scope or higher).284fn is_captured_by_function(285    scope_info: &ScopeInfo,286    binding_scope: ScopeId,287    function_scope: ScopeId,288) -> bool {289    let fn_parent = match scope_info.scopes[function_scope.0 as usize].parent {290        Some(p) => p,291        None => return false,292    };293    if binding_scope == fn_parent {294        return true;295    }296    // Walk up from fn_parent to see if binding_scope is an ancestor297    let mut current = scope_info.scopes[fn_parent.0 as usize].parent;298    while let Some(scope_id) = current {299        if scope_id == binding_scope {300            return true;301        }302        current = scope_info.scopes[scope_id.0 as usize].parent;303    }304    false305}306307/// Build a set of `(BindingId, position)` pairs that are declaration sites308/// in `reference_to_binding`, not true references. Uses node-ID comparison309/// when available (from `ref_node_id_to_binding` + `declaration_node_id`),310/// falling back to position comparison otherwise.311/// Build a set of (BindingId, node_id) pairs for declaration sites in312/// ref_node_id_to_binding. These are entries where the reference's node_id313/// matches the binding's declaration_node_id — i.e., the "reference" is314/// actually the declaration itself.315fn build_declaration_node_ids(scope_info: &ScopeInfo) -> FxHashSet<(BindingId, u32)> {316    let mut result = FxHashSet::default();317    for (&ref_nid, &binding_id) in &scope_info.ref_node_id_to_binding {318        let binding = &scope_info.bindings[binding_id.0 as usize];319        if binding.declaration_node_id == Some(ref_nid) {320            result.insert((binding_id, ref_nid));321        }322    }323    result324}325326/// Find context identifiers for a function: variables that are captured across327/// function boundaries and need StoreContext/LoadContext semantics.328///329/// A binding is a context identifier if:330/// - It is reassigned from inside a nested function (`reassignedByInnerFn`), OR331/// - It is reassigned AND referenced from inside a nested function332///   (`reassigned && referencedByInnerFn`)333///334/// This is the Rust equivalent of the TypeScript `FindContextIdentifiers` pass.335pub fn find_context_identifiers(336    func: &FunctionNode<'_>,337    scope_info: &ScopeInfo,338    env: &mut Environment,339    identifier_locs: &crate::identifier_loc_index::IdentifierLocIndex,340) -> Result<FxHashSet<BindingId>, CompilerError> {341    let func_scope = scope_info342        .resolve_scope_for_node(func.node_id())343        .unwrap_or(scope_info.program_scope);344345    let mut visitor = ContextIdentifierVisitor {346        scope_info,347        env,348        function_stack: Vec::new(),349        binding_info: FxHashMap::default(),350        error: None,351    };352    let mut walker = AstWalker::with_initial_scope(scope_info, func_scope);353354    // Walk params and body (like Babel's func.traverse())355    match func {356        FunctionNode::FunctionDeclaration(d) => {357            for param in &d.params {358                walker.walk_pattern(&mut visitor, param);359            }360            walker.walk_block_statement(&mut visitor, &d.body);361        }362        FunctionNode::FunctionExpression(e) => {363            for param in &e.params {364                walker.walk_pattern(&mut visitor, param);365            }366            walker.walk_block_statement(&mut visitor, &e.body);367        }368        FunctionNode::ArrowFunctionExpression(a) => {369            for param in &a.params {370                walker.walk_pattern(&mut visitor, param);371            }372            match a.body.as_ref() {373                ArrowFunctionBody::BlockStatement(block) => {374                    walker.walk_block_statement(&mut visitor, block);375                }376                ArrowFunctionBody::Expression(expr) => {377                    walker.walk_expression(&mut visitor, expr);378                }379            }380        }381    }382383    if let Some(error) = visitor.error {384        return Err(error);385    }386387    // Supplement the walker-based analysis with referenceToBinding data.388    // The AST walker doesn't visit identifiers inside type annotations,389    // but Babel's traverse (used by TS findContextIdentifiers) does.390    // After scope extraction includes type annotation references,391    // we check if any reassigned binding has references in nested function scopes392    // via referenceToBinding — matching the TS behavior.393    //394    // We must skip declaration sites (e.g., the `x` in `function x() {}`),395    // which are included in reference_to_binding but are not true references.396    // Prefer node-ID comparison (immune to position-0 collisions from synthetic397    // nodes), falling back to position when node-IDs are unavailable.398    let declaration_node_ids = build_declaration_node_ids(scope_info);399    for (&ref_nid, &binding_id) in &scope_info.ref_node_id_to_binding {400        let info = match visitor.binding_info.get(&binding_id) {401            Some(info) if info.reassigned && !info.referenced_by_inner_fn => info,402            _ => continue,403        };404        let _ = info;405        if declaration_node_ids.contains(&(binding_id, ref_nid)) {406            continue;407        }408        // Get the reference's position from identifier_locs for range checks409        let ref_pos = match identifier_locs.get(&ref_nid) {410            Some(entry) => entry.start,411            None => continue,412        };413        let binding = &scope_info.bindings[binding_id.0 as usize];414        // Check if ref_pos is inside a nested function scope415        for (&scope_start, &scope_id) in &scope_info.node_to_scope {416            if scope_start <= ref_pos {417                if let Some(&scope_end) = scope_info.node_to_scope_end.get(&scope_start) {418                    if ref_pos < scope_end419                        && matches!(420                            scope_info.scopes[scope_id.0 as usize].kind,421                            ScopeKind::Function422                        )423                        && is_captured_by_function(scope_info, binding.scope, scope_id)424                    {425                        visitor426                            .binding_info427                            .get_mut(&binding_id)428                            .unwrap()429                            .referenced_by_inner_fn = true;430                        break;431                    }432                }433            }434        }435    }436437    // Collect results438    Ok(visitor439        .binding_info440        .into_iter()441        .filter(|(_, info)| {442            info.reassigned_by_inner_fn || (info.reassigned && info.referenced_by_inner_fn)443        })444        .map(|(id, _)| id)445        .collect())446}