jvmrs 0.1.2

//! Profile-Guided Optimization (PGO) support.
//!
//! Collects runtime profiling data (method hotness, branch frequencies,
//! type profiles) and uses it to guide JIT compilation decisions.

use crate::jit::CompilationLevel;
use std::collections::HashMap;

/// Branch outcome recorded at a conditional branch site
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BranchOutcome {
    Taken,
    NotTaken,
}

/// Type profile entry: receiver class seen at a call site
#[derive(Debug, Clone)]
pub struct TypeProfileEntry {
    /// Receiver class name
    pub class_name: String,
    /// Number of times this type was seen
    pub count: u64,
}

/// Profile data for a single call site
#[derive(Debug, Clone, Default)]
pub struct CallSiteProfile {
    /// Total invocations at this call site
    pub total_calls: u64,
    /// Type profiles (receiver classes seen)
    pub type_profiles: Vec<TypeProfileEntry>,
    /// Whether this site is monomorphic (single receiver type)
    pub is_monomorphic: bool,
    /// Whether this site is megamorphic (many receiver types)
    pub is_megamorphic: bool,
}

impl CallSiteProfile {
    const MEGAMORPHIC_THRESHOLD: usize = 4;

    /// Record a call with the given receiver class
    pub fn record_call(&mut self, receiver_class: &str) {
        self.total_calls += 1;

        if let Some(entry) = self
            .type_profiles
            .iter_mut()
            .find(|e| e.class_name == receiver_class)
        {
            entry.count += 1;
        } else {
            self.type_profiles.push(TypeProfileEntry {
                class_name: receiver_class.to_string(),
                count: 1,
            });
        }

        self.is_monomorphic = self.type_profiles.len() == 1;
        self.is_megamorphic = self.type_profiles.len() >= Self::MEGAMORPHIC_THRESHOLD;
    }

    /// Get the dominant receiver class (if monomorphic or bimorphic)
    pub fn dominant_type(&self) -> Option<&str> {
        self.type_profiles
            .iter()
            .max_by_key(|e| e.count)
            .map(|e| e.class_name.as_str())
    }

    /// Get the frequency of the dominant type (0.0 - 1.0)
    pub fn dominant_type_frequency(&self) -> f64 {
        if self.total_calls == 0 {
            return 0.0;
        }
        self.type_profiles
            .iter()
            .map(|e| e.count)
            .max()
            .map(|max| max as f64 / self.total_calls as f64)
            .unwrap_or(0.0)
    }
}

/// Branch profile at a specific bytecode offset
#[derive(Debug, Clone, Default)]
pub struct BranchProfile {
    /// Number of times the branch was taken
    pub taken_count: u64,
    /// Number of times the branch was not taken
    pub not_taken_count: u64,
}

impl BranchProfile {
    pub fn record(&mut self, outcome: BranchOutcome) {
        match outcome {
            BranchOutcome::Taken => self.taken_count += 1,
            BranchOutcome::NotTaken => self.not_taken_count += 1,
        }
    }

    pub fn total(&self) -> u64 {
        self.taken_count + self.not_taken_count
    }

    /// Probability that the branch is taken (0.0 - 1.0)
    pub fn taken_probability(&self) -> f64 {
        let total = self.total();
        if total == 0 {
            0.5
        } else {
            self.taken_count as f64 / total as f64
        }
    }

    /// True if the branch is strongly biased (>90% one way)
    pub fn is_biased(&self) -> bool {
        let p = self.taken_probability();
        p > 0.9 || p < 0.1
    }
}

/// PGO profile for a single method
#[derive(Debug, Clone, Default)]
pub struct MethodPgoProfile {
    /// Method invocation count
    pub invocation_count: u64,
    /// Branch profiles: bytecode_offset -> BranchProfile
    pub branch_profiles: HashMap<usize, BranchProfile>,
    /// Call site profiles: bytecode_offset -> CallSiteProfile
    pub call_site_profiles: HashMap<usize, CallSiteProfile>,
    /// Estimated "hotness" score (higher = more optimization priority)
    pub hotness_score: f64,
}

impl MethodPgoProfile {
    pub fn record_invocation(&mut self) {
        self.invocation_count += 1;
        self.update_hotness();
    }

    pub fn record_branch(&mut self, pc: usize, outcome: BranchOutcome) {
        self.branch_profiles.entry(pc).or_default().record(outcome);
    }

    pub fn record_call_site(&mut self, pc: usize, receiver_class: &str) {
        self.call_site_profiles
            .entry(pc)
            .or_default()
            .record_call(receiver_class);
    }

    fn update_hotness(&mut self) {
        // Hotness = invocations + weighted branch bias
        let branch_bonus: f64 = self
            .branch_profiles
            .values()
            .filter(|b| b.is_biased())
            .count() as f64
            * 10.0;
        self.hotness_score = self.invocation_count as f64 + branch_bonus;
    }

    /// Recommended compilation level based on profile
    pub fn recommended_level(&self) -> CompilationLevel {
        if self.invocation_count >= 10_000 {
            CompilationLevel::Optimized
        } else if self.invocation_count >= 100 {
            CompilationLevel::Baseline
        } else {
            CompilationLevel::Interpreter
        }
    }
}

/// PGO manager: collects and queries profile data
pub struct PgoManager {
    /// Per-method profiles
    profiles: HashMap<String, MethodPgoProfile>,
    /// Whether PGO collection is enabled
    pub enabled: bool,
    /// Whether to use collected profiles for optimization decisions
    pub use_profiles: bool,
}

impl PgoManager {
    pub fn new() -> Self {
        Self {
            profiles: HashMap::new(),
            enabled: true,
            use_profiles: true,
        }
    }

    fn method_key(class_name: &str, method_name: &str) -> String {
        format!("{}.{}", class_name, method_name)
    }

    /// Record a method invocation
    pub fn record_invocation(&mut self, class_name: &str, method_name: &str) {
        if !self.enabled {
            return;
        }
        let key = Self::method_key(class_name, method_name);
        self.profiles.entry(key).or_default().record_invocation();
    }

    /// Record a branch outcome
    pub fn record_branch(
        &mut self,
        class_name: &str,
        method_name: &str,
        pc: usize,
        outcome: BranchOutcome,
    ) {
        if !self.enabled {
            return;
        }
        let key = Self::method_key(class_name, method_name);
        self.profiles
            .entry(key)
            .or_default()
            .record_branch(pc, outcome);
    }

    /// Record a virtual call site receiver type
    pub fn record_call_site(
        &mut self,
        class_name: &str,
        method_name: &str,
        pc: usize,
        receiver_class: &str,
    ) {
        if !self.enabled {
            return;
        }
        let key = Self::method_key(class_name, method_name);
        self.profiles
            .entry(key)
            .or_default()
            .record_call_site(pc, receiver_class);
    }

    /// Get the profile for a method
    pub fn get_profile(&self, class_name: &str, method_name: &str) -> Option<&MethodPgoProfile> {
        let key = Self::method_key(class_name, method_name);
        self.profiles.get(&key)
    }

    /// Get recommended compilation level for a method
    pub fn recommended_level(
        &self,
        class_name: &str,
        method_name: &str,
    ) -> CompilationLevel {
        if !self.use_profiles {
            return CompilationLevel::Interpreter;
        }
        self.get_profile(class_name, method_name)
            .map(|p| p.recommended_level())
            .unwrap_or(CompilationLevel::Interpreter)
    }

    /// Get the hottest methods (sorted by hotness score, descending)
    pub fn hottest_methods(&self, top_n: usize) -> Vec<(&str, f64)> {
        let mut methods: Vec<(&str, f64)> = self
            .profiles
            .iter()
            .map(|(k, p)| (k.as_str(), p.hotness_score))
            .collect();
        methods.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
        methods.into_iter().take(top_n).collect()
    }

    /// Check if a call site is monomorphic (enables devirtualization)
    pub fn is_monomorphic_call_site(
        &self,
        class_name: &str,
        method_name: &str,
        pc: usize,
    ) -> bool {
        self.get_profile(class_name, method_name)
            .and_then(|p| p.call_site_profiles.get(&pc))
            .map(|cs| cs.is_monomorphic)
            .unwrap_or(false)
    }

    /// Get the dominant receiver type at a call site (for devirtualization)
    pub fn dominant_receiver(
        &self,
        class_name: &str,
        method_name: &str,
        pc: usize,
    ) -> Option<String> {
        self.get_profile(class_name, method_name)
            .and_then(|p| p.call_site_profiles.get(&pc))
            .and_then(|cs| cs.dominant_type())
            .map(|s| s.to_string())
    }

    /// Export profiles to a serializable format (for persistent PGO)
    pub fn export_summary(&self) -> Vec<PgoMethodSummary> {
        self.profiles
            .iter()
            .map(|(key, profile)| PgoMethodSummary {
                method_key: key.clone(),
                invocation_count: profile.invocation_count,
                hotness_score: profile.hotness_score,
                recommended_level: profile.recommended_level(),
                monomorphic_call_sites: profile
                    .call_site_profiles
                    .values()
                    .filter(|cs| cs.is_monomorphic)
                    .count(),
                biased_branches: profile
                    .branch_profiles
                    .values()
                    .filter(|b| b.is_biased())
                    .count(),
            })
            .collect()
    }

    /// Overall PGO statistics
    pub fn stats(&self) -> PgoStats {
        let total_methods = self.profiles.len();
        let optimized = self
            .profiles
            .values()
            .filter(|p| p.recommended_level() == CompilationLevel::Optimized)
            .count();
        let baseline = self
            .profiles
            .values()
            .filter(|p| p.recommended_level() == CompilationLevel::Baseline)
            .count();
        PgoStats {
            total_methods_profiled: total_methods,
            methods_at_optimized: optimized,
            methods_at_baseline: baseline,
            methods_at_interpreter: total_methods - optimized - baseline,
        }
    }
}

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

/// Summary of PGO data for a single method
#[derive(Debug, Clone)]
pub struct PgoMethodSummary {
    pub method_key: String,
    pub invocation_count: u64,
    pub hotness_score: f64,
    pub recommended_level: CompilationLevel,
    pub monomorphic_call_sites: usize,
    pub biased_branches: usize,
}

/// Aggregate PGO statistics
#[derive(Debug, Clone)]
pub struct PgoStats {
    pub total_methods_profiled: usize,
    pub methods_at_optimized: usize,
    pub methods_at_baseline: usize,
    pub methods_at_interpreter: usize,
}

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

    #[test]
    fn test_pgo_manager_creation() {
        let pgo = PgoManager::new();
        assert!(pgo.enabled);
        assert!(pgo.use_profiles);
    }

    #[test]
    fn test_record_invocation() {
        let mut pgo = PgoManager::new();
        pgo.record_invocation("Test", "method");
        let profile = pgo.get_profile("Test", "method").unwrap();
        assert_eq!(profile.invocation_count, 1);
    }

    #[test]
    fn test_recommended_level_interpreter() {
        let mut pgo = PgoManager::new();
        pgo.record_invocation("Test", "cold");
        assert_eq!(
            pgo.recommended_level("Test", "cold"),
            CompilationLevel::Interpreter
        );
    }

    #[test]
    fn test_recommended_level_baseline() {
        let mut pgo = PgoManager::new();
        for _ in 0..200 {
            pgo.record_invocation("Test", "warm");
        }
        assert_eq!(
            pgo.recommended_level("Test", "warm"),
            CompilationLevel::Baseline
        );
    }

    #[test]
    fn test_recommended_level_optimized() {
        let mut pgo = PgoManager::new();
        for _ in 0..15_000 {
            pgo.record_invocation("Test", "hot");
        }
        assert_eq!(
            pgo.recommended_level("Test", "hot"),
            CompilationLevel::Optimized
        );
    }

    #[test]
    fn test_branch_profile() {
        let mut pgo = PgoManager::new();
        for _ in 0..95 {
            pgo.record_branch("Test", "m", 10, BranchOutcome::Taken);
        }
        for _ in 0..5 {
            pgo.record_branch("Test", "m", 10, BranchOutcome::NotTaken);
        }
        let profile = pgo.get_profile("Test", "m").unwrap();
        let branch = &profile.branch_profiles[&10];
        assert!(branch.is_biased());
        assert!(branch.taken_probability() > 0.9);
    }

    #[test]
    fn test_call_site_monomorphic() {
        let mut pgo = PgoManager::new();
        for _ in 0..100 {
            pgo.record_call_site("Test", "m", 20, "java/lang/String");
        }
        assert!(pgo.is_monomorphic_call_site("Test", "m", 20));
        assert_eq!(
            pgo.dominant_receiver("Test", "m", 20),
            Some("java/lang/String".to_string())
        );
    }

    #[test]
    fn test_call_site_megamorphic() {
        let mut pgo = PgoManager::new();
        for i in 0..5 {
            pgo.record_call_site("Test", "m", 30, &format!("Class{}", i));
        }
        let profile = pgo.get_profile("Test", "m").unwrap();
        assert!(profile.call_site_profiles[&30].is_megamorphic);
    }

    #[test]
    fn test_hottest_methods() {
        let mut pgo = PgoManager::new();
        for _ in 0..1000 {
            pgo.record_invocation("Test", "hot");
        }
        for _ in 0..10 {
            pgo.record_invocation("Test", "cold");
        }
        let hottest = pgo.hottest_methods(1);
        assert_eq!(hottest.len(), 1);
        assert_eq!(hottest[0].0, "Test.hot");
    }
}