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//! Removes manual memoization using `useMemo` and `useCallback` APIs.7//!8//! For useMemo: replaces `Call useMemo(fn, deps)` with `Call fn()`9//! For useCallback: replaces `Call useCallback(fn, deps)` with `LoadLocal fn`10//!11//! When validation flags are set, inserts `StartMemoize`/`FinishMemoize` markers.12//!13//! Analogous to TS `Inference/DropManualMemoization.ts`.1415use rustc_hash::{FxHashMap, FxHashSet};1617use react_compiler_diagnostics::CompilerDiagnostic;18use react_compiler_diagnostics::CompilerDiagnosticDetail;19use react_compiler_diagnostics::ErrorCategory;20use react_compiler_hir::ArrayElement;21use react_compiler_hir::DependencyPathEntry;22use react_compiler_hir::Effect;23use react_compiler_hir::EvaluationOrder;24use react_compiler_hir::HirFunction;25use react_compiler_hir::IdentifierId;26use react_compiler_hir::IdentifierName;27use react_compiler_hir::Instruction;28use react_compiler_hir::InstructionId;29use react_compiler_hir::InstructionValue;30use react_compiler_hir::ManualMemoDependency;31use react_compiler_hir::ManualMemoDependencyRoot;32use react_compiler_hir::NonLocalBinding;33use react_compiler_hir::Place;34use react_compiler_hir::PlaceOrSpread;35use react_compiler_hir::PropertyLiteral;36use react_compiler_hir::SourceLocation;37use react_compiler_hir::environment::Environment;38use react_compiler_lowering::create_temporary_place;39use react_compiler_lowering::mark_instruction_ids;4041// =============================================================================42// Types43// =============================================================================4445#[derive(Debug, Clone, Copy, PartialEq, Eq)]46enum ManualMemoKind {47    UseMemo,48    UseCallback,49}5051#[derive(Debug, Clone)]52struct ManualMemoCallee {53    kind: ManualMemoKind,54    /// InstructionId of the LoadGlobal or PropertyLoad that loaded the callee.55    load_instr_id: InstructionId,56}5758struct IdentifierSidemap {59    /// Maps identifier id -> InstructionId of FunctionExpression instructions60    functions: FxHashSet<IdentifierId>,61    /// Maps identifier id -> ManualMemoCallee for useMemo/useCallback callees62    manual_memos: FxHashMap<IdentifierId, ManualMemoCallee>,63    /// Set of identifier ids that loaded 'React' global64    react: FxHashSet<IdentifierId>,65    /// Maps identifier id -> deps list info for array expressions66    maybe_deps_lists: FxHashMap<IdentifierId, MaybeDepsListInfo>,67    /// Maps identifier id -> ManualMemoDependency for dependency tracking68    maybe_deps: FxHashMap<IdentifierId, ManualMemoDependency>,69    /// Set of identifier ids that are results of optional chains70    optionals: FxHashSet<IdentifierId>,71}7273#[derive(Debug, Clone)]74struct MaybeDepsListInfo {75    loc: Option<SourceLocation>,76    deps: Vec<Place>,77}7879struct ExtractedMemoArgs {80    fn_place: Place,81    deps_list: Option<Vec<ManualMemoDependency>>,82    deps_loc: Option<SourceLocation>,83}8485// =============================================================================86// Main pass87// =============================================================================8889/// Drop manual memoization (useMemo/useCallback calls), replacing them90/// with direct invocations/references.91pub fn drop_manual_memoization(92    func: &mut HirFunction,93    env: &mut Environment,94) -> Result<(), CompilerDiagnostic> {95    let is_validation_enabled = env.validate_preserve_existing_memoization_guarantees96        || env.validate_no_set_state_in_render97        || env.enable_preserve_existing_memoization_guarantees;9899    let optionals = find_optional_places(func)?;100    let mut sidemap = IdentifierSidemap {101        functions: FxHashSet::default(),102        manual_memos: FxHashMap::default(),103        react: FxHashSet::default(),104        maybe_deps: FxHashMap::default(),105        maybe_deps_lists: FxHashMap::default(),106        optionals,107    };108    let mut next_manual_memo_id: u32 = 0;109110    // Phase 1:111    // - Overwrite manual memoization CallExpression/MethodCall112    // - (if validation is enabled) collect manual memoization markers113    //114    // queued_inserts maps InstructionId -> new Instruction to insert after that instruction115    let mut queued_inserts: FxHashMap<InstructionId, Instruction> = FxHashMap::default();116117    // Collect all block instruction lists up front to avoid borrowing func immutably118    // while needing to mutate it119    let all_block_instructions: Vec<Vec<InstructionId>> = func120        .body121        .blocks122        .values()123        .map(|block| block.instructions.clone())124        .collect();125126    for block_instructions in &all_block_instructions {127        for &instr_id in block_instructions {128            let instr = &func.instructions[instr_id.0 as usize];129130            // Extract the identifier we need to look up, and whether it's a call/method131            let lookup_id = match &instr.value {132                InstructionValue::CallExpression { callee, .. } => Some(callee.identifier),133                InstructionValue::MethodCall { property, .. } => Some(property.identifier),134                _ => None,135            };136137            let manual_memo = lookup_id.and_then(|id| sidemap.manual_memos.get(&id).cloned());138139            if let Some(manual_memo) = manual_memo {140                process_manual_memo_call(141                    func,142                    env,143                    instr_id,144                    &manual_memo,145                    &mut sidemap,146                    is_validation_enabled,147                    &mut next_manual_memo_id,148                    &mut queued_inserts,149                );150            } else {151                collect_temporaries(func, env, instr_id, &mut sidemap);152            }153        }154    }155156    // Phase 2: Insert manual memoization markers as needed157    if !queued_inserts.is_empty() {158        let mut has_changes = false;159        for block in func.body.blocks.values_mut() {160            let mut next_instructions: Option<Vec<InstructionId>> = None;161            for i in 0..block.instructions.len() {162                let instr_id = block.instructions[i];163                if let Some(insert_instr) = queued_inserts.remove(&instr_id) {164                    if next_instructions.is_none() {165                        next_instructions = Some(block.instructions[..i].to_vec());166                    }167                    let ni = next_instructions.as_mut().unwrap();168                    ni.push(instr_id);169                    // Add the new instruction to the flat table and get its InstructionId170                    let new_instr_id = InstructionId(func.instructions.len() as u32);171                    func.instructions.push(insert_instr);172                    ni.push(new_instr_id);173                } else if let Some(ni) = next_instructions.as_mut() {174                    ni.push(instr_id);175                }176            }177            if let Some(ni) = next_instructions {178                block.instructions = ni;179                has_changes = true;180            }181        }182183        if has_changes {184            mark_instruction_ids(&mut func.body, &mut func.instructions);185        }186    }187188    Ok(())189}190191// =============================================================================192// Phase 1 helpers193// =============================================================================194195#[allow(clippy::too_many_arguments)]196fn process_manual_memo_call(197    func: &mut HirFunction,198    env: &mut Environment,199    instr_id: InstructionId,200    manual_memo: &ManualMemoCallee,201    sidemap: &mut IdentifierSidemap,202    is_validation_enabled: bool,203    next_manual_memo_id: &mut u32,204    queued_inserts: &mut FxHashMap<InstructionId, Instruction>,205) {206    let instr = &func.instructions[instr_id.0 as usize];207208    let memo_details = extract_manual_memoization_args(instr, manual_memo.kind, sidemap, env);209210    let Some(memo_details) = memo_details else {211        return;212    };213214    let ExtractedMemoArgs {215        fn_place,216        deps_list,217        deps_loc,218    } = memo_details;219220    let loc = func.instructions[instr_id.0 as usize].value.loc().cloned();221222    // Replace the instruction value with the memoization replacement223    let replacement = get_manual_memoization_replacement(&fn_place, loc.clone(), manual_memo.kind);224    func.instructions[instr_id.0 as usize].value = replacement;225226    if is_validation_enabled {227        // Bail out when we encounter manual memoization without inline function expressions228        if !sidemap.functions.contains(&fn_place.identifier) {229            let mut diag = CompilerDiagnostic::new(230                ErrorCategory::UseMemo,231                "Expected the first argument to be an inline function expression",232                Some("Expected the first argument to be an inline function expression".to_string()),233            )234            .with_detail(CompilerDiagnosticDetail::Error {235                loc: fn_place.loc.clone(),236                message: Some(237                    "Expected the first argument to be an inline function expression".to_string(),238                ),239                identifier_name: None,240            });241            // Match TS behavior: suggestions is [] (empty array), not null242            diag.suggestions = Some(vec![]);243            env.record_diagnostic(diag);244            return;245        }246247        let memo_decl: Place = if manual_memo.kind == ManualMemoKind::UseMemo {248            func.instructions[instr_id.0 as usize].lvalue.clone()249        } else {250            Place {251                identifier: fn_place.identifier,252                effect: Effect::Unknown,253                reactive: false,254                loc: fn_place.loc.clone(),255            }256        };257258        let manual_memo_id = *next_manual_memo_id;259        *next_manual_memo_id += 1;260261        let (start_marker, finish_marker) = make_manual_memoization_markers(262            &fn_place,263            env,264            deps_list,265            deps_loc,266            &memo_decl,267            manual_memo_id,268        );269270        queued_inserts.insert(manual_memo.load_instr_id, start_marker);271        queued_inserts.insert(instr_id, finish_marker);272    }273}274275fn collect_temporaries(276    func: &HirFunction,277    env: &Environment,278    instr_id: InstructionId,279    sidemap: &mut IdentifierSidemap,280) {281    let instr = &func.instructions[instr_id.0 as usize];282    let lvalue_id = instr.lvalue.identifier;283284    match &instr.value {285        InstructionValue::FunctionExpression { .. } => {286            sidemap.functions.insert(lvalue_id);287        }288        InstructionValue::LoadGlobal { binding, .. } => {289            let hook_name = get_hook_detection_name(binding);290            let mut detected = false;291            if let Some(name) = hook_name {292                if name == "useMemo" {293                    sidemap.manual_memos.insert(294                        lvalue_id,295                        ManualMemoCallee {296                            kind: ManualMemoKind::UseMemo,297                            load_instr_id: instr_id,298                        },299                    );300                    detected = true;301                } else if name == "useCallback" {302                    sidemap.manual_memos.insert(303                        lvalue_id,304                        ManualMemoCallee {305                            kind: ManualMemoKind::UseCallback,306                            load_instr_id: instr_id,307                        },308                    );309                    detected = true;310                }311            }312            if !detected && binding.name() == "React" {313                sidemap.react.insert(lvalue_id);314            }315        }316        InstructionValue::PropertyLoad {317            object, property, ..318        } => {319            if sidemap.react.contains(&object.identifier) {320                if let PropertyLiteral::String(prop_name) = property {321                    if prop_name == "useMemo" {322                        sidemap.manual_memos.insert(323                            lvalue_id,324                            ManualMemoCallee {325                                kind: ManualMemoKind::UseMemo,326                                load_instr_id: instr_id,327                            },328                        );329                    } else if prop_name == "useCallback" {330                        sidemap.manual_memos.insert(331                            lvalue_id,332                            ManualMemoCallee {333                                kind: ManualMemoKind::UseCallback,334                                load_instr_id: instr_id,335                            },336                        );337                    }338                }339            }340        }341        InstructionValue::ArrayExpression { elements, .. } => {342            // Check if all elements are Identifier (Place) - no spreads or holes343            let all_places: Option<Vec<Place>> = elements344                .iter()345                .map(|e| match e {346                    ArrayElement::Place(p) => Some(p.clone()),347                    _ => None,348                })349                .collect();350351            if let Some(deps) = all_places {352                sidemap.maybe_deps_lists.insert(353                    lvalue_id,354                    MaybeDepsListInfo {355                        loc: instr.value.loc().cloned(),356                        deps,357                    },358                );359            }360        }361        _ => {}362    }363364    let is_optional = sidemap.optionals.contains(&lvalue_id);365    let maybe_dep =366        collect_maybe_memo_dependencies(&instr.value, &sidemap.maybe_deps, is_optional, env);367    if let Some(dep) = maybe_dep {368        // For StoreLocal, also insert under the StoreLocal's lvalue place identifier,369        // matching the TS behavior where collectMaybeMemoDependencies inserts into370        // maybeDeps directly for StoreLocal's target variable.371        if let InstructionValue::StoreLocal { lvalue, .. } = &instr.value {372            sidemap373                .maybe_deps374                .insert(lvalue.place.identifier, dep.clone());375        }376        sidemap.maybe_deps.insert(lvalue_id, dep);377    }378}379380// =============================================================================381// collectMaybeMemoDependencies382// =============================================================================383384/// Collect loads from named variables and property reads into `maybe_deps`.385/// Returns the variable + property reads represented by the instruction value.386pub fn collect_maybe_memo_dependencies(387    value: &InstructionValue,388    maybe_deps: &FxHashMap<IdentifierId, ManualMemoDependency>,389    optional: bool,390    env: &Environment,391) -> Option<ManualMemoDependency> {392    match value {393        InstructionValue::LoadGlobal { binding, loc, .. } => Some(ManualMemoDependency {394            root: ManualMemoDependencyRoot::Global {395                identifier_name: binding.name().to_string(),396            },397            path: vec![],398            loc: loc.clone(),399        }),400        InstructionValue::PropertyLoad {401            object,402            property,403            loc,404            ..405        } => {406            if let Some(object_dep) = maybe_deps.get(&object.identifier) {407                Some(ManualMemoDependency {408                    root: object_dep.root.clone(),409                    path: {410                        let mut path = object_dep.path.clone();411                        path.push(DependencyPathEntry {412                            property: property.clone(),413                            optional,414                            loc: loc.clone(),415                        });416                        path417                    },418                    loc: loc.clone(),419                })420            } else {421                None422            }423        }424        InstructionValue::LoadLocal { place, .. } | InstructionValue::LoadContext { place, .. } => {425            if let Some(source) = maybe_deps.get(&place.identifier) {426                Some(source.clone())427            } else if matches!(428                &env.identifiers[place.identifier.0 as usize].name,429                Some(IdentifierName::Named(_))430            ) {431                Some(ManualMemoDependency {432                    root: ManualMemoDependencyRoot::NamedLocal {433                        value: place.clone(),434                        constant: false,435                    },436                    path: vec![],437                    loc: place.loc.clone(),438                })439            } else {440                None441            }442        }443        InstructionValue::StoreLocal {444            lvalue, value: val, ..445        } => {446            // Value blocks rely on StoreLocal to populate their return value.447            // We need to track these as optional property chains are valid in448            // source depslists449            let lvalue_id = lvalue.place.identifier;450            let rvalue_id = val.identifier;451            if let Some(aliased) = maybe_deps.get(&rvalue_id) {452                let lvalue_name = &env.identifiers[lvalue_id.0 as usize].name;453                if !matches!(lvalue_name, Some(IdentifierName::Named(_))) {454                    // Note: we can't insert into maybe_deps here since we only have455                    // a shared reference. The caller handles insertion.456                    return Some(aliased.clone());457                }458            }459            None460        }461        _ => None,462    }463}464465// =============================================================================466// Replacement helpers467// =============================================================================468469fn get_manual_memoization_replacement(470    fn_place: &Place,471    loc: Option<SourceLocation>,472    kind: ManualMemoKind,473) -> InstructionValue {474    if kind == ManualMemoKind::UseMemo {475        // Replace with Call fn() - invoke the memo function directly476        InstructionValue::CallExpression {477            callee: fn_place.clone(),478            args: vec![],479            loc,480        }481    } else {482        // Replace with LoadLocal fn - just reference the function483        InstructionValue::LoadLocal {484            place: Place {485                identifier: fn_place.identifier,486                effect: Effect::Unknown,487                reactive: false,488                loc: loc.clone(),489            },490            loc,491        }492    }493}494495fn make_manual_memoization_markers(496    fn_expr: &Place,497    env: &mut Environment,498    deps_list: Option<Vec<ManualMemoDependency>>,499    deps_loc: Option<SourceLocation>,500    memo_decl: &Place,501    manual_memo_id: u32,502) -> (Instruction, Instruction) {503    let start = Instruction {504        id: EvaluationOrder(0),505        lvalue: create_temporary_place(env, fn_expr.loc.clone()),506        value: InstructionValue::StartMemoize {507            manual_memo_id,508            deps: deps_list,509            deps_loc: Some(deps_loc),510            has_invalid_deps: false,511            loc: fn_expr.loc.clone(),512        },513        loc: fn_expr.loc.clone(),514        effects: None,515    };516    let finish = Instruction {517        id: EvaluationOrder(0),518        lvalue: create_temporary_place(env, fn_expr.loc.clone()),519        value: InstructionValue::FinishMemoize {520            manual_memo_id,521            decl: memo_decl.clone(),522            pruned: false,523            loc: fn_expr.loc.clone(),524        },525        loc: fn_expr.loc.clone(),526        effects: None,527    };528    (start, finish)529}530531fn extract_manual_memoization_args(532    instr: &Instruction,533    kind: ManualMemoKind,534    sidemap: &IdentifierSidemap,535    env: &mut Environment,536) -> Option<ExtractedMemoArgs> {537    let args: &[PlaceOrSpread] = match &instr.value {538        InstructionValue::CallExpression { args, .. } => args,539        InstructionValue::MethodCall { args, .. } => args,540        _ => return None,541    };542543    let kind_name = match kind {544        ManualMemoKind::UseMemo => "useMemo",545        ManualMemoKind::UseCallback => "useCallback",546    };547548    // Get the first arg (fn)549    let fn_place = match args.first() {550        Some(PlaceOrSpread::Place(p)) => p.clone(),551        _ => {552            let loc = instr.value.loc().cloned();553            env.record_diagnostic(554                CompilerDiagnostic::new(555                    ErrorCategory::UseMemo,556                    format!("Expected a callback function to be passed to {kind_name}"),557                    Some(if kind == ManualMemoKind::UseCallback {558                        "The first argument to useCallback() must be a function to cache".to_string()559                    } else {560                        "The first argument to useMemo() must be a function that calculates a result to cache".to_string()561                    }),562                )563                .with_detail(CompilerDiagnosticDetail::Error {564                    loc,565                    message: Some(if kind == ManualMemoKind::UseCallback {566                        "Expected a callback function".to_string()567                    } else {568                        "Expected a memoization function".to_string()569                    }),570                    identifier_name: None,571                }),572            );573            return None;574        }575    };576577    // Get the second arg (deps list), if present578    let deps_list_place = args.get(1);579    if deps_list_place.is_none() {580        return Some(ExtractedMemoArgs {581            fn_place,582            deps_list: None,583            deps_loc: None,584        });585    }586587    let deps_list_id = match deps_list_place {588        Some(PlaceOrSpread::Place(p)) => Some(p.identifier),589        _ => None,590    };591592    let maybe_deps_list = deps_list_id.and_then(|id| sidemap.maybe_deps_lists.get(&id));593594    if maybe_deps_list.is_none() {595        let loc = match deps_list_place {596            Some(PlaceOrSpread::Place(p)) => p.loc.clone(),597            _ => instr.loc.clone(),598        };599        env.record_diagnostic(600            CompilerDiagnostic::new(601                ErrorCategory::UseMemo,602                format!("Expected the dependency list for {kind_name} to be an array literal"),603                Some(format!(604                    "Expected the dependency list for {kind_name} to be an array literal"605                )),606            )607            .with_detail(CompilerDiagnosticDetail::Error {608                loc,609                message: Some(format!(610                    "Expected the dependency list for {kind_name} to be an array literal"611                )),612                identifier_name: None,613            }),614        );615        return None;616    }617618    let deps_info = maybe_deps_list.unwrap();619    let mut deps_list: Vec<ManualMemoDependency> = Vec::new();620    for dep in &deps_info.deps {621        let maybe_dep = sidemap.maybe_deps.get(&dep.identifier);622        if let Some(d) = maybe_dep {623            deps_list.push(d.clone());624        } else {625            env.record_diagnostic(626                CompilerDiagnostic::new(627                    ErrorCategory::UseMemo,628                    "Expected the dependency list to be an array of simple expressions (e.g. `x`, `x.y.z`, `x?.y?.z`)",629                    Some("Expected the dependency list to be an array of simple expressions (e.g. `x`, `x.y.z`, `x?.y?.z`)".to_string()),630                )631                .with_detail(CompilerDiagnosticDetail::Error {632                    loc: dep.loc.clone(),633                    message: Some("Expected the dependency list to be an array of simple expressions (e.g. `x`, `x.y.z`, `x?.y?.z`)".to_string()),634                    identifier_name: None,635                }),636            );637        }638    }639640    Some(ExtractedMemoArgs {641        fn_place,642        deps_list: Some(deps_list),643        deps_loc: deps_info.loc.clone(),644    })645}646647// =============================================================================648// findOptionalPlaces649// =============================================================================650651fn find_optional_places(func: &HirFunction) -> Result<FxHashSet<IdentifierId>, CompilerDiagnostic> {652    use react_compiler_hir::Terminal;653654    let mut optionals = FxHashSet::default();655    for block in func.body.blocks.values() {656        if let Terminal::Optional {657            optional: true,658            test,659            fallthrough,660            ..661        } = &block.terminal662        {663            let optional_fallthrough = *fallthrough;664            let mut test_block_id = *test;665            loop {666                let test_block = &func.body.blocks[&test_block_id];667                match &test_block.terminal {668                    Terminal::Branch {669                        consequent,670                        fallthrough,671                        ..672                    } => {673                        if *fallthrough == optional_fallthrough {674                            // Found it675                            let consequent_block = &func.body.blocks[consequent];676                            if let Some(&last_instr_id) = consequent_block.instructions.last() {677                                let last_instr = &func.instructions[last_instr_id.0 as usize];678                                if let InstructionValue::StoreLocal { value, .. } =679                                    &last_instr.value680                                {681                                    optionals.insert(value.identifier);682                                }683                            }684                            break;685                        } else {686                            test_block_id = *fallthrough;687                        }688                    }689                    Terminal::Optional { fallthrough, .. }690                    | Terminal::Logical { fallthrough, .. }691                    | Terminal::Sequence { fallthrough, .. }692                    | Terminal::Ternary { fallthrough, .. } => {693                        test_block_id = *fallthrough;694                    }695                    Terminal::MaybeThrow { continuation, .. } => {696                        test_block_id = *continuation;697                    }698                    other => {699                        // Invariant: unexpected terminal in optional700                        // In TS this throws CompilerError.invariant701                        return Err(CompilerDiagnostic::new(702                            ErrorCategory::Invariant,703                            format!(704                                "Unexpected terminal kind in optional: {:?}",705                                std::mem::discriminant(other)706                            ),707                            None,708                        ));709                    }710                }711            }712        }713    }714    Ok(optionals)715}716717fn is_known_react_module(module: &str) -> bool {718    let lower = module.to_lowercase();719    lower == "react" || lower == "react-dom"720}721722/// Returns the name to use for useMemo/useCallback detection, matching the TS723/// behavior of `getGlobalDeclaration` + `getHookKindForType`.724///725/// - `Global`: use the binding name (matches globals.get(name) in TS)726/// - `ImportSpecifier` from known React module: use the `imported` name727/// - `ImportSpecifier` from unknown module: return None (TS returns a generic728///   custom hook type with hookKind 'Custom', not 'useMemo'/'useCallback')729/// - `ModuleLocal`: return None (same reason as above)730/// - `ImportDefault`/`ImportNamespace` from known React module: use the local name731/// - `ImportDefault`/`ImportNamespace` from unknown module: return None732fn get_hook_detection_name(binding: &NonLocalBinding) -> Option<&str> {733    match binding {734        NonLocalBinding::Global { name } => Some(name.as_str()),735        NonLocalBinding::ImportSpecifier {736            imported, module, ..737        } => {738            if is_known_react_module(module) {739                Some(imported.as_str())740            } else {741                None742            }743        }744        NonLocalBinding::ImportDefault { name, module }745        | NonLocalBinding::ImportNamespace { name, module } => {746            if is_known_react_module(module) {747                Some(name.as_str())748            } else {749                None750            }751        }752        NonLocalBinding::ModuleLocal { .. } => None,753    }754}