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//! Builds reactive scope terminals in the HIR.7//!8//! Given a function whose reactive scope ranges have been correctly aligned and9//! merged, this pass rewrites blocks to introduce ReactiveScopeTerminals and10//! their fallthrough blocks.11//!12//! Ported from TypeScript `src/HIR/BuildReactiveScopeTerminalsHIR.ts`.1314use indexmap::{IndexMap, IndexSet};15use rustc_hash::{FxBuildHasher, FxHashMap, FxHashSet};1617use react_compiler_hir::BasicBlock;18use react_compiler_hir::BlockId;19use react_compiler_hir::EvaluationOrder;20use react_compiler_hir::GotoVariant;21use react_compiler_hir::HirFunction;22use react_compiler_hir::IdentifierId;23use react_compiler_hir::ScopeId;24use react_compiler_hir::Terminal;25use react_compiler_hir::environment::Environment;26use react_compiler_hir::visitors::each_instruction_lvalue_ids;27use react_compiler_hir::visitors::each_instruction_operand_ids;28use react_compiler_hir::visitors::each_terminal_operand_ids;29use react_compiler_lowering::get_reverse_postordered_blocks;30use react_compiler_lowering::mark_instruction_ids;31use react_compiler_lowering::mark_predecessors;3233// =============================================================================34// getScopes35// =============================================================================3637/// Collect all unique scopes from places in the function that have non-empty ranges.38/// Corresponds to TS `getScopes(fn)`.39fn get_scopes(func: &HirFunction, env: &Environment) -> Vec<ScopeId> {40    let mut scope_ids: FxHashSet<ScopeId> = FxHashSet::default();4142    let mut visit_place = |identifier_id: IdentifierId| {43        if let Some(scope_id) = env.identifiers[identifier_id.0 as usize].scope {44            let range = &env.scopes[scope_id.0 as usize].range;45            if range.start != range.end {46                scope_ids.insert(scope_id);47            }48        }49    };5051    for (_block_id, block) in &func.body.blocks {52        for &instr_id in &block.instructions {53            let instr = &func.instructions[instr_id.0 as usize];54            // lvalues55            for id in each_instruction_lvalue_ids(instr) {56                visit_place(id);57            }58            // operands59            for id in each_instruction_operand_ids(instr, env) {60                visit_place(id);61            }62        }63        // terminal operands64        for id in each_terminal_operand_ids(&block.terminal) {65            visit_place(id);66        }67    }6869    scope_ids.into_iter().collect()70}7172// =============================================================================73// TerminalRewriteInfo74// =============================================================================7576enum TerminalRewriteInfo {77    StartScope {78        block_id: BlockId,79        fallthrough_id: BlockId,80        instr_id: EvaluationOrder,81        scope_id: ScopeId,82    },83    EndScope {84        instr_id: EvaluationOrder,85        fallthrough_id: BlockId,86    },87}8889impl TerminalRewriteInfo {90    fn instr_id(&self) -> EvaluationOrder {91        match self {92            TerminalRewriteInfo::StartScope { instr_id, .. } => *instr_id,93            TerminalRewriteInfo::EndScope { instr_id, .. } => *instr_id,94        }95    }96}9798// =============================================================================99// collectScopeRewrites100// =============================================================================101102/// Collect all scope rewrites by traversing scopes in pre-order.103fn collect_scope_rewrites(func: &HirFunction, env: &mut Environment) -> Vec<TerminalRewriteInfo> {104    let scope_ids = get_scopes(func, env);105106    // Sort: ascending by start, descending by end for ties107    let mut items: Vec<ScopeId> = scope_ids;108    items.sort_by(|a, b| {109        let a_range = &env.scopes[a.0 as usize].range;110        let b_range = &env.scopes[b.0 as usize].range;111        let start_diff = a_range.start.0.cmp(&b_range.start.0);112        if start_diff != std::cmp::Ordering::Equal {113            return start_diff;114        }115        b_range.end.0.cmp(&a_range.end.0)116    });117118    let mut rewrites: Vec<TerminalRewriteInfo> = Vec::new();119    let mut fallthroughs: FxHashMap<ScopeId, BlockId> = FxHashMap::default();120    let mut active_items: Vec<ScopeId> = Vec::new();121122    for i in 0..items.len() {123        let curr = items[i];124        let curr_start = env.scopes[curr.0 as usize].range.start;125        let curr_end = env.scopes[curr.0 as usize].range.end;126127        // Pop active items that are disjoint with current128        let mut j = active_items.len();129        while j > 0 {130            j -= 1;131            let maybe_parent = active_items[j];132            let parent_end = env.scopes[maybe_parent.0 as usize].range.end;133            let disjoint = curr_start >= parent_end;134            let nested = curr_end <= parent_end;135            assert!(136                disjoint || nested,137                "Invalid nesting in program blocks or scopes"138            );139            if disjoint {140                // Exit this scope141                let fallthrough_id = *fallthroughs142                    .get(&maybe_parent)143                    .expect("Expected scope to exist");144                let end_instr_id = env.scopes[maybe_parent.0 as usize].range.end;145                rewrites.push(TerminalRewriteInfo::EndScope {146                    instr_id: end_instr_id,147                    fallthrough_id,148                });149                active_items.truncate(j);150            } else {151                break;152            }153        }154155        // Enter scope156        let block_id = env.next_block_id();157        let fallthrough_id = env.next_block_id();158        let start_instr_id = env.scopes[curr.0 as usize].range.start;159        rewrites.push(TerminalRewriteInfo::StartScope {160            block_id,161            fallthrough_id,162            instr_id: start_instr_id,163            scope_id: curr,164        });165        fallthroughs.insert(curr, fallthrough_id);166        active_items.push(curr);167    }168169    // Exit remaining active items170    while let Some(curr) = active_items.pop() {171        let fallthrough_id = *fallthroughs.get(&curr).expect("Expected scope to exist");172        let end_instr_id = env.scopes[curr.0 as usize].range.end;173        rewrites.push(TerminalRewriteInfo::EndScope {174            instr_id: end_instr_id,175            fallthrough_id,176        });177    }178179    rewrites180}181182// =============================================================================183// handleRewrite184// =============================================================================185186struct RewriteContext {187    next_block_id: BlockId,188    next_preds: Vec<BlockId>,189    instr_slice_idx: usize,190    rewrites: Vec<BasicBlock>,191}192193fn handle_rewrite(194    terminal_info: &TerminalRewriteInfo,195    idx: usize,196    source_block: &BasicBlock,197    context: &mut RewriteContext,198) {199    let terminal: Terminal = match terminal_info {200        TerminalRewriteInfo::StartScope {201            block_id,202            fallthrough_id,203            instr_id,204            scope_id,205        } => Terminal::Scope {206            fallthrough: *fallthrough_id,207            block: *block_id,208            scope: *scope_id,209            id: *instr_id,210            loc: None,211        },212        TerminalRewriteInfo::EndScope {213            instr_id,214            fallthrough_id,215        } => Terminal::Goto {216            variant: GotoVariant::Break,217            block: *fallthrough_id,218            id: *instr_id,219            loc: None,220        },221    };222223    let curr_block_id = context.next_block_id;224    let mut preds = IndexSet::default();225    for &p in &context.next_preds {226        preds.insert(p);227    }228229    context.rewrites.push(BasicBlock {230        kind: source_block.kind,231        id: curr_block_id,232        instructions: source_block.instructions[context.instr_slice_idx..idx].to_vec(),233        preds,234        // Only the first rewrite should reuse source block phis235        phis: if context.rewrites.is_empty() {236            source_block.phis.clone()237        } else {238            Vec::new()239        },240        terminal,241    });242243    context.next_preds = vec![curr_block_id];244    context.next_block_id = match terminal_info {245        TerminalRewriteInfo::StartScope { block_id, .. } => *block_id,246        TerminalRewriteInfo::EndScope { fallthrough_id, .. } => *fallthrough_id,247    };248    context.instr_slice_idx = idx;249}250251// =============================================================================252// Public API253// =============================================================================254255/// Builds reactive scope terminals in the HIR.256///257/// This pass assumes that all program blocks are properly nested with respect258/// to fallthroughs. Given a function whose reactive scope ranges have been259/// correctly aligned and merged, this pass rewrites blocks to introduce260/// ReactiveScopeTerminals and their fallthrough blocks.261pub fn build_reactive_scope_terminals_hir(func: &mut HirFunction, env: &mut Environment) {262    // Step 1: Collect rewrites263    let mut queued_rewrites = collect_scope_rewrites(func, env);264265    // Step 2: Apply rewrites by splitting blocks266    let mut rewritten_final_blocks: FxHashMap<BlockId, BlockId> = FxHashMap::default();267    let mut next_blocks: IndexMap<BlockId, BasicBlock, FxBuildHasher> = IndexMap::default();268269    // Reverse so we can pop from the end while traversing in ascending order270    queued_rewrites.reverse();271272    for (_block_id, block) in &func.body.blocks {273        let preds_vec: Vec<BlockId> = block.preds.iter().copied().collect();274        let mut context = RewriteContext {275            next_block_id: block.id,276            rewrites: Vec::new(),277            next_preds: preds_vec,278            instr_slice_idx: 0,279        };280281        // Handle queued terminal rewrites at their nearest instruction ID282        for i in 0..block.instructions.len() + 1 {283            let instr_id = if i < block.instructions.len() {284                let instr_idx = block.instructions[i];285                func.instructions[instr_idx.0 as usize].id286            } else {287                block.terminal.evaluation_order()288            };289290            while let Some(rewrite) = queued_rewrites.last() {291                if rewrite.instr_id() <= instr_id {292                    // Need to pop before calling handle_rewrite293                    let rewrite = queued_rewrites.pop().unwrap();294                    handle_rewrite(&rewrite, i, block, &mut context);295                } else {296                    break;297                }298            }299        }300301        if !context.rewrites.is_empty() {302            let mut final_preds = IndexSet::default();303            for &p in &context.next_preds {304                final_preds.insert(p);305            }306            let final_block = BasicBlock {307                id: context.next_block_id,308                kind: block.kind,309                preds: final_preds,310                terminal: block.terminal.clone(),311                instructions: block.instructions[context.instr_slice_idx..].to_vec(),312                phis: Vec::new(),313            };314            let final_block_id = final_block.id;315            context.rewrites.push(final_block);316            for b in context.rewrites {317                next_blocks.insert(b.id, b);318            }319            rewritten_final_blocks.insert(block.id, final_block_id);320        } else {321            next_blocks.insert(block.id, block.clone());322        }323    }324325    func.body.blocks = next_blocks;326327    // Step 3: Repoint phis when they refer to a rewritten block328    for block in func.body.blocks.values_mut() {329        for phi in &mut block.phis {330            let updates: Vec<(BlockId, BlockId)> = phi331                .operands332                .keys()333                .filter_map(|original_id| {334                    rewritten_final_blocks335                        .get(original_id)336                        .map(|new_id| (*original_id, *new_id))337                })338                .collect();339            for (old_id, new_id) in updates {340                if let Some(value) = phi.operands.shift_remove(&old_id) {341                    phi.operands.insert(new_id, value);342                }343            }344        }345    }346347    // Step 4: Fixup HIR to restore RPO, correct predecessors, renumber instructions348    func.body.blocks = get_reverse_postordered_blocks(&func.body, &func.instructions);349    mark_predecessors(&mut func.body);350    mark_instruction_ids(&mut func.body, &mut func.instructions);351352    // Step 5: Fix scope and identifier ranges to account for renumbered instructions353    fix_scope_and_identifier_ranges(func, env);354}355356/// Fix scope ranges after instruction renumbering.357/// Scope ranges should always align to start at the 'scope' terminal358/// and end at the first instruction of the fallthrough block.359///360/// In TS, `identifier.mutableRange` and `scope.range` are the same object361/// reference (after InferReactiveScopeVariables). When scope.range is updated,362/// all identifiers with that scope automatically see the new range.363/// BUT: after MergeOverlappingReactiveScopesHIR, repointed identifiers have364/// mutableRange pointing to the OLD scope's range, NOT the root scope's range.365/// So only identifiers whose mutableRange matches their scope's pre-renumbering366/// range should be updated.367///368/// Corresponds to TS `fixScopeAndIdentifierRanges`.369fn fix_scope_and_identifier_ranges(func: &HirFunction, env: &mut Environment) {370    // Save original scope ranges before updating them. In TS,371    // identifier.mutableRange and scope.range may or may not be the same372    // JS object. Only identifiers whose mutableRange shares the same object373    // reference as scope.range see the update automatically. We simulate374    // this by only syncing identifiers whose mutableRange matches the375    // scope's pre-update range.376    let original_scope_ranges: Vec<react_compiler_hir::MutableRange> =377        env.scopes.iter().map(|s| s.range.clone()).collect();378379    for (_block_id, block) in &func.body.blocks {380        match &block.terminal {381            Terminal::Scope {382                fallthrough,383                scope,384                id,385                ..386            }387            | Terminal::PrunedScope {388                fallthrough,389                scope,390                id,391                ..392            } => {393                let fallthrough_block = func.body.blocks.get(fallthrough).unwrap();394                let first_id = if !fallthrough_block.instructions.is_empty() {395                    func.instructions[fallthrough_block.instructions[0].0 as usize].id396                } else {397                    fallthrough_block.terminal.evaluation_order()398                };399                env.scopes[scope.0 as usize].range.start = *id;400                env.scopes[scope.0 as usize].range.end = first_id;401            }402            _ => {}403        }404    }405406    // Sync identifier mutable ranges with their scope ranges, but ONLY407    // for identifiers whose mutableRange has the same identity as their408    // scope's ORIGINAL range (before the updates above). In TS,409    // identifier.mutableRange and scope.range are only the same JS object410    // for identifiers that were the canonical representative when the scope411    // was created. After MergeOverlappingReactiveScopesHIR, repointed412    // identifiers have mutableRange pointing to the OLD scope's range,413    // not the root scope's range — so they should NOT be synced here.414    for ident in &mut env.identifiers {415        if let Some(scope_id) = ident.scope {416            let original = &original_scope_ranges[scope_id.0 as usize];417            if ident.mutable_range.same_range(original) {418                let scope_range = &env.scopes[scope_id.0 as usize].range;419                ident.mutable_range.start = scope_range.start;420                ident.mutable_range.end = scope_range.end;421            }422        }423    }424}