jvmrs 0.1.2

//! Escape analysis for stack allocation optimization.
//!
//! Determines whether objects allocated in a method can be proven to not
//! escape the method's scope, enabling stack allocation instead of heap
//! allocation to reduce GC pressure and improve cache locality.

use crate::class_file::{ClassFile, MethodInfo};
use std::collections::{HashMap, HashSet};

/// Escape state of an object allocation
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum EscapeState {
    /// Object does not escape the method - can be stack-allocated
    NoEscape,
    /// Object escapes to the calling method only (arg escape)
    ArgEscape,
    /// Object escapes globally (heap allocation required)
    GlobalEscape,
}

/// An allocation site identified by bytecode offset
#[derive(Debug, Clone)]
pub struct AllocationSite {
    /// Bytecode offset of the `new` instruction
    pub pc: usize,
    /// Class being allocated
    pub class_name: String,
    /// Escape state determined by analysis
    pub escape_state: EscapeState,
    /// Whether this site is eligible for stack allocation
    pub stack_allocatable: bool,
}

/// Result of escape analysis for a single method
#[derive(Debug, Clone)]
pub struct EscapeAnalysisResult {
    /// Method identifier
    pub method_key: String,
    /// All allocation sites found in the method
    pub allocation_sites: Vec<AllocationSite>,
    /// Number of sites eligible for stack allocation
    pub stack_allocatable_count: usize,
    /// Number of sites that must remain on heap
    pub heap_allocated_count: usize,
}

impl EscapeAnalysisResult {
    /// Returns true if any allocation can be stack-allocated
    pub fn has_stack_allocatable(&self) -> bool {
        self.stack_allocatable_count > 0
    }

    /// Get all stack-allocatable sites
    pub fn stack_allocatable_sites(&self) -> impl Iterator<Item = &AllocationSite> {
        self.allocation_sites
            .iter()
            .filter(|s| s.stack_allocatable)
    }
}

/// Escape analysis engine
///
/// Performs a conservative flow-insensitive escape analysis over JVM bytecode.
/// Objects are assumed to escape unless proven otherwise.
pub struct EscapeAnalyzer {
    /// Cache of analysis results per method
    cache: HashMap<String, EscapeAnalysisResult>,
    /// Whether analysis is enabled
    pub enabled: bool,
}

impl EscapeAnalyzer {
    pub fn new() -> Self {
        Self {
            cache: HashMap::new(),
            enabled: true,
        }
    }

    /// Analyze a method for escape information.
    /// Returns cached result if already analyzed.
    pub fn analyze(
        &mut self,
        class: &ClassFile,
        method: &MethodInfo,
    ) -> &EscapeAnalysisResult {
        let class_name = class
            .get_class_name()
            .unwrap_or_else(|| "Unknown".to_string());
        let method_name = class
            .get_string(method.name_index)
            .unwrap_or_else(|| "unknown".to_string());
        let key = format!("{}.{}", class_name, method_name);

        if !self.cache.contains_key(&key) {
            let result = self.run_analysis(class, method, &key);
            self.cache.insert(key.clone(), result);
        }

        self.cache.get(&key).unwrap()
    }

    fn run_analysis(
        &self,
        class: &ClassFile,
        method: &MethodInfo,
        key: &str,
    ) -> EscapeAnalysisResult {
        let bytecode = Self::extract_bytecode(method);
        let mut allocation_sites: Vec<AllocationSite> = Vec::new();

        // Track which local variable slots hold newly allocated objects
        // and whether those slots are later stored to fields, returned, or passed
        // to other methods (all of which cause escape).
        let mut allocated_locals: HashSet<usize> = HashSet::new();
        let mut escaped_locals: HashSet<usize> = HashSet::new();

        // Stack simulation: track which stack positions hold new allocations
        // (simplified: we track a set of "tainted" stack depth positions)
        let mut tainted_stack: Vec<bool> = Vec::new();

        let mut pc = 0usize;
        while pc < bytecode.len() {
            let opcode = bytecode[pc];
            pc += 1;

            match opcode {
                // new: allocate object - mark top of stack as tainted
                0xbb => {
                    if pc + 1 < bytecode.len() {
                        let cp_index = ((bytecode[pc] as u16) << 8) | (bytecode[pc + 1] as u16);
                        pc += 2;
                        let alloc_class = class
                            .get_class_name_from_index(cp_index)
                            .unwrap_or_else(|| "Unknown".to_string());
                        let site_pc = pc - 3;
                        allocation_sites.push(AllocationSite {
                            pc: site_pc,
                            class_name: alloc_class,
                            escape_state: EscapeState::NoEscape, // optimistic
                            stack_allocatable: true,
                        });
                        tainted_stack.push(true);
                    }
                }

                // astore_0..3: store reference into local
                0x4b..=0x4e => {
                    let local_idx = (opcode - 0x4b) as usize;
                    let tainted = tainted_stack.pop().unwrap_or(false);
                    if tainted {
                        allocated_locals.insert(local_idx);
                    }
                }

                // astore <index>: store reference into local
                0x3a => {
                    if pc < bytecode.len() {
                        let local_idx = bytecode[pc] as usize;
                        pc += 1;
                        let tainted = tainted_stack.pop().unwrap_or(false);
                        if tainted {
                            allocated_locals.insert(local_idx);
                        }
                    }
                }

                // aload_0..3: load reference from local
                0x2a..=0x2d => {
                    let local_idx = (opcode - 0x2a) as usize;
                    tainted_stack.push(allocated_locals.contains(&local_idx));
                }

                // aload <index>
                0x19 => {
                    if pc < bytecode.len() {
                        let local_idx = bytecode[pc] as usize;
                        pc += 1;
                        tainted_stack.push(allocated_locals.contains(&local_idx));
                    }
                }

                // putfield: store to object field - top of stack escapes
                0xb5 => {
                    pc += 2; // skip field ref
                    // value being stored (top) may escape into the object
                    let value_tainted = tainted_stack.pop().unwrap_or(false);
                    let _obj_tainted = tainted_stack.pop().unwrap_or(false);
                    if value_tainted {
                        // Mark all tainted locals as escaped (conservative)
                        for &local in &allocated_locals {
                            escaped_locals.insert(local);
                        }
                    }
                }

                // putstatic: store to static field - always escapes
                0xb3 => {
                    pc += 2;
                    let tainted = tainted_stack.pop().unwrap_or(false);
                    if tainted {
                        for &local in &allocated_locals {
                            escaped_locals.insert(local);
                        }
                    }
                }

                // areturn: return reference - escapes to caller
                0xb0 => {
                    let tainted = tainted_stack.pop().unwrap_or(false);
                    if tainted {
                        for &local in &allocated_locals {
                            escaped_locals.insert(local);
                        }
                    }
                }

                // invokevirtual, invokespecial, invokestatic, invokeinterface:
                // arguments passed to other methods escape (conservative)
                0xb6 | 0xb7 | 0xb8 | 0xb9 => {
                    pc += 2;
                    if opcode == 0xb9 {
                        pc += 2; // invokeinterface has 2 extra bytes
                    }
                    // Mark any tainted values on stack as escaped
                    for tainted in &tainted_stack {
                        if *tainted {
                            for &local in &allocated_locals {
                                escaped_locals.insert(local);
                            }
                            break;
                        }
                    }
                    // Conservatively clear tainted stack after call
                    tainted_stack.clear();
                    // Push non-tainted return value placeholder
                    tainted_stack.push(false);
                }

                // dup: duplicate top of stack
                0x59 => {
                    let top = tainted_stack.last().copied().unwrap_or(false);
                    tainted_stack.push(top);
                }

                // pop: discard top
                0x57 => {
                    tainted_stack.pop();
                }

                // pop2
                0x58 => {
                    tainted_stack.pop();
                    tainted_stack.pop();
                }

                // Wide prefix
                0xc4 => {
                    if pc < bytecode.len() {
                        let wide_op = bytecode[pc];
                        pc += 1;
                        match wide_op {
                            0x19 | 0x3a => pc += 2, // wide aload/astore
                            _ => pc += 2,
                        }
                    }
                }

                // Branch instructions - skip offset bytes
                0x99..=0xa8 => {
                    pc += 2;
                }
                0xc6 | 0xc7 => {
                    pc += 2; // ifnull, ifnonnull
                }
                0xa9 => {
                    pc += 1; // ret
                }
                0xaa => {
                    // tableswitch - variable length
                    let padding = (4 - (pc % 4)) % 4;
                    pc += padding + 12; // default + low + high
                    if pc < bytecode.len() {
                        // skip entries (simplified: just advance past padding)
                    }
                }
                0xab => {
                    // lookupswitch - variable length
                    let padding = (4 - (pc % 4)) % 4;
                    pc += padding + 8; // default + npairs
                }

                // Default: push a non-tainted placeholder for instructions
                // that produce values (simplified)
                _ => {
                    // For simplicity, push false for any value-producing instruction
                    // A full implementation would track each opcode's stack effect
                }
            }
        }

        // Update escape states based on analysis
        for site in &mut allocation_sites {
            // Check if any local holding this allocation escaped
            // (simplified: if any local escaped, mark all sites as escaped)
            if !escaped_locals.is_empty() {
                site.escape_state = EscapeState::GlobalEscape;
                site.stack_allocatable = false;
            }
        }

        let stack_allocatable_count = allocation_sites
            .iter()
            .filter(|s| s.stack_allocatable)
            .count();
        let heap_allocated_count = allocation_sites.len() - stack_allocatable_count;

        EscapeAnalysisResult {
            method_key: key.to_string(),
            allocation_sites,
            stack_allocatable_count,
            heap_allocated_count,
        }
    }

    fn extract_bytecode(method: &MethodInfo) -> Vec<u8> {
        for attr in &method.attributes {
            if attr.info.len() >= 8 {
                let code_len = ((attr.info[4] as usize) << 24)
                    | ((attr.info[5] as usize) << 16)
                    | ((attr.info[6] as usize) << 8)
                    | (attr.info[7] as usize);
                if let Some(code) = attr.info.get(8..8 + code_len) {
                    return code.to_vec();
                }
            }
        }
        Vec::new()
    }

    /// Get cached result for a method key
    pub fn get_cached(&self, key: &str) -> Option<&EscapeAnalysisResult> {
        self.cache.get(key)
    }

    /// Clear the analysis cache
    pub fn clear_cache(&mut self) {
        self.cache.clear();
    }

    /// Statistics across all analyzed methods
    pub fn stats(&self) -> EscapeAnalysisStats {
        let total_sites: usize = self.cache.values().map(|r| r.allocation_sites.len()).sum();
        let stack_allocatable: usize = self
            .cache
            .values()
            .map(|r| r.stack_allocatable_count)
            .sum();
        EscapeAnalysisStats {
            methods_analyzed: self.cache.len(),
            total_allocation_sites: total_sites,
            stack_allocatable_sites: stack_allocatable,
            heap_allocated_sites: total_sites - stack_allocatable,
            stack_allocation_rate: if total_sites > 0 {
                (stack_allocatable as f64 / total_sites as f64) * 100.0
            } else {
                0.0
            },
        }
    }
}

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

/// Aggregate statistics from escape analysis
#[derive(Debug, Clone)]
pub struct EscapeAnalysisStats {
    pub methods_analyzed: usize,
    pub total_allocation_sites: usize,
    pub stack_allocatable_sites: usize,
    pub heap_allocated_sites: usize,
    /// Percentage of allocations that can be stack-allocated
    pub stack_allocation_rate: f64,
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_escape_analyzer_creation() {
        let analyzer = EscapeAnalyzer::new();
        assert!(analyzer.enabled);
        assert_eq!(analyzer.cache.len(), 0);
    }

    #[test]
    fn test_escape_state_ordering() {
        assert_eq!(EscapeState::NoEscape, EscapeState::NoEscape);
        assert_ne!(EscapeState::NoEscape, EscapeState::GlobalEscape);
    }

    #[test]
    fn test_escape_analysis_stats_empty() {
        let analyzer = EscapeAnalyzer::new();
        let stats = analyzer.stats();
        assert_eq!(stats.methods_analyzed, 0);
        assert_eq!(stats.total_allocation_sites, 0);
        assert_eq!(stats.stack_allocation_rate, 0.0);
    }
}