1use rustc_hash::{FxBuildHasher, FxHashMap, FxHashSet};23use indexmap::IndexMap;4use react_compiler_diagnostics::{CompilerDiagnostic, CompilerDiagnosticDetail, ErrorCategory};5use react_compiler_hir::environment::Environment;6use react_compiler_hir::visitors;7use react_compiler_hir::*;89// =============================================================================10// SSABuilder11// =============================================================================1213struct IncompletePhi {14    old_place: Place,15    new_place: Place,16}1718struct State {19    defs: FxHashMap<IdentifierId, IdentifierId>,20    incomplete_phis: Vec<IncompletePhi>,21}2223struct SSABuilder {24    states: FxHashMap<BlockId, State>,25    current: Option<BlockId>,26    unsealed_preds: FxHashMap<BlockId, u32>,27    block_preds: FxHashMap<BlockId, Vec<BlockId>>,28    unknown: FxHashSet<IdentifierId>,29    context: FxHashSet<IdentifierId>,30    pending_phis: FxHashMap<BlockId, Vec<Phi>>,31    processed_functions: Vec<FunctionId>,32}3334impl SSABuilder {35    fn new(blocks: &IndexMap<BlockId, BasicBlock, FxBuildHasher>) -> Self {36        let mut block_preds = FxHashMap::default();37        for (id, block) in blocks {38            block_preds.insert(*id, block.preds.iter().copied().collect());39        }40        SSABuilder {41            states: FxHashMap::default(),42            current: None,43            unsealed_preds: FxHashMap::default(),44            block_preds,45            unknown: FxHashSet::default(),46            context: FxHashSet::default(),47            pending_phis: FxHashMap::default(),48            processed_functions: Vec::new(),49        }50    }5152    fn define_function(&mut self, func: &HirFunction) {53        for (id, block) in &func.body.blocks {54            self.block_preds55                .insert(*id, block.preds.iter().copied().collect());56        }57    }5859    fn state_mut(&mut self) -> &mut State {60        let current = self61            .current62            .expect("we need to be in a block to access state!");63        self.states64            .get_mut(&current)65            .expect("state not found for current block")66    }6768    fn make_id(&mut self, old_id: IdentifierId, env: &mut Environment) -> IdentifierId {69        let new_id = env.next_identifier_id();70        let old = &env.identifiers[old_id.0 as usize];71        let declaration_id = old.declaration_id;72        let name = old.name.clone();73        let loc = old.loc;74        let new_ident = &mut env.identifiers[new_id.0 as usize];75        new_ident.declaration_id = declaration_id;76        new_ident.name = name;77        new_ident.loc = loc;78        new_id79    }8081    fn define_place(82        &mut self,83        old_place: &Place,84        env: &mut Environment,85    ) -> Result<Place, CompilerDiagnostic> {86        let old_id = old_place.identifier;8788        if self.unknown.contains(&old_id) {89            let ident = &env.identifiers[old_id.0 as usize];90            let name = match &ident.name {91                Some(name) => format!("{}${}", name.value(), old_id.0),92                None => format!("${}", old_id.0),93            };94            return Err(CompilerDiagnostic::new(95                ErrorCategory::Todo,96                "[hoisting] EnterSSA: Expected identifier to be defined before being used",97                Some(format!("Identifier {} is undefined", name)),98            )99            .with_detail(CompilerDiagnosticDetail::Error {100                loc: old_place.loc,101                message: None,102                identifier_name: None,103            }));104        }105106        // Do not redefine context references.107        if self.context.contains(&old_id) {108            return Ok(self.get_place(old_place, env));109        }110111        let new_id = self.make_id(old_id, env);112        self.state_mut().defs.insert(old_id, new_id);113        Ok(Place {114            identifier: new_id,115            effect: old_place.effect,116            reactive: old_place.reactive,117            loc: old_place.loc,118        })119    }120121    #[allow(dead_code)]122    fn define_context(123        &mut self,124        old_place: &Place,125        env: &mut Environment,126    ) -> Result<Place, CompilerDiagnostic> {127        let old_id = old_place.identifier;128        let new_place = self.define_place(old_place, env)?;129        self.context.insert(old_id);130        Ok(new_place)131    }132133    /// A function's context places capture a *binding*, not a value: the134    /// variable is only read when the function is later called, so a context135    /// place may reference a binding that is declared after the function136    /// expression itself (eg `const colgroup = useMemo(() => <colgroup>...)`,137    /// where the JSX tag name resolves to the variable being assigned). Unmark138    /// such identifiers so the later declaration doesn't error; if the function139    /// body actually *reads* the variable before it is defined, visiting the140    /// body re-marks it and the hoisting bailout in define_place still applies.141    fn unmark_unknown(&mut self, id: IdentifierId) {142        self.unknown.remove(&id);143    }144145    fn get_place(&mut self, old_place: &Place, env: &mut Environment) -> Place {146        let current_id = self.current.expect("must be in a block");147        let new_id = self.get_id_at(old_place, current_id, env);148        Place {149            identifier: new_id,150            effect: old_place.effect,151            reactive: old_place.reactive,152            loc: old_place.loc,153        }154    }155156    fn get_id_at(157        &mut self,158        old_place: &Place,159        block_id: BlockId,160        env: &mut Environment,161    ) -> IdentifierId {162        if let Some(state) = self.states.get(&block_id) {163            if let Some(&new_id) = state.defs.get(&old_place.identifier) {164                return new_id;165            }166        }167168        let preds = self.block_preds.get(&block_id).cloned().unwrap_or_default();169170        if preds.is_empty() {171            self.unknown.insert(old_place.identifier);172            return old_place.identifier;173        }174175        let unsealed = self.unsealed_preds.get(&block_id).copied().unwrap_or(0);176        if unsealed > 0 {177            let new_id = self.make_id(old_place.identifier, env);178            let new_place = Place {179                identifier: new_id,180                effect: old_place.effect,181                reactive: old_place.reactive,182                loc: old_place.loc,183            };184            let state = self.states.get_mut(&block_id).unwrap();185            state.incomplete_phis.push(IncompletePhi {186                old_place: old_place.clone(),187                new_place,188            });189            state.defs.insert(old_place.identifier, new_id);190            return new_id;191        }192193        if preds.len() == 1 {194            let pred = preds[0];195            let new_id = self.get_id_at(old_place, pred, env);196            self.states197                .get_mut(&block_id)198                .unwrap()199                .defs200                .insert(old_place.identifier, new_id);201            return new_id;202        }203204        let new_id = self.make_id(old_place.identifier, env);205        self.states206            .get_mut(&block_id)207            .unwrap()208            .defs209            .insert(old_place.identifier, new_id);210        let new_place = Place {211            identifier: new_id,212            effect: old_place.effect,213            reactive: old_place.reactive,214            loc: old_place.loc,215        };216        self.add_phi(block_id, old_place, &new_place, env);217        new_id218    }219220    fn add_phi(221        &mut self,222        block_id: BlockId,223        old_place: &Place,224        new_place: &Place,225        env: &mut Environment,226    ) {227        let preds = self.block_preds.get(&block_id).cloned().unwrap_or_default();228229        let mut pred_defs: IndexMap<BlockId, Place, FxBuildHasher> = IndexMap::default();230        for pred_block_id in &preds {231            let pred_id = self.get_id_at(old_place, *pred_block_id, env);232            pred_defs.insert(233                *pred_block_id,234                Place {235                    identifier: pred_id,236                    effect: old_place.effect,237                    reactive: old_place.reactive,238                    loc: old_place.loc,239                },240            );241        }242243        let phi = Phi {244            place: new_place.clone(),245            operands: pred_defs,246        };247248        self.pending_phis.entry(block_id).or_default().push(phi);249    }250251    fn fix_incomplete_phis(&mut self, block_id: BlockId, env: &mut Environment) {252        let incomplete_phis: Vec<IncompletePhi> = self253            .states254            .get_mut(&block_id)255            .unwrap()256            .incomplete_phis257            .drain(..)258            .collect();259        for phi in &incomplete_phis {260            self.add_phi(block_id, &phi.old_place, &phi.new_place, env);261        }262    }263264    fn start_block(&mut self, block_id: BlockId) {265        self.current = Some(block_id);266        self.states.insert(267            block_id,268            State {269                defs: FxHashMap::default(),270                incomplete_phis: Vec::new(),271            },272        );273    }274}275276// =============================================================================277// Public entry point278// =============================================================================279280pub fn enter_ssa(func: &mut HirFunction, env: &mut Environment) -> Result<(), CompilerDiagnostic> {281    let mut builder = SSABuilder::new(&func.body.blocks);282    let root_entry = func.body.entry;283    enter_ssa_impl(func, &mut builder, env, root_entry)?;284285    // Apply all pending phis to the actual blocks286    apply_pending_phis(func, env, &mut builder);287288    Ok(())289}290291fn apply_pending_phis(func: &mut HirFunction, env: &mut Environment, builder: &mut SSABuilder) {292    for (block_id, block) in func.body.blocks.iter_mut() {293        if let Some(phis) = builder.pending_phis.remove(block_id) {294            block.phis.extend(phis);295        }296    }297    for fid in &builder.processed_functions.clone() {298        let inner_func = &mut env.functions[fid.0 as usize];299        for (block_id, block) in inner_func.body.blocks.iter_mut() {300            if let Some(phis) = builder.pending_phis.remove(block_id) {301                block.phis.extend(phis);302            }303        }304    }305}306307fn enter_ssa_impl(308    func: &mut HirFunction,309    builder: &mut SSABuilder,310    env: &mut Environment,311    root_entry: BlockId,312) -> Result<(), CompilerDiagnostic> {313    let mut visited_blocks: FxHashSet<BlockId> = FxHashSet::default();314    let block_ids: Vec<BlockId> = func.body.blocks.keys().copied().collect();315316    for block_id in &block_ids {317        let block_id = *block_id;318319        if visited_blocks.contains(&block_id) {320            return Err(CompilerDiagnostic::new(321                ErrorCategory::Invariant,322                format!("found a cycle! visiting bb{} again", block_id.0),323                None,324            ));325        }326327        visited_blocks.insert(block_id);328        builder.start_block(block_id);329330        // Handle params at the root entry331        if block_id == root_entry {332            if !func.context.is_empty() {333                return Err(CompilerDiagnostic::new(334                    ErrorCategory::Invariant,335                    "Expected function context to be empty for outer function declarations",336                    None,337                ));338            }339            let params = std::mem::take(&mut func.params);340            let mut new_params = Vec::with_capacity(params.len());341            for param in params {342                new_params.push(match param {343                    ParamPattern::Place(p) => ParamPattern::Place(builder.define_place(&p, env)?),344                    ParamPattern::Spread(s) => ParamPattern::Spread(SpreadPattern {345                        place: builder.define_place(&s.place, env)?,346                    }),347                });348            }349            func.params = new_params;350        }351352        // Process instructions353        let instruction_ids: Vec<InstructionId> = func354            .body355            .blocks356            .get(&block_id)357            .unwrap()358            .instructions359            .clone();360361        for instr_id in &instruction_ids {362            let instr_idx = instr_id.0 as usize;363            let instr = &mut func.instructions[instr_idx];364365            // For FunctionExpression/ObjectMethod, we need to handle context366            // mapping specially because env.functions is borrowed by the closure.367            // First, check if this is a FunctionExpression/ObjectMethod and handle368            // context mapping separately.369            let func_expr_id = match &instr.value {370                InstructionValue::FunctionExpression { lowered_func, .. }371                | InstructionValue::ObjectMethod { lowered_func, .. } => Some(lowered_func.func),372                _ => None,373            };374375            // Map context places for function expressions before other operands376            if let Some(fid) = func_expr_id {377                let context = std::mem::take(&mut env.functions[fid.0 as usize].context);378                env.functions[fid.0 as usize].context = context379                    .into_iter()380                    .map(|place| builder.get_place(&place, env))381                    .collect();382            }383384            // Map non-context operands385            visitors::for_each_instruction_value_operand_mut(&mut instr.value, &mut |place| {386                *place = builder.get_place(place, env);387            });388389            // Map lvalues (skip DeclareContext/StoreContext — context variables390            // don't participate in SSA renaming)391            let instr = &mut func.instructions[instr_idx];392            let mut lvalue_err: Option<CompilerDiagnostic> = None;393            visitors::for_each_instruction_lvalue_mut(instr, &mut |place| {394                if lvalue_err.is_none() {395                    match builder.define_place(place, env) {396                        Ok(new_place) => *place = new_place,397                        Err(e) => lvalue_err = Some(e),398                    }399                }400            });401            if let Some(e) = lvalue_err {402                return Err(e);403            }404405            // Handle inner function SSA406            if let Some(fid) = func_expr_id {407                let context_ids: Vec<IdentifierId> = env.functions[fid.0 as usize]408                    .context409                    .iter()410                    .map(|place| place.identifier)411                    .collect();412                for id in context_ids {413                    builder.unmark_unknown(id);414                }415                builder.processed_functions.push(fid);416                let inner_func = &mut env.functions[fid.0 as usize];417                let inner_entry = inner_func.body.entry;418                let entry_block = inner_func.body.blocks.get_mut(&inner_entry).unwrap();419420                if !entry_block.preds.is_empty() {421                    return Err(CompilerDiagnostic::new(422                        ErrorCategory::Invariant,423                        "Expected function expression entry block to have zero predecessors",424                        None,425                    ));426                }427                entry_block.preds.insert(block_id);428429                builder.define_function(inner_func);430431                let saved_current = builder.current;432433                // Map inner function params434                let inner_params = std::mem::take(&mut env.functions[fid.0 as usize].params);435                let mut new_inner_params = Vec::with_capacity(inner_params.len());436                for param in inner_params {437                    new_inner_params.push(match param {438                        ParamPattern::Place(p) => {439                            ParamPattern::Place(builder.define_place(&p, env)?)440                        }441                        ParamPattern::Spread(s) => ParamPattern::Spread(SpreadPattern {442                            place: builder.define_place(&s.place, env)?,443                        }),444                    });445                }446                env.functions[fid.0 as usize].params = new_inner_params;447448                // Take the inner function out of the arena to process it449                let mut inner_func =450                    std::mem::replace(&mut env.functions[fid.0 as usize], placeholder_function());451452                enter_ssa_impl(&mut inner_func, builder, env, root_entry)?;453454                // Put it back455                env.functions[fid.0 as usize] = inner_func;456457                builder.current = saved_current;458459                // Clear entry preds460                env.functions[fid.0 as usize]461                    .body462                    .blocks463                    .get_mut(&inner_entry)464                    .unwrap()465                    .preds466                    .clear();467                builder.block_preds.insert(inner_entry, Vec::new());468            }469        }470471        // Map terminal operands472        let terminal = &mut func.body.blocks.get_mut(&block_id).unwrap().terminal;473        visitors::for_each_terminal_operand_mut(terminal, &mut |place| {474            *place = builder.get_place(place, env);475        });476477        // Handle successors478        let terminal_ref = &func.body.blocks.get(&block_id).unwrap().terminal;479        let successors = visitors::each_terminal_successor(terminal_ref);480        for output_id in successors {481            let output_preds_len = builder482                .block_preds483                .get(&output_id)484                .map(|p| p.len() as u32)485                .unwrap_or(0);486487            let count = if builder.unsealed_preds.contains_key(&output_id) {488                builder.unsealed_preds[&output_id] - 1489            } else {490                output_preds_len - 1491            };492            builder.unsealed_preds.insert(output_id, count);493494            if count == 0 && visited_blocks.contains(&output_id) {495                builder.fix_incomplete_phis(output_id, env);496            }497        }498    }499500    Ok(())501}502503/// Create a placeholder HirFunction for temporarily swapping an inner function504/// out of `env.functions` via `std::mem::replace`. The placeholder is never505/// read — the real function is swapped back immediately after processing.506pub fn placeholder_function() -> HirFunction {507    HirFunction {508        loc: None,509        id: None,510        name_hint: None,511        fn_type: ReactFunctionType::Other,512        params: Vec::new(),513        return_type_annotation: None,514        returns: Place {515            identifier: IdentifierId(0),516            effect: Effect::Unknown,517            reactive: false,518            loc: None,519        },520        context: Vec::new(),521        body: HIR {522            entry: BlockId(0),523            blocks: IndexMap::default(),524        },525        instructions: Vec::new(),526        generator: false,527        is_async: false,528        directives: Vec::new(),529        aliasing_effects: None,530    }531}