lamina 0.0.10

//! Loop optimization transforms for MIR.

use super::{Transform, TransformCategory, TransformLevel};
use crate::mir::{Block, Function, Instruction, IntBinOp, IntCmpOp, Operand, Register};
use std::collections::{HashMap, HashSet};

/// Loop invariant code motion that moves loop-invariant computations outside loops.
#[derive(Default)]
pub struct LoopInvariantCodeMotion;

impl Transform for LoopInvariantCodeMotion {
    fn name(&self) -> &'static str {
        "loop_invariant_code_motion"
    }

    fn description(&self) -> &'static str {
        "Moves loop-invariant computations outside loops"
    }

    fn category(&self) -> TransformCategory {
        TransformCategory::ControlFlowOptimization
    }

    fn level(&self) -> TransformLevel {
        TransformLevel::Stable
    }

    fn apply(&self, func: &mut Function) -> Result<bool, String> {
        self.apply_internal(func)
    }
}

#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
mod tests_licm {
    use super::*;
    use crate::mir::{
        FunctionBuilder, Immediate, IntBinOp, MirType, Operand, ScalarType, VirtualReg,
    };

    #[test]
    fn test_licm_basic() {
        // Create a loop with invariant calculation
        // entry:
        //   jmp loop
        // loop:
        //   v0 = add v0, 1  (induction)
        //   v1 = add 5, 10  (invariant)
        //   v2 = lt v0, v1
        //   br v2, loop, exit
        // exit:
        //   ret

        let mut func = FunctionBuilder::new("test_licm")
            .returns(MirType::Scalar(ScalarType::I64))
            .block("entry")
            .instr(Instruction::IntBinary {
                op: IntBinOp::Add,
                ty: MirType::Scalar(ScalarType::I64),
                dst: VirtualReg::gpr(0).into(),
                lhs: Operand::Immediate(Immediate::I64(0)),
                rhs: Operand::Immediate(Immediate::I64(0)),
            }) // init v0
            .instr(Instruction::Jmp {
                target: "loop".to_string(),
            })
            .block("loop")
            .instr(Instruction::IntBinary {
                op: IntBinOp::Add,
                ty: MirType::Scalar(ScalarType::I64),
                dst: VirtualReg::gpr(0).into(),
                lhs: Operand::Register(VirtualReg::gpr(0).into()),
                rhs: Operand::Immediate(Immediate::I64(1)),
            })
            // Invariant: 5 + 10
            .instr(Instruction::IntBinary {
                op: IntBinOp::Add,
                ty: MirType::Scalar(ScalarType::I64),
                dst: VirtualReg::gpr(1).into(),
                lhs: Operand::Immediate(Immediate::I64(5)),
                rhs: Operand::Immediate(Immediate::I64(10)),
            })
            .instr(Instruction::IntCmp {
                op: IntCmpOp::SLt,
                ty: MirType::Scalar(ScalarType::I1),
                dst: VirtualReg::gpr(2).into(),
                lhs: Operand::Register(VirtualReg::gpr(0).into()),
                rhs: Operand::Register(VirtualReg::gpr(1).into()),
            })
            .instr(Instruction::Br {
                cond: VirtualReg::gpr(2).into(),
                true_target: "loop".to_string(),
                false_target: "exit".to_string(),
            })
            .block("exit")
            .instr(Instruction::Ret { value: None })
            .build();

        let licm = LoopInvariantCodeMotion;
        let changed = licm.apply(&mut func).expect("LICM failed");

        assert!(changed);

        // Loop block should no longer have the 5+10 addition
        let loop_block = func.get_block("loop").expect("loop block exists");
        let has_invariant = loop_block.instructions.iter().any(|i| {
            if let Instruction::IntBinary { lhs, rhs, .. } = i {
                matches!(lhs, Operand::Immediate(Immediate::I64(5)))
                    && matches!(rhs, Operand::Immediate(Immediate::I64(10)))
            } else {
                false
            }
        });
        assert!(
            !has_invariant,
            "Invariant instruction should be moved out of loop"
        );

        let pre_header = func.get_block("loop_pre").expect("pre-header created");
        let has_moved = pre_header.instructions.iter().any(|i| {
            if let Instruction::IntBinary { lhs, rhs, .. } = i {
                matches!(lhs, Operand::Immediate(Immediate::I64(5)))
                    && matches!(rhs, Operand::Immediate(Immediate::I64(10)))
            } else {
                false
            }
        });
        assert!(has_moved, "Invariant instruction should be in pre-header");
    }
}

impl LoopInvariantCodeMotion {
    fn apply_internal(&self, func: &mut Function) -> Result<bool, String> {
        let mut changed = false;

        let loops = self.find_loops(func);

        let max_loops = 10;
        let max_iterations_per_loop = 5;
        let loops_to_process = loops.into_iter().take(max_loops);

        for loop_info in loops_to_process {
            if loop_info.blocks.len() > 50 {
                continue;
            }

            let mut iterations = 0;
            while iterations < max_iterations_per_loop {
                if !self.optimize_loop(func, &loop_info)? {
                    break;
                }
                changed = true;
                iterations += 1;
                if iterations >= max_iterations_per_loop {
                    break;
                }
            }
        }

        Ok(changed)
    }

    fn find_loops(&self, func: &Function) -> Vec<LoopInfo> {
        let mut loops = Vec::new();

        let dominators = self.calculate_dominators(func);

        for block in &func.blocks {
            for instr in &block.instructions {
                if let Instruction::Br {
                    true_target,
                    false_target,
                    ..
                } = instr
                {
                    // Check for back edges: target dominates source
                    if self.is_back_edge(&dominators, &block.label, true_target)
                        && let Some(loop_info) =
                            self.analyze_loop(func, true_target, &block.label, &dominators)
                    {
                        loops.push(loop_info);
                    }
                    if self.is_back_edge(&dominators, &block.label, false_target)
                        && let Some(loop_info) =
                            self.analyze_loop(func, false_target, &block.label, &dominators)
                    {
                        loops.push(loop_info);
                    }
                }
                if let Instruction::Jmp { target } = instr
                    && self.is_back_edge(&dominators, &block.label, target)
                    && let Some(loop_info) =
                        self.analyze_loop(func, target, &block.label, &dominators)
                {
                    loops.push(loop_info);
                }
            }
        }

        // Merge loops with the same header to handle multiple backedges
        let mut merged_loops: HashMap<String, HashSet<String>> = HashMap::new();

        for loop_info in loops {
            match merged_loops.entry(loop_info.header) {
                std::collections::hash_map::Entry::Vacant(e) => {
                    e.insert(loop_info.blocks);
                }
                std::collections::hash_map::Entry::Occupied(mut e) => {
                    e.get_mut().extend(loop_info.blocks);
                }
            }
        }

        merged_loops
            .into_iter()
            .map(|(header, blocks)| LoopInfo { header, blocks })
            .collect()
    }

    fn is_back_edge(
        &self,
        dominators: &HashMap<String, HashSet<String>>,
        from: &str,
        to: &str,
    ) -> bool {
        // A back edge exists if the target (header) dominates the source
        if let Some(doms) = dominators.get(from) {
            doms.contains(to)
        } else {
            false
        }
    }

    fn calculate_dominators(&self, func: &Function) -> HashMap<String, HashSet<String>> {
        let mut dominators: HashMap<String, HashSet<String>> = HashMap::new();
        let all_blocks: HashSet<String> = func.blocks.iter().map(|b| b.label.clone()).collect();

        // Initialize: value for entry is {entry}, others are all blocks
        for block in &func.blocks {
            if block.label == func.entry {
                let mut set = HashSet::new();
                set.insert(block.label.clone());
                dominators.insert(block.label.clone(), set);
            } else {
                dominators.insert(block.label.clone(), all_blocks.clone());
            }
        }

        let mut changed = true;
        while changed {
            changed = false;

            for block in &func.blocks {
                if block.label == func.entry {
                    continue;
                }

                // Intersection of dominators of all predecessors
                let mut new_doms: Option<HashSet<String>> = None;

                // Find predecessors
                let preds: Vec<&Block> = func
                    .blocks
                    .iter()
                    .filter(|pred| self.has_edge_to(func, &pred.label, &block.label))
                    .collect();

                if preds.is_empty() {
                    continue;
                }

                for pred in preds {
                    if let Some(pred_doms) = dominators.get(&pred.label) {
                        if let Some(current_intersect) = &mut new_doms {
                            current_intersect.retain(|d| pred_doms.contains(d));
                        } else {
                            new_doms = Some(pred_doms.clone());
                        }
                    }
                }

                let mut final_doms = new_doms.unwrap_or_default();
                final_doms.insert(block.label.clone());

                if let Some(current_doms) = dominators.get(&block.label)
                    && final_doms != *current_doms
                {
                    dominators.insert(block.label.clone(), final_doms);
                    changed = true;
                }
            }
        }

        dominators
    }

    fn analyze_loop(
        &self,
        func: &Function,
        header: &str,
        back_edge_source: &str,
        _dominators: &HashMap<String, HashSet<String>>,
    ) -> Option<LoopInfo> {
        // Natural loop identification:
        // A natural loop consists of the header, the back edge source,
        // and all nodes that can reach the back edge without going through the header.
        // (Optimized: we can traverse backwards from back_edge_source stopping at header)
        let mut loop_blocks = HashSet::new();
        let mut to_visit = vec![back_edge_source.to_string()];
        let mut visited = HashSet::new();
        const MAX_LOOP_ANALYSIS_ITERATIONS: usize = 1000;
        let mut iterations = 0;

        // First, collect all blocks that can reach the back edge source
        // Initialize queue with back edge source
        if header != back_edge_source {
            to_visit.push(back_edge_source.to_string());
        }
        loop_blocks.insert(header.to_string());
        loop_blocks.insert(back_edge_source.to_string());

        while let Some(block_label) = to_visit.pop() {
            if iterations >= MAX_LOOP_ANALYSIS_ITERATIONS {
                return None;
            }
            iterations += 1;

            if visited.contains(&block_label) {
                continue;
            }
            visited.insert(block_label.clone());
            loop_blocks.insert(block_label.clone());

            // Add predecessors that are not the header
            for block in &func.blocks {
                if self.has_edge_to(func, &block.label, &block_label) && block.label != header {
                    to_visit.push(block.label.clone());
                }
            }
        }

        // Add the header if not already included
        loop_blocks.insert(header.to_string());

        if !loop_blocks.contains(header) || !loop_blocks.contains(back_edge_source) {
            return None;
        }

        if loop_blocks.len() > 50 {
            return None;
        }

        Some(LoopInfo {
            header: header.to_string(),
            blocks: loop_blocks,
        })
    }

    fn has_edge_to(&self, func: &Function, from: &str, to: &str) -> bool {
        if let Some(block) = func.blocks.iter().find(|b| b.label == *from) {
            for instr in &block.instructions {
                match instr {
                    Instruction::Jmp { target } if target == to => return true,
                    Instruction::Br {
                        true_target,
                        false_target,
                        ..
                    } if true_target == to || false_target == to => return true,
                    _ => {}
                }
            }
        }
        false
    }

    fn optimize_loop(&self, func: &mut Function, loop_info: &LoopInfo) -> Result<bool, String> {
        let mut changed = false;

        let invariant_instrs = self.find_invariant_instructions(func, loop_info)?;

        if !invariant_instrs.is_empty() {
            self.move_invariant_instructions(func, loop_info, &invariant_instrs)?;
            changed = true;
        }

        Ok(changed)
    }

    fn find_invariant_instructions(
        &self,
        func: &Function,
        loop_info: &LoopInfo,
    ) -> Result<Vec<(usize, usize)>, String> {
        let mut invariant = Vec::new();
        let defs_in_loop = self.collect_defs_in_loop(func, loop_info);

        // Count how many times each register is defined inside the loop.
        // If a register is defined more than once (e.g. accumulator %sum, induction
        // var %j), it cannot be safely hoisted — one of those assignments is a
        // loop-carried update that must execute every iteration.
        let mut def_counts: HashMap<Register, usize> =
            HashMap::new();
        for block in &func.blocks {
            if !loop_info.blocks.contains(&block.label) {
                continue;
            }
            for instr in &block.instructions {
                if let Some(def) = instr.def_reg() {
                    *def_counts.entry(def.clone()).or_insert(0) += 1;
                }
            }
        }

        let max_invariant_instructions = 20;

        for (block_idx, block) in func.blocks.iter().enumerate() {
            if !loop_info.blocks.contains(&block.label) {
                continue;
            }

            for (instr_idx, instr) in block.instructions.iter().enumerate() {
                // Skip if this register is defined more than once in the loop.
                if let Some(def) = instr.def_reg()
                    && def_counts.get(def).copied().unwrap_or(0) > 1
                {
                    continue;
                }
                if self.is_invariant_instruction(func, loop_info, &defs_in_loop, instr) {
                    invariant.push((block_idx, instr_idx));

                    if invariant.len() >= max_invariant_instructions {
                        break;
                    }
                }
            }

            if invariant.len() >= max_invariant_instructions {
                break;
            }
        }

        Ok(invariant)
    }

    fn is_invariant_instruction(
        &self,
        func: &Function,
        loop_info: &LoopInfo,
        defs_in_loop: &HashSet<Register>,
        instr: &Instruction,
    ) -> bool {
        // An instruction is invariant if:
        // 1. It defines a register
        // 2. All its operands are either constants or defined outside the loop
        // 3. It's not a side-effecting instruction

        if let Some(_def_reg) = instr.def_reg() {
            // Check if all operands are invariant
            let operands_invariant = instr
                .use_regs()
                .iter()
                .all(|reg| self.is_invariant_register(func, loop_info, defs_in_loop, reg));

            // Check if instruction has no side effects
            let no_side_effects = !self.has_side_effects(instr);

            operands_invariant && no_side_effects
        } else {
            false
        }
    }

    fn is_invariant_register(
        &self,
        _func: &Function,
        _loop_info: &LoopInfo,
        defs_in_loop: &HashSet<Register>,
        reg: &Register,
    ) -> bool {
        match reg {
            Register::Physical(_) => false,
            Register::Virtual(_) => !defs_in_loop.contains(reg),
        }
    }

    fn has_side_effects(&self, instr: &Instruction) -> bool {
        matches!(
            instr,
            Instruction::Load { .. }
                | Instruction::Store { .. }
                | Instruction::Call { .. }
                | Instruction::Ret { .. }
        )
    }

    fn move_invariant_instructions(
        &self,
        func: &mut Function,
        loop_info: &LoopInfo,
        invariant_instrs: &[(usize, usize)],
    ) -> Result<(), String> {
        if invariant_instrs.is_empty() {
            return Ok(());
        }

        let header_idx = func
            .blocks
            .iter()
            .position(|b| b.label == loop_info.header)
            .ok_or_else(|| format!("Loop header '{}' not found", loop_info.header))?;

        // Generate unique pre-header label to avoid conflicts
        let mut pre_header_label = format!("{}_pre", loop_info.header);
        let mut counter = 0;
        while func.blocks.iter().any(|b| b.label == pre_header_label) {
            counter += 1;
            pre_header_label = format!("{}_pre_{}", loop_info.header, counter);
        }

        let mut pre_header_block = Block {
            label: pre_header_label.clone(),
            instructions: Vec::new(),
        };

        let mut sorted_invariant = invariant_instrs.to_vec();
        sorted_invariant.sort_by(|a, b| b.1.cmp(&a.1));
        sorted_invariant.sort_by(|a, b| b.0.cmp(&a.0));

        sorted_invariant.dedup();
        let mut instructions_to_move = Vec::new();
        for &(block_idx, instr_idx) in &sorted_invariant {
            if block_idx >= func.blocks.len() {
                continue;
            }
            let block = &func.blocks[block_idx];
            if instr_idx >= block.instructions.len() {
                continue;
            }
            let instr = block.instructions[instr_idx].clone();
            instructions_to_move.push(instr);
        }

        // Move instructions to the pre-header block
        for instr in instructions_to_move {
            pre_header_block.instructions.push(instr);
        }

        // Add jump to the original header
        pre_header_block.instructions.push(Instruction::Jmp {
            target: loop_info.header.clone(),
        });

        // Insert pre-header before the original header
        func.blocks.insert(header_idx, pre_header_block);

        // Remove original instructions from loop blocks
        // Since we sorted by (block_idx, instr_idx) descending, we can remove them safely
        for &(block_idx, instr_idx) in &sorted_invariant {
            // Adjust for the inserted pre-header block
            let adjusted_block_idx = if block_idx >= header_idx {
                block_idx + 1
            } else {
                block_idx
            };
            if adjusted_block_idx < func.blocks.len()
                && instr_idx < func.blocks[adjusted_block_idx].instructions.len()
            {
                func.blocks[adjusted_block_idx]
                    .instructions
                    .remove(instr_idx);
            }
        }

        // Redirect non-loop incoming edges to go through the pre-header.
        // Back-edges from inside the loop are skipped.
        for (block_idx, block) in func.blocks.iter_mut().enumerate() {
            if block_idx == header_idx || loop_info.blocks.contains(&block.label) {
                continue;
            }

            for instr in &mut block.instructions {
                match instr {
                    Instruction::Jmp { target } if *target == loop_info.header => {
                        *target = pre_header_label.clone();
                    }
                    Instruction::Br {
                        true_target,
                        false_target,
                        ..
                    } => {
                        if *true_target == loop_info.header {
                            *true_target = pre_header_label.clone();
                        }
                        if *false_target == loop_info.header {
                            *false_target = pre_header_label.clone();
                        }
                    }
                    _ => {}
                }
            }
        }

        // If the loop header was the function entry point, update it to the pre-header
        if func.entry == loop_info.header {
            func.entry = pre_header_label;
        }

        Ok(())
    }

    fn collect_defs_in_loop(
        &self,
        func: &Function,
        loop_info: &LoopInfo,
    ) -> HashSet<Register> {
        let mut defs = HashSet::new();
        for block in &func.blocks {
            if !loop_info.blocks.contains(&block.label) {
                continue;
            }
            for instr in &block.instructions {
                if let Some(def) = instr.def_reg() {
                    defs.insert(def.clone());
                }
            }
        }
        defs
    }
}

#[derive(Debug)]
struct LoopInfo {
    header: String,
    blocks: HashSet<String>,
}

#[derive(Debug)]
struct UnrollableLoop {
    header: String,
    body_blocks: Vec<String>,
    bound: i64,
    exit_target: String,
}

/// Loop unrolling that unrolls small loops to reduce overhead and improve ILP.
#[derive(Default)]
pub struct LoopUnrolling;

impl Transform for LoopUnrolling {
    fn name(&self) -> &'static str {
        "loop_unrolling"
    }

    fn description(&self) -> &'static str {
        "Unrolls small loops to reduce overhead and improve parallelism"
    }

    fn category(&self) -> TransformCategory {
        TransformCategory::ControlFlowOptimization
    }

    fn level(&self) -> TransformLevel {
        TransformLevel::Experimental
    }

    fn apply(&self, func: &mut Function) -> Result<bool, String> {
        self.apply_internal(func)
    }
}

impl LoopUnrolling {
    fn apply_internal(&self, func: &mut Function) -> Result<bool, String> {
        const MAX_BLOCKS: usize = 1000;
        const MAX_LOOPS_TO_UNROLL: usize = 10;

        if func.blocks.len() > MAX_BLOCKS {
            return Err(format!(
                "Function too large for loop unrolling ({} blocks, max {})",
                func.blocks.len(),
                MAX_BLOCKS
            ));
        }

        let mut changed = false;

        // 1. Identify loops
        // We reuse the logic from LoopInvariantCodeMotion (conceptually) but we need to own the detection here
        // or refactor to share. For now, we implement a detector here.
        let loops = self.find_unrollable_loops(func);

        let loops_to_unroll: Vec<_> = loops.into_iter().take(MAX_LOOPS_TO_UNROLL).collect();

        for loop_info in loops_to_unroll {
            // Apply unrolling
            if self.unroll_loop(func, &loop_info)? {
                changed = true;
            }
        }

        Ok(changed)
    }

    /// Find loops that can be safely unrolled (constant bounds, single entry, single backedge)
    fn find_unrollable_loops(&self, func: &Function) -> Vec<UnrollableLoop> {
        let mut unrollable = Vec::new();
        // Use a simplified CFG analysis to find natural loops
        // We look for back edges: Jmp/Br to an ancestor in the dominance tree (approx via block order)

        // Map label -> index
        let label_to_idx: HashMap<_, _> = func
            .blocks
            .iter()
            .enumerate()
            .map(|(i, b)| (b.label.as_str(), i))
            .collect();

        for (idx, block) in func.blocks.iter().enumerate() {
            if let Some(terminator) = block.instructions.last() {
                let targets = match terminator {
                    Instruction::Jmp { target } => vec![target],
                    Instruction::Br {
                        true_target,
                        false_target,
                        ..
                    } => vec![true_target, false_target],
                    _ => vec![],
                };

                for target in targets {
                    // Check if target is a header (back edge)
                    // Target must be before current block or same (self-loop)
                    if let Some(&target_idx) = label_to_idx.get(target.as_str())
                        && target_idx <= idx
                    {
                        // Potential loop header found.
                        // Now analyze if it's a "simple" counting loop we can unroll.
                        // 1. Must have constant bound
                        // 2. Must be small enough count
                        if let Some((bound, exit_target)) =
                            self.analyze_loop(func, target, &block.label)
                        {
                            // Collect body blocks (basic reachability from header to latch, excluding exit)
                            let body_blocks = self.collect_loop_body(func, target, &block.label);

                            if !body_blocks.is_empty() && body_blocks.len() <= 10 {
                                // Max 10 blocks in body
                                unrollable.push(UnrollableLoop {
                                    header: target.to_string(),
                                    body_blocks,
                                    bound,
                                    exit_target,
                                });
                            }
                        }
                    }
                }
            }
        }

        // Dedup loops by header
        unrollable.sort_by(|a, b| a.header.cmp(&b.header));
        unrollable.dedup_by(|a, b| a.header == b.header);

        unrollable
    }

    /// Analyze loop structure to determine exact invocation count and exit target
    fn analyze_loop(
        &self,
        func: &Function,
        header_label: &str,
        latch_label: &str,
    ) -> Option<(i64, String)> {
        // We look for the condition at the LATCH block (bottom of loop), or HEADER (top of loop).
        // Standard "for" loops often look like:
        // Header: check condition -> Exit or Body
        // ...
        // Latch: Jmp Header
        // Or
        // Header: ...
        // Latch: increment, check condition -> Header or Exit

        // Let's check LATCH block first for condition
        let latch_block = func.get_block(latch_label)?;

        let (cond_reg, true_target, false_target) = match latch_block.instructions.last()? {
            Instruction::Br {
                cond,
                true_target,
                false_target,
                ..
            } => (cond, true_target, false_target),
            Instruction::Jmp { target } if target == header_label => {
                // Unconditional latch. Check header for condition (while loop)
                let header_block = func.get_block(header_label)?;
                return match header_block.instructions.last()? {
                    Instruction::Br {
                        cond,
                        true_target,
                        false_target,
                        ..
                    } => {
                        self.analyze_bound_from_cond(
                            func,
                            header_block,
                            cond,
                            true_target,
                            false_target,
                            header_label,
                        )
                    }
                    _ => None,
                }
            }
            _ => return None,
        };

        self.analyze_bound_from_cond(
            func,
            latch_block,
            cond_reg,
            true_target,
            false_target,
            header_label,
        )
    }

    fn analyze_bound_from_cond(
        &self,
        _func: &Function,
        block: &Block,
        cond_reg: &Register,
        true_target: &str,
        false_target: &str,
        header_label: &str,
    ) -> Option<(i64, String)> {
        // If true_target goes to loop (header or body) and false_target exits, or vice versa
        let (_loop_target, exit_target, _branch_on_true) = if true_target == header_label {
            (true_target, false_target, true)
        } else if false_target == header_label {
            (false_target, true_target, false)
        } else {
            // Maybe points to body? Simplifying assumption: must point to header.
            return None;
        };

        // Find definition of cond_reg in this block
        let cond_def = block
            .instructions
            .iter()
            .find(|i| i.def_reg() == Some(cond_reg))?;

        if let Instruction::IntCmp { op, lhs, rhs, .. } = cond_def {
            // Expect induction variable vs Constant
            let (reg, limit) = match (lhs, rhs) {
                (Operand::Register(r), Operand::Immediate(crate::mir::Immediate::I64(c))) => {
                    (r, *c)
                }
                (Operand::Immediate(crate::mir::Immediate::I64(_c)), Operand::Register(_r)) => {
                    // Reverse op if needed?
                    return None; // Simplify
                }
                _ => return None,
            };

            // Find increment of 'reg'
            let inc_def = block
                .instructions
                .iter()
                .find(|i| i.def_reg() == Some(reg))?;
            if let Instruction::IntBinary {
                op: IntBinOp::Add,
                lhs: Operand::Register(src),
                rhs: Operand::Immediate(crate::mir::Immediate::I64(step)),
                ..
            } = inc_def
                && src == reg
                && *step == 1
            {
                // i++ loop.
                // Standard "while (i < N)"
                // If we are at latch (do-while logic effectively):
                //   I start at 0? We assume 0 for now as per previous task spec, or we only unroll safe small checks.
                //   If condition is "i < N", and we exit when False.
                //   Then we run while i < N.
                //   Count = N.
                // Support bound up to 16-32.
                let count = match op {
                    IntCmpOp::SLt | IntCmpOp::ULt => limit,
                    IntCmpOp::SLe | IntCmpOp::ULe => limit + 1,
                    IntCmpOp::Ne => limit, // i != N
                    _ => return None,
                };

                if count > 0 && count <= 32 {
                    return Some((count, exit_target.to_string()));
                }
            }
        }

        None
    }

    fn collect_loop_body(&self, func: &Function, header: &str, latch: &str) -> Vec<String> {
        // Collect all blocks reachable from header that can reach latch, without passing through exit?
        // Simplified: BFS from header, stop if we see something not dominated by header?
        // Actually simplest is: all nodes N such that Header -> ... -> N -> ... -> Latch
        // For unrolling, we need a topological sort or just a list of blocks part of the loop.
        // We will just do a simple reachability trace within strict limits.

        let mut body = HashSet::new();
        let mut queue = vec![header.to_string()];
        body.insert(header.to_string());

        while let Some(current) = queue.pop() {
            if current == latch {
                continue; // Latch is part of body but has no successors in loop (except header)
            }
            if let Some(block) = func.get_block(&current) {
                let successors = block.successors();
                for succ in successors {
                    // To be in the loop, successor must eventually reach latch (or be latch)
                    // and not be the header (backedge handled elsewhere)
                    if succ != header && !body.contains(&succ) {
                        // Check if succ can reach latch?
                        // This is expensive.
                        // Simplified assumption: All successors of header that are not exit are in loop.
                        // We already know exit_target from analyze_loop. But we don't have it here.
                        // Actually, finding natural loop nodes is standard.
                        // For this task, let's assume if we hit latch, stop.
                        // If we are branching, we follow both.
                        // This is heuristic.
                        body.insert(succ.clone());
                        queue.push(succ);
                    }
                }
            }
        }
        body.insert(latch.to_string());

        // Return sorted list
        let mut list: Vec<_> = body.into_iter().collect();
        // Sort by order in function to preserve approx topological order
        let order: HashMap<_, _> = func
            .blocks
            .iter()
            .enumerate()
            .map(|(i, b)| (b.label.clone(), i))
            .collect();
        list.sort_by_key(|lbl| order.get(lbl).copied().unwrap_or(usize::MAX));
        list
    }

    fn unroll_loop(&self, func: &mut Function, loop_info: &UnrollableLoop) -> Result<bool, String> {
        // 1. Locate blocks
        let mut body_indices = Vec::new();
        for lbl in &loop_info.body_blocks {
            if let Some(idx) = func.blocks.iter().position(|b| &b.label == lbl) {
                body_indices.push(idx);
            } else {
                return Ok(false); // Block not found??
            }
        }

        // 2. Clone bodies
        let mut unrolled_blocks = Vec::new();
        // Original blocks are "iteration 0". We will modify them in place or replace them?
        // Replacing is easier to manage control flow.
        // We will create N copies of the body chain.

        for i in 0..loop_info.bound {
            for &idx in &body_indices {
                let original_block = &func.blocks[idx];
                let mut new_block = original_block.clone();

                // Rename label
                new_block.label = format!("{}_unroll_{}", original_block.label, i);

                // Rewrite targets
                if let Some(term) = new_block.instructions.last_mut() {
                    match term {
                        Instruction::Jmp { target } => {
                            if target == &loop_info.header {
                                // Back edge: Jump to next iteration's header, or EXIT if last iteration
                                if i == loop_info.bound - 1 {
                                    *target = loop_info.exit_target.clone();
                                } else {
                                    *target = format!("{}_unroll_{}", loop_info.header, i + 1);
                                }
                            } else if loop_info.body_blocks.contains(target) {
                                // Internal edge: Jump to same iteration's version
                                *target = format!("{}_unroll_{}", target, i);
                            }
                            // Else: exit edge (break/return), keep as is
                        }
                        Instruction::Br {
                            true_target,
                            false_target,
                            ..
                        } => {
                            // Handle True target
                            if true_target == &loop_info.header {
                                if i == loop_info.bound - 1 {
                                    *true_target = loop_info.exit_target.clone();
                                } else {
                                    *true_target = format!("{}_unroll_{}", loop_info.header, i + 1);
                                }
                            } else if loop_info.body_blocks.contains(true_target) {
                                *true_target = format!("{}_unroll_{}", true_target, i);
                            }

                            // Handle False target
                            if false_target == &loop_info.header {
                                if i == loop_info.bound - 1 {
                                    *false_target = loop_info.exit_target.clone();
                                } else {
                                    *false_target =
                                        format!("{}_unroll_{}", loop_info.header, i + 1);
                                }
                            } else if loop_info.body_blocks.contains(false_target) {
                                *false_target = format!("{}_unroll_{}", false_target, i);
                            }
                        }
                        _ => {}
                    }
                }
                unrolled_blocks.push(new_block);
            }
        }

        // 3. Replace original blocks with unrolled blocks
        // We need to keep the function valid.
        // Identify where to insert. We can remove the old body blocks and insert new ones.
        // We must patch predecessor of header to jump to "header_unroll_0".
        // BUT wait, "header_unroll_0" has same label as original header??
        // No, we renamed it.
        // Predecessors of original header (from outside loop) must now jump to `header_unroll_0`.

        let entry_label = format!("{}_unroll_0", loop_info.header);

        // Fix external predecessors
        for block in &mut func.blocks {
            if loop_info.body_blocks.contains(&block.label) {
                continue;
            }
            for instr in &mut block.instructions {
                match instr {
                    Instruction::Jmp { target } if target == &loop_info.header => {
                        *target = entry_label.clone();
                    }
                    Instruction::Br {
                        true_target,
                        false_target,
                        ..
                    } => {
                        if true_target == &loop_info.header {
                            *true_target = entry_label.clone();
                        }
                        if false_target == &loop_info.header {
                            *false_target = entry_label.clone();
                        }
                    }
                    _ => {}
                }
            }
        }

        let old_body_set: HashSet<_> = loop_info.body_blocks.iter().collect();
        func.blocks.retain(|b| !old_body_set.contains(&b.label));
        func.blocks.extend(unrolled_blocks);

        for block in &mut func.blocks {
            if block
                .label
                .starts_with(&format!("{}_unroll_", loop_info.header))
                && let Some(term) = block.instructions.last_mut()
            {
                // If it's the header/latch that had the condition
                // Wait, did we rename header? Yes.
                // If we are in `header_unroll_i`, and it was `Br`, we want `Jmp`.
                if matches!(term, Instruction::Br { .. }) {
                    // Check if targets match our "Next Iter" or "Exit" pattern
                    // If so, replace with Jmp.
                    // But we need to know which target is taken.
                    // We assumed unrolling is valid, implying we ALWAYS take loop path.
                    // So finding the target that points to "Next Iteration" (or Exit on last) is enough.

                    let target_to_take = match term {
                        Instruction::Br { .. } => {
                            // One of these points to next iter (or exit).
                            // The other points to exit (early) or loop?
                            // Since we verified it's a counting loop 0..N, we take loop path N times.
                            // The "Exit" target of the ORIGINAL loop was identifying the Loop Exit.
                            // The "Loop" target was identifying Loop Body.
                            // We remapped Loop Body -> Next Iter.
                            // We remapped Loop Exit -> Exit.
                            // So we want to jump to (Next Iter).
                            // Which one is it?
                            // We can check if true_target was original exit?
                            // Hard to track original.
                            // But we know the resulting definition:
                            // If `true_target` includes `unroll_` (next iter) or matches `exit_target` (last iter logic).
                            // Actually, `unroll_loop` Step 2 logic modified targets.
                            // We can just look if one target represents "Continue".
                            // Simplified: Leave as Br. Constant Propagation/Folding handles it IF we unroll values too.
                            // But we didn't unroll values (const prop of I).
                            // So explicit `Jmp` is better.
                            // Let's assume the loop-back target is the one we want.
                            // But wait, in the last iteration, we want the EXIT target.
                            None // Skip optimization for safety in this step
                        }
                        _ => None,
                    };

                    if let Some(t) = target_to_take {
                        *term = Instruction::Jmp { target: t };
                    }
                }
            }
        }

        Ok(true)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::mir::{
        FunctionBuilder, Immediate, IntBinOp, MirType, Operand, ScalarType, VirtualReg,
    };

    #[test]
    fn test_loop_unrolling_simple() {
        // Create a simple loop:
        // entry:
        //   v0 = 0
        //   br loop
        // loop:
        //   v0 = add v0, 1
        //   v1 = lt v0, 4
        //   br v1, loop, exit
        // exit:
        //   ret

        let mut func = FunctionBuilder::new("test_loop")
            .returns(MirType::Scalar(ScalarType::I64))
            .block("entry")
            .instr(Instruction::IntBinary {
                op: IntBinOp::Add,
                ty: MirType::Scalar(ScalarType::I64),
                dst: VirtualReg::gpr(0).into(),
                lhs: Operand::Immediate(Immediate::I64(0)),
                rhs: Operand::Immediate(Immediate::I64(0)),
            })
            .instr(Instruction::Jmp {
                target: "loop".to_string(),
            })
            .block("loop")
            // Body: v2 = v2 + 1 (dummy work)
            .instr(Instruction::IntBinary {
                op: IntBinOp::Add,
                ty: MirType::Scalar(ScalarType::I64),
                dst: VirtualReg::gpr(2).into(),
                lhs: Operand::Register(VirtualReg::gpr(2).into()),
                rhs: Operand::Immediate(Immediate::I64(1)),
            })
            // Increment: v0 = v0 + 1
            .instr(Instruction::IntBinary {
                op: IntBinOp::Add,
                ty: MirType::Scalar(ScalarType::I64),
                dst: VirtualReg::gpr(0).into(),
                lhs: Operand::Register(VirtualReg::gpr(0).into()),
                rhs: Operand::Immediate(Immediate::I64(1)),
            })
            // Compare: v1 = v0 < 4
            .instr(Instruction::IntCmp {
                op: IntCmpOp::SLt,
                ty: MirType::Scalar(ScalarType::I1),
                dst: VirtualReg::gpr(1).into(),
                lhs: Operand::Register(VirtualReg::gpr(0).into()),
                rhs: Operand::Immediate(Immediate::I64(4)),
            })
            // Branch
            .instr(Instruction::Br {
                cond: VirtualReg::gpr(1).into(),
                true_target: "loop".to_string(),
                false_target: "exit".to_string(),
            })
            .block("exit")
            .instr(Instruction::Ret { value: None })
            .build();

        let unrolling = LoopUnrolling;
        let changed = unrolling.apply(&mut func).expect("Unrolling failed");

        assert!(changed);

        // With multi-block unrolling, "loop" is replaced by "loop_unroll_0"..."loop_unroll_3"
        // We expect "loop" to be gone.
        assert!(func.get_block("loop").is_none());

        // Check first iteration
        let loop_0 = func
            .get_block("loop_unroll_0")
            .expect("loop_unroll_0 exists");
        assert!(!loop_0.instructions.is_empty());

        match loop_0.instructions.last().unwrap() {
            Instruction::Jmp { target } => assert_eq!(target, "loop_unroll_1"),
            Instruction::Br { true_target, .. } => assert_eq!(true_target, "loop_unroll_1"), // If conversion failed
            _ => panic!("Expected branch/jump to next iter"),
        }

        // Check last iteration
        let loop_3 = func
            .get_block("loop_unroll_3")
            .expect("loop_unroll_3 exists");
        match loop_3.instructions.last().unwrap() {
            Instruction::Jmp { target } => assert_eq!(target, "exit"),
            Instruction::Br { true_target, .. } => assert_eq!(true_target, "exit"), // Last iter jumps to exit
            _ => panic!("Expected branch/jump to exit"),
        }
    }

    #[test]
    fn test_licm_empty_function() {
        let mut func = FunctionBuilder::new("empty")
            .returns(MirType::Scalar(ScalarType::I64))
            .block("entry")
            .instr(Instruction::Ret { value: None })
            .build();

        let licm = LoopInvariantCodeMotion;
        let result = licm.apply(&mut func);
        assert!(result.is_ok());
        assert!(!result.unwrap()); // No changes expected
    }

    #[test]
    fn test_licm_no_loops() {
        // Function with no loops - nothing to optimize
        let mut func = FunctionBuilder::new("no_loop")
            .returns(MirType::Scalar(ScalarType::I64))
            .block("entry")
            .instr(Instruction::IntBinary {
                op: IntBinOp::Add,
                ty: MirType::Scalar(ScalarType::I64),
                dst: VirtualReg::gpr(0).into(),
                lhs: Operand::Immediate(Immediate::I64(1)),
                rhs: Operand::Immediate(Immediate::I64(2)),
            })
            .instr(Instruction::Jmp {
                target: "exit".to_string(),
            })
            .block("exit")
            .instr(Instruction::Ret {
                value: Some(Operand::Register(VirtualReg::gpr(0).into())),
            })
            .build();

        let licm = LoopInvariantCodeMotion;
        let changed = licm.apply(&mut func).expect("should succeed");
        assert!(!changed);
    }

    #[test]
    fn test_licm_side_effect_not_moved() {
        // Store inside loop should NOT be moved out
        let mut func = FunctionBuilder::new("side_effect")
            .param(VirtualReg::gpr(0).into(), MirType::Scalar(ScalarType::I64))
            .returns(MirType::Scalar(ScalarType::I64))
            .block("entry")
            .instr(Instruction::Jmp {
                target: "loop".to_string(),
            })
            .block("loop")
            .instr(Instruction::Store {
                ty: MirType::Scalar(ScalarType::I64),
                src: Operand::Immediate(Immediate::I64(42)),
                addr: crate::mir::AddressMode::BaseOffset {
                    base: VirtualReg::gpr(0).into(),
                    offset: 0,
                },
                attrs: crate::mir::MemoryAttrs::default(),
            })
            .instr(Instruction::Br {
                cond: VirtualReg::gpr(0).into(),
                true_target: "loop".to_string(),
                false_target: "exit".to_string(),
            })
            .block("exit")
            .instr(Instruction::Ret { value: None })
            .build();

        let licm = LoopInvariantCodeMotion;
        let changed = licm.apply(&mut func).expect("should succeed");
        // Store has side effects - should not be moved
        assert!(!changed);
    }

    #[test]
    fn test_loop_unrolling_empty_function() {
        let mut func = FunctionBuilder::new("empty")
            .returns(MirType::Scalar(ScalarType::I64))
            .block("entry")
            .instr(Instruction::Ret { value: None })
            .build();

        let unroll = LoopUnrolling;
        let result = unroll.apply(&mut func);
        assert!(result.is_ok());
        assert!(!result.unwrap());
    }

    #[test]
    fn test_loop_unrolling_no_loops() {
        let mut func = FunctionBuilder::new("no_loop")
            .returns(MirType::Scalar(ScalarType::I64))
            .block("entry")
            .instr(Instruction::IntBinary {
                op: IntBinOp::Add,
                ty: MirType::Scalar(ScalarType::I64),
                dst: VirtualReg::gpr(0).into(),
                lhs: Operand::Immediate(Immediate::I64(1)),
                rhs: Operand::Immediate(Immediate::I64(2)),
            })
            .instr(Instruction::Ret {
                value: Some(Operand::Register(VirtualReg::gpr(0).into())),
            })
            .build();

        let unroll = LoopUnrolling;
        let changed = unroll.apply(&mut func).expect("should succeed");
        assert!(!changed);
    }

    #[test]
    fn test_loop_unrolling_large_bound_not_unrolled() {
        // Loop with bound > 32 should not be unrolled
        let mut func = FunctionBuilder::new("large_bound")
            .returns(MirType::Scalar(ScalarType::I64))
            .block("entry")
            .instr(Instruction::IntBinary {
                op: IntBinOp::Add,
                ty: MirType::Scalar(ScalarType::I64),
                dst: VirtualReg::gpr(0).into(),
                lhs: Operand::Immediate(Immediate::I64(0)),
                rhs: Operand::Immediate(Immediate::I64(0)),
            })
            .instr(Instruction::Jmp {
                target: "loop".to_string(),
            })
            .block("loop")
            .instr(Instruction::IntBinary {
                op: IntBinOp::Add,
                ty: MirType::Scalar(ScalarType::I64),
                dst: VirtualReg::gpr(0).into(),
                lhs: Operand::Register(VirtualReg::gpr(0).into()),
                rhs: Operand::Immediate(Immediate::I64(1)),
            })
            .instr(Instruction::IntCmp {
                op: IntCmpOp::SLt,
                ty: MirType::Scalar(ScalarType::I1),
                dst: VirtualReg::gpr(1).into(),
                lhs: Operand::Register(VirtualReg::gpr(0).into()),
                rhs: Operand::Immediate(Immediate::I64(100)), // > 32
            })
            .instr(Instruction::Br {
                cond: VirtualReg::gpr(1).into(),
                true_target: "loop".to_string(),
                false_target: "exit".to_string(),
            })
            .block("exit")
            .instr(Instruction::Ret { value: None })
            .build();

        let unroll = LoopUnrolling;
        let changed = unroll.apply(&mut func).expect("should succeed");
        // Loop with large bound should not be unrolled
        assert!(!changed);
        // Original loop should still exist
        assert!(func.get_block("loop").is_some());
    }

    #[test]
    fn test_licm_stress_no_infinite_loop() {
        // Create function with nested-looking structure but no actual loop
        let mut func = FunctionBuilder::new("stress")
            .returns(MirType::Scalar(ScalarType::I64))
            .block("entry")
            .build();

        // Add many blocks in linear fashion
        for i in 0..50 {
            let label = format!("block{}", i);
            let next_label = if i < 49 {
                format!("block{}", i + 1)
            } else {
                "exit".to_string()
            };

            let mut block = Block::new(&label);
            block.push(Instruction::IntBinary {
                op: IntBinOp::Add,
                ty: MirType::Scalar(ScalarType::I64),
                dst: VirtualReg::gpr(i).into(),
                lhs: Operand::Immediate(Immediate::I64(i as i64)),
                rhs: Operand::Immediate(Immediate::I64(1)),
            });
            block.push(Instruction::Jmp { target: next_label });
            func.add_block(block);
        }

        let mut exit_block = Block::new("exit");
        exit_block.push(Instruction::Ret { value: None });
        func.add_block(exit_block);

        // Update entry to jump to block0
        func.blocks[0].instructions.push(Instruction::Jmp {
            target: "block0".to_string(),
        });

        let licm = LoopInvariantCodeMotion;
        let result = licm.apply(&mut func);
        assert!(result.is_ok());
    }

    #[test]
    fn test_loop_unrolling_stress_no_infinite_loop() {
        // Many simple blocks should not cause infinite loop
        let mut func = FunctionBuilder::new("stress")
            .returns(MirType::Scalar(ScalarType::I64))
            .block("entry")
            .instr(Instruction::Ret { value: None })
            .build();

        // Add 100 independent blocks
        for i in 0..100 {
            let label = format!("dead_block_{}", i);
            let mut block = Block::new(&label);
            block.push(Instruction::Ret { value: None });
            func.add_block(block);
        }

        let unroll = LoopUnrolling;
        let result = unroll.apply(&mut func);
        assert!(result.is_ok());
    }
}