unluac 1.3.1

//! 这个文件集中承载“StructureFacts 如何消费 Dataflow phi”的共享翻译规则。
//!
//! `loops / branch_values / short_circuit` 都会把 `phi.incoming` 重新整理成更贴近
//! 源码恢复的结构事实。如果每个 pass 都各自维护一套 `incoming -> arm/value defs`
//! 转换，规则一变就会三处平行返工。这里把这层翻译集中成单一 owner，让结构层
//! 共享同一套 phi 语义。
//!
//! 它依赖 Dataflow 已经提供稳定的 `phi_candidates / reaching_defs / def 元数据`，
//! 这里只负责把这些底层 merge 事实改写成 StructureFacts 可直接消费的形状；
//! 它不会越权决定最终 HIR 表达式或语法结构。
//!
//! 例子：
//! - branch merge 会把 `phi.incoming` 直接整理成 `then_arm / else_arm` 两臂 def 集
//! - loop header/exit merge 会整理成 `inside_arm / outside_arm` 或按 predecessor
//!   分组的 incoming facts
//! - short-circuit value merge 会提前带出 `entry_defs / value_incomings`，避免 HIR
//!   再回头拆 phi

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

use crate::structure::{BlockRef, Cfg, DataflowFacts, DefId, GraphFacts, PhiCandidate, PhiId};
use crate::transformer::Reg;

use super::common::{
    BranchValueMergeArm, BranchValueMergeCandidate, BranchValueMergeValue,
    GenericPhiMaterialization, GenericPhiSource, LoopCandidate, LoopValueArm, LoopValueIncoming,
    LoopValueMerge, ShortCircuitCandidate, ShortCircuitValueIncoming,
};

pub(super) struct ShortCircuitPhiFacts {
    pub(super) entry_defs: BTreeSet<DefId>,
    pub(super) value_incomings: Vec<ShortCircuitValueIncoming>,
}

pub(super) fn branch_value_merge_from_phi(
    header: BlockRef,
    dataflow: &DataflowFacts,
    phi: &PhiCandidate,
    then_preds: &BTreeSet<BlockRef>,
    else_preds: &BTreeSet<BlockRef>,
) -> Option<BranchValueMergeValue> {
    let mut then_arm = BranchValueMergeArm {
        preds: BTreeSet::new(),
        defs: BTreeSet::new(),
        non_header_defs: BTreeSet::new(),
    };
    let mut else_arm = BranchValueMergeArm {
        preds: BTreeSet::new(),
        defs: BTreeSet::new(),
        non_header_defs: BTreeSet::new(),
    };

    for incoming in &phi.incoming {
        let pred = incoming.pred?;
        if then_preds.contains(&pred) {
            extend_branch_value_arm(header, dataflow, &mut then_arm, incoming);
        } else if else_preds.contains(&pred) {
            extend_branch_value_arm(header, dataflow, &mut else_arm, incoming);
        } else {
            return None;
        }
    }

    (!then_arm.preds.is_empty() && !else_arm.preds.is_empty()).then_some(BranchValueMergeValue {
        phi_id: phi.id,
        reg: phi.reg,
        then_arm,
        else_arm,
    })
}

pub(super) fn branch_value_merges_in_block(
    header: BlockRef,
    dataflow: &DataflowFacts,
    block: BlockRef,
    then_preds: &BTreeSet<BlockRef>,
    else_preds: &BTreeSet<BlockRef>,
) -> Vec<BranchValueMergeValue> {
    dataflow
        .phi_candidates_in_block(block)
        .iter()
        .filter_map(|phi| {
            branch_value_merge_from_phi(header, dataflow, phi, then_preds, else_preds)
        })
        .collect()
}

pub(super) fn loop_value_merge_from_phi(
    phi: &PhiCandidate,
    loop_blocks: &BTreeSet<BlockRef>,
) -> Option<LoopValueMerge> {
    let mut inside_arm = LoopValueArm::default();
    let mut outside_arm = LoopValueArm::default();

    for incoming in &phi.incoming {
        let arm = if incoming
            .pred
            .is_some_and(|pred| loop_blocks.contains(&pred))
        {
            &mut inside_arm
        } else {
            &mut outside_arm
        };
        arm.incomings.push(LoopValueIncoming {
            pred: incoming.pred,
            defs: incoming.defs.clone(),
        });
    }

    Some(LoopValueMerge {
        phi_id: phi.id,
        reg: phi.reg,
        inside_arm,
        outside_arm,
    })
}

pub(super) fn loop_value_merges_in_block(
    dataflow: &DataflowFacts,
    block: BlockRef,
    loop_blocks: &BTreeSet<BlockRef>,
) -> Vec<LoopValueMerge> {
    dataflow
        .phi_candidates_in_block(block)
        .iter()
        .filter_map(|phi| loop_value_merge_from_phi(phi, loop_blocks))
        .collect()
}

pub(super) fn short_circuit_phi_facts(
    cfg: &Cfg,
    dataflow: &DataflowFacts,
    header: BlockRef,
    reg: Reg,
    phi: &PhiCandidate,
) -> ShortCircuitPhiFacts {
    ShortCircuitPhiFacts {
        entry_defs: value_merge_entry_defs(cfg, dataflow, header, reg),
        value_incomings: phi
            .incoming
            .iter()
            .filter_map(|incoming| {
                let pred = incoming.pred?;
                Some(ShortCircuitValueIncoming {
                    pred,
                    defs: incoming.defs.clone(),
                    latest_local_def: latest_local_incoming_def(dataflow, pred, &incoming.defs),
                })
            })
            .collect(),
    }
}

pub(super) fn analyze_generic_phi_materializations(
    cfg: &Cfg,
    graph_facts: &GraphFacts,
    dataflow: &DataflowFacts,
    branch_value_merge_candidates: &[BranchValueMergeCandidate],
    loop_candidates: &[LoopCandidate],
    short_circuit_candidates: &[ShortCircuitCandidate],
) -> Vec<GenericPhiMaterialization> {
    let mut covered = short_circuit_candidates
        .iter()
        .filter(|candidate| candidate.reducible)
        .filter_map(|candidate| candidate.result_phi_id)
        .collect::<BTreeSet<_>>();
    covered.extend(consumed_branch_value_merge_ids(
        branch_value_merge_candidates,
    ));
    covered.extend(consumed_loop_value_merge_ids(loop_candidates));

    let mut generic = dataflow
        .phi_candidates
        .iter()
        .filter(|phi| !covered.contains(&phi.id))
        .map(|phi| GenericPhiMaterialization {
            block: phi.block,
            phi_id: phi.id,
            reg: phi.reg,
            source: generic_phi_source(cfg, graph_facts, dataflow, phi),
        })
        .collect::<Vec<_>>();
    generic.sort_by_key(|phi| (phi.block, phi.phi_id));
    generic
}

fn consumed_branch_value_merge_ids(
    candidates: &[BranchValueMergeCandidate],
) -> impl Iterator<Item = PhiId> + '_ {
    let mut by_header = BTreeMap::<BlockRef, Option<&BranchValueMergeCandidate>>::new();

    for candidate in candidates {
        by_header
            .entry(candidate.header)
            .and_modify(|entry| *entry = None)
            .or_insert(Some(candidate));
    }

    by_header
        .into_values()
        .flatten()
        .flat_map(|candidate| candidate.values.iter().map(|value| value.phi_id))
}

fn generic_phi_source(
    cfg: &Cfg,
    graph_facts: &GraphFacts,
    dataflow: &DataflowFacts,
    phi: &PhiCandidate,
) -> GenericPhiSource {
    let Some(idom) = graph_facts
        .dominator_tree
        .parent
        .get(phi.block.index())
        .copied()
        .flatten()
    else {
        return GenericPhiSource::Unresolved;
    };
    let Some(idom_defs) = block_exit_defs_for_reg(cfg, dataflow, idom, phi.reg) else {
        return GenericPhiSource::Unresolved;
    };
    if phi
        .incoming
        .iter()
        .all(|incoming| incoming.defs == idom_defs)
    {
        GenericPhiSource::IdomExit(idom)
    } else {
        GenericPhiSource::Unresolved
    }
}

fn block_exit_defs_for_reg(
    cfg: &Cfg,
    dataflow: &DataflowFacts,
    block: BlockRef,
    reg: Reg,
) -> Option<BTreeSet<DefId>> {
    let range = cfg.blocks[block.index()].instrs;
    let Some(last_instr_ref) = range.last() else {
        return Some(BTreeSet::new());
    };

    let effect = &dataflow.instr_effects[last_instr_ref.index()];
    if effect.fixed_must_defs.contains(&reg) {
        return dataflow
            .instr_def_for_reg(last_instr_ref, reg)
            .map(|def| BTreeSet::from([def]));
    }

    let mut defs = dataflow
        .reaching_defs_at(last_instr_ref)
        .fixed
        .get(reg)
        .map(|defs| defs.iter().copied().collect::<BTreeSet<_>>())
        .unwrap_or_default();
    if effect.fixed_may_defs.contains(&reg) {
        defs.insert(dataflow.instr_def_for_reg(last_instr_ref, reg)?);
    }
    Some(defs)
}

fn extend_branch_value_arm(
    header: BlockRef,
    dataflow: &DataflowFacts,
    arm: &mut BranchValueMergeArm,
    incoming: &crate::structure::PhiIncoming,
) {
    let Some(pred) = incoming.pred else {
        return;
    };
    arm.preds.insert(pred);
    for &def in &incoming.defs {
        arm.defs.insert(def);
        if dataflow.def_block(def) != header {
            arm.non_header_defs.insert(def);
        }
    }
}

fn latest_local_incoming_def(
    dataflow: &DataflowFacts,
    block: BlockRef,
    defs: &BTreeSet<DefId>,
) -> Option<DefId> {
    dataflow.latest_local_def_in_block(block, defs.iter().copied())
}

fn value_merge_entry_defs(
    cfg: &Cfg,
    dataflow: &DataflowFacts,
    header: BlockRef,
    reg: Reg,
) -> BTreeSet<DefId> {
    let Some(instr_ref) = cfg.blocks[header.index()].instrs.last() else {
        return BTreeSet::new();
    };

    dataflow
        .reaching_defs_at(instr_ref)
        .fixed
        .get(reg)
        .map(|defs| defs.iter().copied().collect())
        .unwrap_or_default()
}

fn consumed_loop_value_merge_ids(loop_candidates: &[LoopCandidate]) -> BTreeSet<PhiId> {
    let mut ids = BTreeSet::new();

    for candidate in loop_candidates {
        let header_regs = candidate
            .header_value_merges
            .iter()
            .map(|value| value.reg)
            .collect::<BTreeSet<_>>();

        ids.extend(
            candidate
                .header_value_merges
                .iter()
                .map(|value| value.phi_id),
        );
        ids.extend(candidate.exit_value_merges.iter().flat_map(|exit| {
            exit.values
                .iter()
                .filter(|value| header_regs.contains(&value.reg))
                .map(|value| value.phi_id)
        }));
    }

    ids
}