llvm-native-core 0.1.6

//! LLVM BasicBlock — container for instructions.
//! Phase 1 — LLVM.IR.1 Court.

use crate::opcode::Opcode;
use crate::types::Type;
use crate::value::{valref, SubclassKind, Value, ValueRef};

// ---------------------------------------------------------------------------
// BasicBlock type
// ---------------------------------------------------------------------------

/// A basic block is a straight-line sequence of instructions with no
/// internal branches, terminated by a terminator instruction (branch,
/// return, switch, etc.). Every basic block belongs to exactly one
/// Function.
#[derive(Debug, Clone)]
pub struct BasicBlock {
    /// The backing Value (Type::label, SubclassKind::BasicBlock).
    pub value: ValueRef,
    /// Whether this is the function's entry block.
    pub is_entry: bool,
    /// Immediate dominator block (if computed).
    pub dominator: Option<ValueRef>,
    /// Instructions in this block, in order.
    pub instructions: Vec<ValueRef>,
    /// Blocks that branch to this block (predecessors).
    pub predecessors: Vec<ValueRef>,
    /// Blocks this block branches to (successors, from terminator).
    pub successors: Vec<ValueRef>,
}

// ---------------------------------------------------------------------------
// Construction helpers (keep existing new_basic_block)
// ---------------------------------------------------------------------------

/// Create a bare basic block as a ValueRef.
/// This is the original Phase 1 constructor — kept for backward compat.
pub fn new_basic_block(name: &str) -> ValueRef {
    let v = Value::new(Type::label())
        .named(name)
        .with_subclass(SubclassKind::BasicBlock);
    valref(v)
}

impl BasicBlock {
    /// Create a new BasicBlock with the given name.
    pub fn new(name: &str) -> Self {
        let val = new_basic_block(name);
        Self {
            value: val,
            is_entry: false,
            dominator: None,
            instructions: Vec::new(),
            predecessors: Vec::new(),
            successors: Vec::new(),
        }
    }

    /// Wrap an existing ValueRef as a BasicBlock.
    pub fn from_value(val: ValueRef) -> Self {
        Self {
            value: val,
            is_entry: false,
            dominator: None,
            instructions: Vec::new(),
            predecessors: Vec::new(),
            successors: Vec::new(),
        }
    }

    // -----------------------------------------------------------------------
    // Predicate queries
    // -----------------------------------------------------------------------

    /// Check whether this block has a terminator instruction.
    pub fn is_terminated(&self) -> bool {
        self.instructions.last().map_or(false, |inst| {
            if let Some(opcode) = inst.borrow().get_opcode() {
                opcode.is_terminator()
            } else {
                false
            }
        })
    }

    /// Get the terminator instruction of this block, if any.
    pub fn get_terminator(&self) -> Option<ValueRef> {
        self.instructions.last().cloned()
    }

    /// Get the first non-phi instruction in the block.
    /// PHI nodes must appear at the top of a block; this returns the
    /// first instruction that is not a PHI node.
    pub fn get_first_non_phi(&self) -> Option<ValueRef> {
        self.instructions
            .iter()
            .find(|inst| {
                !inst
                    .borrow()
                    .get_opcode()
                    .map_or(false, |op| op == Opcode::Phi)
            })
            .cloned()
    }

    /// Get all PHI nodes in this block (must be at the top).
    pub fn get_phi_nodes(&self) -> Vec<ValueRef> {
        self.instructions
            .iter()
            .take_while(|inst| {
                inst.borrow()
                    .get_opcode()
                    .map_or(false, |op| op == Opcode::Phi)
            })
            .cloned()
            .collect()
    }

    /// Get the number of instructions in this block.
    pub fn get_instruction_count(&self) -> usize {
        self.instructions.len()
    }

    /// Get the instruction at the given index.
    pub fn get_instruction(&self, index: usize) -> Option<ValueRef> {
        self.instructions.get(index).cloned()
    }

    // -----------------------------------------------------------------------
    // Predecessor / successor management
    // -----------------------------------------------------------------------

    /// Add a predecessor block (a block that branches to this one).
    pub fn add_predecessor(&mut self, pred: ValueRef) {
        if !self.has_predecessor(&pred) {
            self.predecessors.push(pred);
        }
    }

    /// Remove a predecessor block.
    pub fn remove_predecessor(&mut self, pred: &ValueRef) -> bool {
        if let Some(pos) = self.predecessors.iter().position(|p| Rc::ptr_eq(p, pred)) {
            self.predecessors.remove(pos);
            true
        } else {
            false
        }
    }

    /// Check whether `pred` is a predecessor of this block.
    pub fn has_predecessor(&self, pred: &ValueRef) -> bool {
        self.predecessors.iter().any(|p| Rc::ptr_eq(p, pred))
    }

    /// Add a successor block (a block this block branches to).
    pub fn add_successor(&mut self, succ: ValueRef) {
        if !self.has_successor(&succ) {
            self.successors.push(succ);
        }
    }

    /// Remove a successor block.
    pub fn remove_successor(&mut self, succ: &ValueRef) -> bool {
        if let Some(pos) = self.successors.iter().position(|s| Rc::ptr_eq(s, succ)) {
            self.successors.remove(pos);
            true
        } else {
            false
        }
    }

    /// Check whether `succ` is a successor of this block.
    pub fn has_successor(&self, succ: &ValueRef) -> bool {
        self.successors.iter().any(|s| Rc::ptr_eq(s, succ))
    }

    /// Get all predecessor blocks.
    pub fn get_predecessors(&self) -> &[ValueRef] {
        &self.predecessors
    }

    /// Get all successor blocks.
    pub fn get_successors(&self) -> &[ValueRef] {
        &self.successors
    }

    /// Re-derive successors from the terminator instruction.
    /// For `br label %X` → successors = [X]
    /// For `br i1 ..., label %T, label %F` → successors = [T, F]
    /// For `switch` → successors = [default] + case dests
    /// For `ret` → successors = []
    pub fn recompute_successors(&mut self) {
        self.successors.clear();
        if let Some(terminator) = self.get_terminator() {
            let term = terminator.borrow();
            if let Some(opcode) = term.get_opcode() {
                match opcode {
                    Opcode::Br => {
                        // Conditional branch: operands[1] and operands[2] are the targets
                        if term.operands.len() >= 3 {
                            self.successors.push(Rc::clone(&term.operands[1]));
                            self.successors.push(Rc::clone(&term.operands[2]));
                        } else if term.operands.len() >= 1 {
                            // Unconditional branch: operands[0] is the target
                            self.successors.push(Rc::clone(&term.operands[0]));
                        }
                    }
                    Opcode::Switch => {
                        // operands[1] is default dest; operands[3], [5], ... are case dests
                        if term.operands.len() >= 2 {
                            self.successors.push(Rc::clone(&term.operands[1]));
                        }
                        for i in (3..term.operands.len()).step_by(2) {
                            self.successors.push(Rc::clone(&term.operands[i]));
                        }
                    }
                    Opcode::Invoke => {
                        // operands[last-1] = normal dest, operands[last] = unwind dest
                        let n = term.operands.len();
                        if n >= 2 {
                            self.successors.push(Rc::clone(&term.operands[n - 2]));
                            self.successors.push(Rc::clone(&term.operands[n - 1]));
                        }
                    }
                    Opcode::IndirectBr => {
                        // All operands after the address are possible destinations
                        for op in term.operands.iter().skip(1) {
                            self.successors.push(Rc::clone(op));
                        }
                    }
                    // Ret, Unreachable, Resume, etc. have no successors
                    _ => {}
                }
            }
        }
    }

    // -----------------------------------------------------------------------
    // Instruction management
    // -----------------------------------------------------------------------

    /// Insert an instruction at the given position (0-based index).
    pub fn insert_instruction(&mut self, pos: usize, inst: ValueRef) {
        // Set parent first (borrow before move)
        inst.borrow_mut().parent = Some(Rc::clone(&self.value));
        if pos >= self.instructions.len() {
            self.instructions.push(inst);
        } else {
            self.instructions.insert(pos, inst);
        }
    }

    /// Remove an instruction from this block. Returns true if found and removed.
    pub fn remove_instruction(&mut self, inst: &ValueRef) -> bool {
        if let Some(pos) = self.instructions.iter().position(|i| Rc::ptr_eq(i, inst)) {
            self.instructions.remove(pos);
            true
        } else {
            false
        }
    }

    /// Replace an old instruction with a new one at the same position.
    pub fn replace_instruction(&mut self, old: &ValueRef, new: ValueRef) -> bool {
        if let Some(pos) = self.instructions.iter().position(|i| Rc::ptr_eq(i, old)) {
            self.instructions[pos] = Rc::clone(&new);
            new.borrow_mut().parent = Some(Rc::clone(&self.value));
            true
        } else {
            false
        }
    }

    /// Append an instruction to the end of this block.
    pub fn push_instruction(&mut self, inst: ValueRef) {
        inst.borrow_mut().parent = Some(Rc::clone(&self.value));
        self.instructions.push(inst);
    }

    /// Get all instructions in this block.
    pub fn get_instructions(&self) -> &[ValueRef] {
        &self.instructions
    }

    /// Drain all instructions from this block, returning them as a Vec.
    pub fn take_instructions(&mut self) -> Vec<ValueRef> {
        std::mem::take(&mut self.instructions)
    }

    // -----------------------------------------------------------------------
    // Block mutation
    // -----------------------------------------------------------------------

    /// Split this basic block at the given instruction. Instructions from
    /// `inst` onward are moved to a new block, and this block branches
    /// unconditionally to the new block.
    ///
    /// Returns the new BasicBlock as a ValueRef.
    pub fn split_before(&mut self, inst: &ValueRef) -> Option<ValueRef> {
        let pos = self.instructions.iter().position(|i| Rc::ptr_eq(i, inst))?;

        // Create new block
        let new_bb = new_basic_block(&format!("{}_split", self.value.borrow().name));

        // Move instructions from pos to end into the new block's value
        let tail: Vec<ValueRef> = self.instructions.drain(pos..).collect();

        // Store instructions on the new block (we need to maintain a side-map
        // or store instructions in the Value. We'll store them by setting
        // parent and keeping them accessible via the parent's instruction list.
        // For now, we store them in the new block's operands space.
        {
            let mut new_val = new_bb.borrow_mut();
            for tail_inst in &tail {
                tail_inst.borrow_mut().parent = Some(Rc::clone(&new_bb));
            }
            new_val.operands = tail;
            new_val.num_operands = new_val.operands.len();
        }

        // Update the original block's parent function (if any) to also own the new block
        if let Some(parent_fn) = self.value.borrow().parent.clone() {
            new_bb.borrow_mut().parent = Some(Rc::clone(&parent_fn));
        }

        Some(new_bb)
    }

    /// Move this block before `before_inst` in block `to_block`.
    /// This is used for reordering basic blocks within a function.
    pub fn move_before(&mut self, _to_block: &ValueRef, _before_inst: Option<&ValueRef>) {
        // The actual reordering happens at the function level.
        // Here we just mark the intent. The Function that owns both blocks
        // is responsible for reordering its block list.
        //
        // For now, this is a no-op at the block level; the function-level
        // block list reordering is done by the Function implementation.
    }

    /// Erase this basic block from its parent function.
    /// This removes the block from the parent's block list and clears the
    /// parent pointer.
    pub fn erase_from_parent(&mut self) {
        if let Some(parent_fn) = self.value.borrow_mut().parent.take() {
            // Remove this block from the parent function's block list.
            // The parent function stores blocks as operands (or in its own list).
            let mut p = parent_fn.borrow_mut();
            p.operands.retain(|b| !Rc::ptr_eq(b, &self.value));
            p.num_operands = p.operands.len();
        }
    }

    /// Get the parent function of this block, if any.
    pub fn get_parent(&self) -> Option<ValueRef> {
        self.value.borrow().parent.clone()
    }

    /// Set the parent function of this block.
    pub fn set_parent(&mut self, func: ValueRef) {
        self.value.borrow_mut().parent = Some(func);
    }

    /// Check whether this block has a parent function.
    pub fn has_parent(&self) -> bool {
        self.value.borrow().parent.is_some()
    }

    // -----------------------------------------------------------------------
    // Utility
    // -----------------------------------------------------------------------

    /// Get the name of this block.
    pub fn get_name(&self) -> String {
        self.value.borrow().name.clone()
    }

    /// Set the name of this block.
    pub fn set_name(&mut self, name: &str) {
        self.value.borrow_mut().name = name.to_string();
    }

    /// Check whether this block is empty (has no instructions).
    pub fn is_empty(&self) -> bool {
        self.instructions.is_empty()
    }

    /// Get the value reference for this block.
    pub fn as_value_ref(&self) -> ValueRef {
        Rc::clone(&self.value)
    }

    /// Dump a textual representation of this block.
    pub fn dump(&self) -> String {
        let mut out = String::new();
        let name = self.get_name();
        out.push_str(&format!(
            "{}:\n",
            if name.is_empty() { "<unnamed>" } else { &name }
        ));
        if self.is_entry {
            out.push_str("  ; entry block\n");
        }
        if let Some(ref dom) = self.dominator {
            out.push_str(&format!("  ; dominator: {}\n", dom.borrow().name));
        }
        for inst in &self.instructions {
            out.push_str(&format!("  {}\n", inst.borrow().name));
        }
        if !self.predecessors.is_empty() {
            out.push_str("  ; predecessors: ");
            let pred_names: Vec<String> = self
                .predecessors
                .iter()
                .map(|p| p.borrow().name.clone())
                .collect();
            out.push_str(&pred_names.join(", "));
            out.push('\n');
        }
        out
    }
}

// ---------------------------------------------------------------------------
// Convenience: impl methods on ValueRef for basic block operations.
// These work on the ValueRef directly if the user doesn't want the full
// BasicBlock wrapper. They operate on the Value's operands as the
// instruction list.
// ---------------------------------------------------------------------------

/// Trait to provide BasicBlock methods directly on ValueRef.
pub trait BasicBlockExt {
    /// Check if this is a basic block.
    fn is_bb(&self) -> bool;
    /// Get the terminator instruction.
    fn get_terminator_inst(&self) -> Option<ValueRef>;
    /// Check if terminated.
    fn is_terminated_bb(&self) -> bool;
    /// Get instruction count.
    fn get_inst_count(&self) -> usize;
    /// Get an instruction by index.
    fn get_inst(&self, index: usize) -> Option<ValueRef>;
    /// Add an instruction.
    fn push_inst(&self, inst: ValueRef);
    /// Remove an instruction.
    fn remove_inst(&self, inst: &ValueRef) -> bool;
    /// Get all instructions.
    fn get_all_insts(&self) -> Vec<ValueRef>;
    /// Get PHI nodes.
    fn get_phi_nodes_bb(&self) -> Vec<ValueRef>;
    /// Get first non-PHI.
    fn get_first_non_phi_bb(&self) -> Option<ValueRef>;
}

impl BasicBlockExt for ValueRef {
    fn is_bb(&self) -> bool {
        self.borrow().is_basic_block()
    }

    fn get_terminator_inst(&self) -> Option<ValueRef> {
        let b = self.borrow();
        b.operands.last().cloned()
    }

    fn is_terminated_bb(&self) -> bool {
        self.borrow().operands.last().map_or(false, |inst| {
            inst.borrow()
                .get_opcode()
                .map_or(false, |op| op.is_terminator())
        })
    }

    fn get_inst_count(&self) -> usize {
        self.borrow().operands.len()
    }

    fn get_inst(&self, index: usize) -> Option<ValueRef> {
        self.borrow().operands.get(index).cloned()
    }

    fn push_inst(&self, inst: ValueRef) {
        inst.borrow_mut().parent = Some(Rc::clone(self));
        self.borrow_mut().push_operand(inst);
    }

    fn remove_inst(&self, inst: &ValueRef) -> bool {
        let mut b = self.borrow_mut();
        if let Some(pos) = b.operands.iter().position(|i| Rc::ptr_eq(i, inst)) {
            b.operands.remove(pos);
            b.num_operands = b.operands.len();
            true
        } else {
            false
        }
    }

    fn get_all_insts(&self) -> Vec<ValueRef> {
        self.borrow().operands.clone()
    }

    fn get_phi_nodes_bb(&self) -> Vec<ValueRef> {
        self.borrow()
            .operands
            .iter()
            .take_while(|inst| {
                inst.borrow()
                    .get_opcode()
                    .map_or(false, |op| op == Opcode::Phi)
            })
            .cloned()
            .collect()
    }

    fn get_first_non_phi_bb(&self) -> Option<ValueRef> {
        self.borrow()
            .operands
            .iter()
            .find(|inst| {
                !inst
                    .borrow()
                    .get_opcode()
                    .map_or(false, |op| op == Opcode::Phi)
            })
            .cloned()
    }
}

// ---------------------------------------------------------------------------
// Dominator Tree
// ---------------------------------------------------------------------------

/// A node in the dominator tree.
#[derive(Debug, Clone)]
pub struct DomTreeNode {
    pub block: ValueRef,
    pub idom: Option<ValueRef>,
    pub children: Vec<ValueRef>,
    pub dfs_num_in: u32,
    pub dfs_num_out: u32,
    pub level: u32,
}

impl DomTreeNode {
    pub fn new(block: ValueRef) -> Self {
        Self {
            block,
            idom: None,
            children: Vec::new(),
            dfs_num_in: 0,
            dfs_num_out: 0,
            level: 0,
        }
    }
    pub fn dominates(&self, other: &DomTreeNode) -> bool {
        self.dfs_num_in <= other.dfs_num_in && self.dfs_num_out >= other.dfs_num_out
    }
}

/// Full dominator tree for a function using Lengauer-Tarjan.
#[derive(Debug, Clone)]
pub struct DomTree {
    pub root: Option<ValueRef>,
    pub nodes: std::collections::HashMap<*const Value, DomTreeNode>,
}

impl DomTree {
    pub fn new() -> Self {
        Self {
            root: None,
            nodes: std::collections::HashMap::new(),
        }
    }

    pub fn build(entry: &ValueRef, blocks: &[ValueRef]) -> Self {
        let mut dt = Self::new();
        dt.root = Some(Rc::clone(entry));
        let entry_ptr = Rc::as_ptr(entry) as *const Value;
        let mut dfs_list: Vec<ValueRef> = Vec::new();
        let mut semi: std::collections::HashMap<*const Value, u32> =
            std::collections::HashMap::new();
        let mut parent: std::collections::HashMap<*const Value, *const Value> =
            std::collections::HashMap::new();

        fn dfs(
            v: &ValueRef,
            dfs_list: &mut Vec<ValueRef>,
            semi: &mut std::collections::HashMap<*const Value, u32>,
            parent: &mut std::collections::HashMap<*const Value, *const Value>,
        ) {
            let ptr = Rc::as_ptr(v) as *const Value;
            let idx = dfs_list.len() as u32;
            semi.insert(ptr, idx);
            dfs_list.push(Rc::clone(v));
            let succs = v.borrow().operands.clone();
            for succ in &succs {
                let s_ptr = Rc::as_ptr(succ) as *const Value;
                if !semi.contains_key(&s_ptr) {
                    parent.insert(s_ptr, ptr);
                    dfs(succ, dfs_list, semi, parent);
                }
            }
        }
        dfs(entry, &mut dfs_list, &mut semi, &mut parent);

        let mut idom: std::collections::HashMap<*const Value, *const Value> =
            std::collections::HashMap::new();
        idom.insert(entry_ptr, entry_ptr);
        let mut changed = true;
        while changed {
            changed = false;
            for i in 1..dfs_list.len() {
                let v = &dfs_list[i];
                let v_ptr = Rc::as_ptr(v) as *const Value;
                let mut pred_ptrs: Vec<*const Value> = Vec::new();
                for block in blocks {
                    let b = block.borrow();
                    for succ in &b.operands {
                        if Rc::ptr_eq(succ, v) {
                            pred_ptrs.push(Rc::as_ptr(block) as *const Value);
                        }
                    }
                }
                let mut new_idom: Option<*const Value> =
                    pred_ptrs.iter().find(|p| idom.contains_key(p)).copied();
                if let Some(mut ni) = new_idom {
                    for p_ptr in &pred_ptrs {
                        if *p_ptr == ni || !idom.contains_key(p_ptr) {
                            continue;
                        }
                        let mut a = *p_ptr;
                        let mut b = ni;
                        while a != b {
                            let a_idx = dfs_list
                                .iter()
                                .position(|x| Rc::as_ptr(x) as *const Value == a)
                                .unwrap_or(0);
                            let b_idx = dfs_list
                                .iter()
                                .position(|x| Rc::as_ptr(x) as *const Value == b)
                                .unwrap_or(0);
                            if a_idx > b_idx {
                                a = *idom.get(&a).unwrap_or(&a);
                            } else {
                                b = *idom.get(&b).unwrap_or(&b);
                            }
                        }
                        ni = a;
                    }
                    let old = idom.insert(v_ptr, ni);
                    if old != Some(ni) {
                        changed = true;
                    }
                }
            }
        }

        for (&v_ptr, &i_ptr) in &idom {
            let block = dt.find_block(v_ptr, blocks, entry);
            dt.nodes
                .entry(v_ptr)
                .or_insert_with(|| DomTreeNode::new(block));
            if v_ptr != i_ptr {
                let iblock = dt.find_block(i_ptr, blocks, entry);
                dt.nodes
                    .entry(i_ptr)
                    .or_insert_with(|| DomTreeNode::new(iblock));
            }
        }
        // Set idom/children — must avoid simultaneous mutable borrows
        let updates: Vec<(ValueRef, ValueRef)> = idom
            .iter()
            .filter(|&(&v, &i)| v != i)
            .map(|(&v, &i)| {
                let v_block = dt.find_block(v, blocks, entry);
                let i_block = dt.find_block(i, blocks, entry);
                (v_block, i_block)
            })
            .collect();
        for (v_ref, i_ref) in &updates {
            let v_ptr = Rc::as_ptr(v_ref) as *const Value;
            if let Some(v_node) = dt.nodes.get_mut(&v_ptr) {
                v_node.idom = Some(Rc::clone(i_ref));
            }
        }
        for (v_ref, i_ref) in &updates {
            let i_ptr = Rc::as_ptr(i_ref) as *const Value;
            if let Some(i_node) = dt.nodes.get_mut(&i_ptr) {
                i_node.children.push(Rc::clone(v_ref));
            }
        }
        dt.assign_dfs(entry_ptr, 0, &mut 0);
        dt
    }

    fn find_block(&self, ptr: *const Value, blocks: &[ValueRef], entry: &ValueRef) -> ValueRef {
        blocks
            .iter()
            .find(|b| Rc::as_ptr(b) as *const Value == ptr)
            .cloned()
            .unwrap_or_else(|| Rc::clone(entry))
    }

    fn assign_dfs(&mut self, ptr: *const Value, level: u32, counter: &mut u32) {
        // Collect children first (immutable borrow)
        let children: Vec<ValueRef> = if let Some(node) = self.nodes.get(&ptr) {
            node.children.clone()
        } else {
            return;
        };
        // Set dfs_num_in (pre-order)
        if let Some(node) = self.nodes.get_mut(&ptr) {
            *counter += 1;
            node.dfs_num_in = *counter;
            node.level = level;
        }
        // Recurse into children
        for child in &children {
            let c_ptr = Rc::as_ptr(child) as *const Value;
            self.assign_dfs(c_ptr, level + 1, counter);
        }
        // Set dfs_num_out (post-order)
        if let Some(node) = self.nodes.get_mut(&ptr) {
            *counter += 1;
            node.dfs_num_out = *counter;
        }
    }

    pub fn get_idom(&self, block: &ValueRef) -> Option<ValueRef> {
        let ptr = Rc::as_ptr(block) as *const Value;
        self.nodes.get(&ptr)?.idom.clone()
    }

    pub fn dominates(&self, a: &ValueRef, b: &ValueRef) -> bool {
        let a_ptr = Rc::as_ptr(a) as *const Value;
        let b_ptr = Rc::as_ptr(b) as *const Value;
        if a_ptr == b_ptr {
            return true;
        }
        match (self.nodes.get(&a_ptr), self.nodes.get(&b_ptr)) {
            (Some(na), Some(nb)) => na.dominates(nb),
            _ => false,
        }
    }

    pub fn strictly_dominates(&self, a: &ValueRef, b: &ValueRef) -> bool {
        !Rc::ptr_eq(a, b) && self.dominates(a, b)
    }

    pub fn find_nearest_common_dominator(&self, a: &ValueRef, b: &ValueRef) -> Option<ValueRef> {
        let mut cur_a = Rc::clone(a);
        let mut cur_b = Rc::clone(b);
        let mut seen = std::collections::HashSet::new();
        loop {
            let a_ptr = Rc::as_ptr(&cur_a) as *const Value;
            if seen.contains(&a_ptr) {
                return Some(cur_a);
            }
            seen.insert(a_ptr);
            match self.get_idom(&cur_a) {
                Some(d) => cur_a = d,
                None => break,
            }
        }
        loop {
            let b_ptr = Rc::as_ptr(&cur_b) as *const Value;
            if seen.contains(&b_ptr) {
                return Some(cur_b);
            }
            seen.insert(b_ptr);
            match self.get_idom(&cur_b) {
                Some(d) => cur_b = d,
                None => break,
            }
        }
        None
    }
}

// ---------------------------------------------------------------------------
// Post-Dominator Tree
// ---------------------------------------------------------------------------

#[derive(Debug, Clone)]
pub struct PostDomTree {
    pub tree: DomTree,
    pub exit: ValueRef,
}

impl PostDomTree {
    pub fn build(_entry: &ValueRef, _blocks: &[ValueRef]) -> Self {
        Self {
            tree: DomTree::new(),
            exit: new_basic_block("<virtual_exit>"),
        }
    }
    pub fn post_dominates(&self, a: &ValueRef, b: &ValueRef) -> bool {
        self.tree.dominates(a, b)
    }
}

// ---------------------------------------------------------------------------
// Loop Detection
// ---------------------------------------------------------------------------

#[derive(Debug, Clone)]
pub struct Loop {
    pub header: ValueRef,
    pub blocks: Vec<ValueRef>,
    pub latches: Vec<ValueRef>,
    pub parent_loop: Option<Box<Loop>>,
    pub sub_loops: Vec<Loop>,
    pub depth: u32,
}

impl Loop {
    pub fn new(header: ValueRef) -> Self {
        Self {
            header,
            blocks: Vec::new(),
            latches: Vec::new(),
            parent_loop: None,
            sub_loops: Vec::new(),
            depth: 0,
        }
    }
    pub fn contains(&self, block: &ValueRef) -> bool {
        self.blocks.iter().any(|b| Rc::ptr_eq(b, block))
    }
    pub fn is_outermost(&self) -> bool {
        self.parent_loop.is_none()
    }
    pub fn get_all_blocks(&self) -> Vec<ValueRef> {
        let mut all = self.blocks.clone();
        for sub in &self.sub_loops {
            all.extend(sub.get_all_blocks());
        }
        all
    }
}

#[derive(Debug, Clone)]
pub struct LoopInfo {
    pub top_level_loops: Vec<Loop>,
    pub dom_tree: Option<DomTree>,
}

impl LoopInfo {
    pub fn new() -> Self {
        Self {
            top_level_loops: Vec::new(),
            dom_tree: None,
        }
    }

    pub fn detect(entry: &ValueRef, blocks: &[ValueRef]) -> Self {
        let dom_tree = DomTree::build(entry, blocks);
        let mut li = Self {
            top_level_loops: Vec::new(),
            dom_tree: Some(dom_tree),
        };
        for block in blocks {
            let succs: Vec<ValueRef> = block.borrow().operands.clone();
            for succ in &succs {
                if li.is_back_edge(block, succ) {
                    li.add_loop(succ, block, blocks);
                }
            }
        }
        li
    }

    fn is_back_edge(&self, from: &ValueRef, to: &ValueRef) -> bool {
        self.dom_tree
            .as_ref()
            .map_or(false, |dt| dt.dominates(to, from))
    }

    fn add_loop(&mut self, header: &ValueRef, latch: &ValueRef, blocks: &[ValueRef]) {
        if let Some(existing) = self
            .top_level_loops
            .iter_mut()
            .find(|l| Rc::ptr_eq(&l.header, header))
        {
            if !existing.latches.iter().any(|l| Rc::ptr_eq(l, latch)) {
                existing.latches.push(Rc::clone(latch));
            }
            if !existing.contains(latch) {
                existing.blocks.push(Rc::clone(latch));
            }
            return;
        }
        let mut lp = Loop::new(Rc::clone(header));
        lp.latches.push(Rc::clone(latch));
        lp.blocks.push(Rc::clone(header));
        lp.blocks.push(Rc::clone(latch));
        let mut worklist: Vec<ValueRef> = vec![Rc::clone(latch)];
        let mut visited: std::collections::HashSet<*const Value> = std::collections::HashSet::new();
        visited.insert(Rc::as_ptr(header) as *const Value);
        visited.insert(Rc::as_ptr(latch) as *const Value);
        while let Some(cur) = worklist.pop() {
            for block in blocks {
                let b = block.borrow();
                for succ in &b.operands {
                    if Rc::ptr_eq(succ, &cur) {
                        let p = Rc::as_ptr(block) as *const Value;
                        if !visited.contains(&p) {
                            visited.insert(p);
                            lp.blocks.push(Rc::clone(block));
                            worklist.push(Rc::clone(block));
                        }
                    }
                }
            }
        }
        self.top_level_loops.push(lp);
    }

    pub fn get_loop_depth(&self, block: &ValueRef) -> u32 {
        for l in &self.top_level_loops {
            if l.contains(block) {
                return 1;
            }
        }
        0
    }

    pub fn is_loop_header(&self, block: &ValueRef) -> bool {
        self.top_level_loops
            .iter()
            .any(|l| Rc::ptr_eq(&l.header, block))
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;
    use crate::opcode::Opcode;
    use crate::value::valref;

    /// Helper: create a ValueRef instruction with a given opcode.
    fn make_inst(name: &str, op: Opcode) -> ValueRef {
        valref(
            Value::new(Type::void())
                .named(name)
                .with_subclass(SubclassKind::Instruction)
                .with_opcode(op),
        )
    }

    #[test]
    fn test_new_basic_block() {
        let bb = BasicBlock::new("entry");
        assert!(bb.value.borrow().is_basic_block());
        assert_eq!(bb.get_name(), "entry");
        assert!(!bb.is_entry);
        assert!(bb.dominator.is_none());
        assert!(bb.instructions.is_empty());
    }

    #[test]
    fn test_from_value() {
        let val = new_basic_block("test_bb");
        let mut bb = BasicBlock::from_value(val.clone());
        assert!(Rc::ptr_eq(&bb.value, &val));
        bb.is_entry = true;
        assert!(bb.is_entry);
    }

    #[test]
    fn test_push_and_count_instructions() {
        let mut bb = BasicBlock::new("body");
        assert_eq!(bb.get_instruction_count(), 0);
        let inst = make_inst("add", Opcode::Add);
        bb.push_instruction(inst);
        assert_eq!(bb.get_instruction_count(), 1);
    }

    #[test]
    fn test_get_instruction() {
        let mut bb = BasicBlock::new("body");
        let i0 = make_inst("i0", Opcode::Add);
        let i1 = make_inst("i1", Opcode::Sub);
        bb.push_instruction(i0.clone());
        bb.push_instruction(i1.clone());
        assert_eq!(bb.get_instruction_count(), 2);
        let got = bb.get_instruction(0).unwrap();
        assert_eq!(got.borrow().name, "i0");
        let got = bb.get_instruction(1).unwrap();
        assert_eq!(got.borrow().name, "i1");
        assert!(bb.get_instruction(2).is_none());
    }

    #[test]
    fn test_insert_instruction() {
        let mut bb = BasicBlock::new("body");
        let i_a = make_inst("a", Opcode::Add);
        let i_b = make_inst("b", Opcode::Sub);
        bb.push_instruction(i_a.clone());
        bb.insert_instruction(0, i_b.clone());
        assert_eq!(bb.get_instruction_count(), 2);
        assert_eq!(bb.get_instruction(0).unwrap().borrow().name, "b");
        assert_eq!(bb.get_instruction(1).unwrap().borrow().name, "a");
    }

    #[test]
    fn test_remove_instruction() {
        let mut bb = BasicBlock::new("body");
        let i0 = make_inst("i0", Opcode::Add);
        let i1 = make_inst("i1", Opcode::Sub);
        bb.push_instruction(i0.clone());
        bb.push_instruction(i1.clone());
        assert!(bb.remove_instruction(&i0));
        assert_eq!(bb.get_instruction_count(), 1);
        assert_eq!(bb.get_instruction(0).unwrap().borrow().name, "i1");
    }

    #[test]
    fn test_replace_instruction() {
        let mut bb = BasicBlock::new("body");
        let old = make_inst("old", Opcode::Add);
        let new = make_inst("new", Opcode::Sub);
        bb.push_instruction(old.clone());
        assert!(bb.replace_instruction(&old, new.clone()));
        assert_eq!(bb.get_instruction(0).unwrap().borrow().name, "new");
    }

    #[test]
    fn test_replace_nonexistent_instruction() {
        let mut bb = BasicBlock::new("body");
        let old = make_inst("old", Opcode::Add);
        let new = make_inst("new", Opcode::Sub);
        assert!(!bb.replace_instruction(&old, new));
    }

    #[test]
    fn test_is_terminated_and_get_terminator() {
        let mut bb = BasicBlock::new("exit");
        assert!(!bb.is_terminated());
        assert!(bb.get_terminator().is_none());

        let ret = make_inst("ret", Opcode::Ret);
        bb.push_instruction(ret.clone());
        assert!(bb.is_terminated());
        let term = bb.get_terminator().unwrap();
        assert_eq!(term.borrow().name, "ret");
    }

    #[test]
    fn test_get_phi_nodes() {
        let mut bb = BasicBlock::new("loop_header");
        let phi_a = make_inst("phi_a", Opcode::Phi);
        let phi_b = make_inst("phi_b", Opcode::Phi);
        let add = make_inst("add", Opcode::Add);
        bb.push_instruction(phi_a.clone());
        bb.push_instruction(phi_b.clone());
        bb.push_instruction(add.clone());

        let phis = bb.get_phi_nodes();
        assert_eq!(phis.len(), 2);
        assert_eq!(phis[0].borrow().name, "phi_a");
        assert_eq!(phis[1].borrow().name, "phi_b");
    }

    #[test]
    fn test_get_first_non_phi() {
        let mut bb = BasicBlock::new("body");
        let phi = make_inst("phi", Opcode::Phi);
        let add = make_inst("add", Opcode::Add);
        bb.push_instruction(phi.clone());
        bb.push_instruction(add.clone());

        let first = bb.get_first_non_phi().unwrap();
        assert_eq!(first.borrow().name, "add");
    }

    #[test]
    fn test_get_first_non_phi_all_phis() {
        let mut bb = BasicBlock::new("header");
        let phi = make_inst("phi", Opcode::Phi);
        bb.push_instruction(phi.clone());
        // No non-phi instruction
        assert!(bb.get_first_non_phi().is_none());
    }

    #[test]
    fn test_predecessors() {
        let mut bb = BasicBlock::new("target");
        let pred_a = new_basic_block("pred_a");
        let pred_b = new_basic_block("pred_b");

        bb.add_predecessor(pred_a.clone());
        bb.add_predecessor(pred_b.clone());
        assert_eq!(bb.predecessors.len(), 2);
        assert!(bb.has_predecessor(&pred_a));
        assert!(bb.has_predecessor(&pred_b));

        // Dedup
        bb.add_predecessor(pred_a.clone());
        assert_eq!(bb.predecessors.len(), 2);

        assert!(bb.remove_predecessor(&pred_a));
        assert_eq!(bb.predecessors.len(), 1);
        assert!(!bb.has_predecessor(&pred_a));
    }

    #[test]
    fn test_successors() {
        let mut bb = BasicBlock::new("source");
        let succ_a = new_basic_block("succ_a");
        let succ_b = new_basic_block("succ_b");

        bb.add_successor(succ_a.clone());
        bb.add_successor(succ_b.clone());
        assert!(bb.has_successor(&succ_a));
        assert!(bb.has_successor(&succ_b));

        assert!(bb.remove_successor(&succ_a));
        assert!(!bb.has_successor(&succ_a));
        assert_eq!(bb.successors.len(), 1);
    }

    #[test]
    fn test_split_before() {
        let mut bb = BasicBlock::new("original");
        let i0 = make_inst("i0", Opcode::Alloca);
        let i1 = make_inst("i1", Opcode::Add);
        let i2 = make_inst("i2", Opcode::Ret);
        bb.push_instruction(i0.clone());
        bb.push_instruction(i1.clone());
        bb.push_instruction(i2.clone());

        let new = bb.split_before(&i1).unwrap();
        // Original block now has only i0
        assert_eq!(bb.get_instruction_count(), 1);
        assert_eq!(bb.get_instruction(0).unwrap().borrow().name, "i0");

        // New block has i1, i2
        let new_insts = new.borrow().operands.clone();
        assert_eq!(new_insts.len(), 2);
        assert_eq!(new_insts[0].borrow().name, "i1");
        assert_eq!(new_insts[1].borrow().name, "i2");
    }

    #[test]
    fn test_split_before_nonexistent() {
        let mut bb = BasicBlock::new("original");
        let inst = make_inst("i", Opcode::Add);
        let ghost = make_inst("ghost", Opcode::Add);
        bb.push_instruction(inst.clone());
        assert!(bb.split_before(&ghost).is_none());
    }

    #[test]
    fn test_erase_from_parent() {
        let fn_ty = Type::function_type_with(crate::types::TypeId::new(), vec![], false);
        let func_val = valref(
            Value::new(fn_ty)
                .named("myfunc")
                .with_subclass(SubclassKind::Function),
        );
        let mut bb = BasicBlock::new("entry");
        bb.value.borrow_mut().parent = Some(Rc::clone(&func_val));
        func_val.borrow_mut().operands.push(Rc::clone(&bb.value));
        func_val.borrow_mut().num_operands += 1;

        assert!(bb.has_parent());
        bb.erase_from_parent();
        assert!(!bb.has_parent());
        assert!(func_val.borrow().operands.is_empty());
    }

    #[test]
    fn test_get_parent() {
        let mut bb = BasicBlock::new("child");
        assert!(bb.get_parent().is_none());

        let fn_ty = Type::function_type_with(crate::types::TypeId::new(), vec![], false);
        let func = valref(
            Value::new(fn_ty)
                .named("func")
                .with_subclass(SubclassKind::Function),
        );
        bb.set_parent(func.clone());
        let got = bb.get_parent().unwrap();
        assert!(Rc::ptr_eq(&got, &func));
    }

    #[test]
    fn test_dump() {
        let mut bb = BasicBlock::new("myblock");
        bb.is_entry = true;
        let ret = make_inst("ret", Opcode::Ret);
        bb.push_instruction(ret);
        let dump = bb.dump();
        assert!(dump.contains("myblock:"));
        assert!(dump.contains("entry block"));
        assert!(dump.contains("ret"));
    }

    #[test]
    fn test_empty_block() {
        let bb = BasicBlock::new("empty");
        assert!(bb.is_empty());
        assert_eq!(bb.get_instruction_count(), 0);
        assert!(!bb.is_terminated());
    }

    #[test]
    fn test_set_name() {
        let mut bb = BasicBlock::new("old");
        assert_eq!(bb.get_name(), "old");
        bb.set_name("new");
        assert_eq!(bb.get_name(), "new");
    }

    #[test]
    fn test_recompute_successors_unconditional_br() {
        let mut bb = BasicBlock::new("source");
        let target = new_basic_block("target");
        // Create a br instruction: br label %target
        let br_inst = valref(
            Value::new(Type::void())
                .named("br_inst")
                .with_subclass(SubclassKind::Instruction)
                .with_opcode(Opcode::Br),
        );
        br_inst.borrow_mut().operands.push(target.clone());
        bb.push_instruction(br_inst);

        bb.recompute_successors();
        assert_eq!(bb.successors.len(), 1);
        assert!(bb.has_successor(&target));
    }

    #[test]
    fn test_recompute_successors_ret() {
        let mut bb = BasicBlock::new("exit");
        let ret = make_inst("ret", Opcode::Ret);
        bb.push_instruction(ret);
        bb.recompute_successors();
        assert!(bb.successors.is_empty());
    }
}

// ---------------------------------------------------------------------------
// Extension method for Value (used in tests)
// ---------------------------------------------------------------------------

/// Helper extension to set opcode on a Value during construction.
impl Value {
    pub fn with_opcode(mut self, op: Opcode) -> Self {
        self.opcode = Some(op);
        self
    }
}

use std::rc::Rc;