1// Copyright (c) Meta Platforms, Inc. and affiliates.2//3// This source code is licensed under the MIT license found in the4// LICENSE file in the root directory of this source tree.56//! Aligns scopes of object method values to that of their enclosing object expressions.7//! To produce a well-formed JS program in Codegen, object methods and object expressions8//! must be in the same ReactiveBlock as object method definitions must be inlined.9//!10//! Ported from TypeScript `src/ReactiveScopes/AlignObjectMethodScopes.ts`.1112use rustc_hash::{FxHashMap, FxHashSet};13use std::cmp;1415use react_compiler_hir::environment::Environment;16use react_compiler_hir::{17    EvaluationOrder, HirFunction, IdentifierId, InstructionValue, ObjectPropertyOrSpread, ScopeId,18};19use react_compiler_utils::DisjointSet;2021// =============================================================================22// findScopesToMerge23// =============================================================================2425/// Identifies ObjectMethod lvalue identifiers and then finds ObjectExpression26/// instructions whose operands reference those methods. Returns a disjoint set27/// of scopes that must be merged.28fn find_scopes_to_merge(func: &HirFunction, env: &Environment) -> DisjointSet<ScopeId> {29    let mut object_method_decls: FxHashSet<IdentifierId> = FxHashSet::default();30    let mut merged_scopes = DisjointSet::<ScopeId>::new();3132    for (_block_id, block) in &func.body.blocks {33        for &instr_id in &block.instructions {34            let instr = &func.instructions[instr_id.0 as usize];35            match &instr.value {36                InstructionValue::ObjectMethod { .. } => {37                    object_method_decls.insert(instr.lvalue.identifier);38                }39                InstructionValue::ObjectExpression { properties, .. } => {40                    for prop_or_spread in properties {41                        let operand_place = match prop_or_spread {42                            ObjectPropertyOrSpread::Property(prop) => &prop.place,43                            ObjectPropertyOrSpread::Spread(spread) => &spread.place,44                        };45                        if object_method_decls.contains(&operand_place.identifier) {46                            let operand_scope =47                                env.identifiers[operand_place.identifier.0 as usize].scope;48                            let lvalue_scope =49                                env.identifiers[instr.lvalue.identifier.0 as usize].scope;5051                            // TS: CompilerError.invariant(operandScope != null && lvalueScope != null, ...)52                            let operand_sid = operand_scope.expect(53                                "Internal error: Expected all ObjectExpressions and ObjectMethods to have non-null scope.",54                            );55                            let lvalue_sid = lvalue_scope.expect(56                                "Internal error: Expected all ObjectExpressions and ObjectMethods to have non-null scope.",57                            );58                            merged_scopes.union(&[operand_sid, lvalue_sid]);59                        }60                    }61                }62                _ => {}63            }64        }65    }6667    merged_scopes68}6970// =============================================================================71// Public API72// =============================================================================7374/// Aligns object method scopes so that ObjectMethod values and their enclosing75/// ObjectExpression share the same scope.76///77/// Corresponds to TS `alignObjectMethodScopes(fn: HIRFunction): void`.78pub fn align_object_method_scopes(func: &mut HirFunction, env: &mut Environment) {79    // Handle inner functions first (TS recurses before processing the outer function)80    for (_block_id, block) in &func.body.blocks {81        for &instr_id in &block.instructions {82            let instr = &func.instructions[instr_id.0 as usize];83            match &instr.value {84                InstructionValue::FunctionExpression { lowered_func, .. }85                | InstructionValue::ObjectMethod { lowered_func, .. } => {86                    let func_id = lowered_func.func;87                    let mut inner_func = std::mem::replace(88                        &mut env.functions[func_id.0 as usize],89                        react_compiler_ssa::enter_ssa::placeholder_function(),90                    );91                    align_object_method_scopes(&mut inner_func, env);92                    env.functions[func_id.0 as usize] = inner_func;93                }94                _ => {}95            }96        }97    }9899    let mut merged_scopes = find_scopes_to_merge(func, env);100101    // Step 1: Merge affected scopes to their canonical root.102    // Use a FxHashMap to accumulate min/max across all scopes mapping to the same root,103    // matching TS behavior where root.range is updated in-place during iteration.104    let mut range_updates: FxHashMap<ScopeId, (EvaluationOrder, EvaluationOrder)> =105        FxHashMap::default();106107    merged_scopes.for_each(|scope_id, root_id| {108        if scope_id == root_id {109            return;110        }111        let scope_range = env.scopes[scope_id.0 as usize].range.clone();112        let root_range = env.scopes[root_id.0 as usize].range.clone();113114        let entry = range_updates115            .entry(root_id)116            .or_insert_with(|| (root_range.start, root_range.end));117        entry.0 = EvaluationOrder(cmp::min(entry.0.0, scope_range.start.0));118        entry.1 = EvaluationOrder(cmp::max(entry.1.0, scope_range.end.0));119    });120121    // Save original scope range IDs before updating122    let original_range_ids: FxHashMap<ScopeId, react_compiler_hir::MutableRangeId> = range_updates123        .keys()124        .map(|&root_id| {125            let range_id = env.scopes[root_id.0 as usize].range.id;126            (root_id, range_id)127        })128        .collect();129130    for (root_id, (new_start, new_end)) in &range_updates {131        env.scopes[root_id.0 as usize].range.start = *new_start;132        env.scopes[root_id.0 as usize].range.end = *new_end;133    }134135    // Sync identifier mutable_ranges that shared the old scope range.136    // Uses MutableRangeId for exact identity matching instead of value comparison.137    for ident in &mut env.identifiers {138        if let Some(scope_id) = ident.scope {139            if let Some(&orig_range_id) = original_range_ids.get(&scope_id) {140                if ident.mutable_range.id == orig_range_id {141                    let new_range = &env.scopes[scope_id.0 as usize].range;142                    ident.mutable_range.start = new_range.start;143                    ident.mutable_range.end = new_range.end;144                }145            }146        }147    }148149    // Step 2: Repoint identifiers whose scopes were merged150    // Build a map from old scope -> root scope for quick lookup151    let mut scope_remap: FxHashMap<ScopeId, ScopeId> = FxHashMap::default();152    merged_scopes.for_each(|scope_id, root_id| {153        if scope_id != root_id {154            scope_remap.insert(scope_id, root_id);155        }156    });157158    for (_block_id, block) in &func.body.blocks {159        for &instr_id in &block.instructions {160            let lvalue_id = func.instructions[instr_id.0 as usize].lvalue.identifier;161162            if let Some(current_scope) = env.identifiers[lvalue_id.0 as usize].scope {163                if let Some(&root) = scope_remap.get(&current_scope) {164                    env.identifiers[lvalue_id.0 as usize].scope = Some(root);165                }166            }167        }168    }169}