llvm-native-core-ext 0.1.0

//! Loop Rotation — Canonicalize loops so the header contains the
//! loop condition check.
//!
//! Clean-room behavioral reconstruction based on compiler optimization
//! literature and the LLVM LoopRotate pass specification.
//!
//! @llvm_behavior: Loop rotation transforms a loop with the condition
//! check in the header by duplicating the header into the preheader.
//! This creates a loop where:
//! - The preheader checks the condition once before entering
//! - The new header is the old body (guaranteed to execute at least once)
//! - The latch branches back to the preheader check
//!
//! This canonicalization enables:
//! - LICM (loop-invariant code motion) to hoist invariant condition checks
//! - Better induction variable recognition
//! - More efficient loop unrolling
//!
//! Algorithm:
//! 1. Find loops with a preheader and a header containing a conditional branch
//! 2. Duplicate the header into the preheader
//! 3. Adjust PHI nodes in the header for the new CFG structure
//! 4. Update the latch to branch to the new preheader instead of the header

use llvm_native_core::SubclassKind;
use llvm_native_core::analysis::LoopInfo;
use llvm_native_core::opcode::Opcode;
use llvm_native_core::value::ValueRef;

// ============================================================================
// Loop Rotate Pass
// ============================================================================

/// Loop rotation pass: canonicalizes loop headers to contain the loop body.
///
/// Transforms:
/// ```text
/// preheader → header(cond, body, exit) → [body → header]
/// ```
/// into:
/// ```text
/// preheader' → header(body, cond', exit) → [body → preheader']
/// ```
pub struct LoopRotatePass {
    /// Number of loops rotated.
    pub rotated: usize,
    /// Maximum number of instructions in the header to allow rotation.
    /// If the header is too large, rotating duplicates too much code.
    pub max_rotation_size: usize,
}

impl LoopRotatePass {
    /// Create a new loop rotation pass.
    pub fn new() -> Self {
        Self {
            rotated: 0,
            max_rotation_size: 50,
        }
    }

    /// Create a loop rotation pass with a custom maximum rotation size.
    pub fn with_max_size(max_rotation_size: usize) -> Self {
        Self {
            rotated: 0,
            max_rotation_size,
        }
    }

    /// Run loop rotation on a function.
    ///
    /// Returns the number of loops successfully rotated.
    pub fn run_on_function(&mut self, func: &ValueRef) -> usize {
        // Step 1: Find all loops using analysis
        let analysis = llvm_native_core::analysis::LoopAnalysis::compute(func);
        let loops = analysis.loops.clone();

        let mut rotated_count = 0;
        for loop_info in &loops {
            if self.can_rotate(loop_info) && self.is_profitable_to_rotate(loop_info) {
                if self.rotate_loop(loop_info, func) {
                    rotated_count += 1;
                }
            }
        }

        self.rotated += rotated_count;
        rotated_count
    }

    /// Check if a loop can be rotated.
    ///
    /// Requirements:
    /// - Loop has a preheader
    /// - Loop has a latch
    /// - The header contains a conditional branch to the body and exit
    /// - The header is not too large
    pub fn can_rotate(&self, loop_info: &LoopInfo) -> bool {
        // Need a preheader to duplicate the header into
        if loop_info.preheader.is_none() {
            return false;
        }

        // Need a latch
        if loop_info.latch.is_none() {
            return false;
        }

        // Check header has a conditional branch
        let header = &loop_info.header;
        let hdr = header.borrow();
        let has_cond_br = hdr.operands.iter().any(|op| {
            let inst = op.borrow();
            inst.get_opcode() == Some(Opcode::Br) && inst.operands.len() == 3
        });

        if !has_cond_br {
            return false;
        }

        // Check header size
        let mut inst_count = 0usize;
        for op in &hdr.operands {
            if op.borrow().is_instruction() {
                inst_count += 1;
            }
        }
        if inst_count > self.max_rotation_size {
            return false;
        }

        true
    }

    /// Check if it's profitable to rotate this loop.
    ///
    /// Rotation is profitable when:
    /// - The loop body will be executed at least once (do-while form)
    /// - The duplicated header code is small
    /// - There's a PHI in the header that benefits from rotation
    pub fn is_profitable_to_rotate(&self, loop_info: &LoopInfo) -> bool {
        let header = &loop_info.header;
        let hdr = header.borrow();

        // Count PHI nodes in the header
        let phi_count = hdr
            .operands
            .iter()
            .filter(|op| {
                let inst = op.borrow();
                inst.get_opcode() == Some(Opcode::Phi)
            })
            .count();

        // Rotation is profitable if there are PHIs (they'll be
        // simplified after rotation) or the header is small.
        let mut inst_count = 0usize;
        for op in &hdr.operands {
            if op.borrow().is_instruction() {
                inst_count += 1;
            }
        }

        phi_count > 0 || inst_count <= 5
    }

    /// Rotate a single loop.
    ///
    /// The transformation:
    /// 1. The preheader gets the condition check and branches to
    ///    the header (if condition) or exit (if not)
    /// 2. The header becomes the loop body entry
    /// 3. The latch branches to the preheader instead of the header
    /// 4. PHI nodes are updated to reflect the new CFG
    pub fn rotate_loop(&mut self, loop_info: &LoopInfo, _func: &ValueRef) -> bool {
        let header = &loop_info.header;
        let preheader = match &loop_info.preheader {
            Some(ph) => ph,
            None => return false,
        };
        let latch = match &loop_info.latch {
            Some(l) => l,
            None => return false,
        };

        // Find the exit block from the header's conditional branch
        let exit_block = {
            let hdr = header.borrow();
            let mut exit: Option<ValueRef> = None;
            let mut _body: Option<ValueRef> = None;

            for op in &hdr.operands {
                let inst = op.borrow();
                if inst.get_opcode() == Some(Opcode::Br) && inst.operands.len() == 3 {
                    // cond, true_block, false_block
                    let t_block = &inst.operands[1];
                    let f_block = &inst.operands[2];

                    // Determine which is the loop body and which is the exit
                    // The block that is in the loop (but not the header) is the body
                    // The other is the exit
                    let t_in_loop = loop_info.blocks.iter().any(|b| {
                        b.borrow().name == t_block.borrow().name
                            && b.borrow().name != header.borrow().name
                    });
                    let f_in_loop = loop_info.blocks.iter().any(|b| {
                        b.borrow().name == f_block.borrow().name
                            && b.borrow().name != header.borrow().name
                    });

                    if t_in_loop {
                        _body = Some(t_block.clone());
                        exit = Some(f_block.clone());
                    } else if f_in_loop {
                        _body = Some(f_block.clone());
                        exit = Some(t_block.clone());
                    }
                    break;
                }
            }

            match exit {
                Some(e) => e,
                None => return false,
            }
        };

        // Step 1: Duplicate the header's condition check into the preheader
        self.duplicate_header_into_preheader(header, preheader, latch, &exit_block);

        // Step 2: Update PHI nodes in the header
        self.update_phi_nodes_after_rotation(header, preheader, latch);

        true
    }

    /// Duplicate the header's instructions (the condition check portion)
    /// into the preheader block.
    ///
    /// The preheader gets:
    /// - All instructions from the header that are needed for the condition
    /// - A conditional branch to (body/header, exit)
    /// - The preheader then branches to either the header (if condition holds)
    ///   or the exit block
    fn duplicate_header_into_preheader(
        &self,
        header: &ValueRef,
        preheader: &ValueRef,
        latch: &ValueRef,
        exit_block: &ValueRef,
    ) {
        let mut pre = preheader.borrow_mut();

        // Remove the unconditional branch at the end of preheader
        // (the preheader originally just branched to the header)
        pre.operands.retain(|op| {
            let inst = op.borrow();
            !(inst.get_opcode() == Some(Opcode::Br) && inst.operands.len() == 1)
        });

        // Clone the header's condition-producing instructions
        let hdr = header.borrow();
        let mut cond_insts: Vec<ValueRef> = Vec::new();
        let mut cond_val: Option<ValueRef> = None;
        let mut _true_dest: Option<ValueRef> = None;
        let mut _false_dest: Option<ValueRef> = None;

        for op in &hdr.operands {
            let inst = op.borrow();
            if !inst.is_instruction() {
                continue;
            }

            let opc = inst.get_opcode();
            if opc == Some(Opcode::Br) && inst.operands.len() == 3 {
                cond_val = Some(inst.operands[0].clone());
                _true_dest = Some(inst.operands[1].clone());
                _false_dest = Some(inst.operands[2].clone());
            } else if opc != Some(Opcode::Phi) {
                // Clone non-phi, non-terminator instructions
                cond_insts.push(op.clone());
            }
        }

        // Add cloned condition instructions to preheader
        for inst in &cond_insts {
            pre.operands.push(inst.clone());
        }

        // Add the conditional branch to preheader
        // br cond, header (body), exit
        if let Some(cond) = cond_val {
            let br = llvm_native_core::instruction::br_cond(
                cond,
                header.clone(),     // true: go to body
                exit_block.clone(), // false: exit the loop
            );
            pre.operands.push(br);
        }

        // Update the latch to branch to the preheader instead of the header
        let mut latch_bb = latch.borrow_mut();
        for op in &mut latch_bb.operands {
            let mut inst = op.borrow_mut();
            if inst.get_opcode() == Some(Opcode::Br) && inst.operands.len() == 1 {
                // Unconditional branch in the latch — change target to preheader
                inst.operands[0] = preheader.clone();
            }
        }
    }

    /// Update PHI nodes in the header after rotation.
    ///
    /// After rotation, the header's incoming edges change:
    /// - Old: incoming from preheader and latch
    /// - New: incoming from preheader (now conditional) and latch
    ///
    /// The PHI values from the preheader need to be updated because
    /// the preheader now contains the duplicated code.
    fn update_phi_nodes_after_rotation(
        &self,
        header: &ValueRef,
        preheader: &ValueRef,
        _latch: &ValueRef,
    ) {
        let mut hdr = header.borrow_mut();

        for op in &mut hdr.operands {
            let inst = op.borrow_mut();
            if inst.get_opcode() != Some(Opcode::Phi) {
                continue;
            }

            // PHI nodes have pairs: [val_from_predecessor1, pred1, val_from_predecessor2, pred2, ...]
            // After rotation, the preheader still feeds the header, but with potentially
            // different values due to instruction cloning.
            // For a simple implementation, we keep the PHI as-is since the cloned
            // instructions in the preheader should produce the same values.

            // Update the preheader value reference if needed
            for i in (0..inst.operands.len()).step_by(2) {
                if i + 1 < inst.operands.len() {
                    let pred = &inst.operands[i + 1];
                    if pred.borrow().name == preheader.borrow().name {
                        // The value from preheader may have changed — keep it for now
                        // A full implementation would map the cloned values
                    }
                }
            }
        }
    }

    /// Get the number of loops rotated.
    pub fn rotated_count(&self) -> usize {
        self.rotated
    }
}

impl Default for LoopRotatePass {
    fn default() -> Self {
        Self::new()
    }
}

// ============================================================================
// Loop Rotation Cost Model
// ============================================================================

/// Cost model for evaluating whether loop rotation is profitable.
#[derive(Debug, Clone)]
pub struct LoopRotationCostModel {
    /// Maximum header size (in instructions) for rotation.
    pub max_header_size: usize,
    /// Cost of duplicating the latch into the header.
    pub latch_duplication_cost: i64,
    /// Benefit of enabling LICM after rotation.
    pub licm_enablement_benefit: i64,
    /// Benefit of simplifying the loop form.
    pub simplification_benefit: i64,
}

impl LoopRotationCostModel {
    /// Create a default cost model.
    pub fn new() -> Self {
        Self {
            max_header_size: 50,
            latch_duplication_cost: 5,
            licm_enablement_benefit: 20,
            simplification_benefit: 10,
        }
    }

    /// Evaluate whether rotating a loop is profitable.
    pub fn is_profitable(&self, header_size: usize, latch_size: usize) -> bool {
        if header_size > self.max_header_size {
            return false;
        }

        let cost = latch_size as i64 * self.latch_duplication_cost;
        let benefit = self.licm_enablement_benefit + self.simplification_benefit;

        benefit > cost
    }

    /// Estimate header size from a block.
    pub fn estimate_header_size(&self, header: &ValueRef) -> usize {
        header
            .borrow()
            .operands
            .iter()
            .filter(|op| op.borrow().subclass == SubclassKind::BasicBlock)
            .count()
    }

    /// Analyze latch duplication cost.
    pub fn analyze_latch(&self, latch: &ValueRef) -> usize {
        latch
            .borrow()
            .operands
            .iter()
            .filter(|op| op.borrow().subclass == SubclassKind::Instruction)
            .count()
    }
}

impl Default for LoopRotationCostModel {
    fn default() -> Self {
        Self::new()
    }
}

// ============================================================================
// Profitability Heuristics for Loop Rotation
// ============================================================================

/// Advanced heuristics for loop rotation profitability.
#[derive(Debug, Default)]
pub struct RotationHeuristics {
    /// Cost model.
    pub cost_model: LoopRotationCostModel,
    /// Rotated loops count.
    pub rotated: usize,
    /// Loops rejected by cost model.
    pub rejected: usize,
}

impl RotationHeuristics {
    /// Create new heuristics.
    pub fn new() -> Self {
        Self::default()
    }

    /// Evaluate whether to rotate a specific loop.
    pub fn should_rotate(&mut self, header: &ValueRef, latch: &ValueRef) -> bool {
        let header_size = self.cost_model.estimate_header_size(header);
        let latch_size = self.cost_model.analyze_latch(latch);

        if self.cost_model.is_profitable(header_size, latch_size) {
            self.rotated += 1;
            true
        } else {
            self.rejected += 1;
            false
        }
    }
}

// ============================================================================
// Top-level entry points
// ============================================================================

/// Run loop rotation with profitability analysis.
pub fn run_loop_rotate_with_cost(func: &ValueRef) -> usize {
    let mut heuristics = RotationHeuristics::new();
    let f = func.borrow();
    for bb in &f.operands {
        if bb.borrow().subclass == SubclassKind::BasicBlock {
            let _ = heuristics.should_rotate(bb, bb);
        }
    }
    heuristics.rotated
}

// ============================================================================
// Tests
// ============================================================================

#[cfg(test)]
mod tests {
    use super::*;
    use llvm_native_core::analysis::LoopAnalysis;
    use llvm_native_core::basic_block::new_basic_block;
    use llvm_native_core::constants;
    use llvm_native_core::function::new_function;
    use llvm_native_core::instruction;
    use llvm_native_core::types::Type;

    /// Build a simple loop with preheader:
    /// entry(preheader) → loop_hdr(cond: loop_body / exit) → loop_body → loop_hdr
    fn build_simple_loop() -> ValueRef {
        let func = new_function("simple_loop", Type::void(), &[]);
        let preheader = new_basic_block("preheader");
        let loop_hdr = new_basic_block("loop_header");
        let loop_body = new_basic_block("loop_body");
        let exit = new_basic_block("exit");

        // preheader: br loop_header
        preheader
            .borrow_mut()
            .push_operand(instruction::br(loop_hdr.clone()));

        // loop_header: br_cond cond, loop_body, exit
        let cond = constants::const_bool(true);
        loop_hdr.borrow_mut().push_operand(instruction::br_cond(
            cond,
            loop_body.clone(),
            exit.clone(),
        ));

        // loop_body: br loop_header (backedge)
        loop_body
            .borrow_mut()
            .push_operand(instruction::br(loop_hdr.clone()));

        // exit: ret void
        exit.borrow_mut().push_operand(instruction::ret_void());

        func.borrow_mut().push_operand(preheader);
        func.borrow_mut().push_operand(loop_hdr);
        func.borrow_mut().push_operand(loop_body);
        func.borrow_mut().push_operand(exit);
        func
    }

    /// Build a loop with PHI node in header:
    /// preheader → loop_hdr(phi, cond: loop_body / exit) → loop_body(inc, br loop_hdr) → exit
    fn build_loop_with_phi() -> ValueRef {
        let func = new_function("phi_loop", Type::void(), &[]);
        let preheader = new_basic_block("preheader");
        let loop_hdr = new_basic_block("loop_header");
        let loop_body = new_basic_block("loop_body");
        let exit = new_basic_block("exit");

        // preheader: br loop_header
        preheader
            .borrow_mut()
            .push_operand(instruction::br(loop_hdr.clone()));

        // loop_header: phi(i, [0, preheader], [inc, loop_body]); icmp(cond); br_cond cond, loop_body, exit
        let zero = constants::const_i32(0);
        let inc_val = constants::const_i32(1); // placeholder — real inc would be in loop_body
        let phi = instruction::phi(
            Type::i32(),
            vec![
                (zero.clone(), preheader.clone()),
                (inc_val.clone(), loop_body.clone()),
            ],
        );
        phi.borrow_mut().name = "i".into();
        loop_hdr.borrow_mut().push_operand(phi);

        let cond = constants::const_bool(true);
        loop_hdr.borrow_mut().push_operand(instruction::br_cond(
            cond,
            loop_body.clone(),
            exit.clone(),
        ));

        // loop_body: br loop_header
        loop_body
            .borrow_mut()
            .push_operand(instruction::br(loop_hdr.clone()));

        // exit: ret void
        exit.borrow_mut().push_operand(instruction::ret_void());

        func.borrow_mut().push_operand(preheader);
        func.borrow_mut().push_operand(loop_hdr);
        func.borrow_mut().push_operand(loop_body);
        func.borrow_mut().push_operand(exit);
        func
    }

    /// Build a loop without preheader (not rotatable):
    /// entry → loop_hdr(cond: loop_body / exit) → loop_body → loop_hdr
    fn build_loop_no_preheader() -> ValueRef {
        let func = new_function("no_preheader", Type::void(), &[]);
        let entry = new_basic_block("entry");
        let loop_hdr = new_basic_block("loop_header");
        let loop_body = new_basic_block("loop_body");
        let exit = new_basic_block("exit");

        entry
            .borrow_mut()
            .push_operand(instruction::br(loop_hdr.clone()));

        let cond = constants::const_bool(true);
        loop_hdr.borrow_mut().push_operand(instruction::br_cond(
            cond,
            loop_body.clone(),
            exit.clone(),
        ));

        loop_body
            .borrow_mut()
            .push_operand(instruction::br(loop_hdr.clone()));
        exit.borrow_mut().push_operand(instruction::ret_void());

        func.borrow_mut().push_operand(entry);
        func.borrow_mut().push_operand(loop_hdr);
        func.borrow_mut().push_operand(loop_body);
        func.borrow_mut().push_operand(exit);
        func
    }

    // === LoopRotatePass::new Tests ===

    #[test]
    fn test_loop_rotate_pass_new() {
        let pass = LoopRotatePass::new();
        assert_eq!(pass.rotated, 0);
        assert_eq!(pass.max_rotation_size, 50);
    }

    #[test]
    fn test_loop_rotate_pass_with_max_size() {
        let pass = LoopRotatePass::with_max_size(10);
        assert_eq!(pass.max_rotation_size, 10);
    }

    #[test]
    fn test_loop_rotate_pass_default() {
        let pass = LoopRotatePass::default();
        assert_eq!(pass.rotated, 0);
    }

    // === can_rotate Tests ===

    #[test]
    fn test_can_rotate_simple_loop() {
        let pass = LoopRotatePass::new();
        let func = build_simple_loop();
        let analysis = LoopAnalysis::compute(&func);
        assert!(!analysis.loops.is_empty());
        assert!(pass.can_rotate(&analysis.loops[0]));
    }

    #[test]
    fn test_can_rotate_no_preheader() {
        let pass = LoopRotatePass::new();
        let func = build_loop_no_preheader();
        let analysis = LoopAnalysis::compute(&func);
        if !analysis.loops.is_empty() {
            assert!(!pass.can_rotate(&analysis.loops[0]));
        }
    }

    #[test]
    fn test_can_rotate_with_phi() {
        let pass = LoopRotatePass::new();
        let func = build_loop_with_phi();
        let analysis = LoopAnalysis::compute(&func);
        assert!(!analysis.loops.is_empty());
        assert!(pass.can_rotate(&analysis.loops[0]));
    }

    // === is_profitable_to_rotate Tests ===

    #[test]
    fn test_is_profitable_to_rotate_simple() {
        let pass = LoopRotatePass::new();
        let func = build_simple_loop();
        let analysis = LoopAnalysis::compute(&func);
        if !analysis.loops.is_empty() {
            assert!(pass.is_profitable_to_rotate(&analysis.loops[0]));
        }
    }

    #[test]
    fn test_is_profitable_to_rotate_with_phi() {
        let pass = LoopRotatePass::new();
        let func = build_loop_with_phi();
        let analysis = LoopAnalysis::compute(&func);
        if !analysis.loops.is_empty() {
            assert!(pass.is_profitable_to_rotate(&analysis.loops[0]));
        }
    }

    // === rotate_loop Tests ===

    #[test]
    fn test_rotate_loop_simple() {
        let mut pass = LoopRotatePass::new();
        let func = build_simple_loop();
        let analysis = LoopAnalysis::compute(&func);
        if !analysis.loops.is_empty() {
            let result = pass.rotate_loop(&analysis.loops[0], &func);
            assert!(result);
        }
    }

    #[test]
    fn test_rotate_loop_with_phi() {
        let mut pass = LoopRotatePass::new();
        let func = build_loop_with_phi();
        let analysis = LoopAnalysis::compute(&func);
        if !analysis.loops.is_empty() {
            let result = pass.rotate_loop(&analysis.loops[0], &func);
            assert!(result);
        }
    }

    // === run_on_function Tests ===

    #[test]
    fn test_run_on_function_simple() {
        let mut pass = LoopRotatePass::new();
        let func = build_simple_loop();
        let count = pass.run_on_function(&func);
        assert!(count >= 0);
    }

    #[test]
    fn test_run_on_function_with_phi() {
        let mut pass = LoopRotatePass::new();
        let func = build_loop_with_phi();
        let count = pass.run_on_function(&func);
        assert!(count >= 0);
    }

    #[test]
    fn test_run_on_function_no_rotatable() {
        let mut pass = LoopRotatePass::new();
        let func = build_loop_no_preheader();
        let count = pass.run_on_function(&func);
        // No rotatable loops found
        assert_eq!(count, 0);
    }

    // === Duplicate Header Tests ===

    #[test]
    fn test_duplicate_header_into_preheader() {
        let pass = LoopRotatePass::new();
        let func = build_simple_loop();
        let f = func.borrow();

        // Find the preheader and header
        let preheader = f.operands[0].clone(); // preheader is first
        let header = f.operands[1].clone(); // loop_header is second
        let latch = f.operands[2].clone(); // loop_body is third
        let exit_block = f.operands[3].clone(); // exit is fourth

        pass.duplicate_header_into_preheader(&header, &preheader, &latch, &exit_block);

        // Preheader should now have more than just an unconditional branch
        let pre = preheader.borrow();
        assert!(
            pre.operands.len() > 1,
            "Preheader should have condition check"
        );
    }

    // === Edge Cases ===

    #[test]
    fn test_rotate_empty_function() {
        let mut pass = LoopRotatePass::new();
        let func = new_function("empty", Type::void(), &[]);
        let count = pass.run_on_function(&func);
        assert_eq!(count, 0);
    }

    #[test]
    fn test_rotate_function_without_loops() {
        let mut pass = LoopRotatePass::new();
        let func = new_function("no_loop", Type::i32(), &[]);
        let entry = new_basic_block("entry");
        entry
            .borrow_mut()
            .push_operand(instruction::ret_val(constants::const_i32(42)));
        func.borrow_mut().push_operand(entry);
        let count = pass.run_on_function(&func);
        assert_eq!(count, 0);
    }
}