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//! PruneHoistedContexts — removes hoisted context variable declarations7//! and transforms references to their original instruction kinds.8//!9//! Corresponds to `src/ReactiveScopes/PruneHoistedContexts.ts`.1011use rustc_hash::FxHashMap;1213use react_compiler_diagnostics::{CompilerError, CompilerErrorDetail, ErrorCategory};14use react_compiler_hir::{15    EvaluationOrder, IdentifierId, InstructionKind, InstructionValue, Place, ReactiveFunction,16    ReactiveInstruction, ReactiveScopeBlock, ReactiveStatement, ReactiveValue,17    environment::Environment,18};1920use crate::visitors::{ReactiveFunctionTransform, Transformed, transform_reactive_function};2122// =============================================================================23// Public entry point24// =============================================================================2526/// Prunes DeclareContexts lowered for HoistedConsts and transforms any27/// references back to their original instruction kind.28/// TS: `pruneHoistedContexts`29pub fn prune_hoisted_contexts(30    func: &mut ReactiveFunction,31    env: &Environment,32) -> Result<(), CompilerError> {33    let mut transform = Transform { env };34    let mut state = VisitorState {35        active_scopes: Vec::new(),36        uninitialized: FxHashMap::default(),37    };38    transform_reactive_function(func, &mut transform, &mut state)39}4041// =============================================================================42// State43// =============================================================================4445#[derive(Debug, Clone)]46enum UninitializedKind {47    UnknownKind,48    Func { definition: Option<IdentifierId> },49}5051struct VisitorState {52    active_scopes: Vec<rustc_hash::FxHashSet<IdentifierId>>,53    uninitialized: FxHashMap<IdentifierId, UninitializedKind>,54}5556impl VisitorState {57    fn find_in_active_scopes(&self, id: IdentifierId) -> bool {58        for scope in &self.active_scopes {59            if scope.contains(&id) {60                return true;61            }62        }63        false64    }65}6667struct Transform<'a> {68    env: &'a Environment,69}7071impl<'a> ReactiveFunctionTransform for Transform<'a> {72    type State = VisitorState;7374    fn env(&self) -> &Environment {75        self.env76    }7778    fn visit_scope(79        &mut self,80        scope: &mut ReactiveScopeBlock,81        state: &mut VisitorState,82    ) -> Result<(), CompilerError> {83        let scope_data = &self.env.scopes[scope.scope.0 as usize];84        let decl_ids: rustc_hash::FxHashSet<IdentifierId> =85            scope_data.declarations.iter().map(|(id, _)| *id).collect();8687        // Add declared but not initialized variables88        for (_, decl) in &scope_data.declarations {89            state90                .uninitialized91                .insert(decl.identifier, UninitializedKind::UnknownKind);92        }9394        state.active_scopes.push(decl_ids);95        self.traverse_scope(scope, state)?;96        state.active_scopes.pop();9798        // Clean up uninitialized after scope99        let scope_data = &self.env.scopes[scope.scope.0 as usize];100        for (_, decl) in &scope_data.declarations {101            state.uninitialized.remove(&decl.identifier);102        }103        Ok(())104    }105106    fn visit_place(107        &mut self,108        _id: EvaluationOrder,109        place: &Place,110        state: &mut VisitorState,111    ) -> Result<(), CompilerError> {112        if let Some(kind) = state.uninitialized.get(&place.identifier) {113            if let UninitializedKind::Func { definition } = kind {114                if *definition != Some(place.identifier) {115                    let mut err = CompilerError::new();116                    err.push_error_detail(117                        CompilerErrorDetail::new(118                            ErrorCategory::Todo,119                            "[PruneHoistedContexts] Rewrite hoisted function references"120                                .to_string(),121                        )122                        .with_loc(place.loc),123                    );124                    return Err(err);125                }126            }127        }128        Ok(())129    }130131    fn transform_instruction(132        &mut self,133        instruction: &mut ReactiveInstruction,134        state: &mut VisitorState,135    ) -> Result<Transformed<ReactiveStatement>, CompilerError> {136        // Remove hoisted declarations to preserve TDZ137        if let ReactiveValue::Instruction(InstructionValue::DeclareContext { lvalue, .. }) =138            &instruction.value139        {140            let maybe_non_hoisted = convert_hoisted_lvalue_kind(lvalue.kind);141            if let Some(non_hoisted) = maybe_non_hoisted {142                if non_hoisted == InstructionKind::Function143                    && state.uninitialized.contains_key(&lvalue.place.identifier)144                {145                    state.uninitialized.insert(146                        lvalue.place.identifier,147                        UninitializedKind::Func { definition: None },148                    );149                }150                return Ok(Transformed::Remove);151            }152        }153154        if let ReactiveValue::Instruction(InstructionValue::StoreContext { lvalue, .. }) =155            &mut instruction.value156        {157            if lvalue.kind != InstructionKind::Reassign {158                let lvalue_id = lvalue.place.identifier;159                let is_declared_by_scope = state.find_in_active_scopes(lvalue_id);160                if is_declared_by_scope {161                    if lvalue.kind == InstructionKind::Let || lvalue.kind == InstructionKind::Const162                    {163                        lvalue.kind = InstructionKind::Reassign;164                    } else if lvalue.kind == InstructionKind::Function {165                        if let Some(kind) = state.uninitialized.get(&lvalue_id) {166                            if !matches!(kind, UninitializedKind::Func { .. }) {167                                let mut err = CompilerError::new();168                                err.push_error_detail(169                                    CompilerErrorDetail::new(170                                        ErrorCategory::Invariant,171                                        "[PruneHoistedContexts] Unexpected hoisted function"172                                            .to_string(),173                                    )174                                    .with_loc(instruction.loc),175                                );176                                return Err(err);177                            }178                            // References to hoisted functions are now "safe" as179                            // variable assignments have finished.180                            state.uninitialized.remove(&lvalue_id);181                        }182                    } else {183                        let mut err = CompilerError::new();184                        err.push_error_detail(185                            CompilerErrorDetail::new(186                                ErrorCategory::Todo,187                                "[PruneHoistedContexts] Unexpected kind".to_string(),188                            )189                            .with_loc(instruction.loc),190                        );191                        return Err(err);192                    }193                }194            }195        }196197        self.visit_instruction(instruction, state)?;198        Ok(Transformed::Keep)199    }200}201202/// Corresponds to TS `convertHoistedLValueKind` — returns None for non-hoisted kinds.203fn convert_hoisted_lvalue_kind(kind: InstructionKind) -> Option<InstructionKind> {204    match kind {205        InstructionKind::HoistedLet => Some(InstructionKind::Let),206        InstructionKind::HoistedConst => Some(InstructionKind::Const),207        InstructionKind::HoistedFunction => Some(InstructionKind::Function),208        _ => None,209    }210}