| 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 | //! Infers which variables belong to reactive scopes. |
| 7 | //! |
| 8 | //! Ported from TypeScript `src/ReactiveScopes/InferReactiveScopeVariables.ts`. |
| 9 | //! |
| 10 | //! This is the 1st of 4 passes that determine how to break a function into |
| 11 | //! discrete reactive scopes (independently memoizable units of code): |
| 12 | //! 1. InferReactiveScopeVariables (this pass, on HIR) determines operands that |
| 13 | //! mutate together and assigns them a unique reactive scope. |
| 14 | //! 2. AlignReactiveScopesToBlockScopes aligns reactive scopes to block scopes. |
| 15 | //! 3. MergeOverlappingReactiveScopes ensures scopes do not overlap. |
| 16 | //! 4. BuildReactiveBlocks groups the statements for each scope. |
| 17 | |
| 18 | use rustc_hash::FxHashMap; |
| 19 | |
| 20 | use react_compiler_diagnostics::{CompilerDiagnostic, ErrorCategory}; |
| 21 | use react_compiler_hir::environment::Environment; |
| 22 | use react_compiler_hir::visitors; |
| 23 | use react_compiler_hir::{ |
| 24 | DeclarationId, EvaluationOrder, HirFunction, IdentifierId, InstructionValue, Pattern, Position, |
| 25 | SourceLocation, |
| 26 | }; |
| 27 | use react_compiler_utils::DisjointSet; |
| 28 | |
| 29 | // ============================================================================= |
| 30 | // Public API |
| 31 | // ============================================================================= |
| 32 | |
| 33 | /// Infer reactive scope variables for a function. |
| 34 | /// |
| 35 | /// For each mutable variable, infers a reactive scope which will construct that |
| 36 | /// variable. Variables that co-mutate are assigned to the same reactive scope. |
| 37 | /// |
| 38 | /// Corresponds to TS `inferReactiveScopeVariables(fn: HIRFunction): void`. |
| 39 | pub fn infer_reactive_scope_variables( |
| 40 | func: &mut HirFunction, |
| 41 | env: &mut Environment, |
| 42 | ) -> Result<(), CompilerDiagnostic> { |
| 43 | // Phase 1: find disjoint sets of co-mutating identifiers |
| 44 | let mut scope_identifiers = find_disjoint_mutable_values(func, env); |
| 45 | |
| 46 | // Phase 2: assign scopes |
| 47 | // Maps each group root identifier to the ScopeId assigned to that group. |
| 48 | let mut scopes: FxHashMap<IdentifierId, ScopeState> = FxHashMap::default(); |
| 49 | |
| 50 | scope_identifiers.for_each(|identifier_id, group_id| { |
| 51 | let ident_range = env.identifiers[identifier_id.0 as usize] |
| 52 | .mutable_range |
| 53 | .clone(); |
| 54 | let ident_loc = env.identifiers[identifier_id.0 as usize].loc; |
| 55 | |
| 56 | let state = scopes.entry(group_id).or_insert_with(|| { |
| 57 | let scope_id = env.next_scope_id(); |
| 58 | // Initialize scope range from the first member |
| 59 | let scope = &mut env.scopes[scope_id.0 as usize]; |
| 60 | scope.range = ident_range.clone(); |
| 61 | ScopeState { |
| 62 | scope_id, |
| 63 | loc: ident_loc, |
| 64 | } |
| 65 | }); |
| 66 | |
| 67 | // Update scope range |
| 68 | let scope = &mut env.scopes[state.scope_id.0 as usize]; |
| 69 | |
| 70 | // If this is not the first identifier (scope was already created), merge ranges |
| 71 | if scope.range.start != ident_range.start || scope.range.end != ident_range.end { |
| 72 | if scope.range.start == EvaluationOrder(0) { |
| 73 | scope.range.start = ident_range.start; |
| 74 | } else if ident_range.start != EvaluationOrder(0) { |
| 75 | scope.range.start = EvaluationOrder(scope.range.start.0.min(ident_range.start.0)); |
| 76 | } |
| 77 | scope.range.end = EvaluationOrder(scope.range.end.0.max(ident_range.end.0)); |
| 78 | } |
| 79 | |
| 80 | // Merge location |
| 81 | state.loc = merge_location(state.loc, ident_loc); |
| 82 | |
| 83 | // Assign the scope to this identifier |
| 84 | let scope_id = state.scope_id; |
| 85 | env.identifiers[identifier_id.0 as usize].scope = Some(scope_id); |
| 86 | }); |
| 87 | |
| 88 | // Set loc on each scope |
| 89 | for (_group_id, state) in &scopes { |
| 90 | env.scopes[state.scope_id.0 as usize].loc = state.loc; |
| 91 | } |
| 92 | |
| 93 | // Update each identifier's mutable_range to match its scope's range |
| 94 | for (&_identifier_id, state) in &scopes { |
| 95 | let scope_range = env.scopes[state.scope_id.0 as usize].range.clone(); |
| 96 | // Find all identifiers with this scope and update their mutable_range |
| 97 | // We iterate through all identifiers and check their scope |
| 98 | for ident in &mut env.identifiers { |
| 99 | if ident.scope == Some(state.scope_id) { |
| 100 | ident.mutable_range = scope_range.clone(); |
| 101 | } |
| 102 | } |
| 103 | } |
| 104 | |
| 105 | // Validate scope ranges |
| 106 | let mut max_instruction = EvaluationOrder(0); |
| 107 | for (_block_id, block) in &func.body.blocks { |
| 108 | for instr_id in &block.instructions { |
| 109 | let instr = &func.instructions[instr_id.0 as usize]; |
| 110 | max_instruction = EvaluationOrder(max_instruction.0.max(instr.id.0)); |
| 111 | } |
| 112 | max_instruction = |
| 113 | EvaluationOrder(max_instruction.0.max(block.terminal.evaluation_order().0)); |
| 114 | } |
| 115 | |
| 116 | for (_group_id, state) in &scopes { |
| 117 | let scope = &env.scopes[state.scope_id.0 as usize]; |
| 118 | if scope.range.start == EvaluationOrder(0) |
| 119 | || scope.range.end == EvaluationOrder(0) |
| 120 | || max_instruction == EvaluationOrder(0) |
| 121 | || scope.range.end.0 > max_instruction.0 + 1 |
| 122 | { |
| 123 | return Err(CompilerDiagnostic::new( |
| 124 | ErrorCategory::Invariant, |
| 125 | &format!( |
| 126 | "Invalid mutable range for scope: Scope @{} has range [{}:{}] but the valid range is [1:{}]", |
| 127 | scope.id.0, |
| 128 | scope.range.start.0, |
| 129 | scope.range.end.0, |
| 130 | max_instruction.0 + 1, |
| 131 | ), |
| 132 | None, |
| 133 | )); |
| 134 | } |
| 135 | } |
| 136 | |
| 137 | Ok(()) |
| 138 | } |
| 139 | |
| 140 | struct ScopeState { |
| 141 | scope_id: react_compiler_hir::ScopeId, |
| 142 | loc: Option<SourceLocation>, |
| 143 | } |
| 144 | |
| 145 | /// Merge two source locations, preferring non-None values. |
| 146 | /// Corresponds to TS `mergeLocation`. |
| 147 | fn merge_location(l: Option<SourceLocation>, r: Option<SourceLocation>) -> Option<SourceLocation> { |
| 148 | match (l, r) { |
| 149 | (None, r) => r, |
| 150 | (l, None) => l, |
| 151 | (Some(l), Some(r)) => Some(SourceLocation { |
| 152 | start: Position { |
| 153 | line: l.start.line.min(r.start.line), |
| 154 | column: l.start.column.min(r.start.column), |
| 155 | index: match (l.start.index, r.start.index) { |
| 156 | (Some(a), Some(b)) => Some(a.min(b)), |
| 157 | (a, b) => a.or(b), |
| 158 | }, |
| 159 | }, |
| 160 | end: Position { |
| 161 | line: l.end.line.max(r.end.line), |
| 162 | column: l.end.column.max(r.end.column), |
| 163 | index: match (l.end.index, r.end.index) { |
| 164 | (Some(a), Some(b)) => Some(a.max(b)), |
| 165 | (a, b) => a.or(b), |
| 166 | }, |
| 167 | }, |
| 168 | }), |
| 169 | } |
| 170 | } |
| 171 | |
| 172 | // ============================================================================= |
| 173 | // is_mutable / in_range helpers |
| 174 | // ============================================================================= |
| 175 | |
| 176 | // ============================================================================= |
| 177 | // may_allocate |
| 178 | // ============================================================================= |
| 179 | |
| 180 | /// Check if an instruction may allocate. Corresponds to TS `mayAllocate`. |
| 181 | fn may_allocate(value: &InstructionValue, lvalue_type_is_primitive: bool) -> bool { |
| 182 | match value { |
| 183 | InstructionValue::Destructure { lvalue, .. } => { |
| 184 | visitors::does_pattern_contain_spread_element(&lvalue.pattern) |
| 185 | } |
| 186 | InstructionValue::PostfixUpdate { .. } |
| 187 | | InstructionValue::PrefixUpdate { .. } |
| 188 | | InstructionValue::Await { .. } |
| 189 | | InstructionValue::DeclareLocal { .. } |
| 190 | | InstructionValue::DeclareContext { .. } |
| 191 | | InstructionValue::StoreLocal { .. } |
| 192 | | InstructionValue::LoadGlobal { .. } |
| 193 | | InstructionValue::MetaProperty { .. } |
| 194 | | InstructionValue::TypeCastExpression { .. } |
| 195 | | InstructionValue::LoadLocal { .. } |
| 196 | | InstructionValue::LoadContext { .. } |
| 197 | | InstructionValue::StoreContext { .. } |
| 198 | | InstructionValue::PropertyDelete { .. } |
| 199 | | InstructionValue::ComputedLoad { .. } |
| 200 | | InstructionValue::ComputedDelete { .. } |
| 201 | | InstructionValue::JSXText { .. } |
| 202 | | InstructionValue::TemplateLiteral { .. } |
| 203 | | InstructionValue::Primitive { .. } |
| 204 | | InstructionValue::GetIterator { .. } |
| 205 | | InstructionValue::IteratorNext { .. } |
| 206 | | InstructionValue::NextPropertyOf { .. } |
| 207 | | InstructionValue::Debugger { .. } |
| 208 | | InstructionValue::StartMemoize { .. } |
| 209 | | InstructionValue::FinishMemoize { .. } |
| 210 | | InstructionValue::UnaryExpression { .. } |
| 211 | | InstructionValue::BinaryExpression { .. } |
| 212 | | InstructionValue::PropertyLoad { .. } |
| 213 | | InstructionValue::StoreGlobal { .. } => false, |
| 214 | |
| 215 | InstructionValue::TaggedTemplateExpression { .. } |
| 216 | | InstructionValue::CallExpression { .. } |
| 217 | | InstructionValue::MethodCall { .. } => !lvalue_type_is_primitive, |
| 218 | |
| 219 | InstructionValue::RegExpLiteral { .. } |
| 220 | | InstructionValue::PropertyStore { .. } |
| 221 | | InstructionValue::ComputedStore { .. } |
| 222 | | InstructionValue::ArrayExpression { .. } |
| 223 | | InstructionValue::JsxExpression { .. } |
| 224 | | InstructionValue::JsxFragment { .. } |
| 225 | | InstructionValue::NewExpression { .. } |
| 226 | | InstructionValue::ObjectExpression { .. } |
| 227 | | InstructionValue::UnsupportedNode { .. } |
| 228 | | InstructionValue::ObjectMethod { .. } |
| 229 | | InstructionValue::FunctionExpression { .. } => true, |
| 230 | } |
| 231 | } |
| 232 | |
| 233 | // ============================================================================= |
| 234 | // Pattern helpers |
| 235 | // ============================================================================= |
| 236 | |
| 237 | /// Collect all Place identifiers from a destructure pattern. |
| 238 | /// Corresponds to TS `eachPatternOperand`. |
| 239 | fn each_pattern_operand(pattern: &Pattern) -> Vec<IdentifierId> { |
| 240 | visitors::each_pattern_operand(pattern) |
| 241 | .into_iter() |
| 242 | .map(|p| p.identifier) |
| 243 | .collect() |
| 244 | } |
| 245 | |
| 246 | /// Collect all operand identifiers from an instruction value. |
| 247 | /// Corresponds to TS `eachInstructionValueOperand`. |
| 248 | fn each_instruction_value_operand( |
| 249 | value: &InstructionValue, |
| 250 | env: &Environment, |
| 251 | ) -> Vec<IdentifierId> { |
| 252 | visitors::each_instruction_value_operand(value, env) |
| 253 | .into_iter() |
| 254 | .map(|p| p.identifier) |
| 255 | .collect() |
| 256 | } |
| 257 | |
| 258 | // ============================================================================= |
| 259 | // findDisjointMutableValues |
| 260 | // ============================================================================= |
| 261 | |
| 262 | /// Find disjoint sets of co-mutating identifier IDs. |
| 263 | /// |
| 264 | /// Corresponds to TS `findDisjointMutableValues(fn: HIRFunction): DisjointSet<Identifier>`. |
| 265 | pub(crate) fn find_disjoint_mutable_values( |
| 266 | func: &HirFunction, |
| 267 | env: &Environment, |
| 268 | ) -> DisjointSet<IdentifierId> { |
| 269 | let mut scope_identifiers = DisjointSet::<IdentifierId>::new(); |
| 270 | let mut declarations: FxHashMap<DeclarationId, IdentifierId> = FxHashMap::default(); |
| 271 | |
| 272 | let enable_forest = env.config.enable_forest; |
| 273 | |
| 274 | for (_block_id, block) in &func.body.blocks { |
| 275 | // Handle phi nodes |
| 276 | for phi in &block.phis { |
| 277 | let phi_id = phi.place.identifier; |
| 278 | let phi_range = &env.identifiers[phi_id.0 as usize].mutable_range; |
| 279 | let phi_decl_id = env.identifiers[phi_id.0 as usize].declaration_id; |
| 280 | |
| 281 | let first_instr_id = block |
| 282 | .instructions |
| 283 | .first() |
| 284 | .map(|iid| func.instructions[iid.0 as usize].id) |
| 285 | .unwrap_or(block.terminal.evaluation_order()); |
| 286 | |
| 287 | let is_phi_mutated_after_creation = |
| 288 | phi_range.start.0 + 1 != phi_range.end.0 && phi_range.end > first_instr_id; |
| 289 | // A phi operand defined at or after the phi's block is a loop |
| 290 | // back-edge: the variable is reassigned within the loop (eg a |
| 291 | // counter `a++` or `a = a + 1`). The reassignment must count as |
| 292 | // the loop's scope reassigning the variable, so union the phi |
| 293 | // with its operands and declaration. Otherwise the variable's |
| 294 | // pre-loop value would become a dependency of the scope even |
| 295 | // though the scope changes the value as it executes, making the |
| 296 | // scope's dependencies unstable (the cached dependency would be |
| 297 | // the post-loop value, which can never match the pre-loop value |
| 298 | // compared at the top of the scope). |
| 299 | let is_loop_carried_reassignment = !is_phi_mutated_after_creation |
| 300 | && phi.operands.iter().any(|(_pred_id, operand)| { |
| 301 | env.identifiers[operand.identifier.0 as usize] |
| 302 | .mutable_range |
| 303 | .start |
| 304 | >= first_instr_id |
| 305 | }); |
| 306 | if is_phi_mutated_after_creation || is_loop_carried_reassignment { |
| 307 | let mut operands = vec![phi_id]; |
| 308 | if let Some(&decl_id) = declarations.get(&phi_decl_id) { |
| 309 | operands.push(decl_id); |
| 310 | } |
| 311 | for (_pred_id, phi_operand) in &phi.operands { |
| 312 | operands.push(phi_operand.identifier); |
| 313 | } |
| 314 | scope_identifiers.union(&operands); |
| 315 | } else if enable_forest { |
| 316 | for (_pred_id, phi_operand) in &phi.operands { |
| 317 | scope_identifiers.union(&[phi_id, phi_operand.identifier]); |
| 318 | } |
| 319 | } |
| 320 | } |
| 321 | |
| 322 | // Handle instructions |
| 323 | for instr_id in &block.instructions { |
| 324 | let instr = &func.instructions[instr_id.0 as usize]; |
| 325 | let mut operands: Vec<IdentifierId> = Vec::new(); |
| 326 | |
| 327 | let lvalue_id = instr.lvalue.identifier; |
| 328 | let lvalue_range = &env.identifiers[lvalue_id.0 as usize].mutable_range; |
| 329 | let lvalue_type = &env.types[env.identifiers[lvalue_id.0 as usize].type_.0 as usize]; |
| 330 | let lvalue_type_is_primitive = react_compiler_hir::is_primitive_type(lvalue_type); |
| 331 | |
| 332 | if lvalue_range.end.0 > lvalue_range.start.0 + 1 |
| 333 | || may_allocate(&instr.value, lvalue_type_is_primitive) |
| 334 | { |
| 335 | operands.push(lvalue_id); |
| 336 | } |
| 337 | |
| 338 | match &instr.value { |
| 339 | InstructionValue::DeclareLocal { lvalue, .. } |
| 340 | | InstructionValue::DeclareContext { lvalue, .. } => { |
| 341 | let place_id = lvalue.place.identifier; |
| 342 | let decl_id = env.identifiers[place_id.0 as usize].declaration_id; |
| 343 | declarations.entry(decl_id).or_insert(place_id); |
| 344 | } |
| 345 | InstructionValue::StoreLocal { lvalue, value, .. } |
| 346 | | InstructionValue::StoreContext { lvalue, value, .. } => { |
| 347 | let place_id = lvalue.place.identifier; |
| 348 | let decl_id = env.identifiers[place_id.0 as usize].declaration_id; |
| 349 | declarations.entry(decl_id).or_insert(place_id); |
| 350 | |
| 351 | let place_range = &env.identifiers[place_id.0 as usize].mutable_range; |
| 352 | if place_range.end.0 > place_range.start.0 + 1 { |
| 353 | operands.push(place_id); |
| 354 | } |
| 355 | |
| 356 | let value_range = &env.identifiers[value.identifier.0 as usize].mutable_range; |
| 357 | if value_range.contains(instr.id) && value_range.start.0 > 0 { |
| 358 | operands.push(value.identifier); |
| 359 | } |
| 360 | } |
| 361 | InstructionValue::Destructure { lvalue, value, .. } => { |
| 362 | let pattern_places = each_pattern_operand(&lvalue.pattern); |
| 363 | for place_id in &pattern_places { |
| 364 | let decl_id = env.identifiers[place_id.0 as usize].declaration_id; |
| 365 | declarations.entry(decl_id).or_insert(*place_id); |
| 366 | |
| 367 | let place_range = &env.identifiers[place_id.0 as usize].mutable_range; |
| 368 | if place_range.end.0 > place_range.start.0 + 1 { |
| 369 | operands.push(*place_id); |
| 370 | } |
| 371 | } |
| 372 | |
| 373 | let value_range = &env.identifiers[value.identifier.0 as usize].mutable_range; |
| 374 | if value_range.contains(instr.id) && value_range.start.0 > 0 { |
| 375 | operands.push(value.identifier); |
| 376 | } |
| 377 | } |
| 378 | InstructionValue::MethodCall { property, .. } => { |
| 379 | // For MethodCall: include all mutable operands plus the computed property |
| 380 | let all_operands = each_instruction_value_operand(&instr.value, env); |
| 381 | for op_id in &all_operands { |
| 382 | let op_range = &env.identifiers[op_id.0 as usize].mutable_range; |
| 383 | if op_range.contains(instr.id) && op_range.start.0 > 0 { |
| 384 | operands.push(*op_id); |
| 385 | } |
| 386 | } |
| 387 | // Ensure method property is in the same scope as the call |
| 388 | operands.push(property.identifier); |
| 389 | } |
| 390 | _ => { |
| 391 | // For all other instructions: include mutable operands |
| 392 | let all_operands = each_instruction_value_operand(&instr.value, env); |
| 393 | for op_id in &all_operands { |
| 394 | let op_range = &env.identifiers[op_id.0 as usize].mutable_range; |
| 395 | if op_range.contains(instr.id) && op_range.start.0 > 0 { |
| 396 | operands.push(*op_id); |
| 397 | } |
| 398 | } |
| 399 | } |
| 400 | } |
| 401 | |
| 402 | if !operands.is_empty() { |
| 403 | scope_identifiers.union(&operands); |
| 404 | } |
| 405 | } |
| 406 | } |
| 407 | scope_identifiers |
| 408 | } |