torsh_jit/

speculative_optimization.rs

//! Speculative Optimization for ToRSh JIT
//!
//! This module implements speculative optimization techniques that make optimistic
//! assumptions about runtime behavior and provide deoptimization mechanisms when
//! those assumptions are violated.

use crate::{ComputationGraph, JitError, JitResult, Node, NodeId};
use petgraph::graph::NodeIndex;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::{
    atomic::{AtomicBool, AtomicU64, Ordering},
    Arc, RwLock,
};

/// Speculative optimization manager
pub struct SpeculativeOptimizer {
    config: SpeculativeConfig,
    assumptions: Arc<RwLock<HashMap<AssumptionId, Assumption>>>,
    guards: Arc<RwLock<HashMap<NodeId, Vec<Guard>>>>,
    deopt_counter: AtomicU64,
    enabled: AtomicBool,
}

/// Configuration for speculative optimization
#[derive(Debug, Clone)]
pub struct SpeculativeConfig {
    /// Maximum number of active assumptions
    pub max_assumptions: usize,

    /// Deoptimization threshold - disable after this many failures
    pub deopt_threshold: u64,

    /// Confidence threshold for applying speculation
    pub confidence_threshold: f64,

    /// Enable type speculation
    pub enable_type_speculation: bool,

    /// Enable shape speculation
    pub enable_shape_speculation: bool,

    /// Enable value speculation
    pub enable_value_speculation: bool,

    /// Enable nullability speculation
    pub enable_nullability_speculation: bool,

    /// Enable branch speculation
    pub enable_branch_speculation: bool,

    /// Enable loop iteration speculation
    pub enable_loop_speculation: bool,

    /// Speculation aggressiveness (0.0 to 1.0)
    pub aggressiveness: f64,
}

impl Default for SpeculativeConfig {
    fn default() -> Self {
        Self {
            max_assumptions: 1000,
            deopt_threshold: 100,
            confidence_threshold: 0.8,
            enable_type_speculation: true,
            enable_shape_speculation: true,
            enable_value_speculation: false, // More risky
            enable_nullability_speculation: true,
            enable_branch_speculation: true,
            enable_loop_speculation: true,
            aggressiveness: 0.7,
        }
    }
}

/// Unique identifier for assumptions
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct AssumptionId(pub u64);

/// Speculation assumption
#[derive(Debug, Clone)]
pub struct Assumption {
    pub id: AssumptionId,
    pub assumption_type: AssumptionType,
    pub node_id: NodeId,
    pub confidence: f64,
    pub success_count: u64,
    pub failure_count: u64,
    pub created_at: std::time::SystemTime,
    pub metadata: HashMap<String, String>,
}

/// Types of speculative assumptions
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum AssumptionType {
    /// Assume tensor has specific data type
    TypeSpeculation { expected_type: String },

    /// Assume tensor has specific shape
    ShapeSpeculation { expected_shape: Vec<usize> },

    /// Assume value is constant
    ValueSpeculation { expected_value: f64, tolerance: f64 },

    /// Assume value is not null/NaN
    NullabilitySpeculation,

    /// Assume branch is usually taken/not taken
    BranchSpeculation {
        usually_taken: bool,
        probability: f64,
    },

    /// Assume loop iterates specific number of times
    LoopSpeculation {
        expected_iterations: u64,
        tolerance: u64,
    },

    /// Assume memory access pattern
    MemorySpeculation { access_pattern: MemoryAccessPattern },
}

/// Memory access patterns for speculation
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum MemoryAccessPattern {
    Sequential,
    Random,
    Strided { stride: usize },
    Clustered { cluster_size: usize },
}

/// Runtime guard for checking assumptions
#[derive(Debug, Clone)]
pub struct Guard {
    pub assumption_id: AssumptionId,
    pub guard_type: GuardType,
    pub check_frequency: GuardFrequency,
}

/// Types of runtime guards
#[derive(Debug, Clone, PartialEq)]
pub enum GuardType {
    /// Check tensor data type
    TypeCheck,

    /// Check tensor shape
    ShapeCheck,

    /// Check value against expected
    ValueCheck,

    /// Check for null/NaN values
    NullabilityCheck,

    /// Check branch outcome
    BranchCheck,

    /// Check loop iteration count
    LoopCheck,

    /// Check memory access pattern
    MemoryCheck,
}

/// Guard check frequency
#[derive(Debug, Clone, PartialEq)]
pub enum GuardFrequency {
    /// Check every execution
    Always,

    /// Check probabilistically
    Probabilistic(f64),

    /// Check after N executions
    Periodic(u64),

    /// Check only first N executions
    InitialOnly(u64),
}

/// Result of a speculation attempt
#[derive(Debug, Clone)]
pub struct SpeculationResult {
    pub assumptions_made: Vec<AssumptionId>,
    pub optimizations_applied: Vec<SpeculativeOptimization>,
    pub guards_installed: Vec<Guard>,
    pub estimated_speedup: f64,
}

/// Speculative optimizations that can be applied
#[derive(Debug, Clone)]
pub struct SpeculativeOptimization {
    pub optimization_type: SpeculativeOptimizationType,
    pub node_id: NodeId,
    pub description: String,
    pub estimated_benefit: f64,
}

/// Types of speculative optimizations
#[derive(Debug, Clone, PartialEq)]
pub enum SpeculativeOptimizationType {
    /// Remove type checks
    TypeCheckElimination,

    /// Remove bounds checks
    BoundsCheckElimination,

    /// Optimize for specific shape
    ShapeSpecialization,

    /// Constant propagation
    ConstantPropagation,

    /// Dead code elimination
    DeadCodeElimination,

    /// Loop unrolling
    LoopUnrolling,

    /// Branch elimination
    BranchElimination,

    /// Memory prefetching
    MemoryPrefetching,

    /// Vectorization
    VectorizationOptimization,
}

/// Deoptimization event
#[derive(Debug, Clone)]
pub struct DeoptimizationEvent {
    pub assumption_id: AssumptionId,
    pub node_id: NodeId,
    pub reason: String,
    pub timestamp: std::time::SystemTime,
    pub execution_count: u64,
}

impl SpeculativeOptimizer {
    /// Create a new speculative optimizer
    pub fn new(config: SpeculativeConfig) -> Self {
        Self {
            config,
            assumptions: Arc::new(RwLock::new(HashMap::new())),
            guards: Arc::new(RwLock::new(HashMap::new())),
            deopt_counter: AtomicU64::new(0),
            enabled: AtomicBool::new(true),
        }
    }

    /// Analyze graph and generate speculative optimizations
    pub fn analyze_and_speculate(
        &self,
        graph: &ComputationGraph,
        execution_history: &ExecutionHistory,
    ) -> JitResult<SpeculationResult> {
        if !self.enabled.load(Ordering::Relaxed) {
            return Ok(SpeculationResult {
                assumptions_made: Vec::new(),
                optimizations_applied: Vec::new(),
                guards_installed: Vec::new(),
                estimated_speedup: 1.0,
            });
        }

        let mut assumptions_made = Vec::new();
        let mut optimizations = Vec::new();
        let mut guards = Vec::new();
        let mut total_speedup = 1.0;

        for (node_id, node) in graph.nodes() {
            // Analyze node history for speculation opportunities
            if let Some(node_history) = execution_history.get_node_history(node_id) {
                // Type speculation
                if self.config.enable_type_speculation {
                    if let Some(spec_result) =
                        self.analyze_type_speculation(node_id, node, node_history)?
                    {
                        assumptions_made.extend(spec_result.assumptions_made);
                        optimizations.extend(spec_result.optimizations_applied);
                        guards.extend(spec_result.guards_installed);
                        total_speedup *= spec_result.estimated_speedup;
                    }
                }

                // Shape speculation
                if self.config.enable_shape_speculation {
                    if let Some(spec_result) =
                        self.analyze_shape_speculation(node_id, node, node_history)?
                    {
                        assumptions_made.extend(spec_result.assumptions_made);
                        optimizations.extend(spec_result.optimizations_applied);
                        guards.extend(spec_result.guards_installed);
                        total_speedup *= spec_result.estimated_speedup;
                    }
                }

                // Value speculation
                if self.config.enable_value_speculation {
                    if let Some(spec_result) =
                        self.analyze_value_speculation(node_id, node, node_history)?
                    {
                        assumptions_made.extend(spec_result.assumptions_made);
                        optimizations.extend(spec_result.optimizations_applied);
                        guards.extend(spec_result.guards_installed);
                        total_speedup *= spec_result.estimated_speedup;
                    }
                }

                // Branch speculation
                if self.config.enable_branch_speculation {
                    if let Some(spec_result) =
                        self.analyze_branch_speculation(node_id, node, node_history)?
                    {
                        assumptions_made.extend(spec_result.assumptions_made);
                        optimizations.extend(spec_result.optimizations_applied);
                        guards.extend(spec_result.guards_installed);
                        total_speedup *= spec_result.estimated_speedup;
                    }
                }
            }
        }

        // Install guards
        if let Ok(mut guard_map) = self.guards.write() {
            for guard in &guards {
                guard_map
                    .entry(NodeIndex::new(guard.assumption_id.0 as usize))
                    .or_insert_with(Vec::new)
                    .push(guard.clone());
            }
        }

        // Record assumptions — pair each AssumptionId with the optimization that produced
        // it.  The two vecs are grown in lock-step (one optimization per assumption), so
        // zip gives correct context for every entry.
        if let Ok(mut assumption_map) = self.assumptions.write() {
            for (assumption_id, optimization) in assumptions_made.iter().zip(optimizations.iter()) {
                let assumption = self.build_assumption(*assumption_id, optimization);
                assumption_map.insert(*assumption_id, assumption);
            }
        }

        Ok(SpeculationResult {
            assumptions_made,
            optimizations_applied: optimizations,
            guards_installed: guards,
            estimated_speedup: total_speedup,
        })
    }

    /// Apply speculative optimizations to the graph
    pub fn apply_speculative_optimizations(
        &self,
        graph: &mut ComputationGraph,
        result: &SpeculationResult,
    ) -> JitResult<usize> {
        let mut applied_count = 0;

        for optimization in &result.optimizations_applied {
            match optimization.optimization_type {
                SpeculativeOptimizationType::TypeCheckElimination => {
                    if self.apply_type_check_elimination(graph, optimization)? {
                        applied_count += 1;
                    }
                }
                SpeculativeOptimizationType::ShapeSpecialization => {
                    if self.apply_shape_specialization(graph, optimization)? {
                        applied_count += 1;
                    }
                }
                SpeculativeOptimizationType::ConstantPropagation => {
                    if self.apply_constant_propagation(graph, optimization)? {
                        applied_count += 1;
                    }
                }
                SpeculativeOptimizationType::BranchElimination => {
                    if self.apply_branch_elimination(graph, optimization)? {
                        applied_count += 1;
                    }
                }
                _ => {
                    // Other optimizations can be implemented as needed
                }
            }
        }

        Ok(applied_count)
    }

    /// Check guards during execution and handle deoptimization
    pub fn check_guards(&self, node_id: NodeId, runtime_info: &RuntimeInfo) -> JitResult<bool> {
        let guard_map = self
            .guards
            .read()
            .map_err(|_| JitError::RuntimeError("Failed to read guards".to_string()))?;

        if let Some(node_guards) = guard_map.get(&node_id) {
            for guard in node_guards {
                if self.should_check_guard(guard, runtime_info.execution_count) {
                    let check_passed = self.execute_guard_check(guard, runtime_info)?;

                    if !check_passed {
                        self.handle_deoptimization(
                            guard.assumption_id,
                            node_id,
                            "Guard check failed",
                        )?;
                        return Ok(false);
                    }
                }
            }
        }

        Ok(true)
    }

    /// Record successful execution (reinforces assumptions)
    pub fn record_success(&self, assumption_id: AssumptionId) {
        if let Ok(mut assumptions) = self.assumptions.write() {
            if let Some(assumption) = assumptions.get_mut(&assumption_id) {
                assumption.success_count += 1;
                assumption.confidence = self.calculate_confidence(assumption);
            }
        }
    }

    /// Handle deoptimization when assumptions fail
    pub fn handle_deoptimization(
        &self,
        assumption_id: AssumptionId,
        node_id: NodeId,
        reason: &str,
    ) -> JitResult<()> {
        let deopt_count = self.deopt_counter.fetch_add(1, Ordering::Relaxed);

        // Update assumption failure count
        if let Ok(mut assumptions) = self.assumptions.write() {
            if let Some(assumption) = assumptions.get_mut(&assumption_id) {
                assumption.failure_count += 1;
                assumption.confidence = self.calculate_confidence(assumption);

                // Remove assumption if confidence drops too low
                if assumption.confidence < 0.3 {
                    assumptions.remove(&assumption_id);
                }
            }
        }

        // Disable speculative optimization if too many failures
        if deopt_count > self.config.deopt_threshold {
            self.enabled.store(false, Ordering::Relaxed);
        }

        // Log deoptimization event
        let event = DeoptimizationEvent {
            assumption_id,
            node_id,
            reason: reason.to_string(),
            timestamp: std::time::SystemTime::now(),
            execution_count: deopt_count,
        };

        self.log_deoptimization_event(&event);

        Ok(())
    }

    /// Get speculation statistics
    pub fn get_statistics(&self) -> JitResult<SpeculationStatistics> {
        let assumptions = self
            .assumptions
            .read()
            .map_err(|_| JitError::RuntimeError("Failed to read assumptions".to_string()))?;

        let active_assumptions = assumptions.len();
        let total_successes = assumptions.values().map(|a| a.success_count).sum();
        let total_failures = assumptions.values().map(|a| a.failure_count).sum();
        let avg_confidence = if !assumptions.is_empty() {
            assumptions.values().map(|a| a.confidence).sum::<f64>() / assumptions.len() as f64
        } else {
            0.0
        };

        let deopt_count = self.deopt_counter.load(Ordering::Relaxed);
        let enabled = self.enabled.load(Ordering::Relaxed);

        Ok(SpeculationStatistics {
            active_assumptions,
            total_successes,
            total_failures,
            avg_confidence,
            deoptimization_count: deopt_count,
            enabled,
        })
    }

    // Helper methods
    fn analyze_type_speculation(
        &self,
        node_id: NodeId,
        _node: &Node,
        history: &NodeExecutionHistory,
    ) -> JitResult<Option<SpeculationResult>> {
        // Analyze type patterns in execution history
        if let Some(dominant_type) = history.get_dominant_type(self.config.confidence_threshold) {
            let assumption_id = self.generate_assumption_id();

            let optimization = SpeculativeOptimization {
                optimization_type: SpeculativeOptimizationType::TypeCheckElimination,
                node_id,
                description: format!("Assume type is always {}", dominant_type),
                estimated_benefit: 0.05, // 5% speedup from eliminating type checks
            };

            let guard = Guard {
                assumption_id,
                guard_type: GuardType::TypeCheck,
                check_frequency: GuardFrequency::Probabilistic(0.1), // Check 10% of the time
            };

            return Ok(Some(SpeculationResult {
                assumptions_made: vec![assumption_id],
                optimizations_applied: vec![optimization],
                guards_installed: vec![guard],
                estimated_speedup: 1.05,
            }));
        }

        Ok(None)
    }

    fn analyze_shape_speculation(
        &self,
        node_id: NodeId,
        _node: &Node,
        history: &NodeExecutionHistory,
    ) -> JitResult<Option<SpeculationResult>> {
        // Analyze shape patterns in execution history
        if let Some(dominant_shape) = history.get_dominant_shape(self.config.confidence_threshold) {
            let assumption_id = self.generate_assumption_id();

            let optimization = SpeculativeOptimization {
                optimization_type: SpeculativeOptimizationType::ShapeSpecialization,
                node_id,
                description: format!("Specialize for shape {:?}", dominant_shape),
                estimated_benefit: 0.15, // 15% speedup from shape specialization
            };

            let guard = Guard {
                assumption_id,
                guard_type: GuardType::ShapeCheck,
                check_frequency: GuardFrequency::Always, // Shape changes are critical
            };

            return Ok(Some(SpeculationResult {
                assumptions_made: vec![assumption_id],
                optimizations_applied: vec![optimization],
                guards_installed: vec![guard],
                estimated_speedup: 1.15,
            }));
        }

        Ok(None)
    }

    fn analyze_value_speculation(
        &self,
        node_id: NodeId,
        _node: &Node,
        history: &NodeExecutionHistory,
    ) -> JitResult<Option<SpeculationResult>> {
        // Analyze value patterns - look for constants
        if let Some(constant_value) = history.get_constant_value(self.config.confidence_threshold) {
            let assumption_id = self.generate_assumption_id();

            let optimization = SpeculativeOptimization {
                optimization_type: SpeculativeOptimizationType::ConstantPropagation,
                node_id,
                description: format!("Assume constant value {}", constant_value),
                estimated_benefit: 0.20, // 20% speedup from constant propagation
            };

            let guard = Guard {
                assumption_id,
                guard_type: GuardType::ValueCheck,
                check_frequency: GuardFrequency::Periodic(100), // Check every 100 executions
            };

            return Ok(Some(SpeculationResult {
                assumptions_made: vec![assumption_id],
                optimizations_applied: vec![optimization],
                guards_installed: vec![guard],
                estimated_speedup: 1.20,
            }));
        }

        Ok(None)
    }

    fn analyze_branch_speculation(
        &self,
        node_id: NodeId,
        _node: &Node,
        history: &NodeExecutionHistory,
    ) -> JitResult<Option<SpeculationResult>> {
        // Analyze branch patterns
        if let Some(branch_bias) = history.get_branch_bias(self.config.confidence_threshold) {
            let assumption_id = self.generate_assumption_id();

            let optimization = SpeculativeOptimization {
                optimization_type: SpeculativeOptimizationType::BranchElimination,
                node_id,
                description: format!(
                    "Assume branch is usually {}",
                    if branch_bias > 0.5 {
                        "taken"
                    } else {
                        "not taken"
                    }
                ),
                estimated_benefit: 0.10, // 10% speedup from branch elimination
            };

            let guard = Guard {
                assumption_id,
                guard_type: GuardType::BranchCheck,
                check_frequency: GuardFrequency::Probabilistic(0.05), // Check 5% of the time
            };

            return Ok(Some(SpeculationResult {
                assumptions_made: vec![assumption_id],
                optimizations_applied: vec![optimization],
                guards_installed: vec![guard],
                estimated_speedup: 1.10,
            }));
        }

        Ok(None)
    }

    fn apply_type_check_elimination(
        &self,
        graph: &mut ComputationGraph,
        optimization: &SpeculativeOptimization,
    ) -> JitResult<bool> {
        let node_id = optimization.node_id;

        if let Some(node) = graph.node_mut(node_id) {
            // Remove redundant type checks for nodes with stable types
            node.set_optimization_hint("eliminate_type_checks", "true")?;
            node.set_optimization_hint("assumed_type_stable", "true")?;

            // Add guard to verify type assumption at runtime
            node.set_optimization_hint("add_type_guard", "true")?;
            node.set_optimization_hint("guard_frequency", "low")?;

            return Ok(true);
        }

        Ok(false)
    }

    fn apply_shape_specialization(
        &self,
        graph: &mut ComputationGraph,
        optimization: &SpeculativeOptimization,
    ) -> JitResult<bool> {
        let node_id = optimization.node_id;

        if let Some(node) = graph.node_mut(node_id) {
            // Specialize operations for the most common shape
            node.set_optimization_hint("shape_specialized", "true")?;
            node.set_optimization_hint("eliminate_shape_checks", "true")?;

            // Extract assumed shape from optimization description
            if optimization.description.contains("shape") {
                node.set_optimization_hint("specialized_shape_source", "speculation")?;
                node.set_optimization_hint("add_shape_guard", "true")?;
            }

            return Ok(true);
        }

        Ok(false)
    }

    fn apply_constant_propagation(
        &self,
        graph: &mut ComputationGraph,
        optimization: &SpeculativeOptimization,
    ) -> JitResult<bool> {
        let node_id = optimization.node_id;

        if let Some(node) = graph.node_mut(node_id) {
            // Mark node for constant propagation optimization
            node.set_optimization_hint("constant_propagation", "true")?;
            node.set_optimization_hint("assumed_constant", "true")?;

            // Extract assumed constant value from description
            if let Some(start) = optimization.description.find("value ") {
                if let Some(end) = optimization.description[start + 6..].find(' ') {
                    let value_str = &optimization.description[start + 6..start + 6 + end];
                    node.set_optimization_hint("assumed_constant_value", value_str)?;
                }
            }

            // Add value guard for verification
            node.set_optimization_hint("add_value_guard", "true")?;
            node.set_optimization_hint("guard_tolerance", "1e-10")?;

            return Ok(true);
        }

        Ok(false)
    }

    fn apply_branch_elimination(
        &self,
        graph: &mut ComputationGraph,
        optimization: &SpeculativeOptimization,
    ) -> JitResult<bool> {
        let node_id = optimization.node_id;

        if let Some(node) = graph.node_mut(node_id) {
            // Determine branch bias from description
            let usually_taken = optimization.description.contains("usually taken");

            if usually_taken {
                node.set_optimization_hint("branch_likely", "true")?;
                node.set_optimization_hint("optimize_taken_path", "true")?;
            } else {
                node.set_optimization_hint("branch_unlikely", "true")?;
                node.set_optimization_hint("optimize_not_taken_path", "true")?;
            }

            // For highly predictable branches, consider elimination
            if optimization.estimated_benefit > 0.08 {
                // 8% or higher benefit
                node.set_optimization_hint("branch_elimination_candidate", "true")?;
                node.set_optimization_hint("speculative_branch_elimination", "true")?;
            }

            // Add branch guard
            node.set_optimization_hint("add_branch_guard", "true")?;

            return Ok(true);
        }

        Ok(false)
    }

    fn should_check_guard(&self, guard: &Guard, execution_count: u64) -> bool {
        match guard.check_frequency {
            GuardFrequency::Always => true,
            GuardFrequency::Probabilistic(probability) => {
                use std::collections::hash_map::DefaultHasher;
                use std::hash::{Hash, Hasher};

                let mut hasher = DefaultHasher::new();
                execution_count.hash(&mut hasher);
                let hash = hasher.finish();
                (hash as f64 / u64::MAX as f64) < probability
            }
            GuardFrequency::Periodic(period) => execution_count % period == 0,
            GuardFrequency::InitialOnly(limit) => execution_count < limit,
        }
    }

    fn execute_guard_check(&self, guard: &Guard, runtime_info: &RuntimeInfo) -> JitResult<bool> {
        match guard.guard_type {
            GuardType::TypeCheck => {
                // Check if actual type matches expected type
                Ok(runtime_info.actual_type == runtime_info.expected_type)
            }
            GuardType::ShapeCheck => {
                // Check if actual shape matches expected shape
                Ok(runtime_info.actual_shape == runtime_info.expected_shape)
            }
            GuardType::ValueCheck => {
                // Check if actual value matches expected value within tolerance
                Ok(
                    (runtime_info.actual_value - runtime_info.expected_value).abs()
                        < runtime_info.tolerance,
                )
            }
            GuardType::NullabilityCheck => {
                // Check if value is not null/NaN
                Ok(!runtime_info.actual_value.is_nan() && runtime_info.actual_value.is_finite())
            }
            GuardType::BranchCheck => {
                // Check if branch outcome matches prediction
                Ok(runtime_info.branch_taken == runtime_info.expected_branch_taken)
            }
            GuardType::LoopCheck => {
                // Check if loop iterations match expectation
                let diff = (runtime_info.actual_iterations as i64
                    - runtime_info.expected_iterations as i64)
                    .abs();
                Ok(diff <= runtime_info.iteration_tolerance as i64)
            }
            GuardType::MemoryCheck => {
                // Check if memory access pattern matches expectation
                Ok(runtime_info.memory_pattern == runtime_info.expected_memory_pattern)
            }
        }
    }

    fn generate_assumption_id(&self) -> AssumptionId {
        use std::sync::atomic::{AtomicU64, Ordering};
        static COUNTER: AtomicU64 = AtomicU64::new(0);
        AssumptionId(COUNTER.fetch_add(1, Ordering::Relaxed))
    }

    /// Build an `Assumption` from a known `(id, optimization)` pair so that
    /// every field reflects the actual speculation context rather than a
    /// hardcoded placeholder.
    fn build_assumption(
        &self,
        id: AssumptionId,
        optimization: &SpeculativeOptimization,
    ) -> Assumption {
        let assumption_type = match &optimization.optimization_type {
            SpeculativeOptimizationType::TypeCheckElimination => {
                // Extract the expected type from the optimization description.
                // Description format: "Assume type is always <type_name>"
                let expected_type = optimization
                    .description
                    .strip_prefix("Assume type is always ")
                    .unwrap_or("unknown")
                    .to_string();
                AssumptionType::TypeSpeculation { expected_type }
            }
            SpeculativeOptimizationType::ShapeSpecialization => {
                // Description format: "Specialize for shape [d0, d1, ...]"
                // Parse the shape from the description; fall back to empty on parse error.
                let expected_shape = optimization
                    .description
                    .strip_prefix("Specialize for shape ")
                    .and_then(|s| {
                        // Remove brackets and split on ", "
                        let inner = s.trim_start_matches('[').trim_end_matches(']');
                        if inner.is_empty() {
                            Some(Vec::new())
                        } else {
                            inner
                                .split(", ")
                                .map(|tok| tok.trim().parse::<usize>().ok())
                                .collect::<Option<Vec<usize>>>()
                        }
                    })
                    .unwrap_or_default();
                AssumptionType::ShapeSpeculation { expected_shape }
            }
            SpeculativeOptimizationType::ConstantPropagation => {
                // Description format: "Assume constant value <f64>"
                let (expected_value, tolerance) = optimization
                    .description
                    .strip_prefix("Assume constant value ")
                    .and_then(|s| s.split_whitespace().next())
                    .and_then(|tok| tok.parse::<f64>().ok())
                    .map(|v| (v, 1e-10))
                    .unwrap_or((0.0, 1e-10));
                AssumptionType::ValueSpeculation {
                    expected_value,
                    tolerance,
                }
            }
            SpeculativeOptimizationType::BranchElimination => {
                let usually_taken = optimization.description.contains("usually taken");
                // Use the optimization's estimated benefit as a proxy for confidence.
                let probability = (optimization.estimated_benefit * 10.0).clamp(0.5, 0.99);
                AssumptionType::BranchSpeculation {
                    usually_taken,
                    probability,
                }
            }
            // For other optimization types the most conservative assumption is
            // NullabilitySpeculation (guards against NaN/null inputs).
            _ => AssumptionType::NullabilitySpeculation,
        };

        let initial_confidence = self.config.confidence_threshold;
        let mut metadata = HashMap::new();
        metadata.insert(
            "optimization_type".to_string(),
            format!("{:?}", optimization.optimization_type),
        );
        metadata.insert(
            "estimated_benefit".to_string(),
            optimization.estimated_benefit.to_string(),
        );

        Assumption {
            id,
            assumption_type,
            node_id: optimization.node_id,
            confidence: initial_confidence,
            success_count: 0,
            failure_count: 0,
            created_at: std::time::SystemTime::now(),
            metadata,
        }
    }

    fn calculate_confidence(&self, assumption: &Assumption) -> f64 {
        let total = assumption.success_count + assumption.failure_count;
        if total == 0 {
            return 0.5; // No data, neutral confidence
        }

        assumption.success_count as f64 / total as f64
    }

    fn log_deoptimization_event(&self, event: &DeoptimizationEvent) {
        // Log the deoptimization event for debugging and analysis
        eprintln!("Deoptimization: {:?}", event);
    }
}

/// Runtime information for guard checks
#[derive(Debug, Clone)]
pub struct RuntimeInfo {
    pub execution_count: u64,
    pub actual_type: String,
    pub expected_type: String,
    pub actual_shape: Vec<usize>,
    pub expected_shape: Vec<usize>,
    pub actual_value: f64,
    pub expected_value: f64,
    pub tolerance: f64,
    pub branch_taken: bool,
    pub expected_branch_taken: bool,
    pub actual_iterations: u64,
    pub expected_iterations: u64,
    pub iteration_tolerance: u64,
    pub memory_pattern: MemoryAccessPattern,
    pub expected_memory_pattern: MemoryAccessPattern,
}

/// Execution history for a node
#[derive(Debug, Clone)]
pub struct NodeExecutionHistory {
    types: Vec<String>,
    shapes: Vec<Vec<usize>>,
    values: Vec<f64>,
    branch_outcomes: Vec<bool>,
    loop_iterations: Vec<u64>,
}

impl NodeExecutionHistory {
    pub fn get_dominant_type(&self, threshold: f64) -> Option<String> {
        let mut type_counts = HashMap::new();
        for type_name in &self.types {
            *type_counts.entry(type_name.clone()).or_insert(0) += 1;
        }

        if let Some((dominant_type, count)) = type_counts.iter().max_by_key(|(_, &count)| count) {
            if *count as f64 / self.types.len() as f64 >= threshold {
                return Some(dominant_type.clone());
            }
        }

        None
    }

    pub fn get_dominant_shape(&self, threshold: f64) -> Option<Vec<usize>> {
        let mut shape_counts = HashMap::new();
        for shape in &self.shapes {
            *shape_counts.entry(shape.clone()).or_insert(0) += 1;
        }

        if let Some((dominant_shape, count)) = shape_counts.iter().max_by_key(|(_, &count)| count) {
            if *count as f64 / self.shapes.len() as f64 >= threshold {
                return Some(dominant_shape.clone());
            }
        }

        None
    }

    pub fn get_constant_value(&self, threshold: f64) -> Option<f64> {
        if self.values.is_empty() {
            return None;
        }

        // Check if all values are approximately the same
        let first_value = self.values[0];
        let tolerance = 1e-10;
        let constant_count = self
            .values
            .iter()
            .filter(|&&v| (v - first_value).abs() < tolerance)
            .count();

        if constant_count as f64 / self.values.len() as f64 >= threshold {
            Some(first_value)
        } else {
            None
        }
    }

    pub fn get_branch_bias(&self, threshold: f64) -> Option<f64> {
        if self.branch_outcomes.is_empty() {
            return None;
        }

        let taken_count = self.branch_outcomes.iter().filter(|&&taken| taken).count();
        let bias = taken_count as f64 / self.branch_outcomes.len() as f64;

        // Return bias if it's significantly different from 50/50
        if (bias - 0.5).abs() >= (threshold - 0.5) {
            Some(bias)
        } else {
            None
        }
    }
}

/// Execution history for the entire graph
#[derive(Debug, Clone)]
pub struct ExecutionHistory {
    node_histories: HashMap<NodeId, NodeExecutionHistory>,
}

impl ExecutionHistory {
    pub fn new() -> Self {
        Self {
            node_histories: HashMap::new(),
        }
    }

    pub fn get_node_history(&self, node_id: NodeId) -> Option<&NodeExecutionHistory> {
        self.node_histories.get(&node_id)
    }

    pub fn record_execution(&mut self, node_id: NodeId, info: NodeExecutionInfo) {
        let history = self
            .node_histories
            .entry(node_id)
            .or_insert_with(|| NodeExecutionHistory {
                types: Vec::new(),
                shapes: Vec::new(),
                values: Vec::new(),
                branch_outcomes: Vec::new(),
                loop_iterations: Vec::new(),
            });

        if let Some(type_name) = info.type_name {
            history.types.push(type_name);
        }
        if let Some(shape) = info.shape {
            history.shapes.push(shape);
        }
        if let Some(value) = info.value {
            history.values.push(value);
        }
        if let Some(branch_taken) = info.branch_taken {
            history.branch_outcomes.push(branch_taken);
        }
        if let Some(iterations) = info.loop_iterations {
            history.loop_iterations.push(iterations);
        }
    }
}

/// Information about a single node execution
#[derive(Debug, Clone)]
pub struct NodeExecutionInfo {
    pub type_name: Option<String>,
    pub shape: Option<Vec<usize>>,
    pub value: Option<f64>,
    pub branch_taken: Option<bool>,
    pub loop_iterations: Option<u64>,
}

/// Statistics about speculative optimization
#[derive(Debug, Clone)]
pub struct SpeculationStatistics {
    pub active_assumptions: usize,
    pub total_successes: u64,
    pub total_failures: u64,
    pub avg_confidence: f64,
    pub deoptimization_count: u64,
    pub enabled: bool,
}

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

    #[test]
    fn test_speculative_optimizer_creation() {
        let config = SpeculativeConfig::default();
        let optimizer = SpeculativeOptimizer::new(config);
        assert!(optimizer.enabled.load(Ordering::Relaxed));
        assert_eq!(optimizer.deopt_counter.load(Ordering::Relaxed), 0);
    }

    #[test]
    fn test_assumption_id_generation() {
        let optimizer = SpeculativeOptimizer::new(SpeculativeConfig::default());
        let id1 = optimizer.generate_assumption_id();
        let id2 = optimizer.generate_assumption_id();
        assert_ne!(id1, id2);
    }

    #[test]
    fn test_guard_frequency_checking() {
        let optimizer = SpeculativeOptimizer::new(SpeculativeConfig::default());

        let always_guard = Guard {
            assumption_id: AssumptionId(1),
            guard_type: GuardType::TypeCheck,
            check_frequency: GuardFrequency::Always,
        };
        assert!(optimizer.should_check_guard(&always_guard, 100));

        let periodic_guard = Guard {
            assumption_id: AssumptionId(2),
            guard_type: GuardType::TypeCheck,
            check_frequency: GuardFrequency::Periodic(10),
        };
        assert!(optimizer.should_check_guard(&periodic_guard, 100));
        assert!(!optimizer.should_check_guard(&periodic_guard, 101));
    }

    #[test]
    fn test_execution_history() {
        let mut history = ExecutionHistory::new();
        let node_id = NodeId::new(1);

        // Record some executions
        history.record_execution(
            node_id,
            NodeExecutionInfo {
                type_name: Some("f32".to_string()),
                shape: Some(vec![10, 20]),
                value: Some(1.0),
                branch_taken: Some(true),
                loop_iterations: Some(5),
            },
        );

        history.record_execution(
            node_id,
            NodeExecutionInfo {
                type_name: Some("f32".to_string()),
                shape: Some(vec![10, 20]),
                value: Some(1.0),
                branch_taken: Some(true),
                loop_iterations: Some(5),
            },
        );

        let node_history = history.get_node_history(node_id).unwrap();
        assert_eq!(node_history.get_dominant_type(0.8), Some("f32".to_string()));
        assert_eq!(node_history.get_dominant_shape(0.8), Some(vec![10, 20]));
        assert_eq!(node_history.get_constant_value(0.8), Some(1.0));
        assert_eq!(node_history.get_branch_bias(0.8), Some(1.0));
    }
}