1/**2 * Copyright (c) Meta Platforms, Inc. and affiliates.3 *4 * This source code is licensed under the MIT license found in the5 * LICENSE file in the root directory of this source tree.6 */7use rustc_hash::{FxHashMap, FxHashSet};89use react_compiler_ast::common::BaseNode;10use react_compiler_ast::declarations::{11    ImportDeclaration, ImportKind, ImportSpecifier, ImportSpecifierData, ModuleExportName,12};13use react_compiler_ast::expressions::{CallExpression, Expression, Identifier};14use react_compiler_ast::literals::StringLiteral;15use react_compiler_ast::patterns::{16    ObjectPattern, ObjectPatternProp, ObjectPatternProperty, PatternLike,17};18use react_compiler_ast::scope::ScopeInfo;19use react_compiler_ast::statements::{20    Statement, VariableDeclaration, VariableDeclarationKind, VariableDeclarator,21};22use react_compiler_ast::{Program, SourceType};23use react_compiler_diagnostics::{24    CompilerError, CompilerErrorDetail, ErrorCategory, Position, SourceLocation,25};2627use super::compile_result::{DebugLogEntry, LoggerEvent, OrderedLogItem};28use super::plugin_options::{CompilerTarget, PluginOptions};29use super::suppression::SuppressionRange;30use crate::timing::TimingData;3132/// An import specifier tracked by ProgramContext.33/// Corresponds to NonLocalImportSpecifier in the TS compiler.34#[derive(Debug, Clone)]35pub struct NonLocalImportSpecifier {36    pub name: String,37    pub module: String,38    pub imported: String,39}4041/// Context for the program being compiled.42/// Tracks compiled functions, generated names, and import requirements.43/// Equivalent to ProgramContext class in Imports.ts.44pub struct ProgramContext {45    pub opts: PluginOptions,46    pub filename: Option<String>,47    /// The source filename from the parser's sourceFilename option.48    /// This is the filename stored on AST node `loc.filename` fields,49    /// which may differ from `filename` (e.g., no path prefix).50    source_filename: Option<String>,51    pub code: Option<String>,52    pub react_runtime_module: String,53    pub suppressions: Vec<SuppressionRange>,54    pub has_module_scope_opt_out: bool,55    pub events: Vec<LoggerEvent>,56    /// Unified ordered log that interleaves events and debug entries57    /// in the order they were emitted during compilation.58    pub ordered_log: Vec<OrderedLogItem>,5960    // Pre-resolved import local names for codegen61    pub instrument_fn_name: Option<String>,62    pub instrument_gating_name: Option<String>,63    pub hook_guard_name: Option<String>,6465    // Variable renames from lowering, to be applied back to the Babel AST66    pub renames: Vec<react_compiler_hir::environment::BindingRename>,6768    /// Timing data for profiling. Accumulates across all function compilations.69    pub timing: TimingData,7071    /// Whether debug logging is enabled (HIR formatting after each pass).72    pub debug_enabled: bool,7374    // Internal state75    already_compiled: FxHashSet<u32>,76    known_referenced_names: FxHashSet<String>,77    imports: FxHashMap<String, FxHashMap<String, NonLocalImportSpecifier>>,78}7980impl ProgramContext {81    pub fn new(82        opts: PluginOptions,83        filename: Option<String>,84        code: Option<String>,85        suppressions: Vec<SuppressionRange>,86        has_module_scope_opt_out: bool,87    ) -> Self {88        let react_runtime_module = get_react_compiler_runtime_module(&opts.target);89        let profiling = opts.profiling;90        let debug_enabled = opts.debug;91        Self {92            opts,93            filename,94            source_filename: None,95            code,96            react_runtime_module,97            suppressions,98            has_module_scope_opt_out,99            events: Vec::new(),100            ordered_log: Vec::new(),101            instrument_fn_name: None,102            instrument_gating_name: None,103            hook_guard_name: None,104            renames: Vec::new(),105            timing: TimingData::new(profiling),106            debug_enabled,107            already_compiled: FxHashSet::default(),108            known_referenced_names: FxHashSet::default(),109            imports: FxHashMap::default(),110        }111    }112113    /// Set the source filename (from AST node loc.filename).114    pub fn set_source_filename(&mut self, filename: Option<String>) {115        if self.source_filename.is_none() {116            self.source_filename = filename;117        }118    }119120    /// Get the source filename for logger events.121    pub fn source_filename(&self) -> Option<String> {122        self.source_filename.clone()123    }124125    /// Check if a function at the given start position has already been compiled.126    /// This is a workaround for Babel not consistently respecting skip().127    pub fn is_already_compiled(&self, start: u32) -> bool {128        self.already_compiled.contains(&start)129    }130131    /// Mark a function at the given start position as compiled.132    pub fn mark_compiled(&mut self, start: u32) {133        self.already_compiled.insert(start);134    }135136    /// Initialize known referenced names from scope bindings.137    /// Call this after construction to seed conflict detection with program scope bindings.138    pub fn init_from_scope(&mut self, scope: &ScopeInfo) {139        // Register ALL bindings (not just program-scope) so that UID generation140        // avoids name conflicts with any binding in the file. This matches141        // Babel's generateUid() which checks all scopes.142        for binding in &scope.bindings {143            self.known_referenced_names.insert(binding.name.clone());144        }145    }146147    /// Check if a name conflicts with known references.148    pub fn has_reference(&self, name: &str) -> bool {149        self.known_referenced_names.contains(name)150    }151152    /// Generate a unique identifier name that doesn't conflict with existing bindings.153    ///154    /// For hook names (use*), preserves the original name to avoid breaking155    /// hook-name-based type inference. For other names, prefixes with underscore156    /// similar to Babel's generateUid.157    pub fn new_uid(&mut self, name: &str) -> String {158        if is_hook_name(name) {159            // Don't prefix hooks with underscore, since InferTypes might160            // type HookKind based on callee naming convention.161            let mut uid = name.to_string();162            let mut i = 0;163            while self.has_reference(&uid) {164                uid = format!("{}_{}", name, i);165                i += 1;166            }167            self.known_referenced_names.insert(uid.clone());168            uid169        } else if !self.has_reference(name) {170            self.known_referenced_names.insert(name.to_string());171            name.to_string()172        } else {173            // Generate unique name with underscore prefix (similar to Babel's generateUid).174            // Babel strips leading underscores before prefixing, so:175            //   generateUid("_c") → strips to "c" → generates "_c", "_c2", "_c3", ...176            //   generateUid("foo") → generates "_foo", "_foo2", "_foo3", ...177            let base = name.trim_start_matches('_');178            let mut uid = format!("_{}", base);179            let mut i = 2;180            while self.has_reference(&uid) {181                uid = format!("_{}{}", base, i);182                i += 1;183            }184            self.known_referenced_names.insert(uid.clone());185            uid186        }187    }188189    /// Add the memo cache import (the `c` function from the compiler runtime).190    pub fn add_memo_cache_import(&mut self) -> NonLocalImportSpecifier {191        let module = self.react_runtime_module.clone();192        self.add_import_specifier(&module, "c", Some("_c"))193    }194195    /// Add an import specifier, reusing an existing one if it was already added.196    ///197    /// If `name_hint` is provided, it will be used as the basis for the local198    /// name; otherwise `specifier` is used.199    pub fn add_import_specifier(200        &mut self,201        module: &str,202        specifier: &str,203        name_hint: Option<&str>,204    ) -> NonLocalImportSpecifier {205        // Check if already imported206        if let Some(module_imports) = self.imports.get(module) {207            if let Some(existing) = module_imports.get(specifier) {208                return existing.clone();209            }210        }211212        let name = self.new_uid(name_hint.unwrap_or(specifier));213        let binding = NonLocalImportSpecifier {214            name,215            module: module.to_string(),216            imported: specifier.to_string(),217        };218219        self.imports220            .entry(module.to_string())221            .or_default()222            .insert(specifier.to_string(), binding.clone());223224        binding225    }226227    /// Register a name as referenced so future uid generation avoids it.228    pub fn add_new_reference(&mut self, name: String) {229        self.known_referenced_names.insert(name);230    }231232    /// Get the set of known referenced names for seeding per-function Environment UID generation.233    pub fn known_referenced_names(&self) -> &FxHashSet<String> {234        &self.known_referenced_names235    }236237    /// Merge UID names generated during a function compilation back into the program context,238    /// so subsequent function compilations avoid collisions.239    pub fn merge_uid_known_names(&mut self, names: &FxHashSet<String>) {240        self.known_referenced_names.extend(names.iter().cloned());241    }242243    /// Log a compilation event.244    pub fn log_event(&mut self, event: LoggerEvent) {245        self.ordered_log.push(OrderedLogItem::Event {246            event: event.clone(),247        });248        self.events.push(event);249    }250251    /// Log a debug entry (for debugLogIRs support).252    pub fn log_debug(&mut self, entry: DebugLogEntry) {253        self.ordered_log.push(OrderedLogItem::Debug { entry });254    }255256    /// Check if there are any pending imports to add to the program.257    pub fn has_pending_imports(&self) -> bool {258        !self.imports.is_empty()259    }260261    /// Get an immutable view of the generated imports.262    pub fn imports(&self) -> &FxHashMap<String, FxHashMap<String, NonLocalImportSpecifier>> {263        &self.imports264    }265}266267/// Check for blocklisted import modules.268/// Returns a CompilerError if any blocklisted imports are found.269pub fn validate_restricted_imports(270    program: &Program,271    blocklisted: &Option<Vec<String>>,272) -> Option<CompilerError> {273    let blocklisted = match blocklisted {274        Some(b) if !b.is_empty() => b,275        _ => return None,276    };277    let restricted: FxHashSet<&str> = blocklisted.iter().map(|s| s.as_str()).collect();278    let mut error = CompilerError::new();279280    for stmt in &program.body {281        if let Statement::ImportDeclaration(import) = stmt {282            if import283                .source284                .value285                .as_str()286                .is_some_and(|v| restricted.contains(v))287            {288                let mut detail = CompilerErrorDetail::new(289                    ErrorCategory::Todo,290                    "Bailing out due to blocklisted import",291                )292                .with_description(format!("Import from module {}", import.source.value));293                detail.loc = import.base.loc.as_ref().map(|loc| SourceLocation {294                    start: Position {295                        line: loc.start.line,296                        column: loc.start.column,297                        index: loc.start.index,298                    },299                    end: Position {300                        line: loc.end.line,301                        column: loc.end.column,302                        index: loc.end.index,303                    },304                });305                error.push_error_detail(detail);306            }307        }308    }309310    if error.has_any_errors() {311        Some(error)312    } else {313        None314    }315}316317/// Insert import declarations into the program body.318/// Handles both ESM imports and CommonJS require.319///320/// For existing imports of the same module (non-namespaced, value imports),321/// new specifiers are merged into the existing declaration. Otherwise,322/// new import/require statements are prepended to the program body.323pub fn add_imports_to_program(program: &mut Program, context: &ProgramContext) {324    if context.imports.is_empty() {325        return;326    }327328    // Collect existing non-namespaced imports by module name329    let existing_import_indices: FxHashMap<String, usize> = program330        .body331        .iter()332        .enumerate()333        .filter_map(|(idx, stmt)| {334            if let Statement::ImportDeclaration(import) = stmt {335                if is_non_namespaced_import(import) {336                    return Some((import.source.value.to_marker_string(), idx));337                }338            }339            None340        })341        .collect();342343    let mut stmts: Vec<Statement> = Vec::new();344    let mut sorted_modules: Vec<_> = context.imports.iter().collect();345    sorted_modules.sort_by(|(a, _), (b, _)| a.to_lowercase().cmp(&b.to_lowercase()));346347    for (module_name, imports_map) in sorted_modules {348        let sorted_imports = {349            let mut sorted: Vec<_> = imports_map.values().collect();350            sorted.sort_by_key(|s| &s.imported);351            sorted352        };353354        let import_specifiers: Vec<ImportSpecifier> = sorted_imports355            .iter()356            .map(|spec| make_import_specifier(spec))357            .collect();358359        // If an existing import of this module exists, merge into it360        if let Some(&idx) = existing_import_indices.get(module_name.as_str()) {361            if let Statement::ImportDeclaration(ref mut import) = program.body[idx] {362                import.specifiers.extend(import_specifiers);363            }364        } else if matches!(program.source_type, SourceType::Module) {365            // ESM: import { ... } from 'module'366            stmts.push(Statement::ImportDeclaration(ImportDeclaration {367                base: BaseNode::typed("ImportDeclaration"),368                specifiers: import_specifiers,369                source: StringLiteral {370                    base: BaseNode::typed("StringLiteral"),371                    value: module_name.clone().into(),372                },373                import_kind: None,374                assertions: None,375                attributes: None,376            }));377        } else {378            // CommonJS: const { imported: local, ... } = require('module')379            let properties: Vec<ObjectPatternProperty> = sorted_imports380                .iter()381                .map(|spec| {382                    ObjectPatternProperty::ObjectProperty(ObjectPatternProp {383                        base: BaseNode::typed("ObjectProperty"),384                        key: Box::new(Expression::Identifier(Identifier {385                            base: BaseNode::typed("Identifier"),386                            name: spec.imported.clone(),387                            type_annotation: None,388                            optional: None,389                            decorators: None,390                        })),391                        value: Box::new(PatternLike::Identifier(Identifier {392                            base: BaseNode::typed("Identifier"),393                            name: spec.name.clone(),394                            type_annotation: None,395                            optional: None,396                            decorators: None,397                        })),398                        computed: false,399                        shorthand: false,400                        decorators: None,401                        method: None,402                    })403                })404                .collect();405406            stmts.push(Statement::VariableDeclaration(VariableDeclaration {407                base: BaseNode::typed("VariableDeclaration"),408                kind: VariableDeclarationKind::Const,409                declarations: vec![VariableDeclarator {410                    base: BaseNode::typed("VariableDeclarator"),411                    id: PatternLike::ObjectPattern(ObjectPattern {412                        base: BaseNode::typed("ObjectPattern"),413                        properties,414                        type_annotation: None,415                        decorators: None,416                    }),417                    init: Some(Box::new(Expression::CallExpression(CallExpression {418                        base: BaseNode::typed("CallExpression"),419                        callee: Box::new(Expression::Identifier(Identifier {420                            base: BaseNode::typed("Identifier"),421                            name: "require".to_string(),422                            type_annotation: None,423                            optional: None,424                            decorators: None,425                        })),426                        arguments: vec![Expression::StringLiteral(StringLiteral {427                            base: BaseNode::typed("StringLiteral"),428                            value: module_name.clone().into(),429                        })],430                        type_parameters: None,431                        type_arguments: None,432                        optional: None,433                    }))),434                    definite: None,435                }],436                declare: None,437            }));438        }439    }440441    // Prepend new import statements to the program body442    if !stmts.is_empty() {443        let mut new_body = stmts;444        new_body.append(&mut program.body);445        program.body = new_body;446    }447}448449/// Create an ImportSpecifier AST node from a NonLocalImportSpecifier.450fn make_import_specifier(spec: &NonLocalImportSpecifier) -> ImportSpecifier {451    ImportSpecifier::ImportSpecifier(ImportSpecifierData {452        base: BaseNode::typed("ImportSpecifier"),453        local: Identifier {454            base: BaseNode::typed("Identifier"),455            name: spec.name.clone(),456            type_annotation: None,457            optional: None,458            decorators: None,459        },460        imported: ModuleExportName::Identifier(Identifier {461            base: BaseNode::typed("Identifier"),462            name: spec.imported.clone(),463            type_annotation: None,464            optional: None,465            decorators: None,466        }),467        import_kind: None,468    })469}470471/// Check if an import declaration is a non-namespaced value import.472/// Matches `import { ... } from 'module'` but NOT:473///   - `import * as Foo from 'module'` (namespace)474///   - `import type { Foo } from 'module'` (type import)475///   - `import typeof { Foo } from 'module'` (typeof import)476fn is_non_namespaced_import(import: &ImportDeclaration) -> bool {477    import478        .specifiers479        .iter()480        .all(|s| matches!(s, ImportSpecifier::ImportSpecifier(_)))481        && import482            .import_kind483            .as_ref()484            .map_or(true, |k| matches!(k, ImportKind::Value))485}486487/// Check if a name follows the React hook naming convention (use[A-Z0-9]...).488fn is_hook_name(name: &str) -> bool {489    let bytes = name.as_bytes();490    bytes.len() >= 4491        && bytes[0] == b'u'492        && bytes[1] == b's'493        && bytes[2] == b'e'494        && bytes495            .get(3)496            .map_or(false, |c| c.is_ascii_uppercase() || c.is_ascii_digit())497}498499/// Get the runtime module name based on the compiler target.500pub fn get_react_compiler_runtime_module(target: &CompilerTarget) -> String {501    match target {502        CompilerTarget::Version(v) if v == "19" => "react/compiler-runtime".to_string(),503        CompilerTarget::Version(v) if v == "17" || v == "18" => {504            "react-compiler-runtime".to_string()505        }506        CompilerTarget::MetaInternal { runtime_module, .. } => runtime_module.clone(),507        // Default to React 19 runtime for unrecognized versions508        CompilerTarget::Version(_) => "react/compiler-runtime".to_string(),509    }510}