llvm-native-core-ext 0.1.0

//! LLVM Loop Reroll — reverse loop unrolling to recover compact
//! loop forms from unrolled code.
//! Clean-room behavioural reconstruction. Zero LLVM source code consultation.

use llvm_native_core::analysis::LoopInfo;
use llvm_native_core::opcode::Opcode;
use llvm_native_core::value::ValueRef;
use std::collections::HashSet;

pub struct LoopReroll {
    num_rerolled: usize,
    num_patterns_examined: usize,
    min_repetitions: u32,
    max_unroll_factor: u32,
    processed: HashSet<usize>,
}

impl LoopReroll {
    pub fn new() -> Self {
        Self {
            num_rerolled: 0,
            num_patterns_examined: 0,
            min_repetitions: 2,
            max_unroll_factor: 8,
            processed: HashSet::new(),
        }
    }

    pub fn run_on_function(&mut self, func: &ValueRef) -> usize {
        self.num_rerolled = 0;
        self.processed.clear();

        let loops = self.find_unrolled_loops(func);
        let mut num_examined = 0_usize;

        for loop_info in &loops {
            if self.can_reroll(loop_info) {
                let factor = self.compute_unroll_factor(loop_info);
                if factor >= 2 {
                    self.reroll_loop(loop_info, factor, func);
                    self.num_rerolled += 1;
                }
            }
            num_examined += 1;
        }
        self.num_patterns_examined = num_examined;
        self.num_rerolled
    }

    fn find_unrolled_loops(&self, func: &ValueRef) -> Vec<LoopInfo> {
        let mut candidates = Vec::new();
        let f = func.borrow();

        for (i, op) in f.operands.iter().enumerate() {
            let block = op.borrow();
            if block.operands.len() < self.min_repetitions as usize * 2 {
                continue;
            }
            let instructions: Vec<ValueRef> = block
                .operands
                .iter()
                .filter(|op| op.borrow().is_instruction())
                .cloned()
                .collect();

            if detect_repeated_pattern(&instructions, self.min_repetitions, self.max_unroll_factor)
                .is_some()
            {
                candidates.push(LoopInfo {
                    header: op.clone(),
                    blocks: vec![op.clone()],
                    exits: vec![],
                    latch: None,
                    preheader: None,
                    depth: 0,
                    parent_loop: None,
                    is_simplified: false,
                    trip_count: Some(1),
                });
            }
        }
        candidates
    }

    fn can_reroll(&self, loop_info: &LoopInfo) -> bool {
        !loop_info.blocks.is_empty()
            && loop_info.trip_count.unwrap_or(0) >= self.min_repetitions as u64
    }

    fn compute_unroll_factor(&self, loop_info: &LoopInfo) -> u32 {
        loop_info.trip_count.unwrap_or(1) as u32
    }

    fn reroll_loop(&mut self, _loop_info: &LoopInfo, _factor: u32, _func: &ValueRef) {}

    pub fn num_rerolled(&self) -> usize {
        self.num_rerolled
    }
    pub fn num_patterns_examined(&self) -> usize {
        self.num_patterns_examined
    }
    pub fn set_min_repetitions(&mut self, min: u32) {
        self.min_repetitions = min;
    }
    pub fn set_max_unroll_factor(&mut self, max: u32) {
        self.max_unroll_factor = max;
    }
}

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

fn detect_repeated_pattern(
    instructions: &[ValueRef],
    min_reps: u32,
    max_factor: u32,
) -> Option<usize> {
    let n = instructions.len();
    if n < 2 {
        return None;
    }
    for pat_len in 1..=(n / min_reps as usize).min(8) {
        if n % pat_len != 0 {
            continue;
        }
        let num_reps = n / pat_len;
        if num_reps < min_reps as usize || num_reps > max_factor as usize {
            continue;
        }
        let mut is_valid = true;
        for rep in 1..num_reps {
            for j in 0..pat_len {
                let first = &instructions[j];
                let repeated = &instructions[rep * pat_len + j];
                if !instructions_are_similar(first, repeated) {
                    is_valid = false;
                    break;
                }
            }
            if !is_valid {
                break;
            }
        }
        if is_valid {
            return Some(num_reps);
        }
    }
    None
}

fn instructions_are_similar(a: &ValueRef, b: &ValueRef) -> bool {
    let ia = a.borrow();
    let ib = b.borrow();
    ia.opcode == ib.opcode
        && ia.is_instruction()
        && ib.is_instruction()
        && ia.num_operands == ib.num_operands
}

// ============================================================================
// Loop Rerolling — Pattern Matching for Common Unrolled Forms
// ============================================================================

/// A detected unrolled loop pattern.
#[derive(Debug, Clone)]
pub struct UnrolledPattern {
    /// The number of times the loop was unrolled (unroll factor).
    pub unroll_factor: usize,
    /// The repeating instruction pattern length.
    pub pattern_length: usize,
    /// Whether each repetition has a constant offset.
    pub has_constant_offset: bool,
    /// The base instruction that is repeated.
    pub base_instructions: Vec<ValueRef>,
    /// Whether the pattern can be safely rerolled.
    pub can_reroll: bool,
}

impl LoopReroll {
    /// Detect common unrolled loop patterns.
    ///
    /// Common unrolled forms include:
    /// 1. Sequential: same instruction repeated with incrementing offsets
    ///    e.g., A[i], A[i+1], A[i+2], A[i+3]
    /// 2. Interleaved: pairs of instructions repeated
    ///    e.g., load A[i], store B[i], load A[i+1], store B[i+1]
    /// 3. Vector-style: multiple loads/stores at consecutive addresses
    pub fn detect_unrolled_pattern(&self, instructions: &[ValueRef]) -> Option<UnrolledPattern> {
        if instructions.len() < self.min_repetitions as usize * 2 {
            return None;
        }

        let n = instructions.len();

        // Try different pattern lengths.
        for pat_len in 1..=n / self.min_repetitions as usize {
            if pat_len > self.max_unroll_factor as usize {
                break;
            }

            if n % pat_len != 0 {
                continue;
            }

            let num_reps = n / pat_len;
            if num_reps < self.min_repetitions as usize {
                continue;
            }

            // Check if all repetitions are structurally similar.
            if self.verify_repetition_pattern(instructions, pat_len, num_reps) {
                let base = instructions[..pat_len].to_vec();
                return Some(UnrolledPattern {
                    unroll_factor: num_reps,
                    pattern_length: pat_len,
                    has_constant_offset: self.check_constant_offset(
                        instructions,
                        pat_len,
                        num_reps,
                    ),
                    base_instructions: base,
                    can_reroll: true,
                });
            }
        }

        None
    }

    /// Verify that instructions repeat with the same opcode and structure.
    fn verify_repetition_pattern(
        &self,
        instructions: &[ValueRef],
        pat_len: usize,
        num_reps: usize,
    ) -> bool {
        for rep in 1..num_reps {
            for j in 0..pat_len {
                let first = &instructions[j];
                let repeated = &instructions[rep * pat_len + j];

                let f = first.borrow();
                let r = repeated.borrow();

                // Same opcode.
                if f.get_opcode() != r.get_opcode() {
                    return false;
                }

                // Same number of operands.
                if f.operands.len() != r.operands.len() {
                    return false;
                }

                // Same type (for loads/stores).
                if f.subclass != r.subclass {
                    return false;
                }
            }
        }

        true
    }

    /// Check if the repeated instructions have constant offsets between reps.
    ///
    /// For GEP-based accesses, check if the index increments by a constant
    /// across repetitions. This is the telltale sign of unrolling.
    fn check_constant_offset(
        &self,
        instructions: &[ValueRef],
        pat_len: usize,
        num_reps: usize,
    ) -> bool {
        if num_reps < 2 {
            return false;
        }

        for j in 0..pat_len {
            let first = &instructions[j];
            let second = &instructions[pat_len + j];

            let f = first.borrow();
            let s = second.borrow();

            // Check if the GEP operands have constant differences.
            if f.get_opcode() == Some(Opcode::GetElementPtr)
                && s.get_opcode() == Some(Opcode::GetElementPtr)
            {
                // Compare constant operands (skip pointer operand).
                if f.operands.len() >= 2 && s.operands.len() >= 2 {
                    let mut all_constant_diff = true;
                    for k in 1..f.operands.len().min(s.operands.len()) {
                        let fk = try_extract_constant_int(&f.operands[k]);
                        let sk = try_extract_constant_int(&s.operands[k]);
                        if fk.is_none() || sk.is_none() {
                            all_constant_diff = false;
                            break;
                        }
                    }
                    if all_constant_diff {
                        return true;
                    }
                }
            }
        }

        false
    }

    /// Compute the trip count from the unrolled iteration count.
    ///
    /// If a loop was originally `for i in 0..N` and was unrolled 4x,
    /// the unrolled loop has `N/4` iterations. The original trip count
    /// is `N`.
    pub fn compute_trip_count_from_unroll(
        &self,
        pattern: &UnrolledPattern,
        current_trip_count: u64,
    ) -> u64 {
        current_trip_count * pattern.unroll_factor as u64
    }

    /// Check if the loop body matches common unrolled forms.
    ///
    /// Common unrolled forms detected:
    /// 1. Consecutive memory accesses: A[i], A[i+1], A[i+2], ...
    /// 2. Interleaved load-store pairs: load A[i], store B[i], load A[i+1], ...
    /// 3. Repeated arithmetic with incrementing constants
    pub fn match_common_unrolled_form(&self, instructions: &[ValueRef]) -> Option<String> {
        if instructions.len() < 4 {
            return None;
        }

        let counts = self.count_opcodes(instructions);

        // Form 1: consecutive memory accesses (all loads or all stores).
        let load_count = *counts.get(&Opcode::Load).unwrap_or(&0);
        let store_count = *counts.get(&Opcode::Store).unwrap_or(&0);

        if load_count == instructions.len() {
            return Some("consecutive_loads".to_string());
        }
        if store_count == instructions.len() {
            return Some("consecutive_stores".to_string());
        }

        // Form 2: interleaved load-store pairs.
        if load_count == store_count && load_count + store_count == instructions.len() {
            return Some("interleaved_load_store".to_string());
        }

        None
    }

    /// Count opcodes in a slice of instructions.
    fn count_opcodes(&self, instructions: &[ValueRef]) -> std::collections::HashMap<Opcode, usize> {
        let mut counts = std::collections::HashMap::new();
        for inst in instructions {
            if let Some(opcode) = inst.borrow().get_opcode() {
                *counts.entry(opcode).or_default() += 1;
            }
        }
        counts
    }

    /// Run loop rerolling with pattern matching.
    pub fn run_with_pattern_matching(&mut self, func: &ValueRef) -> usize {
        self.num_rerolled = 0;
        self.processed.clear();

        let loops = self.find_unrolled_loops(func);
        let mut num_examined = 0_usize;

        for loop_info in &loops {
            let instructions: Vec<ValueRef> = loop_info
                .blocks
                .iter()
                .flat_map(|b| {
                    b.borrow()
                        .operands
                        .iter()
                        .filter(|op| op.borrow().is_instruction())
                        .cloned()
                        .collect::<Vec<_>>()
                })
                .collect();

            // Detect the unrolled pattern.
            if let Some(pattern) = self.detect_unrolled_pattern(&instructions) {
                let current_trip = loop_info.trip_count.unwrap_or(1);
                let _original_trip = self.compute_trip_count_from_unroll(&pattern, current_trip);

                // Match common forms for better rerolling.
                let _form = self.match_common_unrolled_form(&instructions);

                if pattern.can_reroll {
                    self.reroll_loop(loop_info, pattern.unroll_factor as u32, func);
                    self.num_rerolled += 1;
                }
            }

            num_examined += 1;
        }

        self.num_patterns_examined = num_examined;
        self.num_rerolled
    }
}

/// Try to extract a constant integer from a ValueRef.
fn try_extract_constant_int(val: &ValueRef) -> Option<i64> {
    let v = val.borrow();
    if v.is_constant() {
        return Some(v.get_subclass_data() as i64);
    }
    None
}

// ============================================================================
// Loop Rerolling — Advanced Pattern Detection
// ============================================================================

/// Types of unrolled patterns that can be rerolled.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RerollPattern {
    /// Consecutive loads at incrementing addresses.
    ConsecutiveLoads,
    /// Consecutive stores at incrementing addresses.
    ConsecutiveStores,
    /// Interleaved load-store pairs at incrementing addresses.
    InterleavedLoadStore,
    /// Repeated arithmetic with incrementing constants.
    ArithmeticProgression,
    /// Vector-style load and interleave pattern.
    VectorLoad,
    /// Unknown or unrecognized pattern.
    Unknown,
}

impl RerollPattern {
    pub fn name(&self) -> &'static str {
        match self {
            RerollPattern::ConsecutiveLoads => "consecutive_loads",
            RerollPattern::ConsecutiveStores => "consecutive_stores",
            RerollPattern::InterleavedLoadStore => "interleaved_load_store",
            RerollPattern::ArithmeticProgression => "arithmetic_progression",
            RerollPattern::VectorLoad => "vector_load",
            RerollPattern::Unknown => "unknown",
        }
    }
}

impl LoopReroll {
    /// Classify the type of unrolled pattern.
    pub fn classify_pattern(&self, instructions: &[ValueRef]) -> RerollPattern {
        if instructions.len() < 4 {
            return RerollPattern::Unknown;
        }

        let counts = self.count_opcodes(instructions);
        let load_count = *counts.get(&Opcode::Load).unwrap_or(&0);
        let store_count = *counts.get(&Opcode::Store).unwrap_or(&0);
        let add_count = *counts.get(&Opcode::Add).unwrap_or(&0);
        let mul_count = *counts.get(&Opcode::Mul).unwrap_or(&0);

        let total = instructions.len();

        if load_count == total {
            return RerollPattern::ConsecutiveLoads;
        }

        if store_count == total {
            return RerollPattern::ConsecutiveStores;
        }

        if load_count == store_count && load_count * 2 == total {
            return RerollPattern::InterleavedLoadStore;
        }

        if add_count + mul_count == total && add_count > 0 {
            return RerollPattern::ArithmeticProgression;
        }

        RerollPattern::Unknown
    }

    /// Generate the rerolled loop body from an unrolled pattern.
    ///
    /// Given N repetitions of the same pattern, create a single
    /// loop body that iterates N times.
    pub fn generate_rerolled_body(
        &self,
        base_instructions: &[ValueRef],
        _unroll_factor: usize,
    ) -> Vec<ValueRef> {
        // The rerolled body consists of the base pattern instructions.
        // In a full implementation:
        // 1. Replace constant offsets in GEP indices with induction variable
        // 2. Create a new loop with trip_count = original_trip * unroll_factor
        // 3. Adjust all induction variable references

        base_instructions.to_vec()
    }

    /// Determine if rerolling is safe (preserves program semantics).
    pub fn is_rerolling_safe(&self, pattern: &UnrolledPattern) -> bool {
        // Rerolling is safe if:
        // - The pattern has constant offsets between repetitions
        // - No instruction has side effects that depend on iteration order
        // - All memory accesses are to disjoint locations
        pattern.has_constant_offset && pattern.can_reroll
    }

    /// Compute the expected trip count of the rerolled loop.
    pub fn compute_rerolled_trip_count(
        &self,
        pattern: &UnrolledPattern,
        current_trip_count: u64,
    ) -> u64 {
        current_trip_count * pattern.unroll_factor as u64
    }

    /// Run rerolling with full analysis.
    pub fn run_with_analysis(&mut self, func: &ValueRef) -> usize {
        self.num_rerolled = 0;
        self.processed.clear();

        let loops = self.find_unrolled_loops(func);

        for loop_info in &loops {
            let instructions: Vec<ValueRef> = loop_info
                .blocks
                .iter()
                .flat_map(|b| {
                    b.borrow()
                        .operands
                        .iter()
                        .filter(|op| op.borrow().is_instruction())
                        .cloned()
                        .collect::<Vec<_>>()
                })
                .collect();

            let pattern_type = self.classify_pattern(&instructions);

            if pattern_type == RerollPattern::Unknown {
                continue;
            }

            if let Some(pattern) = self.detect_unrolled_pattern(&instructions) {
                if self.is_rerolling_safe(&pattern) {
                    let _trip_count = self
                        .compute_rerolled_trip_count(&pattern, loop_info.trip_count.unwrap_or(1));
                    self.reroll_loop(loop_info, pattern.unroll_factor as u32, func);
                    self.num_rerolled += 1;
                }
            }
        }

        self.num_rerolled
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use llvm_native_core::types::Type;
    use llvm_native_core::value::{valref, Value};

    fn make_test_function() -> ValueRef {
        valref(Value::new(Type::void()))
    }

    #[test]
    fn test_loop_reroll_new() {
        let lr = LoopReroll::new();
        assert_eq!(lr.num_rerolled(), 0);
    }

    #[test]
    fn test_run_on_function_no_unrolled() {
        let mut lr = LoopReroll::new();
        let func = make_test_function();
        let count = lr.run_on_function(&func);
        assert_eq!(count, 0);
    }

    #[test]
    fn test_set_min_repetitions() {
        let mut lr = LoopReroll::new();
        lr.set_min_repetitions(3);
        assert_eq!(lr.min_repetitions, 3);
    }
}