unluac 1.1.1

//! 这个文件实现共享循环候选提取。
//!
//! 这个 pass 只消费 CFG / GraphFacts / Dataflow / low-IR terminator，产出“循环形态 hint +
//! 可直接复用的源码绑定证据 + loop merge incoming 事实”，不会越权决定最终
//! `while/repeat/for` 语法。
//!
//! 例子：
//! - `NumericForInit/Loop` 会产出 `LoopKindHint::NumericForLike`，并把源码绑定寄存器
//!   记录成 `LoopSourceBindings::Numeric`
//! - `GenericForCall/Loop` 会产出 `LoopKindHint::GenericForLike`，并把源码绑定区间
//!   记录成 `LoopSourceBindings::Generic`
//! - `while ... do ... end` 的 header/exit phi 会被整理成 `inside/outside` 两臂的
//!   incoming facts，后续 HIR 直接消费这些结构事实，不再自己回头拆 `phi.incoming`
//! - 普通 `while/repeat` 只保留形态 hint，不会伪造额外 binding 证据

use std::collections::{BTreeMap, BTreeSet};

use crate::cfg::{BlockRef, Cfg, DataflowFacts, EdgeRef, GraphFacts};
use crate::transformer::{LowInstr, LoweredProto, Reg, ResultPack};

use super::common::{
    LoopCandidate, LoopExitValueMergeCandidate, LoopKindHint, LoopSourceBindings, LoopValueMerge,
};
use super::helpers::{collect_region_exits, is_reducible_region};
use super::phi_facts::loop_value_merges_in_block;

pub(super) fn analyze_loops(
    proto: &LoweredProto,
    cfg: &Cfg,
    graph_facts: &GraphFacts,
    dataflow: &DataflowFacts,
) -> Vec<LoopCandidate> {
    let mut grouped_loops = BTreeMap::<BlockRef, (BTreeSet<BlockRef>, Vec<EdgeRef>)>::new();
    for natural_loop in &graph_facts.natural_loops {
        let entry = grouped_loops
            .entry(natural_loop.header)
            .or_insert_with(|| (BTreeSet::new(), Vec::new()));
        entry.0.extend(natural_loop.blocks.iter().copied());
        entry.1.push(natural_loop.backedge);
    }

    let mut loop_candidates = grouped_loops
        .into_iter()
        .map(|(header, (blocks, mut backedges))| {
            backedges.sort();
            backedges.dedup();
            let preheader = unique_loop_preheader(cfg, header, &blocks);
            let exits = collect_region_exits(cfg, &blocks);
            let reducible = is_reducible_region(cfg, header, &blocks);
            let header_value_merges = analyze_loop_header_value_merges(dataflow, header, &blocks);
            let (kind_hint, continue_target, source_bindings) = infer_loop_shape(
                proto,
                cfg,
                header,
                &blocks,
                &backedges,
                preheader,
                &header_value_merges,
            );
            let exit_value_merges = analyze_loop_exit_value_merges(dataflow, &exits, &blocks);

            LoopCandidate {
                header,
                preheader,
                blocks,
                backedges,
                exits,
                continue_target,
                kind_hint,
                source_bindings,
                header_value_merges,
                exit_value_merges,
                reducible,
            }
        })
        .collect::<Vec<_>>();

    loop_candidates.sort_by_key(|candidate| candidate.header);
    loop_candidates
}

fn infer_loop_shape(
    proto: &LoweredProto,
    cfg: &Cfg,
    header: BlockRef,
    blocks: &BTreeSet<BlockRef>,
    backedges: &[EdgeRef],
    preheader: Option<BlockRef>,
    header_value_merges: &[LoopValueMerge],
) -> (LoopKindHint, Option<BlockRef>, Option<LoopSourceBindings>) {
    let backedge_sources = backedges
        .iter()
        .map(|edge_ref| cfg.edges[edge_ref.index()].from)
        .collect::<BTreeSet<_>>();

    if backedge_sources.len() == 1 {
        let source = *backedge_sources
            .iter()
            .next()
            .expect("set length already checked");
        if let Some(terminator) = cfg.terminator(&proto.instrs, source)
            && matches!(terminator, LowInstr::NumericForLoop(_instr))
        {
            return (
                LoopKindHint::NumericForLike,
                Some(source),
                numeric_for_source_bindings(proto, cfg, preheader),
            );
        }
    }

    // generic-for 的 header 本身就携带了比普通回边更强的形状证据。
    // 如果这里先按“回边源是 branch”去判断，很容易把正常的 generic-for
    // 误认成 repeat-like，后面 HIR 就只能回到 unresolved 的 VM 级控制块。
    if matches!(
        cfg.terminator(&proto.instrs, header),
        Some(LowInstr::GenericForLoop(instr))
            if generic_for_has_loop_body_and_exit(proto, cfg, header, instr, blocks)
    ) {
        return (
            LoopKindHint::GenericForLike,
            Some(header),
            generic_for_source_bindings(proto, cfg, header),
        );
    }

    // dialect lowering 往往会把 while 条件需要的临时准备也塞进 header block，再接 branch。
    // 这些前缀仍然属于“每轮先算条件、再决定进不进 body”的源码语义；如果这里只接受
    // 纯常量加载，像 `while i <= #values do`、`while (x & mask) ~= 0 do` 这类最普通的
    // 条件都会被误打成 repeat/unknown，后面整片 loop state 恢复就只能回退成 label/goto。
    if block_is_while_header_like(proto, cfg, header, header_value_merges)
        && branch_has_loop_body_and_exit(cfg, header, blocks)
    {
        return (LoopKindHint::WhileLike, Some(header), None);
    }

    if backedge_sources.len() == 1 {
        let source = *backedge_sources
            .iter()
            .next()
            .expect("set length already checked");
        if matches!(
            cfg.terminator(&proto.instrs, source),
            Some(LowInstr::Branch(_instr)) if branch_has_header_and_exit(cfg, source, header, blocks)
        ) {
            return (LoopKindHint::RepeatLike, Some(source), None);
        }

        if matches!(
            cfg.terminator(&proto.instrs, source),
            Some(LowInstr::Jump(jump))
                if cfg.instr_to_block[jump.target.index()] == header
                    && repeat_continue_target_via_backedge_pad(proto, cfg, source, blocks).is_some()
        ) {
            return (
                LoopKindHint::RepeatLike,
                repeat_continue_target_via_backedge_pad(proto, cfg, source, blocks),
                None,
            );
        }
    }

    let continue_target = if backedge_sources.len() == 1 {
        backedge_sources.iter().next().copied()
    } else {
        None
    };

    (LoopKindHint::Unknown, continue_target, None)
}

fn numeric_for_source_bindings(
    proto: &LoweredProto,
    cfg: &Cfg,
    preheader: Option<BlockRef>,
) -> Option<LoopSourceBindings> {
    let preheader = preheader?;
    let instr_ref = cfg.blocks[preheader.index()].instrs.last()?;

    match proto.instrs.get(instr_ref.index())? {
        LowInstr::NumericForInit(instr) => Some(LoopSourceBindings::Numeric(instr.binding)),
        _ => None,
    }
}

fn generic_for_source_bindings(
    proto: &LoweredProto,
    cfg: &Cfg,
    header: BlockRef,
) -> Option<LoopSourceBindings> {
    let instr_ref = cfg.blocks[header.index()].instrs.last()?;

    match proto.instrs.get(instr_ref.index())? {
        LowInstr::GenericForLoop(instr) => Some(LoopSourceBindings::Generic(instr.bindings)),
        _ => None,
    }
}

fn analyze_loop_header_value_merges(
    dataflow: &DataflowFacts,
    header: BlockRef,
    loop_blocks: &BTreeSet<BlockRef>,
) -> Vec<LoopValueMerge> {
    loop_value_merges_in_block(dataflow, header, loop_blocks)
        .into_iter()
        .filter(loop_value_has_inside_and_outside_incoming)
        .collect()
}

fn analyze_loop_exit_value_merges(
    dataflow: &DataflowFacts,
    exits: &BTreeSet<BlockRef>,
    loop_blocks: &BTreeSet<BlockRef>,
) -> Vec<LoopExitValueMergeCandidate> {
    exits
        .iter()
        .copied()
        .filter_map(|exit| {
            let values = loop_value_merges_in_block(dataflow, exit, loop_blocks)
                .into_iter()
                .filter(|value| !value.inside_arm.is_empty())
                .collect::<Vec<_>>();
            (!values.is_empty()).then_some(LoopExitValueMergeCandidate { exit, values })
        })
        .collect()
}

fn loop_value_has_inside_and_outside_incoming(value: &LoopValueMerge) -> bool {
    !value.inside_arm.is_empty() && !value.outside_arm.is_empty()
}

fn unique_loop_preheader(
    cfg: &Cfg,
    header: BlockRef,
    loop_blocks: &BTreeSet<BlockRef>,
) -> Option<BlockRef> {
    let preds = cfg
        .reachable_predecessors(header)
        .into_iter()
        .filter(|pred| !loop_blocks.contains(pred))
        .collect::<Vec<_>>();
    let [preheader] = preds.as_slice() else {
        return None;
    };
    Some(*preheader)
}

fn branch_has_loop_body_and_exit(cfg: &Cfg, header: BlockRef, blocks: &BTreeSet<BlockRef>) -> bool {
    let Some((then_edge_ref, else_edge_ref)) = cfg.branch_edges(header) else {
        return false;
    };
    let then_block = cfg.edges[then_edge_ref.index()].to;
    let else_block = cfg.edges[else_edge_ref.index()].to;

    (blocks.contains(&then_block) && !blocks.contains(&else_block))
        || (!blocks.contains(&then_block) && blocks.contains(&else_block))
}

fn branch_has_header_and_exit(
    cfg: &Cfg,
    block: BlockRef,
    header: BlockRef,
    blocks: &BTreeSet<BlockRef>,
) -> bool {
    let Some((then_edge_ref, else_edge_ref)) = cfg.branch_edges(block) else {
        return false;
    };
    let then_block = cfg.edges[then_edge_ref.index()].to;
    let else_block = cfg.edges[else_edge_ref.index()].to;

    (then_block == header && !blocks.contains(&else_block))
        || (else_block == header && !blocks.contains(&then_block))
}

fn block_is_while_header_like(
    proto: &LoweredProto,
    cfg: &Cfg,
    block: BlockRef,
    header_value_merges: &[LoopValueMerge],
) -> bool {
    let range = cfg.blocks[block.index()].instrs;
    if !matches!(
        cfg.terminator(&proto.instrs, block),
        Some(LowInstr::Branch(_))
    ) {
        return false;
    }
    if range.len == 1 {
        return true;
    }

    let carried_regs = header_value_merges
        .iter()
        .map(|value| value.reg)
        .collect::<BTreeSet<_>>();
    (range.start.index()..range.end() - 1).all(|instr_index| {
        let instr = &proto.instrs[instr_index];
        instr_is_while_header_prefix(instr) && !instr_writes_any_reg(instr, &carried_regs)
    })
}

/// while 条件求值中不可能出现的指令。这些指令要么是控制流终结指令（已由
/// terminator 位置单独处理），要么是纯副作用写出指令（SetUpvalue / SetTable /
/// SetList），要么是 scope 管理指令（Close / Tbc），不可能出现在 Lua 编译器
/// 为 expression 上下文生成的条件求值序列里。
///
/// 使用排除列表而非允许列表，避免新增 LowInstr 变体时遗漏导致合法的 while
/// 条件被误判为 repeat/unknown。
fn instr_is_while_header_prefix(instr: &LowInstr) -> bool {
    !matches!(
        instr,
        LowInstr::SetUpvalue(_)
            | LowInstr::SetTable(_)
            | LowInstr::SetList(_)
            | LowInstr::TailCall(_)
            | LowInstr::Return(_)
            | LowInstr::Close(_)
            | LowInstr::Tbc(_)
            | LowInstr::NumericForInit(_)
            | LowInstr::NumericForLoop(_)
            | LowInstr::GenericForCall(_)
            | LowInstr::GenericForLoop(_)
            | LowInstr::Jump(_)
            | LowInstr::Branch(_)
    )
}

fn instr_writes_any_reg(instr: &LowInstr, regs: &BTreeSet<Reg>) -> bool {
    match instr {
        LowInstr::Move(instr) => regs.contains(&instr.dst),
        LowInstr::LoadNil(instr) => (0..instr.dst.len)
            .map(|offset| Reg(instr.dst.start.index() + offset))
            .any(|reg| regs.contains(&reg)),
        LowInstr::LoadBool(instr) => regs.contains(&instr.dst),
        LowInstr::LoadConst(instr) => regs.contains(&instr.dst),
        LowInstr::LoadInteger(instr) => regs.contains(&instr.dst),
        LowInstr::LoadNumber(instr) => regs.contains(&instr.dst),
        LowInstr::UnaryOp(instr) => regs.contains(&instr.dst),
        LowInstr::BinaryOp(instr) => regs.contains(&instr.dst),
        LowInstr::Concat(instr) => regs.contains(&instr.dst),
        LowInstr::GetUpvalue(instr) => regs.contains(&instr.dst),
        LowInstr::GetTable(instr) => regs.contains(&instr.dst),
        LowInstr::NewTable(instr) => regs.contains(&instr.dst),
        LowInstr::Closure(instr) => regs.contains(&instr.dst),
        LowInstr::Call(instr) => result_pack_writes_any_reg(&instr.results, regs),
        LowInstr::VarArg(instr) => result_pack_writes_any_reg(&instr.results, regs),
        // ErrNil 只读 subject 不写寄存器。
        LowInstr::ErrNil(_) => false,
        // 以下指令要么不写寄存器，要么已被 instr_is_while_header_prefix 排除。
        LowInstr::SetUpvalue(_)
        | LowInstr::SetTable(_)
        | LowInstr::SetList(_)
        | LowInstr::TailCall(_)
        | LowInstr::Return(_)
        | LowInstr::Close(_)
        | LowInstr::Tbc(_)
        | LowInstr::NumericForInit(_)
        | LowInstr::NumericForLoop(_)
        | LowInstr::GenericForCall(_)
        | LowInstr::GenericForLoop(_)
        | LowInstr::Jump(_)
        | LowInstr::Branch(_) => false,
    }
}

fn result_pack_writes_any_reg(results: &ResultPack, regs: &BTreeSet<Reg>) -> bool {
    match results {
        ResultPack::Fixed(range) => (0..range.len)
            .map(|offset| Reg(range.start.index() + offset))
            .any(|reg| regs.contains(&reg)),
        // Open result pack 从 start 开始向高位扩展，保守地检查 start 本身。
        // 实际 while 条件求值中 open pack 极少出现，但如果出现则 start 之后的
        // 所有寄存器理论上都可能被写入——这里保守处理即可，因为 carried_regs
        // 通常只包含少量低位循环变量。
        ResultPack::Open(start) => regs.iter().any(|reg| reg.index() >= start.index()),
        ResultPack::Ignore => false,
    }
}

fn repeat_continue_target_via_backedge_pad(
    proto: &LoweredProto,
    cfg: &Cfg,
    backedge_source: BlockRef,
    blocks: &BTreeSet<BlockRef>,
) -> Option<BlockRef> {
    let preds = cfg
        .reachable_predecessors(backedge_source)
        .into_iter()
        .filter(|pred| blocks.contains(pred))
        .collect::<Vec<_>>();
    let [continue_target] = preds.as_slice() else {
        return None;
    };

    if !matches!(
        cfg.terminator(&proto.instrs, *continue_target),
        Some(LowInstr::Branch(_))
    ) {
        return None;
    }

    let (then_edge_ref, else_edge_ref) = cfg.branch_edges(*continue_target)?;
    let then_block = cfg.edges[then_edge_ref.index()].to;
    let else_block = cfg.edges[else_edge_ref.index()].to;

    if (then_block == backedge_source && !blocks.contains(&else_block))
        || (else_block == backedge_source && !blocks.contains(&then_block))
    {
        Some(*continue_target)
    } else {
        None
    }
}

fn generic_for_has_loop_body_and_exit(
    proto: &LoweredProto,
    cfg: &Cfg,
    header: BlockRef,
    instr: &crate::transformer::GenericForLoopInstr,
    blocks: &BTreeSet<BlockRef>,
) -> bool {
    let range = cfg.blocks[header.index()].instrs;
    if range.len < 2 {
        return false;
    }
    let Some(call_instr_index) = range.end().checked_sub(2) else {
        return false;
    };
    let Some(LowInstr::GenericForCall(call)) = proto.instrs.get(call_instr_index) else {
        return false;
    };
    let body_block = cfg.instr_to_block[instr.body_target.index()];
    let exit_block = cfg.instr_to_block[instr.exit_target.index()];

    matches!(call.results, crate::transformer::ResultPack::Fixed(range) if range == instr.bindings)
        && blocks.contains(&body_block)
        && !blocks.contains(&exit_block)
}