hakorune/src/runner/json_v0_bridge/lowering/loop_.rs

use super::{lower_stmt_list_with_vars, new_block, BridgeEnv, LoopContext};
use crate::mir::{BasicBlockId, MirFunction, MirInstruction, ValueId};
use std::collections::HashMap;
use super::super::ast::StmtV0;
use super::super::ast::ExprV0;

pub(super) fn lower_loop_stmt(
    f: &mut MirFunction,
    cur_bb: BasicBlockId,
    cond: &ExprV0,
    body: &[StmtV0],
    vars: &mut HashMap<String, ValueId>,
    loop_stack: &mut Vec<LoopContext>,
    env: &BridgeEnv,
) -> Result<BasicBlockId, String> {
    let cond_bb = new_block(f);
    let body_bb = new_block(f);
    let exit_bb = new_block(f);

    // SSOT（phi_core）への統一: preheader Copy → header Phi seed を集中実装に委譲
    // 1) preheader スナップショット
    let base_vars = vars.clone();
    let mut block_var_maps: HashMap<BasicBlockId, HashMap<String, ValueId>> = HashMap::new();
    block_var_maps.insert(cur_bb, base_vars.clone());

    // 2) preheader → header へ Jump（ヘッダで Phi を組む前に制御を渡す）
    if let Some(bb) = f.get_block_mut(cur_bb) {
        if !bb.is_terminated() {
            bb.add_instruction(MirInstruction::Jump { target: cond_bb });
        }
    }

    // 3) LoopPhiOps アダプタ（準備用: preheader seed 専用）
    struct Ops<'a> {
        f: &'a mut MirFunction,
        vars: &'a mut HashMap<String, ValueId>,
        block_var_maps: &'a HashMap<BasicBlockId, HashMap<String, ValueId>>,
    }
    impl crate::mir::phi_core::loop_phi::LoopPhiOps for Ops<'_> {
        fn new_value(&mut self) -> ValueId { self.f.next_value_id() }
        fn emit_phi_at_block_start(
            &mut self,
            block: BasicBlockId,
            dst: ValueId,
            mut inputs: Vec<(BasicBlockId, ValueId)>,
        ) -> Result<(), String> {
            if let Some(bb) = self.f.get_block_mut(block) {
                // If a PHI for the same dst already exists at block start, merge inputs instead of inserting
                let mut found = false;
                // Count PHIs at head and iterate by index to allow mutation
                let phi_count = bb.phi_instructions().count();
                for i in 0..phi_count {
                    if let Some(MirInstruction::Phi { dst: existing_dst, inputs: existing_inputs }) = bb.instructions.get_mut(i) {
                        if *existing_dst == dst {
                            // Merge: add only missing predecessors (build a set first to avoid borrow conflicts)
                            let mut pred_set: std::collections::HashSet<BasicBlockId> = existing_inputs.iter().map(|(bbid, _)| *bbid).collect();
                            for (pred, val) in inputs.drain(..) {
                                if !pred_set.contains(&pred) {
                                    existing_inputs.push((pred, val));
                                    pred_set.insert(pred);
                                }
                            }
                            found = true;
                            break;
                        }
                    }
                }
                if !found {
                    bb.insert_instruction_after_phis(MirInstruction::Phi { dst, inputs });
                }
                Ok(())
            } else {
                Err(format!("block {} not found", block.0))
            }
        }
        fn update_var(&mut self, name: String, value: ValueId) { self.vars.insert(name, value); }
        fn get_variable_at_block(&mut self, name: &str, block: BasicBlockId) -> Option<ValueId> {
            self.block_var_maps.get(&block).and_then(|m| m.get(name).copied())
        }
        fn debug_verify_phi_inputs(&mut self, merge_bb: BasicBlockId, inputs: &[(BasicBlockId, ValueId)]) {
            #[cfg(debug_assertions)]
            if let Some(func) = Some(&*self.f) {
                crate::mir::phi_core::common::debug_verify_phi_inputs(func, merge_bb, inputs);
            }
        }
        fn emit_copy_at_preheader(
            &mut self,
            preheader_id: BasicBlockId,
            dst: ValueId,
            src: ValueId,
        ) -> Result<(), String> {
            if let Some(bb) = self.f.get_block_mut(preheader_id) {
                bb.add_instruction(MirInstruction::Copy { dst, src });
                Ok(())
            } else {
                Err(format!("preheader block {} not found", preheader_id.0))
            }
        }
    }

    let mut ops = Ops { f, vars, block_var_maps: &block_var_maps };
    // 4) header の不完全PHIを宣言（preheaderのコピー値でseed）し、変数マップをPHIへ再束縛
    let mut incomplete = crate::mir::phi_core::loop_phi::prepare_loop_variables_with(
        &mut ops, cond_bb, cur_bb, &base_vars,
    )?;
    // Header snapshot for exit PHIs
    let header_vars_snapshot = ops.vars.clone();

    let (cval, _cend) = super::expr::lower_expr_with_vars(env, ops.f, cond_bb, cond, ops.vars)?;
    if let Some(bb) = ops.f.get_block_mut(cond_bb) {
        bb.set_terminator(MirInstruction::Branch {
            condition: cval,
            then_bb: body_bb,
            else_bb: exit_bb,
        });
    }
    let mut body_vars = ops.vars.clone();
    // open snapshot frames for nested break/continue
    super::push_loop_snapshot_frames();
    loop_stack.push(LoopContext { cond_bb, exit_bb });
    let bend_res = lower_stmt_list_with_vars(ops.f, body_bb, body, &mut body_vars, loop_stack, env);
    loop_stack.pop();
    let bend = bend_res?;
    // collect snapshots for this loop level
    let continue_snaps = super::pop_continue_snapshots();
    let exit_snaps = super::pop_exit_snapshots();
    // latch（body末尾）スナップショットを別マップに構築
    let mut block_var_maps2: HashMap<BasicBlockId, HashMap<String, ValueId>> = HashMap::new();
    block_var_maps2.insert(cur_bb, base_vars.clone());
    block_var_maps2.insert(bend, body_vars.clone());
    if let Some(bb) = ops.f.get_block_mut(bend) {
        if !bb.is_terminated() {
            bb.set_terminator(MirInstruction::Jump { target: cond_bb });
        }
    }
    let backedge_to_cond = matches!(
        ops.f.blocks
            .get(&bend)
            .and_then(|bb| bb.terminator.as_ref()),
        Some(MirInstruction::Jump { target, .. }) if *target == cond_bb
    );
    if backedge_to_cond {
        // 5) header の不完全PHIを完成（latch + continue スナップショット集約）
        let mut ops_seal = Ops { f: ops.f, vars: ops.vars, block_var_maps: &block_var_maps2 };
        crate::mir::phi_core::loop_phi::seal_incomplete_phis_with(
            &mut ops_seal,
            cond_bb,
            bend,
            std::mem::take(&mut incomplete),
            &continue_snaps,
        )?;
    }
    // 6) exit PHIs（breakの合流 + headerからのフォールスルー）
    let mut ops_exit = Ops { f: ops.f, vars: ops.vars, block_var_maps: &block_var_maps2 };
    crate::mir::phi_core::loop_phi::build_exit_phis_with(
        &mut ops_exit,
        cond_bb,
        exit_bb,
        &header_vars_snapshot,
        &exit_snaps,
    )?;
    Ok(exit_bb)
}