lambdust 0.1.1

//! Profile-guided optimization for adaptive JIT compilation
//!
//! This module implements a sophisticated profile-guided optimization system that
//! uses runtime profiling data to make intelligent optimization decisions. The system
//! adapts to changing execution patterns and optimizes code based on actual usage.

use crate::ast::Expr;
use crate::diagnostics::{Result, Error};
use crate::jit::ExecutionProfile;
use std::collections::HashMap;
use std::time::Duration;

/// Configuration for profile-guided optimization
#[derive(Debug, Clone)]
pub struct PgoConfig {
    /// Enable adaptive optimization based on runtime profiles
    pub adaptive_optimization: bool,
    
    /// Minimum profile data required before optimization
    pub min_profile_samples: u64,
    
    /// Profile data retention period
    pub profile_retention_period: Duration,
    
    /// Enable type feedback collection
    pub type_feedback: bool,
    
    /// Enable branch profiling
    pub branch_profiling: bool,
    
    /// Enable memory access pattern analysis
    pub memory_access_profiling: bool,
}

impl From<crate::jit::config::ProfileGuidedOptimizerConfig> for PgoConfig {
    fn from(config: crate::jit::config::ProfileGuidedOptimizerConfig) -> Self {
        Self {
            adaptive_optimization: config.enabled,
            min_profile_samples: config.min_profile_executions,
            profile_retention_period: config.profile_retention_time,
            type_feedback: config.enable_type_profiling,
            branch_profiling: config.enable_branch_profiling,
            memory_access_profiling: config.enable_memory_profiling,
        }
    }
}

impl Default for PgoConfig {
    fn default() -> Self {
        Self {
            adaptive_optimization: true,
            min_profile_samples: 50,
            profile_retention_period: Duration::from_secs(3600), // 1 hour
            type_feedback: true,
            branch_profiling: true,
            memory_access_profiling: true,
        }
    }
}

/// Runtime profile data for expressions
#[derive(Debug, Clone)]
pub struct RuntimeProfile {
    /// Type feedback data
    pub type_feedback: TypeFeedback,
    
    /// Branch prediction data
    pub branch_data: BranchProfile,
    
    /// Memory access patterns
    pub memory_access: MemoryAccessProfile,
    
    /// Performance counters
    pub performance_counters: PerformanceCounters,
}

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

impl RuntimeProfile {
    /// Creates a new runtime profile with default values
    pub fn new() -> Self {
        Self {
            type_feedback: TypeFeedback::new(),
            branch_data: BranchProfile::new(),
            memory_access: MemoryAccessProfile::new(),
            performance_counters: PerformanceCounters::new(),
        }
    }
}

/// Type feedback for dynamic optimization
#[derive(Debug, Clone)]
pub struct TypeFeedback {
    /// Observed types for variable accesses
    pub variable_types: HashMap<String, Vec<TypeObservation>>,
    
    /// Observed types for function arguments
    pub argument_types: HashMap<String, Vec<Vec<TypeObservation>>>,
    
    /// Observed return types
    pub return_types: HashMap<String, Vec<TypeObservation>>,
}

impl TypeFeedback {
    fn new() -> Self {
        Self {
            variable_types: HashMap::new(),
            argument_types: HashMap::new(),
            return_types: HashMap::new(),
        }
    }
}

/// Type observation with frequency data
#[derive(Debug, Clone)]
pub struct TypeObservation {
    /// The observed type
    pub type_info: crate::jit::code_generator::SchemeType,
    
    /// Number of times this type was observed
    pub frequency: u64,
    
    /// Percentage of total observations
    pub percentage: f64,
}

/// Branch profiling data for optimization
#[derive(Debug, Clone)]
pub struct BranchProfile {
    /// Branch taken/not-taken statistics
    pub branch_stats: HashMap<String, BranchStatistics>,
    
    /// Most frequently taken branches
    pub hot_branches: Vec<String>,
    
    /// Branch prediction accuracy
    pub prediction_accuracy: f64,
}

impl BranchProfile {
    fn new() -> Self {
        Self {
            branch_stats: HashMap::new(),
            hot_branches: Vec::new(),
            prediction_accuracy: 0.0,
        }
    }
}

/// Statistics for individual branches
#[derive(Debug, Clone)]
pub struct BranchStatistics {
    /// Number of times branch was taken
    pub taken_count: u64,
    
    /// Number of times branch was not taken
    pub not_taken_count: u64,
    
    /// Branch taken percentage
    pub taken_percentage: f64,
}

/// Memory access pattern profiling
#[derive(Debug, Clone)]
pub struct MemoryAccessProfile {
    /// Frequently accessed memory regions
    pub hot_memory_regions: Vec<MemoryRegion>,
    
    /// Cache miss statistics
    pub cache_miss_data: CacheMissData,
    
    /// Memory allocation patterns
    pub allocation_patterns: AllocationPatterns,
}

impl MemoryAccessProfile {
    fn new() -> Self {
        Self {
            hot_memory_regions: Vec::new(),
            cache_miss_data: CacheMissData::new(),
            allocation_patterns: AllocationPatterns::new(),
        }
    }
}

/// Memory region access information
#[derive(Debug, Clone)]
pub struct MemoryRegion {
    /// Start address of the region
    pub start_address: usize,
    
    /// Size of the region
    pub size: usize,
    
    /// Access frequency
    pub access_count: u64,
    
    /// Access pattern (sequential, random, etc.)
    pub access_pattern: AccessPattern,
}

/// Memory access patterns
#[derive(Debug, Clone, PartialEq)]
pub enum AccessPattern {
    /// Sequential memory access pattern
    Sequential,
    /// Random memory access pattern
    Random,
    /// Strided memory access pattern with fixed stride
    Strided { 
        /// The stride size in bytes
        stride: usize 
    },
    /// Clustered memory access pattern
    Clustered,
}

/// Cache miss statistics
#[derive(Debug, Clone)]
pub struct CacheMissData {
    /// L1 cache miss rate
    pub l1_miss_rate: f64,
    
    /// L2 cache miss rate
    pub l2_miss_rate: f64,
    
    /// TLB miss rate
    pub tlb_miss_rate: f64,
}

impl CacheMissData {
    fn new() -> Self {
        Self {
            l1_miss_rate: 0.0,
            l2_miss_rate: 0.0,
            tlb_miss_rate: 0.0,
        }
    }
}

/// Memory allocation pattern analysis
#[derive(Debug, Clone)]
pub struct AllocationPatterns {
    /// Frequently allocated object sizes
    pub common_sizes: Vec<(usize, u64)>,
    
    /// Allocation frequency
    pub allocation_rate: f64,
    
    /// Average object lifetime
    pub avg_object_lifetime: Duration,
}

impl AllocationPatterns {
    fn new() -> Self {
        Self {
            common_sizes: Vec::new(),
            allocation_rate: 0.0,
            avg_object_lifetime: Duration::ZERO,
        }
    }
}

/// Performance counters from CPU
#[derive(Debug, Clone)]
pub struct PerformanceCounters {
    /// Instructions per cycle
    pub ipc: f64,
    
    /// Branch misprediction rate
    pub branch_misprediction_rate: f64,
    
    /// Cache miss rates
    pub cache_miss_rates: HashMap<String, f64>,
    
    /// CPU utilization
    pub cpu_utilization: f64,
}

impl PerformanceCounters {
    fn new() -> Self {
        Self {
            ipc: 0.0,
            branch_misprediction_rate: 0.0,
            cache_miss_rates: HashMap::new(),
            cpu_utilization: 0.0,
        }
    }
}

/// Adaptive optimization decisions
#[derive(Debug, Clone)]
pub struct AdaptiveOptimization {
    /// Optimizations to apply
    pub optimizations: Vec<OptimizationDecision>,
    
    /// Confidence level in these decisions
    pub confidence: f64,
    
    /// Expected performance improvement
    pub expected_improvement: f64,
}

/// Individual optimization decision
#[derive(Debug, Clone)]
pub struct OptimizationDecision {
    /// Type of optimization
    pub optimization_type: OptimizationType,
    
    /// Target code location
    pub target: String,
    
    /// Optimization parameters
    pub parameters: HashMap<String, OptimizationParameter>,
    
    /// Expected benefit
    pub expected_benefit: f64,
}

/// Types of adaptive optimizations
#[derive(Debug, Clone, PartialEq)]
pub enum OptimizationType {
    /// Type specialization based on feedback
    TypeSpecialization,
    
    /// Branch layout optimization
    BranchOptimization,
    
    /// Memory prefetching
    MemoryPrefetching,
    
    /// Loop optimization
    LoopOptimization,
    
    /// Function inlining
    FunctionInlining,
    
    /// SIMD vectorization
    SIMDVectorization,
}

/// Optimization parameters
#[derive(Debug, Clone)]
pub enum OptimizationParameter {
    /// Integer parameter value
    Integer(i64),
    /// Floating-point parameter value
    Float(f64),
    /// String parameter value
    String(String),
    /// Boolean parameter value
    Boolean(bool),
}

/// Profile-guided optimizer
pub struct ProfileGuidedOptimizer {
    /// Configuration
    config: PgoConfig,
    
    /// Runtime profiles by expression
    profiles: HashMap<String, RuntimeProfile>,
    
    /// Optimization history
    optimization_history: Vec<OptimizationHistory>,
    
    /// Statistics
    stats: PgoStats,
}

impl ProfileGuidedOptimizer {
    /// Creates a new profile-guided optimizer
    pub fn new(config: PgoConfig) -> Result<Self> {
        Ok(Self {
            config,
            profiles: HashMap::new(),
            optimization_history: Vec::new(),
            stats: PgoStats::default(),
        })
    }
    
    /// Optimizes an expression based on runtime profile
    pub fn optimize_expression(&mut self, expr: &Expr, profile: &ExecutionProfile) -> Result<Expr> {
        if !self.config.adaptive_optimization {
            return Ok(expr.clone());
        }
        
        // Check if we have sufficient profile data
        if profile.execution_count < self.config.min_profile_samples {
            return Ok(expr.clone());
        }
        
        let expr_key = self.expression_key(expr);
        
        // Get or create runtime profile
        let runtime_profile = self.profiles.entry(expr_key.clone())
            .or_default();
        
        // Clone the runtime profile to avoid borrow issues
        let runtime_profile_clone = runtime_profile.clone();
        
        // Analyze profile and make optimization decisions
        let optimization_decisions = self.analyze_profile(expr, profile, &runtime_profile_clone)?;
        
        // Apply optimizations
        let optimized_expr = self.apply_optimizations(expr, &optimization_decisions)?;
        
        // Record optimization history
        self.record_optimization_history(expr_key, optimization_decisions);
        
        self.stats.expressions_optimized += 1;
        
        Ok(optimized_expr)
    }
    
    /// Analyzes runtime profile to make optimization decisions
    fn analyze_profile(&self, expr: &Expr, execution_profile: &ExecutionProfile, 
                      runtime_profile: &RuntimeProfile) -> Result<Vec<OptimizationDecision>> {
        let mut decisions = Vec::new();
        
        // Type specialization analysis
        if self.config.type_feedback {
            if let Some(type_decision) = self.analyze_type_feedback(expr, runtime_profile)? {
                decisions.push(type_decision);
            }
        }
        
        // Branch optimization analysis
        if self.config.branch_profiling {
            if let Some(branch_decision) = self.analyze_branch_profile(expr, &runtime_profile.branch_data)? {
                decisions.push(branch_decision);
            }
        }
        
        // Memory access optimization analysis
        if self.config.memory_access_profiling {
            if let Some(memory_decision) = self.analyze_memory_access(expr, &runtime_profile.memory_access)? {
                decisions.push(memory_decision);
            }
        }
        
        // Performance counter analysis
        if let Some(perf_decision) = self.analyze_performance_counters(expr, &runtime_profile.performance_counters)? {
            decisions.push(perf_decision);
        }
        
        Ok(decisions)
    }
    
    /// Analyzes type feedback for specialization opportunities
    fn analyze_type_feedback(&self, expr: &Expr, runtime_profile: &RuntimeProfile) -> Result<Option<OptimizationDecision>> {
        // Look for opportunities to specialize based on observed types
        // For example, if a variable is always an integer, generate specialized integer code
        
        if let Expr::Symbol(var_name) = expr {
            if let Some(type_observations) = runtime_profile.type_feedback.variable_types.get(var_name) {
                // Check if we have a dominant type (>80% of observations)
                for observation in type_observations {
                    if observation.percentage > 0.8 {
                        let mut parameters = HashMap::new();
                        parameters.insert("target_type".to_string(), 
                                        OptimizationParameter::String(format!("{:?}", observation.type_info)));
                        
                        return Ok(Some(OptimizationDecision {
                            optimization_type: OptimizationType::TypeSpecialization,
                            target: var_name.clone(),
                            parameters,
                            expected_benefit: observation.percentage * 2.0, // Simplified benefit calculation
                        }));
                    }
                }
            }
        }
        
        Ok(None)
    }
    
    /// Analyzes branch profile for layout optimization
    fn analyze_branch_profile(&self, expr: &Expr, branch_profile: &BranchProfile) -> Result<Option<OptimizationDecision>> {
        let expr_key = self.expression_key(expr);
        
        if let Some(branch_stats) = branch_profile.branch_stats.get(&expr_key) {
            // If branch is heavily biased, optimize for the common case
            if branch_stats.taken_percentage > 0.9 || branch_stats.taken_percentage < 0.1 {
                let mut parameters = HashMap::new();
                parameters.insert("likely_taken".to_string(), 
                                OptimizationParameter::Boolean(branch_stats.taken_percentage > 0.5));
                
                return Ok(Some(OptimizationDecision {
                    optimization_type: OptimizationType::BranchOptimization,
                    target: expr_key,
                    parameters,
                    expected_benefit: (branch_stats.taken_percentage - 0.5).abs() * 4.0,
                }));
            }
        }
        
        Ok(None)
    }
    
    /// Analyzes memory access patterns for prefetching opportunities
    fn analyze_memory_access(&self, expr: &Expr, memory_profile: &MemoryAccessProfile) -> Result<Option<OptimizationDecision>> {
        // Look for sequential access patterns that would benefit from prefetching
        for region in &memory_profile.hot_memory_regions {
            if matches!(region.access_pattern, AccessPattern::Sequential) && region.access_count > 100 {
                let mut parameters = HashMap::new();
                parameters.insert("prefetch_distance".to_string(), OptimizationParameter::Integer(64));
                
                return Ok(Some(OptimizationDecision {
                    optimization_type: OptimizationType::MemoryPrefetching,
                    target: format!("memory_region_{}", region.start_address),
                    parameters,
                    expected_benefit: 1.5, // Estimated speedup from prefetching
                }));
            }
        }
        
        Ok(None)
    }
    
    /// Analyzes performance counters for optimization opportunities
    fn analyze_performance_counters(&self, expr: &Expr, perf_counters: &PerformanceCounters) -> Result<Option<OptimizationDecision>> {
        // Low IPC might indicate opportunities for better instruction scheduling or SIMD
        if perf_counters.ipc < 1.5 {
            // Consider SIMD vectorization if we're not saturating execution units
            return Ok(Some(OptimizationDecision {
                optimization_type: OptimizationType::SIMDVectorization,
                target: self.expression_key(expr),
                parameters: HashMap::new(),
                expected_benefit: 2.5, // Expected benefit from vectorization
            }));
        }
        
        Ok(None)
    }
    
    /// Applies optimization decisions to an expression
    fn apply_optimizations(&self, expr: &Expr, decisions: &[OptimizationDecision]) -> Result<Expr> {
        let mut optimized_expr = expr.clone();
        
        for decision in decisions {
            optimized_expr = self.apply_single_optimization(optimized_expr, decision)?;
        }
        
        Ok(optimized_expr)
    }
    
    /// Applies a single optimization decision
    fn apply_single_optimization(&self, expr: Expr, decision: &OptimizationDecision) -> Result<Expr> {
        match decision.optimization_type {
            OptimizationType::TypeSpecialization => {
                // In a real implementation, this would transform the expression
                // to use type-specialized operations
                Ok(expr)
            }
            OptimizationType::BranchOptimization => {
                // Reorder branches based on likely taken path
                Ok(expr)
            }
            OptimizationType::MemoryPrefetching => {
                // Insert prefetch instructions
                Ok(expr)
            }
            OptimizationType::SIMDVectorization => {
                // Transform to use SIMD operations
                Ok(expr)
            }
            _ => Ok(expr),
        }
    }
    
    /// Records optimization history for analysis
    fn record_optimization_history(&mut self, expr_key: String, decisions: Vec<OptimizationDecision>) {
        let history = OptimizationHistory {
            expr_key,
            decisions,
            timestamp: std::time::Instant::now(),
            performance_before: 0.0, // Would be measured
            performance_after: 0.0,  // Would be measured
        };
        
        self.optimization_history.push(history);
        
        // Keep bounded history
        if self.optimization_history.len() > 1000 {
            self.optimization_history.remove(0);
        }
    }
    
    /// Generates expression key for profiling
    fn expression_key(&self, expr: &Expr) -> String {
        format!("{expr:?}")
    }
    
    /// Returns optimization statistics
    pub fn stats(&self) -> &PgoStats {
        &self.stats
    }
}

/// Optimization history entry
#[derive(Debug, Clone)]
pub struct OptimizationHistory {
    /// Expression key
    expr_key: String,
    
    /// Optimization decisions applied
    decisions: Vec<OptimizationDecision>,
    
    /// When the optimization was applied
    timestamp: std::time::Instant,
    
    /// Performance before optimization
    performance_before: f64,
    
    /// Performance after optimization
    performance_after: f64,
}

/// Profile-guided optimization statistics
#[derive(Debug, Clone, Default)]
pub struct PgoStats {
    /// Total expressions optimized
    pub expressions_optimized: u64,
    
    /// Total optimizations applied
    pub optimizations_applied: u64,
    
    /// Type specializations performed
    pub type_specializations: u64,
    
    /// Branch optimizations performed
    pub branch_optimizations: u64,
    
    /// Memory optimizations performed
    pub memory_optimizations: u64,
    
    /// SIMD optimizations performed
    pub simd_optimizations: u64,
    
    /// Average performance improvement
    pub avg_performance_improvement: f64,
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ast::Literal;
    use crate::jit::ExecutionProfile;
    use crate::jit::code_generator::SchemeType;
    
    #[test]
    fn test_pgo_config_default() {
        let config = PgoConfig::default();
        assert!(config.adaptive_optimization);
        assert_eq!(config.min_profile_samples, 50);
    }
    
    #[test]
    fn test_runtime_profile_creation() {
        let profile = RuntimeProfile::new();
        assert!(profile.type_feedback.variable_types.is_empty());
        assert!(profile.branch_data.branch_stats.is_empty());
    }
    
    #[test]
    fn test_optimizer_creation() {
        let config = PgoConfig::default();
        let optimizer = ProfileGuidedOptimizer::new(config);
        assert!(optimizer.is_ok());
    }
    
    #[test]
    fn test_type_observation() {
        let observation = TypeObservation {
            type_info: SchemeType::Integer,
            frequency: 100,
            percentage: 0.85,
        };
        
        assert_eq!(observation.frequency, 100);
        assert_eq!(observation.percentage, 0.85);
    }
}