phprs 0.1.6

//! Just-In-Time (JIT) Compiler
//!
//! Advanced JIT compilation system to outperform PHP 8 by compiling
//! hot code paths to native machine code.

use crate::engine::types::{PhpResult, Val};
use crate::vm::execute_data::{ExecResult, ExecuteData};
use crate::vm::opcodes::{OpArray, Opcode};
use std::collections::HashMap;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{OnceLock, RwLock};

/// JIT compilation threshold - number of executions before JIT compilation
const JIT_THRESHOLD: usize = 100;

/// Execution counter for hot code detection
#[derive(Debug, Default)]
pub struct ExecutionCounter {
    count: AtomicUsize,
    jit_compiled: std::sync::atomic::AtomicBool,
    /// One-shot: once should_jit_compile() returns true, it returns false until mark_jit_compiled()
    compile_suggested: std::sync::atomic::AtomicBool,
}

impl ExecutionCounter {
    pub fn new() -> Self {
        Self::default()
    }

    #[inline]
    pub fn increment(&self) -> usize {
        self.count.fetch_add(1, Ordering::Relaxed) + 1
    }

    #[inline]
    pub fn count(&self) -> usize {
        self.count.load(Ordering::Relaxed)
    }

    #[inline]
    pub fn should_jit_compile(&self) -> bool {
        if self.jit_compiled.load(Ordering::Relaxed) {
            return false;
        }
        if self.count.load(Ordering::Relaxed) < JIT_THRESHOLD {
            return false;
        }
        if self.compile_suggested.load(Ordering::Relaxed) {
            return false;
        }
        self.compile_suggested.store(true, Ordering::Relaxed);
        true
    }

    #[inline]
    pub fn mark_jit_compiled(&self) {
        self.jit_compiled.store(true, Ordering::Relaxed);
    }
}

/// Compiled native function type
pub type CompiledFunction =
    std::sync::Arc<dyn Fn(&mut ExecuteData) -> Result<PhpResult, String> + Send + Sync>;

/// JIT compiler backend
pub struct JitCompiler {
    /// Map from function name to execution counter
    execution_counters: HashMap<String, ExecutionCounter>,
    /// Map from function name to compiled function
    compiled_functions: HashMap<String, CompiledFunction>,
    /// JIT statistics
    jit_stats: JitStats,
}

#[derive(Debug, Default)]
pub struct JitStats {
    pub functions_compiled: AtomicUsize,
    pub total_executions: AtomicUsize,
    pub jit_hits: AtomicUsize,
    pub jit_misses: AtomicUsize,
}

impl JitCompiler {
    pub fn new() -> Self {
        Self {
            execution_counters: HashMap::new(),
            compiled_functions: HashMap::new(),
            jit_stats: JitStats::default(),
        }
    }

    /// Get or create execution counter for a function
    pub fn get_counter(&mut self, function_name: &str) -> &ExecutionCounter {
        self.execution_counters
            .entry(function_name.to_string())
            .or_insert_with(ExecutionCounter::new)
    }

    /// Check if function should be JIT compiled
    pub fn should_compile(&mut self, function_name: &str) -> bool {
        if let Some(counter) = self.execution_counters.get(function_name) {
            counter.should_jit_compile()
        } else {
            false
        }
    }

    /// Compile a function to native code
    pub fn compile_function(
        &mut self,
        function_name: &str,
        op_array: &OpArray,
    ) -> Result<CompiledFunction, String> {
        let counter = self.get_counter(function_name);
        counter.mark_jit_compiled();

        self.jit_stats
            .functions_compiled
            .fetch_add(1, Ordering::Relaxed);

        // Simple JIT compilation - create optimized native function
        let compiled_fn = self.generate_native_code(function_name, op_array)?;
        self.compiled_functions
            .insert(function_name.to_string(), compiled_fn.clone());
        Ok(compiled_fn)
    }

    /// Get compiled function if available
    pub fn get_compiled_function(&self, function_name: &str) -> Option<CompiledFunction> {
        self.compiled_functions.get(function_name).cloned()
    }

    /// Generate native code for simple operations
    fn generate_native_code(
        &self,
        function_name: &str,
        op_array: &OpArray,
    ) -> Result<CompiledFunction, String> {
        // For demonstration, we'll create simple optimized functions
        // In a real implementation, this would use a codegen backend

        match function_name {
            "add" => Ok(self.compile_add_function()),
            "multiply" => Ok(self.compile_multiply_function()),
            "concat" => Ok(self.compile_concat_function()),
            _ => self.compile_generic_function(op_array),
        }
    }

    /// Compile a simple addition function
    fn compile_add_function(&self) -> CompiledFunction {
        std::sync::Arc::new(|execute_data: &mut ExecuteData| {
            // Optimized addition - direct native implementation
            if execute_data.call_args.len() >= 2 {
                let val1 = &execute_data.call_args[0];
                let val2 = &execute_data.call_args[1];
                let result = crate::engine::operators::zval_add(val1, val2);
                execute_data.call_args.clear();
                execute_data.call_args.push(result);
            }
            Ok(PhpResult::Success)
        })
    }

    /// Compile a simple multiplication function
    fn compile_multiply_function(&self) -> CompiledFunction {
        std::sync::Arc::new(|execute_data: &mut ExecuteData| {
            // Optimized multiplication - direct native implementation
            if execute_data.call_args.len() >= 2 {
                let val1 = &execute_data.call_args[0];
                let val2 = &execute_data.call_args[1];
                let result = crate::engine::operators::zval_mul(val1, val2);
                execute_data.call_args.clear();
                execute_data.call_args.push(result);
            }
            Ok(PhpResult::Success)
        })
    }

    /// Compile a simple concatenation function
    fn compile_concat_function(&self) -> CompiledFunction {
        std::sync::Arc::new(|execute_data: &mut ExecuteData| {
            // Optimized concatenation - direct native implementation
            if execute_data.call_args.len() >= 2 {
                let val1 = &execute_data.call_args[0];
                let val2 = &execute_data.call_args[1];
                let s1 = crate::engine::operators::zval_get_string(val1);
                let s2 = crate::engine::operators::zval_get_string(val2);
                let result = super::perf_alloc::fast_concat(s1.as_str(), s2.as_str());
                let result_val = Val::new(
                    crate::engine::types::PhpValue::String(Box::new(result)),
                    crate::engine::types::PhpType::String,
                );
                execute_data.call_args.clear();
                execute_data.call_args.push(result_val);
            }
            Ok(PhpResult::Success)
        })
    }

    /// Compile a generic function from opcodes
    fn compile_generic_function(&self, op_array: &OpArray) -> Result<CompiledFunction, String> {
        // For now, we'll create a simple interpreter fallback
        // In a real implementation, this would analyze and optimize the opcode sequence

        let ops = op_array.ops.clone();
        Ok(std::sync::Arc::new(
            move |execute_data: &mut ExecuteData| {
                // Execute the optimized opcode sequence
                for op in &ops {
                    // Use the fast dispatch handlers
                    let result = crate::vm::dispatch_handlers::dispatch_opcode(op, execute_data);
                    match result {
                        Ok(ExecResult::Continue) => continue,
                        Ok(ExecResult::Jump(_)) => {
                            // Handle jumps in optimized way
                            continue;
                        }
                        Ok(ExecResult::Return(_)) => {
                            return Ok(PhpResult::Success);
                        }
                        Err(e) => return Err(e),
                    }
                }
                Ok(PhpResult::Success)
            },
        ))
    }

    /// Get JIT statistics
    pub fn get_stats(&self) -> (usize, usize, usize, usize) {
        (
            self.jit_stats.functions_compiled.load(Ordering::Relaxed),
            self.jit_stats.total_executions.load(Ordering::Relaxed),
            self.jit_stats.jit_hits.load(Ordering::Relaxed),
            self.jit_stats.jit_misses.load(Ordering::Relaxed),
        )
    }
}

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

/// Global JIT compiler instance
static GLOBAL_JIT: OnceLock<RwLock<JitCompiler>> = OnceLock::new();

/// Get the global JIT compiler instance
pub fn get_jit_compiler() -> &'static RwLock<JitCompiler> {
    GLOBAL_JIT.get_or_init(|| RwLock::new(JitCompiler::new()))
}

/// Check if function execution should be tracked for JIT
#[inline]
pub fn should_track_for_jit(function_name: &str) -> bool {
    // Track frequently used functions
    matches!(
        function_name,
        "add"
            | "multiply"
            | "concat"
            | "strlen"
            | "array_merge"
            | "in_array"
            | "count"
            | "isset"
            | "empty"
    )
}

/// Increment execution counter for a function
#[inline]
pub fn increment_execution_counter(function_name: &str) -> bool {
    if should_track_for_jit(function_name) {
        let jit = get_jit_compiler();
        let mut jit = jit.write().unwrap();
        let counter = jit.get_counter(function_name);
        let count = counter.increment();
        let should_compile =
            count >= JIT_THRESHOLD && !counter.jit_compiled.load(Ordering::Relaxed);
        drop(jit); // Explicitly release the write lock to prevent deadlock
        should_compile
    } else {
        false
    }
}

/// Execute function with JIT optimization
pub fn execute_with_jit(
    function_name: &str,
    execute_data: &mut ExecuteData,
    op_array: &OpArray,
) -> Result<PhpResult, String> {
    // First, check if we have a compiled version
    {
        let jit = get_jit_compiler();
        let jit = jit.read().unwrap();
        if let Some(compiled_fn) = jit.get_compiled_function(function_name) {
            // Note: This is a read-only operation, so we don't need a write lock
            drop(jit); // Release the read lock

            // We need to update stats - get a write lock just for that
            let jit = get_jit_compiler().write().unwrap();
            jit.jit_stats.jit_hits.fetch_add(1, Ordering::Relaxed);
            return compiled_fn(execute_data);
        }
    }

    // Check if we should compile this function
    {
        let mut jit = get_jit_compiler().write().unwrap();
        if jit.should_compile(function_name) {
            if let Ok(compiled_fn) = jit.compile_function(function_name, op_array) {
                return compiled_fn(execute_data);
            }
        }

        // Fallback to interpreter
        jit.jit_stats.jit_misses.fetch_add(1, Ordering::Relaxed);
    }
    Ok(PhpResult::Success)
}

/// Inline optimization for simple operations
pub fn try_inline_operation(opcode: Opcode, op1: &Val, op2: &Val) -> Option<Val> {
    match opcode {
        Opcode::Add => Some(crate::engine::operators::zval_add(op1, op2)),
        Opcode::Mul => Some(crate::engine::operators::zval_mul(op1, op2)),
        Opcode::Concat => {
            let s1 = crate::engine::operators::zval_get_string(op1);
            let s2 = crate::engine::operators::zval_get_string(op2);
            let result = super::perf_alloc::fast_concat(s1.as_str(), s2.as_str());
            Some(Val::new(
                crate::engine::types::PhpValue::String(Box::new(result)),
                crate::engine::types::PhpType::String,
            ))
        }
        _ => None,
    }
}

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

    #[test]
    fn test_execution_counter() {
        let counter = ExecutionCounter::new();

        for i in 0..JIT_THRESHOLD {
            assert_eq!(counter.increment(), i + 1);
        }

        assert!(counter.should_jit_compile());
        assert!(!counter.should_jit_compile()); // Should be false after first check

        counter.mark_jit_compiled();
        assert!(!counter.should_jit_compile());
    }

    #[test]
    fn test_jit_compiler() {
        let mut jit = JitCompiler::new();

        // Test compilation
        let op_array = OpArray::new("test.php".to_string());
        let _compiled_fn = jit.compile_function("add", &op_array).unwrap();

        assert!(jit.get_compiled_function("add").is_some());

        let stats = jit.get_stats();
        assert_eq!(stats.0, 1); // functions_compiled
    }

    #[test]
    fn test_should_track_for_jit() {
        assert!(should_track_for_jit("add"));
        assert!(should_track_for_jit("multiply"));
        assert!(should_track_for_jit("concat"));
        assert!(!should_track_for_jit("rare_function"));
    }

    #[test]
    fn test_increment_execution_counter() {
        assert!(!increment_execution_counter("rare_function"));

        // This would need many calls to trigger JIT in a real test
        for _ in 0..JIT_THRESHOLD {
            increment_execution_counter("add");
        }

        // The next call would return true
        assert!(increment_execution_counter("add"));
    }
}