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//! Rewrites InstructionKind of instructions which declare/assign variables,7//! converting the first declaration to Const/Let depending on whether it is8//! subsequently reassigned, and ensuring that subsequent reassignments are9//! marked as Reassign.10//!11//! Ported from TypeScript `src/SSA/RewriteInstructionKindsBasedOnReassignment.ts`.12//!13//! Note that declarations which were const in the original program cannot become14//! `let`, but the inverse is not true: a `let` which was reassigned in the source15//! may be converted to a `const` if the reassignment is not used and was removed16//! by dead code elimination.1718use rustc_hash::FxHashMap;1920use react_compiler_diagnostics::{21    CompilerDiagnostic, CompilerDiagnosticDetail, CompilerError, ErrorCategory, SourceLocation,22};23use react_compiler_hir::visitors::each_pattern_operand;24use react_compiler_hir::{25    BlockKind, DeclarationId, HirFunction, InstructionKind, InstructionValue, ParamPattern, Place,26};2728use react_compiler_hir::environment::Environment;2930/// Create an invariant CompilerError (matches TS CompilerError.invariant).31/// When a loc is provided, creates a CompilerDiagnostic with an error detail item32/// (matching TS CompilerError.invariant which uses .withDetails()).33fn invariant_error(reason: &str, description: Option<String>) -> CompilerError {34    invariant_error_with_loc(reason, description, None)35}3637fn invariant_error_with_loc(38    reason: &str,39    description: Option<String>,40    loc: Option<SourceLocation>,41) -> CompilerError {42    let mut err = CompilerError::new();43    let diagnostic = CompilerDiagnostic::new(ErrorCategory::Invariant, reason, description)44        .with_detail(CompilerDiagnosticDetail::Error {45            loc,46            message: Some(reason.to_string()),47            identifier_name: None,48        });49    err.push_diagnostic(diagnostic);50    err51}5253/// Format an InstructionKind variant name (matches TS `${kind}` interpolation).54fn format_kind(kind: Option<InstructionKind>) -> String {55    match kind {56        Some(InstructionKind::Const) => "Const".to_string(),57        Some(InstructionKind::Let) => "Let".to_string(),58        Some(InstructionKind::Reassign) => "Reassign".to_string(),59        Some(InstructionKind::Catch) => "Catch".to_string(),60        Some(InstructionKind::HoistedConst) => "HoistedConst".to_string(),61        Some(InstructionKind::HoistedLet) => "HoistedLet".to_string(),62        Some(InstructionKind::HoistedFunction) => "HoistedFunction".to_string(),63        Some(InstructionKind::Function) => "Function".to_string(),64        None => "null".to_string(),65    }66}6768/// Format a Place like TS `printPlace()`: `<effect> <name>$<id>[<range>]{reactive}`69fn format_place(place: &Place, env: &Environment) -> String {70    let ident = &env.identifiers[place.identifier.0 as usize];71    let name = match &ident.name {72        Some(n) => n.value().to_string(),73        None => String::new(),74    };75    let scope = match ident.scope {76        Some(scope_id) => format!("_@{}", scope_id.0),77        None => String::new(),78    };79    let mutable_range = if ident.mutable_range.end.0 > ident.mutable_range.start.0 + 1 {80        format!(81            "[{}:{}]",82            ident.mutable_range.start.0, ident.mutable_range.end.083        )84    } else {85        String::new()86    };87    let reactive = if place.reactive { "{reactive}" } else { "" };88    format!(89        "{} {}${}{}{}{}",90        place.effect, name, place.identifier.0, scope, mutable_range, reactive91    )92}9394/// Index into a collected list of declaration mutations to apply.95///96/// We use a two-phase approach: first collect which declarations exist,97/// then apply mutations. This is because in the TS code, `declarations`98/// map stores references to LValue/LValuePattern and mutates `kind` through them.99/// In Rust, we track instruction indices and apply changes in a second pass.100enum DeclarationLoc {101    /// An LValue from DeclareLocal or StoreLocal — identified by (block_index, instr_index_in_block)102    Instruction {103        block_index: usize,104        instr_local_index: usize,105    },106    /// A parameter or context variable (seeded as Let, may be upgraded to Let on reassignment — already Let)107    ParamOrContext,108}109110pub fn rewrite_instruction_kinds_based_on_reassignment(111    func: &mut HirFunction,112    env: &Environment,113) -> Result<(), CompilerError> {114    // Phase 1: Collect all information about which declarations need updates.115    //116    // Track: for each DeclarationId, the location of its first declaration,117    // and whether it needs to be changed to Let (because of reassignment).118    let mut declarations: FxHashMap<DeclarationId, DeclarationLoc> = FxHashMap::default();119    // Track which (block_index, instr_local_index) should have their lvalue.kind set to Reassign120    let mut reassign_locs: Vec<(usize, usize)> = Vec::new();121    // Track which declaration locations need to be set to Let122    let mut let_locs: Vec<(usize, usize)> = Vec::new();123    // Track which (block_index, instr_local_index) should have their lvalue.kind set to Const124    let mut const_locs: Vec<(usize, usize)> = Vec::new();125    // Track which (block_index, instr_local_index) Destructure instructions get a specific kind126    let mut destructure_kind_locs: Vec<(usize, usize, InstructionKind)> = Vec::new();127128    // Seed with parameters129    for param in &func.params {130        let place: &Place = match param {131            ParamPattern::Place(p) => p,132            ParamPattern::Spread(s) => &s.place,133        };134        let ident = &env.identifiers[place.identifier.0 as usize];135        if ident.name.is_some() {136            declarations.insert(ident.declaration_id, DeclarationLoc::ParamOrContext);137        }138    }139140    // Seed with context variables141    for place in &func.context {142        let ident = &env.identifiers[place.identifier.0 as usize];143        if ident.name.is_some() {144            declarations.insert(ident.declaration_id, DeclarationLoc::ParamOrContext);145        }146    }147148    // Process all blocks149    let block_keys: Vec<_> = func.body.blocks.keys().cloned().collect();150    for (block_index, block_id) in block_keys.iter().enumerate() {151        let block = &func.body.blocks[block_id];152        let block_kind = block.kind;153        for (local_idx, instr_id) in block.instructions.iter().enumerate() {154            let instr = &func.instructions[instr_id.0 as usize];155            match &instr.value {156                InstructionValue::DeclareLocal { lvalue, .. } => {157                    let decl_id =158                        env.identifiers[lvalue.place.identifier.0 as usize].declaration_id;159                    if declarations.contains_key(&decl_id) {160                        return Err(invariant_error_with_loc(161                            "Expected variable not to be defined prior to declaration",162                            Some(format!(163                                "{} was already defined",164                                format_place(&lvalue.place, env),165                            )),166                            lvalue.place.loc,167                        ));168                    }169                    declarations.insert(170                        decl_id,171                        DeclarationLoc::Instruction {172                            block_index,173                            instr_local_index: local_idx,174                        },175                    );176                }177                InstructionValue::StoreLocal { lvalue, .. } => {178                    let ident = &env.identifiers[lvalue.place.identifier.0 as usize];179                    if ident.name.is_some() {180                        let decl_id = ident.declaration_id;181                        if let Some(existing) = declarations.get(&decl_id) {182                            // Reassignment: mark existing declaration as Let, current as Reassign183                            match existing {184                                DeclarationLoc::Instruction {185                                    block_index: bi,186                                    instr_local_index: ili,187                                } => {188                                    let_locs.push((*bi, *ili));189                                }190                                DeclarationLoc::ParamOrContext => {191                                    // Already Let, no-op192                                }193                            }194                            reassign_locs.push((block_index, local_idx));195                        } else {196                            // First store — mark as Const197                            // Mirrors TS: CompilerError.invariant(!declarations.has(...))198                            if declarations.contains_key(&decl_id) {199                                return Err(invariant_error_with_loc(200                                    "Expected variable not to be defined prior to declaration",201                                    Some(format!(202                                        "{} was already defined",203                                        format_place(&lvalue.place, env),204                                    )),205                                    lvalue.place.loc,206                                ));207                            }208                            declarations.insert(209                                decl_id,210                                DeclarationLoc::Instruction {211                                    block_index,212                                    instr_local_index: local_idx,213                                },214                            );215                            const_locs.push((block_index, local_idx));216                        }217                    }218                }219                InstructionValue::Destructure { lvalue, .. } => {220                    let mut kind: Option<InstructionKind> = None;221                    for place in each_pattern_operand(&lvalue.pattern) {222                        let ident = &env.identifiers[place.identifier.0 as usize];223                        if ident.name.is_none() {224                            if !(kind.is_none() || kind == Some(InstructionKind::Const)) {225                                return Err(invariant_error_with_loc(226                                    "Expected consistent kind for destructuring",227                                    Some(format!(228                                        "other places were `{}` but '{}' is const",229                                        format_kind(kind),230                                        format_place(&place, env),231                                    )),232                                    place.loc,233                                ));234                            }235                            kind = Some(InstructionKind::Const);236                        } else {237                            let decl_id = ident.declaration_id;238                            if let Some(existing) = declarations.get(&decl_id) {239                                // Reassignment240                                if !(kind.is_none() || kind == Some(InstructionKind::Reassign)) {241                                    return Err(invariant_error_with_loc(242                                        "Expected consistent kind for destructuring",243                                        Some(format!(244                                            "Other places were `{}` but '{}' is reassigned",245                                            format_kind(kind),246                                            format_place(&place, env),247                                        )),248                                        place.loc,249                                    ));250                                }251                                kind = Some(InstructionKind::Reassign);252                                match existing {253                                    DeclarationLoc::Instruction {254                                        block_index: bi,255                                        instr_local_index: ili,256                                    } => {257                                        let_locs.push((*bi, *ili));258                                    }259                                    DeclarationLoc::ParamOrContext => {260                                        // Already Let261                                    }262                                }263                            } else {264                                // New declaration265                                if block_kind == BlockKind::Value {266                                    return Err(invariant_error_with_loc(267                                        "TODO: Handle reassignment in a value block where the original declaration was removed by dead code elimination (DCE)",268                                        None,269                                        place.loc,270                                    ));271                                }272                                declarations.insert(273                                    decl_id,274                                    DeclarationLoc::Instruction {275                                        block_index,276                                        instr_local_index: local_idx,277                                    },278                                );279                                if !(kind.is_none() || kind == Some(InstructionKind::Const)) {280                                    return Err(invariant_error_with_loc(281                                        "Expected consistent kind for destructuring",282                                        Some(format!(283                                            "Other places were `{}` but '{}' is const",284                                            format_kind(kind),285                                            format_place(&place, env),286                                        )),287                                        place.loc,288                                    ));289                                }290                                kind = Some(InstructionKind::Const);291                            }292                        }293                    }294                    let kind =295                        kind.ok_or_else(|| invariant_error("Expected at least one operand", None))?;296                    destructure_kind_locs.push((block_index, local_idx, kind));297                }298                InstructionValue::PostfixUpdate { lvalue, .. }299                | InstructionValue::PrefixUpdate { lvalue, .. } => {300                    let ident = &env.identifiers[lvalue.identifier.0 as usize];301                    let decl_id = ident.declaration_id;302                    let Some(existing) = declarations.get(&decl_id) else {303                        return Err(invariant_error_with_loc(304                            "Expected variable to have been defined",305                            Some(format!("No declaration for {}", format_place(lvalue, env),)),306                            lvalue.loc,307                        ));308                    };309                    match existing {310                        DeclarationLoc::Instruction {311                            block_index: bi,312                            instr_local_index: ili,313                        } => {314                            let_locs.push((*bi, *ili));315                        }316                        DeclarationLoc::ParamOrContext => {317                            // Already Let318                        }319                    }320                }321                _ => {}322            }323        }324    }325326    // Phase 2: Apply all collected mutations.327328    // Helper: given (block_index, instr_local_index), get the InstructionId329    // and mutate the instruction's lvalue kind.330    for (bi, ili) in const_locs {331        let block_id = &block_keys[bi];332        let instr_id = func.body.blocks[block_id].instructions[ili];333        let instr = &mut func.instructions[instr_id.0 as usize];334        match &mut instr.value {335            InstructionValue::StoreLocal { lvalue, .. } => {336                lvalue.kind = InstructionKind::Const;337            }338            _ => {}339        }340    }341342    for (bi, ili) in reassign_locs {343        let block_id = &block_keys[bi];344        let instr_id = func.body.blocks[block_id].instructions[ili];345        let instr = &mut func.instructions[instr_id.0 as usize];346        match &mut instr.value {347            InstructionValue::StoreLocal { lvalue, .. } => {348                lvalue.kind = InstructionKind::Reassign;349            }350            _ => {}351        }352    }353354    // Apply destructure_kind_locs BEFORE let_locs: a Destructure that first355    // declares a variable gets kind=Const here, but if a later instruction356    // reassigns that variable the Destructure must become Let.  Applying357    // let_locs afterwards allows it to override the Const set here, matching358    // the TS behaviour where `declaration.kind = Let` mutates the original359    // lvalue reference after the Destructure's own `lvalue.kind = kind`.360    for (bi, ili, kind) in destructure_kind_locs {361        let block_id = &block_keys[bi];362        let instr_id = func.body.blocks[block_id].instructions[ili];363        let instr = &mut func.instructions[instr_id.0 as usize];364        match &mut instr.value {365            InstructionValue::Destructure { lvalue, .. } => {366                lvalue.kind = kind;367            }368            _ => {}369        }370    }371372    for (bi, ili) in let_locs {373        let block_id = &block_keys[bi];374        let instr_id = func.body.blocks[block_id].instructions[ili];375        let instr = &mut func.instructions[instr_id.0 as usize];376        match &mut instr.value {377            InstructionValue::DeclareLocal { lvalue, .. }378            | InstructionValue::StoreLocal { lvalue, .. } => {379                lvalue.kind = InstructionKind::Let;380            }381            InstructionValue::Destructure { lvalue, .. } => {382                lvalue.kind = InstructionKind::Let;383            }384            _ => {}385        }386    }387388    Ok(())389}