react_compiler_inference/
align_object_method_scopes.rs

1// Copyright (c) Meta Platforms, Inc. and affiliates.
2//
3// This source code is licensed under the MIT license found in the
4// LICENSE file in the root directory of this source tree.
5
6//! 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 expressions
8//! must be in the same ReactiveBlock as object method definitions must be inlined.
9//!
10//! Ported from TypeScript `src/ReactiveScopes/AlignObjectMethodScopes.ts`.
11
12use std::cmp;
13use std::collections::{HashMap, HashSet};
14
15use react_compiler_hir::environment::Environment;
16use react_compiler_hir::{
17    EvaluationOrder, HirFunction, IdentifierId, InstructionValue, ObjectPropertyOrSpread, ScopeId,
18};
19use react_compiler_utils::DisjointSet;
20
21// =============================================================================
22// findScopesToMerge
23// =============================================================================
24
25/// Identifies ObjectMethod lvalue identifiers and then finds ObjectExpression
26/// instructions whose operands reference those methods. Returns a disjoint set
27/// of scopes that must be merged.
28fn find_scopes_to_merge(func: &HirFunction, env: &Environment) -> DisjointSet<ScopeId> {
29    let mut object_method_decls: HashSet<IdentifierId> = HashSet::new();
30    let mut merged_scopes = DisjointSet::<ScopeId>::new();
31
32    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;
50
51                            // 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    }
66
67    merged_scopes
68}
69
70// =============================================================================
71// Public API
72// =============================================================================
73
74/// Aligns object method scopes so that ObjectMethod values and their enclosing
75/// 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    }
98
99    let mut merged_scopes = find_scopes_to_merge(func, env);
100
101    // Step 1: Merge affected scopes to their canonical root.
102    // Use a HashMap 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: HashMap<ScopeId, (EvaluationOrder, EvaluationOrder)> = HashMap::new();
105
106    merged_scopes.for_each(|scope_id, root_id| {
107        if scope_id == root_id {
108            return;
109        }
110        let scope_range = env.scopes[scope_id.0 as usize].range.clone();
111        let root_range = env.scopes[root_id.0 as usize].range.clone();
112
113        let entry = range_updates.entry(root_id).or_insert_with(|| {
114            (root_range.start, root_range.end)
115        });
116        entry.0 = EvaluationOrder(cmp::min(entry.0.0, scope_range.start.0));
117        entry.1 = EvaluationOrder(cmp::max(entry.1.0, scope_range.end.0));
118    });
119
120    // Save original scope range IDs before updating
121    let original_range_ids: HashMap<ScopeId, react_compiler_hir::MutableRangeId> = range_updates
122        .keys()
123        .map(|&root_id| {
124            let range_id = env.scopes[root_id.0 as usize].range.id;
125            (root_id, range_id)
126        })
127        .collect();
128
129    for (root_id, (new_start, new_end)) in &range_updates {
130        env.scopes[root_id.0 as usize].range.start = *new_start;
131        env.scopes[root_id.0 as usize].range.end = *new_end;
132    }
133
134    // Sync identifier mutable_ranges that shared the old scope range.
135    // Uses MutableRangeId for exact identity matching instead of value comparison.
136    for ident in &mut env.identifiers {
137        if let Some(scope_id) = ident.scope {
138            if let Some(&orig_range_id) = original_range_ids.get(&scope_id) {
139                if ident.mutable_range.id == orig_range_id {
140                    let new_range = &env.scopes[scope_id.0 as usize].range;
141                    ident.mutable_range.start = new_range.start;
142                    ident.mutable_range.end = new_range.end;
143                }
144            }
145        }
146    }
147
148    // Step 2: Repoint identifiers whose scopes were merged
149    // Build a map from old scope -> root scope for quick lookup
150    let mut scope_remap: HashMap<ScopeId, ScopeId> = HashMap::new();
151    merged_scopes.for_each(|scope_id, root_id| {
152        if scope_id != root_id {
153            scope_remap.insert(scope_id, root_id);
154        }
155    });
156
157    for (_block_id, block) in &func.body.blocks {
158        for &instr_id in &block.instructions {
159            let lvalue_id = func.instructions[instr_id.0 as usize].lvalue.identifier;
160
161            if let Some(current_scope) = env.identifiers[lvalue_id.0 as usize].scope {
162                if let Some(&root) = scope_remap.get(&current_scope) {
163                    env.identifiers[lvalue_id.0 as usize].scope = Some(root);
164                }
165            }
166        }
167    }
168}
react_compiler_inference/align_object_method_scopes.rs

react_compiler_inference/
align_object_method_scopes.rs
