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//! Port of OutlineJsx from TypeScript.7//!8//! Outlines JSX expressions in callbacks into separate component functions.9//! This pass is conditional on `env.config.enable_jsx_outlining` (defaults to false).1011use indexmap::{IndexMap, IndexSet};12use rustc_hash::{FxBuildHasher, FxHashMap, FxHashSet};1314use react_compiler_hir::environment::Environment;15use react_compiler_hir::{16    BasicBlock, BlockId, BlockKind, EvaluationOrder, FunctionId, HIR, HirFunction, IdentifierId,17    IdentifierName, Instruction, InstructionId, InstructionKind, InstructionValue, JsxAttribute,18    JsxTag, LValuePattern, NonLocalBinding, ObjectPattern, ObjectProperty, ObjectPropertyKey,19    ObjectPropertyOrSpread, ObjectPropertyType, ParamPattern, Pattern, Place, ReactFunctionType,20    ReturnVariant, Terminal,21};2223/// Outline JSX expressions in inner functions into separate outlined components.24///25/// Ported from TS `outlineJSX` in `Optimization/OutlineJsx.ts`.26pub fn outline_jsx(func: &mut HirFunction, env: &mut Environment) {27    let mut outlined_fns: Vec<HirFunction> = Vec::new();28    outline_jsx_impl(func, env, &mut outlined_fns);2930    for outlined_fn in outlined_fns {31        env.outline_function(outlined_fn, Some(ReactFunctionType::Component));32    }33}3435/// Data about a JSX instruction for outlining36struct JsxInstrInfo {37    instr_idx: usize, // index into func.instructions38    #[allow(dead_code)]39    instr_id: InstructionId, // the InstructionId40    lvalue_id: IdentifierId,41    eval_order: EvaluationOrder,42}4344struct OutlinedJsxAttribute {45    original_name: String,46    new_name: String,47    place: Place,48}4950struct OutlinedResult {51    instrs: Vec<Instruction>,52    func: HirFunction,53}5455fn outline_jsx_impl(56    func: &mut HirFunction,57    env: &mut Environment,58    outlined_fns: &mut Vec<HirFunction>,59) {60    // Collect LoadGlobal instructions (tag -> instr)61    let mut globals: FxHashMap<IdentifierId, usize> = FxHashMap::default(); // id -> instr_idx6263    // Process each block64    let block_ids: Vec<BlockId> = func.body.blocks.keys().copied().collect();65    for block_id in &block_ids {66        let block = &func.body.blocks[block_id];67        let instr_ids = block.instructions.clone();6869        let mut rewrite_instr: FxHashMap<EvaluationOrder, Vec<Instruction>> = FxHashMap::default();70        let mut jsx_group: Vec<JsxInstrInfo> = Vec::new();71        let mut children_ids: FxHashSet<IdentifierId> = FxHashSet::default();7273        // First pass: collect all instruction info without borrowing func mutably74        enum InstrAction {75            LoadGlobal {76                lvalue_id: IdentifierId,77                instr_idx: usize,78            },79            FunctionExpr {80                func_id: FunctionId,81            },82            JsxExpr {83                lvalue_id: IdentifierId,84                instr_idx: usize,85                eval_order: EvaluationOrder,86                child_ids: Vec<IdentifierId>,87            },88            Other,89        }9091        let mut actions: Vec<InstrAction> = Vec::new();92        for i in (0..instr_ids.len()).rev() {93            let iid = instr_ids[i];94            let instr = &func.instructions[iid.0 as usize];95            let lvalue_id = instr.lvalue.identifier;9697            match &instr.value {98                InstructionValue::LoadGlobal { .. } => {99                    actions.push(InstrAction::LoadGlobal {100                        lvalue_id,101                        instr_idx: iid.0 as usize,102                    });103                }104                InstructionValue::FunctionExpression { lowered_func, .. } => {105                    actions.push(InstrAction::FunctionExpr {106                        func_id: lowered_func.func,107                    });108                }109                InstructionValue::JsxExpression { children, .. } => {110                    let child_ids = children111                        .as_ref()112                        .map(|kids| kids.iter().map(|c| c.identifier).collect())113                        .unwrap_or_default();114                    actions.push(InstrAction::JsxExpr {115                        lvalue_id,116                        instr_idx: iid.0 as usize,117                        eval_order: instr.id,118                        child_ids,119                    });120                }121                _ => {122                    actions.push(InstrAction::Other);123                }124            }125        }126127        // Second pass: process actions128        for action in actions {129            match action {130                InstrAction::LoadGlobal {131                    lvalue_id,132                    instr_idx,133                } => {134                    globals.insert(lvalue_id, instr_idx);135                }136                InstrAction::FunctionExpr { func_id } => {137                    let mut inner_func = std::mem::replace(138                        &mut env.functions[func_id.0 as usize],139                        react_compiler_ssa::enter_ssa::placeholder_function(),140                    );141                    outline_jsx_impl(&mut inner_func, env, outlined_fns);142                    env.functions[func_id.0 as usize] = inner_func;143                }144                InstrAction::JsxExpr {145                    lvalue_id,146                    instr_idx,147                    eval_order,148                    child_ids,149                } => {150                    if !children_ids.contains(&lvalue_id) {151                        process_and_outline_jsx(152                            func,153                            env,154                            &mut jsx_group,155                            &globals,156                            &mut rewrite_instr,157                            outlined_fns,158                        );159                        jsx_group.clear();160                        children_ids.clear();161                    }162                    jsx_group.push(JsxInstrInfo {163                        instr_idx,164                        instr_id: InstructionId(instr_idx as u32),165                        lvalue_id,166                        eval_order,167                    });168                    for child_id in child_ids {169                        children_ids.insert(child_id);170                    }171                }172                InstrAction::Other => {}173            }174        }175        // Process remaining JSX group after the loop176        process_and_outline_jsx(177            func,178            env,179            &mut jsx_group,180            &globals,181            &mut rewrite_instr,182            outlined_fns,183        );184        if !rewrite_instr.is_empty() {185            let block = func.body.blocks.get_mut(block_id).unwrap();186            let old_instr_ids = block.instructions.clone();187            let mut new_instr_ids = Vec::new();188            for &iid in &old_instr_ids {189                let eval_order = func.instructions[iid.0 as usize].id;190                if let Some(replacement_instrs) = rewrite_instr.get(&eval_order) {191                    // Add replacement instructions to the instruction table and reference them192                    for new_instr in replacement_instrs {193                        let new_idx = func.instructions.len();194                        func.instructions.push(new_instr.clone());195                        new_instr_ids.push(InstructionId(new_idx as u32));196                    }197                } else {198                    new_instr_ids.push(iid);199                }200            }201            let block = func.body.blocks.get_mut(block_id).unwrap();202            block.instructions = new_instr_ids;203204            // Run dead code elimination after rewriting205            super::dead_code_elimination(func, env);206        }207    }208}209210fn process_and_outline_jsx(211    func: &mut HirFunction,212    env: &mut Environment,213    jsx_group: &mut Vec<JsxInstrInfo>,214    globals: &FxHashMap<IdentifierId, usize>,215    rewrite_instr: &mut FxHashMap<EvaluationOrder, Vec<Instruction>>,216    outlined_fns: &mut Vec<HirFunction>,217) {218    if jsx_group.len() <= 1 {219        return;220    }221    // Sort by eval order ascending (TS: sort by a.id - b.id)222    jsx_group.sort_by_key(|j| j.eval_order);223224    let result = process_jsx_group(func, env, jsx_group, globals);225    if let Some(result) = result {226        outlined_fns.push(result.func);227        // Map from the LAST JSX instruction's eval order to the replacement instructions228        // In the TS code, `state.jsx.at(0)` is the first element pushed during reverse iteration,229        // which is the last JSX in forward block order (highest eval order).230        // After sorting by eval_order ascending, that's jsx_group.last().231        let last_eval_order = jsx_group.last().unwrap().eval_order;232        rewrite_instr.insert(last_eval_order, result.instrs);233    }234}235236fn process_jsx_group(237    func: &HirFunction,238    env: &mut Environment,239    jsx_group: &[JsxInstrInfo],240    globals: &FxHashMap<IdentifierId, usize>,241) -> Option<OutlinedResult> {242    // Only outline in callbacks, not top-level components243    if func.fn_type == ReactFunctionType::Component {244        return None;245    }246247    let props = collect_props(func, env, jsx_group)?;248249    let outlined_tag = env.generate_globally_unique_identifier_name(None);250    let new_instrs = emit_outlined_jsx(func, env, jsx_group, &props, &outlined_tag)?;251    let outlined_fn = emit_outlined_fn(func, env, jsx_group, &props, globals)?;252253    // Set the outlined function's id254    let mut outlined_fn = outlined_fn;255    outlined_fn.id = Some(outlined_tag);256257    Some(OutlinedResult {258        instrs: new_instrs,259        func: outlined_fn,260    })261}262263fn collect_props(264    func: &HirFunction,265    env: &mut Environment,266    jsx_group: &[JsxInstrInfo],267) -> Option<Vec<OutlinedJsxAttribute>> {268    let mut id_counter = 1u32;269    let mut seen: FxHashSet<String> = FxHashSet::default();270    let mut attributes = Vec::new();271    let jsx_ids: FxHashSet<IdentifierId> = jsx_group.iter().map(|j| j.lvalue_id).collect();272273    let mut generate_name = |old_name: &str, _env: &mut Environment| -> String {274        let mut new_name = old_name.to_string();275        while seen.contains(&new_name) {276            new_name = format!("{}{}", old_name, id_counter);277            id_counter += 1;278        }279        seen.insert(new_name.clone());280        // TS: env.programContext.addNewReference(newName)281        // We don't have programContext in Rust, but this is needed for unique name tracking282        new_name283    };284285    for info in jsx_group {286        let instr = &func.instructions[info.instr_idx];287        if let InstructionValue::JsxExpression {288            props, children, ..289        } = &instr.value290        {291            for attr in props {292                match attr {293                    JsxAttribute::SpreadAttribute { .. } => return None,294                    JsxAttribute::Attribute { name, place } => {295                        let new_name = generate_name(name, env);296                        attributes.push(OutlinedJsxAttribute {297                            original_name: name.clone(),298                            new_name,299                            place: place.clone(),300                        });301                    }302                }303            }304305            if let Some(kids) = children {306                for child in kids {307                    if jsx_ids.contains(&child.identifier) {308                        continue;309                    }310                    // Promote the child's identifier to a named temporary311                    let child_id = child.identifier;312                    let decl_id = env.identifiers[child_id.0 as usize].declaration_id;313                    if env.identifiers[child_id.0 as usize].name.is_none() {314                        env.identifiers[child_id.0 as usize].name =315                            Some(IdentifierName::Promoted(format!("#t{}", decl_id.0)));316                    }317318                    let child_name = match &env.identifiers[child_id.0 as usize].name {319                        Some(IdentifierName::Named(n)) => n.clone(),320                        Some(IdentifierName::Promoted(n)) => n.clone(),321                        None => format!("#t{}", decl_id.0),322                    };323                    let new_name = generate_name("t", env);324                    attributes.push(OutlinedJsxAttribute {325                        original_name: child_name,326                        new_name,327                        place: child.clone(),328                    });329                }330            }331        }332    }333334    Some(attributes)335}336337fn emit_outlined_jsx(338    func: &HirFunction,339    env: &mut Environment,340    jsx_group: &[JsxInstrInfo],341    outlined_props: &[OutlinedJsxAttribute],342    outlined_tag: &str,343) -> Option<Vec<Instruction>> {344    let props: Vec<JsxAttribute> = outlined_props345        .iter()346        .map(|p| JsxAttribute::Attribute {347            name: p.new_name.clone(),348            place: p.place.clone(),349        })350        .collect();351352    // Create LoadGlobal for the outlined component353    let load_id = env.next_identifier_id();354    // Promote it as a JSX tag temporary355    let decl_id = env.identifiers[load_id.0 as usize].declaration_id;356    env.identifiers[load_id.0 as usize].name =357        Some(IdentifierName::Promoted(format!("#T{}", decl_id.0)));358359    let load_place = Place {360        identifier: load_id,361        effect: react_compiler_hir::Effect::Unknown,362        reactive: false,363        loc: None,364    };365366    let load_jsx = Instruction {367        id: EvaluationOrder(0),368        lvalue: load_place.clone(),369        value: InstructionValue::LoadGlobal {370            binding: NonLocalBinding::ModuleLocal {371                name: outlined_tag.to_string(),372            },373            loc: None,374        },375        loc: None,376        effects: None,377    };378379    // Create the replacement JsxExpression using the last JSX instruction's lvalue380    let last_info = jsx_group.last().unwrap();381    let last_instr = &func.instructions[last_info.instr_idx];382    let jsx_expr = Instruction {383        id: EvaluationOrder(0),384        lvalue: last_instr.lvalue.clone(),385        value: InstructionValue::JsxExpression {386            tag: JsxTag::Place(load_place),387            props,388            children: None,389            loc: None,390            opening_loc: None,391            closing_loc: None,392        },393        loc: None,394        effects: None,395    };396397    Some(vec![load_jsx, jsx_expr])398}399400fn emit_outlined_fn(401    func: &HirFunction,402    env: &mut Environment,403    jsx_group: &[JsxInstrInfo],404    old_props: &[OutlinedJsxAttribute],405    globals: &FxHashMap<IdentifierId, usize>,406) -> Option<HirFunction> {407    let old_to_new_props = create_old_to_new_props_mapping(env, old_props);408409    // Create props parameter410    let props_obj_id = env.next_identifier_id();411    let decl_id = env.identifiers[props_obj_id.0 as usize].declaration_id;412    env.identifiers[props_obj_id.0 as usize].name =413        Some(IdentifierName::Promoted(format!("#t{}", decl_id.0)));414    let props_obj = Place {415        identifier: props_obj_id,416        effect: react_compiler_hir::Effect::Unknown,417        reactive: false,418        loc: None,419    };420421    // Create destructure instruction422    let destructure_instr = emit_destructure_props(env, &props_obj, &old_to_new_props);423424    // Emit load globals for JSX tags425    let load_global_instrs = emit_load_globals(func, jsx_group, globals)?;426427    // Emit updated JSX instructions428    let updated_jsx_instrs = emit_updated_jsx(func, jsx_group, &old_to_new_props);429430    // Build instructions list431    let mut instructions = Vec::new();432    instructions.push(destructure_instr);433    instructions.extend(load_global_instrs);434    instructions.extend(updated_jsx_instrs);435436    // Build instruction table and instruction IDs437    let mut instr_table = Vec::new();438    let mut instr_ids = Vec::new();439    for instr in instructions {440        let idx = instr_table.len();441        instr_table.push(instr);442        instr_ids.push(InstructionId(idx as u32));443    }444445    // Return terminal uses the last instruction's lvalue446    let last_lvalue = instr_table.last().unwrap().lvalue.clone();447448    // Create return place449    let returns_id = env.next_identifier_id();450    let returns_place = Place {451        identifier: returns_id,452        effect: react_compiler_hir::Effect::Unknown,453        reactive: false,454        loc: None,455    };456457    let block = BasicBlock {458        kind: BlockKind::Block,459        id: BlockId(0),460        instructions: instr_ids,461        preds: IndexSet::default(),462        terminal: Terminal::Return {463            value: last_lvalue,464            return_variant: ReturnVariant::Explicit,465            id: EvaluationOrder(0),466            loc: None,467            effects: None,468        },469        phis: Vec::new(),470    };471472    let mut blocks = IndexMap::default();473    blocks.insert(BlockId(0), block);474475    let outlined_fn = HirFunction {476        id: None,477        name_hint: None,478        fn_type: ReactFunctionType::Other,479        params: vec![ParamPattern::Place(props_obj)],480        return_type_annotation: None,481        returns: returns_place,482        context: Vec::new(),483        body: HIR {484            entry: BlockId(0),485            blocks,486        },487        instructions: instr_table,488        generator: false,489        is_async: false,490        directives: Vec::new(),491        aliasing_effects: Some(vec![]),492        loc: None,493    };494495    Some(outlined_fn)496}497498fn emit_load_globals(499    func: &HirFunction,500    jsx_group: &[JsxInstrInfo],501    globals: &FxHashMap<IdentifierId, usize>,502) -> Option<Vec<Instruction>> {503    let mut instructions = Vec::new();504    for info in jsx_group {505        let instr = &func.instructions[info.instr_idx];506        if let InstructionValue::JsxExpression { tag, .. } = &instr.value {507            if let JsxTag::Place(tag_place) = tag {508                let global_instr_idx = globals.get(&tag_place.identifier)?;509                instructions.push(func.instructions[*global_instr_idx].clone());510            }511        }512    }513    Some(instructions)514}515516fn emit_updated_jsx(517    func: &HirFunction,518    jsx_group: &[JsxInstrInfo],519    old_to_new_props: &IndexMap<IdentifierId, OutlinedJsxAttribute, FxBuildHasher>,520) -> Vec<Instruction> {521    let jsx_ids: FxHashSet<IdentifierId> = jsx_group.iter().map(|j| j.lvalue_id).collect();522    let mut new_instrs = Vec::new();523524    for info in jsx_group {525        let instr = &func.instructions[info.instr_idx];526        if let InstructionValue::JsxExpression {527            tag,528            props,529            children,530            loc,531            opening_loc,532            closing_loc,533        } = &instr.value534        {535            let mut new_props = Vec::new();536            for prop in props {537                // TS: invariant(prop.kind === 'JsxAttribute', ...)538                // Spread attributes would have caused collectProps to return null earlier539                let (name, place) = match prop {540                    JsxAttribute::Attribute { name, place } => (name, place),541                    JsxAttribute::SpreadAttribute { .. } => {542                        unreachable!("Expected only JsxAttribute, not spread")543                    }544                };545                if name == "key" {546                    continue;547                }548                // TS: invariant(newProp !== undefined, ...)549                let new_prop = old_to_new_props550                    .get(&place.identifier)551                    .expect("Expected a new property for identifier");552                new_props.push(JsxAttribute::Attribute {553                    name: new_prop.original_name.clone(),554                    place: new_prop.place.clone(),555                });556            }557558            let new_children = children.as_ref().map(|kids| {559                kids.iter()560                    .map(|child| {561                        if jsx_ids.contains(&child.identifier) {562                            child.clone()563                        } else {564                            // TS: invariant(newChild !== undefined, ...)565                            let new_prop = old_to_new_props566                                .get(&child.identifier)567                                .expect("Expected a new prop for child identifier");568                            new_prop.place.clone()569                        }570                    })571                    .collect()572            });573574            new_instrs.push(Instruction {575                id: instr.id,576                lvalue: instr.lvalue.clone(),577                value: InstructionValue::JsxExpression {578                    tag: tag.clone(),579                    props: new_props,580                    children: new_children,581                    loc: *loc,582                    opening_loc: *opening_loc,583                    closing_loc: *closing_loc,584                },585                loc: instr.loc,586                effects: instr.effects.clone(),587            });588        }589    }590591    new_instrs592}593594fn create_old_to_new_props_mapping(595    env: &mut Environment,596    old_props: &[OutlinedJsxAttribute],597) -> IndexMap<IdentifierId, OutlinedJsxAttribute, FxBuildHasher> {598    let mut old_to_new = IndexMap::default();599600    for old_prop in old_props {601        if old_prop.original_name == "key" {602            continue;603        }604605        let new_id = env.next_identifier_id();606        env.identifiers[new_id.0 as usize].name =607            Some(IdentifierName::Named(old_prop.new_name.clone()));608609        let new_place = Place {610            identifier: new_id,611            effect: react_compiler_hir::Effect::Unknown,612            reactive: false,613            loc: None,614        };615616        old_to_new.insert(617            old_prop.place.identifier,618            OutlinedJsxAttribute {619                original_name: old_prop.original_name.clone(),620                new_name: old_prop.new_name.clone(),621                place: new_place,622            },623        );624    }625626    old_to_new627}628629fn emit_destructure_props(630    env: &mut Environment,631    props_obj: &Place,632    old_to_new_props: &IndexMap<IdentifierId, OutlinedJsxAttribute, FxBuildHasher>,633) -> Instruction {634    let mut properties = Vec::new();635    for prop in old_to_new_props.values() {636        properties.push(ObjectPropertyOrSpread::Property(ObjectProperty {637            key: ObjectPropertyKey::String {638                name: prop.new_name.clone(),639            },640            property_type: ObjectPropertyType::Property,641            place: prop.place.clone(),642        }));643    }644645    let lvalue_id = env.next_identifier_id();646    let lvalue = Place {647        identifier: lvalue_id,648        effect: react_compiler_hir::Effect::Unknown,649        reactive: false,650        loc: None,651    };652653    Instruction {654        id: EvaluationOrder(0),655        lvalue,656        value: InstructionValue::Destructure {657            lvalue: LValuePattern {658                pattern: Pattern::Object(ObjectPattern {659                    properties,660                    loc: None,661                }),662                kind: InstructionKind::Let,663            },664            value: props_obj.clone(),665            loc: None,666        },667        loc: None,668        effects: None,669    }670}