hakorune/src/mir/builder/control_flow/joinir/merge/mod.rs

//! JoinIR MIR Block Merging Coordinator
//!
//! This module coordinates the merging of JoinIR-generated MIR functions
//! into the host MIR builder. The process is broken into 6 phases:
//!
//! 1. Block ID allocation (block_allocator.rs)
//! 2. Value collection (value_collector.rs)
//! 3. ValueId remapping (uses JoinIrIdRemapper)
//! 4. Instruction rewriting (instruction_rewriter.rs)
//! 5. Exit PHI construction (exit_phi_builder.rs)
//! 6. Boundary reconnection (inline in this file)
//!
//! Phase 4 Refactoring: Breaking down 714-line merge_joinir_mir_blocks() into focused modules

mod block_allocator;
mod value_collector;
mod instruction_rewriter;
mod exit_phi_builder;

use crate::mir::{MirModule, ValueId};
use crate::mir::join_ir::lowering::inline_boundary::JoinInlineBoundary;

/// Phase 49-3.2: Merge JoinIR-generated MIR blocks into current_function
///
/// # Phase 189: Multi-Function MIR Merge
///
/// This merges JoinIR-generated blocks by:
/// 1. Remapping all block IDs across ALL functions to avoid conflicts
/// 2. Remapping all value IDs across ALL functions to avoid conflicts
/// 3. Adding all blocks from all functions to current_function
/// 4. Jumping from current_block to the entry block
/// 5. Converting Return → Jump to exit block for all functions
///
/// **Multi-Function Support** (Phase 189):
/// - Pattern 1 (Simple While) generates 3 functions: entry + loop_step + k_exit
/// - All functions are flattened into current_function with global ID remapping
/// - Single exit block receives all Return instructions from all functions
///
/// # Phase 188-Impl-3: JoinInlineBoundary Support
///
/// When `boundary` is provided, injects Copy instructions at the entry block
/// to connect host ValueIds to JoinIR local ValueIds:
///
/// ```text
/// entry_block:
///   // Injected by boundary
///   ValueId(100) = Copy ValueId(4)  // join_input → host_input
///   // Original JoinIR instructions follow...
/// ```
///
/// This enables clean separation: JoinIR uses local IDs (0,1,2...),
/// host uses its own IDs, and Copy instructions bridge the gap.
///
/// # Returns
///
/// Returns `Ok(Some(exit_phi_id))` if the merged JoinIR functions have return values
/// that were collected into an exit block PHI. さらに、`boundary` に
/// host_outputs が指定されている場合は、exit PHI の結果をホスト側の
/// SSA スロットへ再接続する（variable_map 内の ValueId を更新する）。
pub(in crate::mir::builder) fn merge_joinir_mir_blocks(
    builder: &mut crate::mir::builder::MirBuilder,
    mir_module: &MirModule,
    boundary: Option<&JoinInlineBoundary>,
    debug: bool,
) -> Result<Option<ValueId>, String> {
    if debug {
        eprintln!(
            "[cf_loop/joinir] merge_joinir_mir_blocks called with {} functions",
            mir_module.functions.len()
        );
    }

    // Phase 1: Allocate block IDs for all functions
    let mut remapper = block_allocator::allocate_blocks(builder, mir_module, debug)?;

    // Phase 2: Collect values from all functions
    let (mut used_values, value_to_func_name, function_params) =
        value_collector::collect_values(mir_module, &remapper, debug)?;

    // Phase 171-fix: Add condition_bindings' join_values to used_values for remapping
    if let Some(boundary) = boundary {
        for binding in &boundary.condition_bindings {
            if debug {
                eprintln!(
                    "[cf_loop/joinir] Phase 171-fix: Adding condition binding '{}' JoinIR {:?} to used_values",
                    binding.name, binding.join_value
                );
            }
            used_values.insert(binding.join_value);
        }
    }

    // Phase 3: Remap ValueIds
    remap_values(builder, &used_values, &mut remapper, debug)?;

    // Phase 4: Merge blocks and rewrite instructions
    let (exit_block_id, exit_phi_inputs) = instruction_rewriter::merge_and_rewrite(
        builder,
        mir_module,
        &mut remapper,
        &value_to_func_name,
        &function_params,
        boundary,
        debug,
    )?;

    // Phase 5: Build exit PHI
    let exit_phi_result_id = exit_phi_builder::build_exit_phi(
        builder,
        exit_block_id,
        &exit_phi_inputs,
        debug,
    )?;

    // Phase 6: Reconnect boundary (if specified)
    // Phase 197-B: Pass remapper to enable per-carrier exit value lookup
    if let Some(boundary) = boundary {
        reconnect_boundary(builder, boundary, &remapper, exit_phi_result_id, debug)?;
    }

    // Jump from current block to entry function's entry block
    let (entry_func_name, entry_func) = mir_module
        .functions
        .iter()
        .next()
        .ok_or("JoinIR module has no functions")?;
    let entry_block = remapper
        .get_block(entry_func_name, entry_func.entry_block)
        .ok_or_else(|| format!("Entry block not found for {}", entry_func_name))?;

    if debug {
        eprintln!("[cf_loop/joinir]   Entry function name: {}", entry_func_name);
        eprintln!(
            "[cf_loop/joinir]   Entry function's entry_block (JoinIR local): {:?}",
            entry_func.entry_block
        );
        eprintln!("[cf_loop/joinir]   Entry block (remapped): {:?}", entry_block);
        eprintln!(
            "[cf_loop/joinir]   Current block before emit_jump: {:?}",
            builder.current_block
        );
        eprintln!("[cf_loop/joinir]   Jumping to entry block: {:?}", entry_block);
    }

    crate::mir::builder::emission::branch::emit_jump(builder, entry_block)?;

    if debug {
        eprintln!(
            "[cf_loop/joinir]   After emit_jump, current_block: {:?}",
            builder.current_block
        );
    }

    // Switch to exit block for subsequent code
    builder.start_new_block(exit_block_id)?;

    if debug {
        eprintln!(
            "[cf_loop/joinir] Phase 189: Merge complete: {} functions merged, continuing from {:?}",
            mir_module.functions.len(),
            exit_block_id
        );
    }

    Ok(exit_phi_result_id)
}

/// Phase 3: Allocate new ValueIds for all collected values
fn remap_values(
    builder: &mut crate::mir::builder::MirBuilder,
    used_values: &std::collections::BTreeSet<ValueId>,
    remapper: &mut crate::mir::builder::joinir_id_remapper::JoinIrIdRemapper,
    debug: bool,
) -> Result<(), String> {
    if debug {
        eprintln!("[cf_loop/joinir] Phase 3: Remapping {} ValueIds", used_values.len());
    }

    for old_value in used_values {
        let new_value = builder.next_value_id();
        remapper.set_value(*old_value, new_value);
        if debug {
            eprintln!(
                "[cf_loop/joinir]   Value remap: {:?} → {:?}",
                old_value, new_value
            );
        }
    }

    Ok(())
}

/// Phase 6: Reconnect boundary to update host variable_map
///
/// Phase 197-B: Multi-carrier support
/// Previously, a single exit_phi_result was applied to all carriers, causing
/// all carriers to get the same value (e.g., both sum and count = 5).
///
/// Now, we use each binding's join_exit_value and the remapper to find
/// the actual exit value for each carrier.
fn reconnect_boundary(
    builder: &mut crate::mir::builder::MirBuilder,
    boundary: &JoinInlineBoundary,
    remapper: &crate::mir::builder::joinir_id_remapper::JoinIrIdRemapper,
    _exit_phi_result: Option<ValueId>,
    debug: bool,
) -> Result<(), String> {
    // Phase 197-B: Multi-carrier exit binding using join_exit_value directly
    //
    // Problem: Previously used single exit_phi_result for all carriers.
    // Solution: Each carrier's exit value is in k_exit parameters.
    //           We use remapper to find the remapped ValueId for each carrier.
    //
    // For Pattern 4 with 2 carriers (sum, count):
    // - k_exit(sum_exit, count_exit) receives values via Jump args
    // - sum_exit = ValueId(N) in JoinIR → remapped to ValueId(M) in host
    // - count_exit = ValueId(N+1) in JoinIR → remapped to ValueId(M+1) in host
    //
    // Each carrier's variable_map entry is updated with its specific remapped value.

    if boundary.exit_bindings.is_empty() {
        if debug {
            eprintln!("[cf_loop/joinir] Phase 197-B: No exit bindings, skipping reconnect");
        }
        return Ok(());
    }

    if debug {
        eprintln!(
            "[cf_loop/joinir] Phase 197-B: Reconnecting {} exit bindings",
            boundary.exit_bindings.len()
        );
    }

    // For each binding, look up the remapped exit value and update variable_map
    for binding in &boundary.exit_bindings {
        // Phase 197-B: Use remapper to find the remapped exit value
        let remapped_exit = remapper.get_value(binding.join_exit_value);

        if debug {
            eprintln!(
                "[cf_loop/joinir] Phase 197-B: Carrier '{}' join_exit={:?} → remapped={:?}",
                binding.carrier_name, binding.join_exit_value, remapped_exit
            );
        }

        if let Some(remapped_value) = remapped_exit {
            // Update variable_map with the remapped exit value
            if let Some(var_vid) = builder.variable_map.get_mut(&binding.carrier_name) {
                if debug {
                    eprintln!(
                        "[cf_loop/joinir] Phase 197-B: Updated variable_map['{}'] {:?} → {:?}",
                        binding.carrier_name, var_vid, remapped_value
                    );
                }
                *var_vid = remapped_value;
            } else if debug {
                eprintln!(
                    "[cf_loop/joinir] Phase 197-B WARNING: Carrier '{}' not found in variable_map",
                    binding.carrier_name
                );
            }
        } else if debug {
            eprintln!(
                "[cf_loop/joinir] Phase 197-B WARNING: No remapped value for join_exit={:?}",
                binding.join_exit_value
            );
        }
    }

    // Phase 190: Backward compatibility - also check deprecated host_outputs
    #[allow(deprecated)]
    if !boundary.host_outputs.is_empty() && debug {
        eprintln!(
            "[cf_loop/joinir] WARNING: Using deprecated host_outputs. Migrate to exit_bindings."
        );
    }

    Ok(())
}