react_compiler_inference/
align_method_call_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//! Ensures that method call instructions have scopes such that either:
7//! - Both the MethodCall and its property have the same scope
8//! - OR neither has a scope
9//!
10//! Ported from TypeScript `src/ReactiveScopes/AlignMethodCallScopes.ts`.
11
12use std::collections::HashMap;
13
14use react_compiler_hir::environment::Environment;
15use react_compiler_hir::{EvaluationOrder, HirFunction, IdentifierId, InstructionValue, ScopeId};
16use react_compiler_utils::DisjointSet;
17
18// =============================================================================
19// Public API
20// =============================================================================
21
22/// Aligns method call scopes so that either both the MethodCall result and its
23/// property operand share the same scope, or neither has a scope.
24///
25/// Corresponds to TS `alignMethodCallScopes(fn: HIRFunction): void`.
26pub fn align_method_call_scopes(func: &mut HirFunction, env: &mut Environment) {
27    // Maps an identifier to the scope it should be assigned to (or None to remove scope)
28    let mut scope_mapping: HashMap<IdentifierId, Option<ScopeId>> = HashMap::new();
29    let mut merged_scopes = DisjointSet::<ScopeId>::new();
30
31    // Phase 1: Walk instructions and collect scope relationships
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::MethodCall { property, .. } => {
37                    let lvalue_scope =
38                        env.identifiers[instr.lvalue.identifier.0 as usize].scope;
39                    let property_scope =
40                        env.identifiers[property.identifier.0 as usize].scope;
41
42                    match (lvalue_scope, property_scope) {
43                        (Some(lvalue_sid), Some(property_sid)) => {
44                            // Both have a scope: merge the scopes
45                            merged_scopes.union(&[lvalue_sid, property_sid]);
46                        }
47                        (Some(lvalue_sid), None) => {
48                            // Call has a scope but not the property:
49                            // record that this property should be in this scope
50                            scope_mapping
51                                .insert(property.identifier, Some(lvalue_sid));
52                        }
53                        (None, Some(_)) => {
54                            // Property has a scope but call doesn't:
55                            // this property does not need a scope
56                            scope_mapping.insert(property.identifier, None);
57                        }
58                        (None, None) => {
59                            // Neither has a scope, nothing to do
60                        }
61                    }
62                }
63                InstructionValue::FunctionExpression { lowered_func, .. }
64                | InstructionValue::ObjectMethod { lowered_func, .. } => {
65                    // Recurse into inner functions
66                    let func_id = lowered_func.func;
67                    let mut inner_func = std::mem::replace(
68                        &mut env.functions[func_id.0 as usize],
69                        react_compiler_ssa::enter_ssa::placeholder_function(),
70                    );
71                    align_method_call_scopes(&mut inner_func, env);
72                    env.functions[func_id.0 as usize] = inner_func;
73                }
74                _ => {}
75            }
76        }
77    }
78
79    // Phase 2: Merge scope ranges for unioned scopes.
80    // Use a HashMap to accumulate min/max across all scopes mapping to the same root,
81    // matching TS behavior where root.range is updated in-place during iteration.
82    let mut range_updates: HashMap<ScopeId, (EvaluationOrder, EvaluationOrder)> = HashMap::new();
83
84    merged_scopes.for_each(|scope_id, root_id| {
85        if scope_id == root_id {
86            return;
87        }
88        let scope_range = env.scopes[scope_id.0 as usize].range.clone();
89        let root_range = env.scopes[root_id.0 as usize].range.clone();
90
91        let entry = range_updates
92            .entry(root_id)
93            .or_insert_with(|| (root_range.start, root_range.end));
94        entry.0 = EvaluationOrder(std::cmp::min(entry.0 .0, scope_range.start.0));
95        entry.1 = EvaluationOrder(std::cmp::max(entry.1 .0, scope_range.end.0));
96    });
97
98    // Save original scope range IDs before updating
99    let original_range_ids: HashMap<ScopeId, react_compiler_hir::MutableRangeId> = range_updates
100        .keys()
101        .map(|&root_id| {
102            let range_id = env.scopes[root_id.0 as usize].range.id;
103            (root_id, range_id)
104        })
105        .collect();
106
107    for (root_id, (new_start, new_end)) in &range_updates {
108        env.scopes[root_id.0 as usize].range.start = *new_start;
109        env.scopes[root_id.0 as usize].range.end = *new_end;
110    }
111
112    // Sync identifier mutable_ranges that shared the old scope range.
113    // Uses MutableRangeId for exact identity matching instead of value comparison.
114    for ident in &mut env.identifiers {
115        if let Some(scope_id) = ident.scope {
116            if let Some(&orig_range_id) = original_range_ids.get(&scope_id) {
117                if ident.mutable_range.id == orig_range_id {
118                    let new_range = &env.scopes[scope_id.0 as usize].range;
119                    ident.mutable_range.start = new_range.start;
120                    ident.mutable_range.end = new_range.end;
121                }
122            }
123        }
124    }
125
126    // Phase 3: Apply scope mappings and merged scope reassignments
127    for (_block_id, block) in &func.body.blocks {
128        for &instr_id in &block.instructions {
129            let lvalue_id = func.instructions[instr_id.0 as usize].lvalue.identifier;
130
131            if let Some(mapped_scope) = scope_mapping.get(&lvalue_id) {
132                env.identifiers[lvalue_id.0 as usize].scope = *mapped_scope;
133            } else if let Some(current_scope) =
134                env.identifiers[lvalue_id.0 as usize].scope
135            {
136                // TS: mergedScopes.find() returns null if not in the set
137                if let Some(merged) = merged_scopes.find_opt(current_scope) {
138                    env.identifiers[lvalue_id.0 as usize].scope = Some(merged);
139                }
140            }
141        }
142    }
143}
react_compiler_inference/align_method_call_scopes.rs

react_compiler_inference/
align_method_call_scopes.rs
