// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is licensed under the MIT license found in the
// LICENSE file in the root directory of this source tree.

//! Inlines immediately invoked function expressions (IIFEs) to allow more
//! fine-grained memoization of the values they produce.
//!
//! Example:
//! ```text
//! const x = (() => {
//!    const x = [];
//!    x.push(foo());
//!    return x;
//! })();
//!
//! =>
//!
//! bb0:
//!     // placeholder for the result, all return statements will assign here
//!    let t0;
//!    // Label allows using a goto (break) to exit out of the body
//!    Label block=bb1 fallthrough=bb2
//! bb1:
//!    // code within the function expression
//!    const x0 = [];
//!    x0.push(foo());
//!    // return is replaced by assignment to the result variable...
//!    t0 = x0;
//!    // ...and a goto to the code after the function expression invocation
//!    Goto bb2
//! bb2:
//!    // code after the IIFE call
//!    const x = t0;
//! ```
//!
//! If the inlined function has only one return, we avoid the labeled block
//! and fully inline the code. The original return is replaced with an assignment
//! to the IIFE's call expression lvalue.
//!
//! Analogous to TS `Inference/InlineImmediatelyInvokedFunctionExpressions.ts`.

use crate::react_compiler_utils::FxIndexSet;
use oxc_str::format_ident;
use rustc_hash::{FxHashMap, FxHashSet};

use crate::react_compiler_hir::environment::Environment;
use crate::react_compiler_hir::visitors;
use crate::react_compiler_hir::{
    BasicBlock, BlockId, BlockKind, EvaluationOrder, FunctionId, GENERATED_SOURCE, GotoVariant,
    HirFunction, IdentifierId, IdentifierName, Instruction, InstructionId, InstructionKind,
    InstructionValue, LValue, Place, Terminal,
};
use crate::react_compiler_lowering::{
    create_temporary_place, get_reverse_postordered_blocks, mark_instruction_ids, mark_predecessors,
};

use crate::react_compiler_optimization::merge_consecutive_blocks::merge_consecutive_blocks;

/// Inline immediately invoked function expressions into the enclosing function's
/// control flow graph.
pub fn inline_immediately_invoked_function_expressions<'a>(
    func: &mut HirFunction<'a>,
    env: &mut Environment<'a>,
) {
    // Track all function expressions that are assigned to a temporary
    let mut functions: FxHashMap<IdentifierId, FunctionId> = FxHashMap::default();
    // Functions that are inlined (by identifier id of the callee)
    let mut inlined_functions: FxHashSet<IdentifierId> = FxHashSet::default();

    // Iterate the *existing* blocks from the outer component to find IIFEs
    // and inline them. During iteration we will modify `func` (by inlining the CFG
    // of IIFEs) so we explicitly copy references to just the original
    // function's block IDs first. As blocks are split to make room for IIFE calls,
    // the split portions of the blocks will be added to this queue.
    let mut queue: Vec<BlockId> = func.body.blocks.keys().copied().collect();
    let mut queue_idx = 0;

    'queue: while queue_idx < queue.len() {
        let block_id = queue[queue_idx];
        queue_idx += 1;

        let block = match func.body.blocks.get(&block_id) {
            Some(b) => b,
            None => continue,
        };

        // We can't handle labels inside expressions yet, so we don't inline IIFEs
        // if they are in an expression block.
        if !is_statement_block_kind(block.kind) {
            continue;
        }

        let num_instructions = block.instructions.len();
        for ii in 0..num_instructions {
            let instr_id = func.body.blocks[&block_id].instructions[ii];
            let instr = &func.instructions[instr_id.index()];

            match &instr.value {
                InstructionValue::FunctionExpression { lowered_func, .. } => {
                    let identifier_id = instr.lvalue.identifier;
                    if env.identifiers[identifier_id].name.is_none() {
                        functions.insert(identifier_id, lowered_func.func);
                    }
                    continue;
                }
                InstructionValue::CallExpression { callee, args, .. } => {
                    if !args.is_empty() {
                        // We don't support inlining when there are arguments
                        continue;
                    }

                    let callee_id = callee.identifier;
                    let inner_func_id = match functions.get(&callee_id) {
                        Some(id) => *id,
                        None => continue, // Not invoking a local function expression
                    };

                    let inner_func = &env.functions[inner_func_id];
                    if !inner_func.params.is_empty() || inner_func.is_async || inner_func.generator
                    {
                        // Can't inline functions with params, or async/generator functions
                        continue;
                    }

                    // We know this function is used for an IIFE and can prune it later
                    inlined_functions.insert(callee_id);

                    // Capture the lvalue from the call instruction
                    let call_lvalue = func.instructions[instr_id.index()].lvalue.clone();
                    let block_terminal_id = func.body.blocks[&block_id].terminal.evaluation_order();
                    let block_terminal_span = func.body.blocks[&block_id].terminal.span().cloned();
                    let block_kind = func.body.blocks[&block_id].kind;

                    // Create a new block which will contain code following the IIFE call
                    let continuation_block_id = env.next_block_id();
                    let continuation_instructions: Vec<InstructionId> =
                        func.body.blocks[&block_id].instructions[ii + 1..].to_vec();
                    let continuation_terminal = func.body.blocks[&block_id].terminal.clone();
                    let continuation_block = BasicBlock {
                        id: continuation_block_id,
                        instructions: continuation_instructions,
                        kind: block_kind,
                        phis: Vec::new(),
                        preds: FxIndexSet::default(),
                        terminal: continuation_terminal,
                    };
                    func.body.blocks.insert(continuation_block_id, continuation_block);

                    // Trim the original block to contain instructions up to (but not including)
                    // the IIFE
                    func.body.blocks.get_mut(&block_id).unwrap().instructions.truncate(ii);

                    let has_single_return =
                        has_single_exit_return_terminal(&env.functions[inner_func_id]);
                    let inner_entry = env.functions[inner_func_id].body.entry;

                    if has_single_return {
                        // Single-return path: simple goto replacement
                        func.body.blocks.get_mut(&block_id).unwrap().terminal = Terminal::Goto {
                            block: inner_entry,
                            id: block_terminal_id,
                            span: block_terminal_span,
                            variant: GotoVariant::Break,
                        };

                        // Take blocks and instructions from inner function
                        let inner_func = &mut env.functions[inner_func_id];
                        let inner_blocks: Vec<(BlockId, BasicBlock)> =
                            inner_func.body.blocks.drain(..).collect();
                        let inner_instructions: Vec<Instruction> =
                            inner_func.instructions.drain(..).collect();

                        // Append inner instructions first, then remap block instruction IDs
                        let instr_offset = func.instructions.len() as u32;
                        func.instructions.extend(inner_instructions);

                        for (_, mut inner_block) in inner_blocks {
                            // Remap instruction IDs in the block
                            for iid in &mut inner_block.instructions {
                                *iid += instr_offset as usize;
                            }
                            inner_block.preds.clear();

                            if let Terminal::Return { value, id: ret_id, span: ret_span, .. } =
                                &inner_block.terminal
                            {
                                // Replace return with LoadLocal + goto
                                let load_instr = Instruction {
                                    id: EvaluationOrder::UNSET,
                                    span: *ret_span,
                                    lvalue: call_lvalue.clone(),
                                    value: InstructionValue::LoadLocal {
                                        place: value.clone(),
                                        span: *ret_span,
                                    },
                                    effects: None,
                                };
                                let load_instr_id =
                                    InstructionId::from_usize(func.instructions.len());
                                func.instructions.push(load_instr);
                                inner_block.instructions.push(load_instr_id);

                                let ret_id = *ret_id;
                                let ret_span = *ret_span;
                                inner_block.terminal = Terminal::Goto {
                                    block: continuation_block_id,
                                    id: ret_id,
                                    span: ret_span,
                                    variant: GotoVariant::Break,
                                };
                            }

                            func.body.blocks.insert(inner_block.id, inner_block);
                        }
                    } else {
                        // Multi-return path: uses LabelTerminal
                        let result = call_lvalue.clone();

                        // Set block terminal to Label
                        func.body.blocks.get_mut(&block_id).unwrap().terminal = Terminal::Label {
                            block: inner_entry,
                            id: EvaluationOrder::UNSET,
                            fallthrough: continuation_block_id,
                            span: block_terminal_span,
                        };

                        // Declare the IIFE temporary
                        declare_temporary(env, func, block_id, &result);

                        // Promote the temporary with a name as we require this to persist
                        let identifier_id = result.identifier;
                        if env.identifiers[identifier_id].name.is_none() {
                            promote_temporary(env, identifier_id);
                        }

                        // Take blocks and instructions from inner function
                        let inner_func = &mut env.functions[inner_func_id];
                        let inner_blocks: Vec<(BlockId, BasicBlock)> =
                            inner_func.body.blocks.drain(..).collect();
                        let inner_instructions: Vec<Instruction> =
                            inner_func.instructions.drain(..).collect();

                        // Append inner instructions first, then remap block instruction IDs
                        let instr_offset = func.instructions.len() as u32;
                        func.instructions.extend(inner_instructions);

                        for (_, mut inner_block) in inner_blocks {
                            for iid in &mut inner_block.instructions {
                                *iid += instr_offset as usize;
                            }
                            inner_block.preds.clear();

                            // Rewrite return terminals to StoreLocal + goto
                            if matches!(inner_block.terminal, Terminal::Return { .. }) {
                                rewrite_block(
                                    env,
                                    &mut func.instructions,
                                    &mut inner_block,
                                    continuation_block_id,
                                    &result,
                                );
                            }

                            func.body.blocks.insert(inner_block.id, inner_block);
                        }
                    }

                    // Ensure we visit the continuation block, since there may have been
                    // sequential IIFEs that need to be visited.
                    queue.push(continuation_block_id);
                    continue 'queue;
                }
                _ => {
                    // Any other use of a function expression means it isn't an IIFE
                    for id in visitors::each_instruction_value_operand_ids(&instr.value, env) {
                        functions.remove(&id);
                    }
                }
            }
        }
    }

    if !inlined_functions.is_empty() {
        // Remove instructions that define lambdas which we inlined
        for block in func.body.blocks.values_mut() {
            block.instructions.retain(|instr_id| {
                let instr = &func.instructions[instr_id.index()];
                !inlined_functions.contains(&instr.lvalue.identifier)
            });
        }

        // If terminals have changed then blocks may have become newly unreachable.
        // Re-run minification of the graph (incl reordering instruction ids).
        func.body.blocks = get_reverse_postordered_blocks(&func.body);
        mark_instruction_ids(&mut func.body, &mut func.instructions);
        mark_predecessors(&mut func.body);
        merge_consecutive_blocks(func, &mut env.functions);
    }
}

/// Returns true for "block" and "catch" block kinds which correspond to statements
/// in the source.
fn is_statement_block_kind(kind: BlockKind) -> bool {
    matches!(kind, BlockKind::Block | BlockKind::Catch)
}

/// Returns true if the function has a single exit terminal (throw/return) which is a return.
fn has_single_exit_return_terminal(func: &HirFunction<'_>) -> bool {
    let mut has_return = false;
    let mut exit_count = 0;
    for block in func.body.blocks.values() {
        match &block.terminal {
            Terminal::Return { .. } => {
                has_return = true;
                exit_count += 1;
            }
            Terminal::Throw { .. } => {
                exit_count += 1;
            }
            _ => {}
        }
    }
    exit_count == 1 && has_return
}

/// Rewrites the block so that all `return` terminals are replaced:
/// * Add a StoreLocal <return_value> = <terminal.value>
/// * Replace the terminal with a Goto to <return_target>
fn rewrite_block<'a>(
    env: &mut Environment<'a>,
    instructions: &mut Vec<Instruction<'a>>,
    block: &mut BasicBlock,
    return_target: BlockId,
    return_value: &Place,
) {
    if let Terminal::Return { value, span: ret_span, .. } = &block.terminal {
        let store_lvalue = create_temporary_place(env, *ret_span);
        let store_instr = Instruction {
            id: EvaluationOrder::UNSET,
            span: *ret_span,
            lvalue: store_lvalue,
            value: InstructionValue::StoreLocal {
                lvalue: LValue { kind: InstructionKind::Reassign, place: return_value.clone() },
                value: value.clone(),
                span: *ret_span,
            },
            effects: None,
        };
        let store_instr_id = InstructionId::from_usize(instructions.len());
        instructions.push(store_instr);
        block.instructions.push(store_instr_id);

        let ret_span = *ret_span;
        block.terminal = Terminal::Goto {
            block: return_target,
            id: EvaluationOrder::UNSET,
            variant: GotoVariant::Break,
            span: ret_span,
        };
    }
}

/// Emits a DeclareLocal instruction for the result temporary.
fn declare_temporary<'a>(
    env: &mut Environment<'a>,
    func: &mut HirFunction<'a>,
    block_id: BlockId,
    result: &Place,
) {
    let declare_lvalue = create_temporary_place(env, result.span);
    let declare_instr = Instruction {
        id: EvaluationOrder::UNSET,
        span: GENERATED_SOURCE,
        lvalue: declare_lvalue,
        value: InstructionValue::DeclareLocal {
            lvalue: LValue { place: result.clone(), kind: InstructionKind::Let },
            span: result.span,
        },
        effects: None,
    };
    let instr_id = InstructionId::from_usize(func.instructions.len());
    func.instructions.push(declare_instr);
    func.body.blocks.get_mut(&block_id).unwrap().instructions.push(instr_id);
}

/// Promote a temporary identifier to a named identifier.
fn promote_temporary(env: &mut Environment<'_>, identifier_id: IdentifierId) {
    let decl_id = env.identifiers[identifier_id].declaration_id;
    env.identifiers[identifier_id].name =
        Some(IdentifierName::Promoted(format_ident!(env.allocator, "#t{}", decl_id.index())));
}
