hakorune/src/mir/builder/phi.rs

use super::MirBuilder;
use crate::ast::ASTNode;
use crate::mir::{BasicBlockId, MirFunction, MirInstruction, MirType, ValueId};
use std::collections::HashMap;

// PHI-based return type inference helper
pub(super) fn infer_type_from_phi(
    function: &MirFunction,
    ret_val: ValueId,
    types: &HashMap<ValueId, MirType>,
) -> Option<MirType> {
    for (_bid, bb) in function.blocks.iter() {
        for inst in bb.instructions.iter() {
            if let MirInstruction::Phi { dst, inputs } = inst {
                if *dst == ret_val {
                    let mut it = inputs.iter().filter_map(|(_, v)| types.get(v));
                    if let Some(first) = it.next() {
                        if it.all(|mt| mt == first) {
                            return Some(first.clone());
                        }
                    }
                }
            }
        }
    }
    None
}

// Local helper for if-statement analysis (moved from stmts.rs)
pub(super) fn extract_assigned_var(ast: &ASTNode) -> Option<String> {
    match ast {
        ASTNode::Assignment { target, .. } => {
            if let ASTNode::Variable { name, .. } = target.as_ref() {
                Some(name.clone())
            } else {
                None
            }
        }
        ASTNode::Program { statements, .. } => {
            statements.last().and_then(|st| extract_assigned_var(st))
        }
        ASTNode::If {
            then_body,
            else_body,
            ..
        } => {
            // Look into nested if: if both sides assign the same variable, propagate that name upward.
            let then_prog = ASTNode::Program {
                statements: then_body.clone(),
                span: crate::ast::Span::unknown(),
            };
            let tvar = extract_assigned_var(&then_prog);
            let evar = else_body.as_ref().and_then(|eb| {
                let ep = ASTNode::Program {
                    statements: eb.clone(),
                    span: crate::ast::Span::unknown(),
                };
                extract_assigned_var(&ep)
            });
            match (tvar, evar) {
                (Some(tv), Some(ev)) if tv == ev => Some(tv),
                _ => None,
            }
        }
        _ => None,
    }
}

impl MirBuilder {
    /// Merge all variables modified in then/else relative to pre_if_snapshot.
    /// In PHI-off mode inserts edge copies from branch exits to merge. In PHI-on mode emits Phi.
    /// `skip_var` allows skipping a variable already merged elsewhere (e.g., bound to an expression result).
    pub(super) fn merge_modified_vars(
        &mut self,
        then_block: super::BasicBlockId,
        else_block: super::BasicBlockId,
        then_exit_block: super::BasicBlockId,
        else_exit_block_opt: Option<super::BasicBlockId>,
        pre_if_snapshot: &std::collections::HashMap<String, super::ValueId>,
        then_map_end: &std::collections::HashMap<String, super::ValueId>,
        else_map_end_opt: &Option<std::collections::HashMap<String, super::ValueId>>,
        skip_var: Option<&str>,
    ) -> Result<(), String> {
        use std::collections::HashSet;
        let mut names: HashSet<&str> = HashSet::new();
        for k in then_map_end.keys() { names.insert(k.as_str()); }
        if let Some(emap) = else_map_end_opt.as_ref() {
            for k in emap.keys() { names.insert(k.as_str()); }
        }
        // Only variables that changed against pre_if_snapshot
        let mut changed: Vec<&str> = Vec::new();
        for &name in &names {
            let pre = pre_if_snapshot.get(name);
            let t = then_map_end.get(name);
            let e = else_map_end_opt.as_ref().and_then(|m| m.get(name));
            // changed when either branch value differs from pre
            if (t.is_some() && Some(t.copied().unwrap()) != pre.copied())
                || (e.is_some() && Some(e.copied().unwrap()) != pre.copied())
            {
                changed.push(name);
            }
        }
        for name in changed {
            if skip_var.map(|s| s == name).unwrap_or(false) {
                continue;
            }
            let pre = match pre_if_snapshot.get(name) {
                Some(v) => *v,
                None => continue, // unknown before-if; skip
            };
            let then_v = then_map_end.get(name).copied().unwrap_or(pre);
            let else_v = else_map_end_opt
                .as_ref()
                .and_then(|m| m.get(name).copied())
                .unwrap_or(pre);
            if self.is_no_phi_mode() {
                let merged = self.value_gen.next();
                // Insert edge copies from then/else exits into merge
                self.insert_edge_copy(then_exit_block, merged, then_v)?;
                if let Some(else_exit_block) = else_exit_block_opt {
                    self.insert_edge_copy(else_exit_block, merged, else_v)?;
                } else {
                    // Fallback: if else missing, copy pre value from then as both inputs already cover
                    self.insert_edge_copy(then_exit_block, merged, then_v)?;
                }
                self.variable_map.insert(name.to_string(), merged);
            } else {
                let merged = self.value_gen.next();
                self.emit_instruction(
                    MirInstruction::Phi {
                        dst: merged,
                        inputs: vec![(then_block, then_v), (else_block, else_v)],
                    }
                )?;
                self.variable_map.insert(name.to_string(), merged);
            }
        }
        Ok(())
    }
    /// Normalize Phi creation for if/else constructs.
    /// This handles variable reassignment patterns and ensures a single exit value.
    pub(super) fn normalize_if_else_phi(
        &mut self,
        then_block: BasicBlockId,
        else_block: BasicBlockId,
        then_exit_block_opt: Option<BasicBlockId>,
        else_exit_block_opt: Option<BasicBlockId>,
        then_value_raw: ValueId,
        else_value_raw: ValueId,
        pre_if_var_map: &HashMap<String, ValueId>,
        then_ast_for_analysis: &ASTNode,
        else_ast_for_analysis: &Option<ASTNode>,
        then_var_map_end: &HashMap<String, ValueId>,
        else_var_map_end_opt: &Option<HashMap<String, ValueId>>,
        pre_then_var_value: Option<ValueId>,
    ) -> Result<ValueId, String> {
        // If only the then-branch assigns a variable (e.g., `if c { x = ... }`) and the else
        // does not assign the same variable, bind that variable to a Phi of (then_value, pre_if_value).
        let assigned_var_then = extract_assigned_var(then_ast_for_analysis);
        let assigned_var_else = else_ast_for_analysis
            .as_ref()
            .and_then(|a| extract_assigned_var(a));
        let result_val = self.value_gen.next();

        if self.is_no_phi_mode() {
            // In PHI-off mode, emit per-predecessor copies into the actual predecessors
            // of the current (merge) block instead of the entry blocks. This correctly
            // handles nested conditionals where the then-branch fans out and merges later.
            let merge_block = self
                .current_block
                .ok_or_else(|| "normalize_if_else_phi: no current (merge) block".to_string())?;
            let preds: Vec<crate::mir::BasicBlockId> = if let Some(ref fun_ro) = self.current_function {
                if let Some(bb) = fun_ro.get_block(merge_block) {
                    bb.predecessors.iter().copied().collect()
                } else {
                    Vec::new()
                }
            } else {
                Vec::new()
            };
            // Prefer explicit exit blocks if provided; fall back to predecessor scan
            let then_exits: Vec<crate::mir::BasicBlockId> = if let Some(b) = then_exit_block_opt {
                vec![b]
            } else {
                preds.iter().copied().filter(|p| *p != else_block).collect()
            };
            let else_exits: Vec<crate::mir::BasicBlockId> = if let Some(b) = else_exit_block_opt {
                vec![b]
            } else {
                preds.iter().copied().filter(|p| *p == else_block).collect()
            };

            if let Some(var_name) = assigned_var_then.clone() {
                let else_assigns_same = assigned_var_else
                    .as_ref()
                    .map(|s| s == &var_name)
                    .unwrap_or(false);
                let then_value_for_var = then_var_map_end
                    .get(&var_name)
                    .copied()
                    .unwrap_or(then_value_raw);
                let else_value_for_var = if else_assigns_same {
                    else_var_map_end_opt
                        .as_ref()
                        .and_then(|m| m.get(&var_name).copied())
                        .unwrap_or(else_value_raw)
                } else {
                    pre_then_var_value.unwrap_or(else_value_raw)
                };
                // Map predecessors: else_block retains else value; others take then value
                for p in then_exits.iter().copied() {
                    self.insert_edge_copy(p, result_val, then_value_for_var)?;
                }
                for p in else_exits.iter().copied() {
                    self.insert_edge_copy(p, result_val, else_value_for_var)?;
                }
                self.variable_map = pre_if_var_map.clone();
                self.variable_map.insert(var_name, result_val);
            } else {
                for p in then_exits.iter().copied() {
                    self.insert_edge_copy(p, result_val, then_value_raw)?;
                }
                for p in else_exits.iter().copied() {
                    self.insert_edge_copy(p, result_val, else_value_raw)?;
                }
                self.variable_map = pre_if_var_map.clone();
            }
            return Ok(result_val);
        }

        if let Some(var_name) = assigned_var_then.clone() {
            let else_assigns_same = assigned_var_else
                .as_ref()
                .map(|s| s == &var_name)
                .unwrap_or(false);
            // Resolve branch-end values for the assigned variable
            let then_value_for_var = then_var_map_end
                .get(&var_name)
                .copied()
                .unwrap_or(then_value_raw);
            let else_value_for_var = if else_assigns_same {
                else_var_map_end_opt
                    .as_ref()
                    .and_then(|m| m.get(&var_name).copied())
                    .unwrap_or(else_value_raw)
            } else {
                // Else doesn't assign: use pre-if value if available
                pre_then_var_value.unwrap_or(else_value_raw)
            };
            // Emit Phi for the assigned variable and bind it
            self.emit_instruction(MirInstruction::Phi {
                dst: result_val,
                inputs: vec![
                    (then_block, then_value_for_var),
                    (else_block, else_value_for_var),
                ],
            })?;
            self.variable_map = pre_if_var_map.clone();
            self.variable_map.insert(var_name, result_val);
        } else {
            // No variable assignment pattern detected – just emit Phi for expression result
            self.emit_instruction(MirInstruction::Phi {
                dst: result_val,
                inputs: vec![(then_block, then_value_raw), (else_block, else_value_raw)],
            })?;
            // Merge variable map conservatively to pre-if snapshot (no new bindings)
            self.variable_map = pre_if_var_map.clone();
        }

        Ok(result_val)
    }
}