dotscope 0.7.0

//! Loop Invariant Code Motion (LICM) Pass.
//!
//! This pass moves computations that produce the same value on every iteration
//! out of loops. This is useful for:
//!
//! - Performance optimization
//! - Cleaning up loop-based obfuscation patterns
//!
//! # Algorithm
//!
//! An instruction is loop-invariant if:
//! 1. All its operands are defined outside the loop, OR
//! 2. All its operands are defined by loop-invariant instructions
//!
//! An instruction can be hoisted if:
//! 1. It is loop-invariant
//! 2. It has no side effects (pure computation)
//! 3. The loop has a preheader where we can place the hoisted code
//!
//! # Example
//!
//! ```text
//! // Before LICM
//! preheader:
//!     a = 5
//!     b = 10
//!     jump header
//!
//! header:
//!     i = phi(0, i')
//!     x = a + b        // Loop invariant!
//!     use(x)
//!     i' = i + 1
//!     branch (i < 10), header, exit
//!
//! // After LICM
//! preheader:
//!     a = 5
//!     b = 10
//!     x = a + b        // Hoisted
//!     jump header
//!
//! header:
//!     i = phi(0, i')
//!     use(x)
//!     i' = i + 1
//!     branch (i < 10), header, exit
//! ```

use std::collections::{HashMap, HashSet, VecDeque};

use crate::{
    analysis::{LoopAnalyzer, LoopInfo, SsaFunction, SsaInstruction, SsaOp, SsaVarId},
    compiler::{
        pass::{ModificationScope, SsaPass},
        CompilerContext, EventKind,
    },
    metadata::token::Token,
    utils::BitSet,
    CilObject, Result,
};

/// Loop Invariant Code Motion Pass.
///
/// Moves loop-invariant computations to the loop preheader.
pub struct LicmPass;

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

impl LicmPass {
    /// Creates a new LICM pass.
    #[must_use]
    pub fn new() -> Self {
        Self
    }
}

impl SsaPass for LicmPass {
    fn name(&self) -> &'static str {
        "licm"
    }

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

    fn modification_scope(&self) -> ModificationScope {
        ModificationScope::InstructionsOnly
    }

    fn run_on_method(
        &self,
        ssa: &mut SsaFunction,
        method_token: Token,
        ctx: &CompilerContext,
        _assembly: &CilObject,
    ) -> Result<bool> {
        let forest = LoopAnalyzer::new(ssa).analyze();

        if forest.is_empty() {
            return Ok(false);
        }

        let mut total_hoisted = 0;

        // Process loops from innermost to outermost.
        // This naturally propagates hoists through nesting levels: inner hoists
        // move instructions to the inner preheader (inside the outer loop body),
        // and outer loop processing then hoists them further out if invariant.
        // The Nop skip in find_loop_invariants prevents exponential blowup by
        // ignoring the Nops left behind by inner hoists.
        for loop_info in forest.by_depth_descending() {
            // Skip loops without preheaders - we need somewhere to hoist to
            let Some(preheader) = loop_info.preheader else {
                continue;
            };

            // Validate that the preheader is an actual immediate predecessor of
            // the loop header. If it isn't (e.g., for loops inside CFF dispatchers
            // where the loop analyzer may pick a distant dominator as preheader),
            // hoisting would move definitions to a block that isn't a CFG
            // predecessor of the header, making phi operand updates invalid.
            let header_idx = loop_info.header.index();
            let preheader_is_pred = ssa
                .block(preheader.index())
                .map(|b| {
                    b.instructions()
                        .last()
                        .map(|i| i.op().successors().contains(&header_idx))
                        .unwrap_or(false)
                })
                .unwrap_or(false);
            if !preheader_is_pred {
                continue;
            }

            // Skip loops whose header is a switch-based dispatcher (CFF pattern).
            // Hoisting into the preheader of such loops adds variable definitions
            // that change the stack layout at the preheader's exit. When a
            // subsequent CfgModifying pass triggers rebuild_ssa(), the phi nodes
            // at the switch header are reconstructed from scratch using stack-slot
            // analysis. The extra definitions shift which variable occupies each
            // stack slot, causing the switch to read the wrong phi — e.g., a
            // hoisted constant instead of the CFF state variable. This misroutes
            // the CFF tracer and drops entire code phases.
            let header_has_switch = ssa
                .block(header_idx)
                .and_then(|b| b.terminator_op())
                .is_some_and(|op| matches!(op, SsaOp::Switch { .. }));
            if header_has_switch {
                continue;
            }

            // Find invariant instructions
            let invariants = find_loop_invariants(ssa, loop_info);

            if invariants.is_empty() {
                continue;
            }

            // Filter to hoistable instructions
            let mut hoistable: Vec<_> = invariants
                .into_iter()
                .filter(|(block_idx, instr_idx)| can_hoist(ssa, loop_info, *block_idx, *instr_idx))
                .collect();

            // Second filter: ensure all operands of hoistable instructions are
            // either defined outside the loop or by other hoistable instructions.
            // Without this, an instruction like Conv(v10) would be hoisted but its
            // operand v10 (from ArrayLength, which is invariant but not hoistable
            // due to side effects) would remain in the loop body, producing a
            // use-before-def.
            let mut outside_defs = BitSet::new(ssa.var_id_capacity());
            for v in ssa.variables() {
                if !loop_info.body.contains(v.def_site().block) {
                    outside_defs.insert(v.id().index());
                }
            }

            loop {
                let mut hoistable_defs = BitSet::new(ssa.var_id_capacity());
                for (block_idx, instr_idx) in hoistable.iter() {
                    if let Some(def) = ssa
                        .block(*block_idx)
                        .and_then(|b| b.instruction(*instr_idx))
                        .and_then(|i| i.def())
                    {
                        hoistable_defs.insert(def.index());
                    }
                }

                let before = hoistable.len();
                hoistable.retain(|(block_idx, instr_idx)| {
                    let Some(block) = ssa.block(*block_idx) else {
                        return false;
                    };
                    let Some(instr) = block.instruction(*instr_idx) else {
                        return false;
                    };
                    instr.op().uses().iter().all(|operand| {
                        outside_defs.contains(operand.index())
                            || hoistable_defs.contains(operand.index())
                    })
                });

                if hoistable.len() == before {
                    break;
                }
            }

            // Guard: if hoisting ALL non-terminator instructions from a block
            // would make it a trampoline AND that block feeds phis at a successor,
            // skip hoisting from that block entirely. Making such blocks trampolines
            // causes block-merging to clear them, and subsequent rebuild_ssa may
            // not correctly reconnect the phi with the preheader's definitions.
            {
                // Count hoistable instructions per block
                let mut hoist_count_per_block: HashMap<usize, usize> = HashMap::new();
                for (block_idx, _) in &hoistable {
                    *hoist_count_per_block.entry(*block_idx).or_insert(0) += 1;
                }

                // Find blocks that would become trampolines
                let mut trampoline_blocks = BitSet::new(ssa.block_count());
                for (&block_idx, &hoist_count) in &hoist_count_per_block {
                    if let Some(block) = ssa.block(block_idx) {
                        let non_term = block
                            .instructions()
                            .iter()
                            .filter(|i| !i.is_terminator() && !matches!(i.op(), SsaOp::Nop))
                            .count();
                        if hoist_count >= non_term {
                            // This block would become a trampoline — check if it feeds phis
                            if let Some(term) = block.terminator_op() {
                                for succ in term.successors() {
                                    if let Some(succ_block) = ssa.block(succ) {
                                        if !succ_block.phi_nodes().is_empty() {
                                            trampoline_blocks.insert(block_idx);
                                        }
                                    }
                                }
                            }
                        }
                    }
                }

                if !trampoline_blocks.is_empty() {
                    hoistable.retain(|(block_idx, _)| !trampoline_blocks.contains(*block_idx));
                }
            }

            if hoistable.is_empty() {
                continue;
            }

            // Collect instructions to hoist (we need to clone them before mutation)
            let mut to_hoist: Vec<(usize, usize, SsaOp)> = Vec::new();
            for (block_idx, instr_idx) in &hoistable {
                if let Some(block) = ssa.block(*block_idx) {
                    if let Some(instr) = block.instruction(*instr_idx) {
                        to_hoist.push((*block_idx, *instr_idx, instr.op().clone()));
                    }
                }
            }

            // Sort hoistable instructions by their dependency order.
            // Instructions must be hoisted in the order they were originally defined
            // to maintain correct dependencies. Sort by (block_idx, instr_idx).
            to_hoist.sort_by_key(|(block_idx, instr_idx, _)| (*block_idx, *instr_idx));

            // Find the insertion point in the preheader (before the terminator)
            let insert_base = if let Some(preheader_block) = ssa.block(preheader.index()) {
                let instrs = preheader_block.instructions();
                if instrs.is_empty() {
                    0
                } else if instrs.last().is_some_and(SsaInstruction::is_terminator) {
                    instrs.len().saturating_sub(1)
                } else {
                    instrs.len()
                }
            } else {
                0
            };

            // Track which source blocks had ALL non-terminator instructions hoisted.
            // These blocks become trampolines, and their successor phis need
            // predecessor updates from the source block to the preheader.
            let mut hoisted_from = BitSet::new(ssa.block_count());

            // Apply hoisting - insert all at once to maintain order
            for (i, (block_idx, instr_idx, op)) in to_hoist.iter().enumerate() {
                hoisted_from.insert(*block_idx);

                // Add to preheader
                if let Some(preheader_block) = ssa.block_mut(preheader.index()) {
                    let new_instr = SsaInstruction::synthetic(op.clone());
                    let instrs = preheader_block.instructions_mut();
                    instrs.insert(insert_base + i, new_instr);
                }

                // Remove from original location (replace with Nop)
                if let Some(block) = ssa.block_mut(*block_idx) {
                    if let Some(instr) = block.instructions_mut().get_mut(*instr_idx) {
                        instr.set_op(SsaOp::Nop);
                    }
                }

                total_hoisted += 1;
            }

            // Update phi operands at successor blocks. When all non-terminator
            // instructions were hoisted from a block to the preheader, the block
            // becomes a trampoline (just a Jump). Phi operands at successors that
            // referenced the source block now need to reference the preheader
            // (where the definitions live after hoisting).
            let preheader_idx = preheader.index();
            for source_block in hoisted_from.iter() {
                // Check if all non-terminator instructions were hoisted (block is now a trampoline)
                let is_trampoline = ssa.block(source_block).is_some_and(|b| {
                    b.instructions()
                        .iter()
                        .all(|i| i.is_terminator() || matches!(i.op(), SsaOp::Nop))
                });

                if !is_trampoline {
                    continue;
                }

                // Get successors of the source block
                let successors: Vec<usize> = ssa
                    .block(source_block)
                    .map(|b| {
                        b.instructions()
                            .last()
                            .map(|i| i.op().successors())
                            .unwrap_or_default()
                    })
                    .unwrap_or_default();

                for succ in successors {
                    if let Some(succ_block) = ssa.block_mut(succ) {
                        for phi in succ_block.phi_nodes_mut() {
                            for operand in phi.operands_mut() {
                                if operand.predecessor() == source_block {
                                    operand.set_predecessor(preheader_idx);
                                }
                            }
                        }
                    }
                }
            }
        }

        if total_hoisted > 0 {
            ctx.events
                .record(EventKind::InstructionRemoved)
                .at(method_token, 0)
                .message(format!(
                    "LICM: hoisted {total_hoisted} loop-invariant instructions"
                ));
        }

        Ok(total_hoisted > 0)
    }
}

/// Finds all loop-invariant instructions in a loop.
///
/// An instruction is loop-invariant if all its operands are:
/// - Defined outside the loop, OR
/// - Defined by loop-invariant instructions
///
/// IMPORTANT: PHI nodes at the loop HEADER define induction variables that change
/// each iteration. Instructions using these values are NOT loop-invariant.
fn find_loop_invariants(ssa: &SsaFunction, loop_info: &LoopInfo) -> Vec<(usize, usize)> {
    let mut invariants: HashSet<(usize, usize)> = HashSet::new();
    let mut invariant_defs = BitSet::new(ssa.var_id_capacity());

    // Collect PHI-defined variables from the loop HEADER only.
    // These are loop induction variables that change each iteration.
    // PHIs at other loop body blocks are path merge points and don't affect invariance.
    let mut header_phi_defs = BitSet::new(ssa.var_id_capacity());
    if let Some(header_block) = ssa.block(loop_info.header.index()) {
        for phi in header_block.phi_nodes() {
            header_phi_defs.insert(phi.result().index());
        }
    }

    // Build map of variables defined outside the loop
    let mut outside_defs = BitSet::new(ssa.var_id_capacity());
    for var in ssa.variables() {
        let def_site = var.def_site();
        if !loop_info.body.contains(def_site.block) {
            outside_defs.insert(var.id().index());
        }
    }

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

        for block_idx in loop_info.body.iter() {
            if let Some(block) = ssa.block(block_idx) {
                for (instr_idx, instr) in block.instructions().iter().enumerate() {
                    // Skip if already marked invariant
                    if invariants.contains(&(block_idx, instr_idx)) {
                        continue;
                    }

                    // Skip terminators
                    if instr.is_terminator() {
                        continue;
                    }

                    // Skip Nop instructions - they have no effect and hoisting them
                    // causes exponential blowup when processing nested loops
                    // (inner hoists create Nops which outer loops then re-hoist)
                    if matches!(instr.op(), SsaOp::Nop) {
                        continue;
                    }

                    // Check if instruction is invariant
                    if is_instruction_invariant(
                        instr,
                        &outside_defs,
                        &invariant_defs,
                        &header_phi_defs,
                    ) {
                        invariants.insert((block_idx, instr_idx));
                        if let Some(def) = instr.def() {
                            invariant_defs.insert(def.index());
                        }
                        changed = true;
                    }
                }
            }
        }
    }

    invariants.into_iter().collect()
}

/// Checks if an instruction is loop-invariant.
///
/// An instruction is NOT loop-invariant if it uses any loop header PHI-defined variable,
/// since those represent induction variables that change each iteration.
fn is_instruction_invariant(
    instr: &SsaInstruction,
    outside_defs: &BitSet,
    invariant_defs: &BitSet,
    header_phi_defs: &BitSet,
) -> bool {
    // Use the built-in uses() method to get all operands
    for operand in instr.op().uses() {
        // If the operand is defined by a PHI at the loop header, it's loop-varying
        if header_phi_defs.contains(operand.index()) {
            return false;
        }
        // Otherwise check if it's defined outside the loop or by an invariant instruction
        if !outside_defs.contains(operand.index()) && !invariant_defs.contains(operand.index()) {
            return false;
        }
    }

    true
}

/// Checks if an instruction can be safely hoisted.
fn can_hoist(ssa: &SsaFunction, loop_info: &LoopInfo, block_idx: usize, instr_idx: usize) -> bool {
    let Some(block) = ssa.block(block_idx) else {
        return false;
    };

    let Some(instr) = block.instruction(instr_idx) else {
        return false;
    };

    // Only hoist instructions that define a value - hoisting effectless
    // instructions (like Nop) is pointless and causes exponential blowup
    if instr.def().is_none() {
        return false;
    }

    // Only hoist pure computations (is_pure is defined on SsaOp)
    if !instr.op().is_pure() {
        return false;
    }

    // Don't hoist if there's no preheader
    if loop_info.preheader.is_none() {
        return false;
    }

    // CRITICAL: Don't hoist if this instruction's result feeds a PHI's back-edge operand.
    // Hoisting such instructions would make the PHI's back-edge operand orphaned or
    // self-referential, breaking the loop structure.
    if let Some(dest) = instr.def() {
        if feeds_phi_back_edge(ssa, loop_info, dest) {
            return false;
        }
    }

    true
}

/// Checks if a variable (directly or indirectly) feeds a PHI operand on an
/// intra-loop edge — i.e. a phi at any block in the loop body whose operand
/// comes from another loop body block.
///
/// Hoisting a def in this category is unsafe because the def's value is
/// attributed to a specific CFG edge (pred → phi-block). Moving it to the
/// shared preheader makes it dominate every intra-loop edge; when
/// `rebuild_ssa` recomputes phi operands by reaching definitions, every
/// such edge sees the hoisted value, collapsing per-edge attribution.
///
/// Typical break case: CFF state-machine dispatcher. Case blocks push
/// different `Const` values (state 5, 4, 0, 2, ...) into a phi at the
/// dispatcher. Without this guard LICM hoists all of them into the
/// single preheader, and SSA rebuild rewrites every phi operand to point
/// at whichever hoisted const happens to be on top of the version stack,
/// destroying the state machine.
fn feeds_phi_back_edge(ssa: &SsaFunction, loop_info: &LoopInfo, var: SsaVarId) -> bool {
    let mut worklist: VecDeque<SsaVarId> = VecDeque::new();
    let mut visited = BitSet::new(ssa.var_id_capacity());

    worklist.push_back(var);
    visited.insert(var.index());

    while let Some(current) = worklist.pop_front() {
        // Check phis at any block in the loop body (including the header).
        // An intra-loop edge is one where both the phi's block and the
        // operand's predecessor are in the loop body.
        for phi_block_idx in loop_info.body.iter() {
            let Some(phi_block) = ssa.block(phi_block_idx) else {
                continue;
            };
            for phi in phi_block.phi_nodes() {
                for operand in phi.operands() {
                    if operand.value() == current && loop_info.body.contains(operand.predecessor())
                    {
                        return true;
                    }
                }
            }
        }

        // Find instructions that use this variable and add their dests to the worklist
        for body_block_idx in loop_info.body.iter() {
            if let Some(body_block) = ssa.block(body_block_idx) {
                for instr in body_block.instructions() {
                    if instr.op().uses().contains(&current) {
                        if let Some(dest) = instr.def() {
                            if visited.insert(dest.index()) {
                                worklist.push_back(dest);
                            }
                        }
                    }
                }
            }
        }
    }

    false
}

#[cfg(test)]
mod tests {
    use crate::{
        analysis::{ConstValue, LoopAnalyzer, MethodRef, SsaFunctionBuilder, SsaOp, SsaVarId},
        compiler::{LicmPass, SsaPass},
        metadata::token::Token,
    };

    #[test]
    fn test_pass_metadata() {
        let pass = LicmPass::new();
        assert_eq!(pass.name(), "licm");
        assert!(!pass.description().is_empty());
    }

    #[test]
    fn test_op_is_pure() {
        let add_op = SsaOp::Add {
            dest: SsaVarId::from_index(0),
            left: SsaVarId::from_index(1),
            right: SsaVarId::from_index(2),
        };
        assert!(add_op.is_pure());

        let const_op = SsaOp::Const {
            dest: SsaVarId::from_index(3),
            value: ConstValue::I32(42),
        };
        assert!(const_op.is_pure());

        let call_op = SsaOp::Call {
            dest: Some(SsaVarId::from_index(4)),
            method: MethodRef::new(Token::new(0x06000001)),
            args: vec![],
        };
        assert!(!call_op.is_pure());
    }

    #[test]
    fn test_op_uses() {
        let v1 = SsaVarId::from_index(0);
        let v2 = SsaVarId::from_index(1);
        let dest = SsaVarId::from_index(2);

        let op = SsaOp::Add {
            dest,
            left: v1,
            right: v2,
        };
        let uses = op.uses();
        assert_eq!(uses.len(), 2);
        assert!(uses.contains(&v1));
        assert!(uses.contains(&v2));

        let const_op = SsaOp::Const {
            dest,
            value: ConstValue::I32(42),
        };
        assert!(const_op.uses().is_empty());
    }

    #[test]
    fn test_no_loops() {
        // Function with no loops should return false
        let ssa = SsaFunctionBuilder::new(0, 0)
            .build_with(|f| {
                f.block(0, |b| {
                    let _ = b.const_i32(42);
                    b.ret();
                });
            })
            .unwrap();

        let forest = LoopAnalyzer::new(&ssa).analyze();
        assert!(forest.is_empty());
    }

    #[test]
    fn test_loop_without_preheader() {
        // Loop without preheader (multiple entry edges) can't be optimized
        // This creates a function where the loop header has multiple predecessors
        let ssa = SsaFunctionBuilder::new(0, 0)
            .build_with(|f| {
                // B0: entry with branch to different blocks
                f.block(0, |b| {
                    let cond = b.const_true();
                    b.branch(cond, 1, 2);
                });
                // B1: goes to loop header
                f.block(1, |b| b.jump(3));
                // B2: also goes to loop header (no single preheader)
                f.block(2, |b| b.jump(3));
                // B3: loop header
                f.block(3, |b| {
                    let cond = b.const_true();
                    b.branch(cond, 3, 4); // self-loop
                });
                // B4: exit
                f.block(4, |b| b.ret());
            })
            .unwrap();

        let forest = LoopAnalyzer::new(&ssa).analyze();
        assert!(!forest.is_empty());

        let loop_info = &forest.loops()[0];
        // This loop has multiple entry edges so no preheader
        assert!(!loop_info.has_preheader());
    }

    #[test]
    fn test_simple_loop_has_preheader() {
        // Create a loop with a single preheader
        let ssa = SsaFunctionBuilder::new(0, 0)
            .build_with(|f| {
                // B0: preheader
                f.block(0, |b| b.jump(1));
                // B1: header with self-loop
                f.block(1, |b| {
                    let cond = b.const_true();
                    b.branch(cond, 1, 2);
                });
                // B2: exit
                f.block(2, |b| b.ret());
            })
            .unwrap();

        let forest = LoopAnalyzer::new(&ssa).analyze();
        assert_eq!(forest.len(), 1);

        let loop_info = &forest.loops()[0];
        assert!(loop_info.has_preheader());
    }
}