oxigdal-gpu-advanced 0.1.4

//! Shader optimization passes.

use crate::error::Result;
use naga::{Literal, Module};
use std::collections::HashSet;

/// Shader optimizer with various optimization passes
pub struct ShaderOptimizer {
    /// Enabled optimization passes
    enabled_passes: HashSet<OptimizationPass>,
}

/// Available optimization passes
#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)]
pub enum OptimizationPass {
    /// Dead code elimination
    DeadCodeElimination,
    /// Constant folding
    ConstantFolding,
    /// Loop unrolling
    LoopUnrolling,
    /// Common subexpression elimination
    CommonSubexpressionElimination,
    /// Register allocation optimization
    RegisterAllocation,
    /// Instruction combining
    InstructionCombining,
}

impl ShaderOptimizer {
    /// Create a new optimizer with default passes
    pub fn new() -> Self {
        let mut enabled_passes = HashSet::new();
        enabled_passes.insert(OptimizationPass::DeadCodeElimination);
        enabled_passes.insert(OptimizationPass::ConstantFolding);

        Self { enabled_passes }
    }

    /// Create optimizer with all passes enabled
    pub fn new_aggressive() -> Self {
        let mut enabled_passes = HashSet::new();
        enabled_passes.insert(OptimizationPass::DeadCodeElimination);
        enabled_passes.insert(OptimizationPass::ConstantFolding);
        enabled_passes.insert(OptimizationPass::LoopUnrolling);
        enabled_passes.insert(OptimizationPass::CommonSubexpressionElimination);
        enabled_passes.insert(OptimizationPass::RegisterAllocation);
        enabled_passes.insert(OptimizationPass::InstructionCombining);

        Self { enabled_passes }
    }

    /// Enable an optimization pass
    pub fn enable_pass(&mut self, pass: OptimizationPass) {
        self.enabled_passes.insert(pass);
    }

    /// Disable an optimization pass
    pub fn disable_pass(&mut self, pass: OptimizationPass) {
        self.enabled_passes.remove(&pass);
    }

    /// Check if a pass is enabled
    pub fn is_pass_enabled(&self, pass: OptimizationPass) -> bool {
        self.enabled_passes.contains(&pass)
    }

    /// Optimize a shader module
    pub fn optimize(&self, module: &Module) -> Result<Module> {
        let mut optimized = module.clone();

        // Apply enabled optimization passes
        if self.is_pass_enabled(OptimizationPass::DeadCodeElimination) {
            optimized = self.eliminate_dead_code(&optimized)?;
        }

        if self.is_pass_enabled(OptimizationPass::ConstantFolding) {
            optimized = self.fold_constants(&optimized)?;
        }

        if self.is_pass_enabled(OptimizationPass::LoopUnrolling) {
            optimized = self.unroll_loops(&optimized)?;
        }

        if self.is_pass_enabled(OptimizationPass::CommonSubexpressionElimination) {
            optimized = self.eliminate_common_subexpressions(&optimized)?;
        }

        if self.is_pass_enabled(OptimizationPass::InstructionCombining) {
            optimized = self.combine_instructions(&optimized)?;
        }

        Ok(optimized)
    }

    /// Dead code elimination pass
    fn eliminate_dead_code(&self, module: &Module) -> Result<Module> {
        let optimized = module.clone();

        // Track which functions are reachable from entry points
        let mut reachable_functions = HashSet::new();

        // Mark all entry points as reachable
        for entry in optimized.entry_points.iter() {
            // Entry points are always reachable
            // Collect called functions from entry point
            self.collect_called_functions(&optimized, &entry.function, &mut reachable_functions);
        }

        // Remove unreachable functions (keep functions called from entry points)
        let mut functions_to_remove = Vec::new();
        for (handle, _func) in optimized.functions.iter() {
            if !reachable_functions.contains(&handle) {
                functions_to_remove.push(handle);
            }
        }

        // Note: Actual removal would require rebuilding the Arena
        // For safety, we keep the module structure intact but mark optimization
        // This avoids handle invalidation issues

        Ok(optimized)
    }

    /// Collect all functions called from a given function
    fn collect_called_functions(
        &self,
        module: &Module,
        function: &naga::Function,
        reachable: &mut HashSet<naga::Handle<naga::Function>>,
    ) {
        use naga::Expression;

        // Walk through function body and collect function calls
        for statement in function.body.iter() {
            self.collect_calls_from_statement(module, statement, reachable);
        }

        // Also check expressions for function calls
        for (_handle, expr) in function.expressions.iter() {
            if let Expression::CallResult(_call_handle) = expr {
                // Track the called function
                // Note: In naga, function calls are tracked differently
                // This is a simplified implementation
            }
        }
    }

    /// Collect function calls from a statement
    fn collect_calls_from_statement(
        &self,
        module: &Module,
        statement: &naga::Statement,
        reachable: &mut HashSet<naga::Handle<naga::Function>>,
    ) {
        use naga::Statement;

        match statement {
            Statement::Block(block) => {
                for stmt in block.iter() {
                    self.collect_calls_from_statement(module, stmt, reachable);
                }
            }
            Statement::If { accept, reject, .. } => {
                for stmt in accept.iter() {
                    self.collect_calls_from_statement(module, stmt, reachable);
                }
                for stmt in reject.iter() {
                    self.collect_calls_from_statement(module, stmt, reachable);
                }
            }
            Statement::Loop {
                body, continuing, ..
            } => {
                for stmt in body.iter() {
                    self.collect_calls_from_statement(module, stmt, reachable);
                }
                for stmt in continuing.iter() {
                    self.collect_calls_from_statement(module, stmt, reachable);
                }
            }
            Statement::Switch { cases, .. } => {
                for case in cases {
                    for stmt in case.body.iter() {
                        self.collect_calls_from_statement(module, stmt, reachable);
                    }
                }
            }
            Statement::Call { function, .. } => {
                reachable.insert(*function);
                // Recursively collect from called function
                if let Ok(func) = module.functions.try_get(*function) {
                    self.collect_called_functions(module, func, reachable);
                }
            }
            _ => {}
        }
    }

    /// Constant folding pass
    fn fold_constants(&self, module: &Module) -> Result<Module> {
        use naga::Expression;

        let mut optimized = module.clone();

        // Walk through all functions and fold constant expressions
        for (_handle, function) in optimized.functions.iter_mut() {
            // Collect modifications first to avoid borrowing conflicts
            let mut modifications = Vec::new();

            for (expr_handle, expr) in function.expressions.iter() {
                // Fold binary operations with constant operands
                if let Expression::Binary { op, left, right } = expr {
                    let left_val = function.expressions.try_get(*left);
                    let right_val = function.expressions.try_get(*right);

                    // Check if both operands are literals
                    if let (Ok(Expression::Literal(left_lit)), Ok(Expression::Literal(right_lit))) =
                        (left_val, right_val)
                    {
                        // Fold arithmetic operations
                        let folded = self.fold_binary_op(*op, left_lit, right_lit);
                        if let Some(result) = folded {
                            modifications.push((expr_handle, Expression::Literal(result)));
                        }
                    }
                }

                // Fold unary operations
                if let Expression::Unary { op, expr: operand } = expr {
                    if let Ok(Expression::Literal(lit)) = function.expressions.try_get(*operand) {
                        let folded = self.fold_unary_op(*op, lit);
                        if let Some(result) = folded {
                            modifications.push((expr_handle, Expression::Literal(result)));
                        }
                    }
                }
            }

            // Apply modifications
            for (handle, new_expr) in modifications {
                function.expressions[handle] = new_expr;
            }
        }

        // Also fold constants in entry points
        for entry in optimized.entry_points.iter_mut() {
            // Collect modifications first to avoid borrowing conflicts
            let mut modifications = Vec::new();

            for (expr_handle, expr) in entry.function.expressions.iter() {
                if let Expression::Binary { op, left, right } = expr {
                    let left_val = entry.function.expressions.try_get(*left);
                    let right_val = entry.function.expressions.try_get(*right);

                    if let (Ok(Expression::Literal(left_lit)), Ok(Expression::Literal(right_lit))) =
                        (left_val, right_val)
                    {
                        let folded = self.fold_binary_op(*op, left_lit, right_lit);
                        if let Some(result) = folded {
                            modifications.push((expr_handle, Expression::Literal(result)));
                        }
                    }
                }

                // Fold unary operations
                if let Expression::Unary { op, expr: operand } = expr {
                    if let Ok(Expression::Literal(lit)) =
                        entry.function.expressions.try_get(*operand)
                    {
                        let folded = self.fold_unary_op(*op, lit);
                        if let Some(result) = folded {
                            modifications.push((expr_handle, Expression::Literal(result)));
                        }
                    }
                }
            }

            // Apply modifications
            for (handle, new_expr) in modifications {
                entry.function.expressions[handle] = new_expr;
            }
        }

        Ok(optimized)
    }

    /// Fold a binary operation on constant literals
    fn fold_binary_op(
        &self,
        op: naga::BinaryOperator,
        left: &Literal,
        right: &Literal,
    ) -> Option<Literal> {
        use naga::{BinaryOperator, Literal};

        match (left, right) {
            (Literal::I32(a), Literal::I32(b)) => match op {
                BinaryOperator::Add => Some(Literal::I32(a.wrapping_add(*b))),
                BinaryOperator::Subtract => Some(Literal::I32(a.wrapping_sub(*b))),
                BinaryOperator::Multiply => Some(Literal::I32(a.wrapping_mul(*b))),
                BinaryOperator::Divide => {
                    if *b != 0 {
                        a.checked_div(*b).map(Literal::I32)
                    } else {
                        None
                    }
                }
                _ => None,
            },
            (Literal::F32(a), Literal::F32(b)) => match op {
                BinaryOperator::Add => Some(Literal::F32(a + b)),
                BinaryOperator::Subtract => Some(Literal::F32(a - b)),
                BinaryOperator::Multiply => Some(Literal::F32(a * b)),
                BinaryOperator::Divide => Some(Literal::F32(a / b)),
                _ => None,
            },
            _ => None,
        }
    }

    /// Fold a unary operation on constant literal
    fn fold_unary_op(&self, op: naga::UnaryOperator, operand: &Literal) -> Option<Literal> {
        use naga::{Literal, UnaryOperator};

        match operand {
            Literal::I32(val) => match op {
                UnaryOperator::Negate => Some(Literal::I32(-val)),
                UnaryOperator::LogicalNot => Some(Literal::Bool(*val == 0)),
                _ => None,
            },
            Literal::F32(val) => match op {
                UnaryOperator::Negate => Some(Literal::F32(-val)),
                _ => None,
            },
            Literal::Bool(val) => match op {
                UnaryOperator::LogicalNot => Some(Literal::Bool(!val)),
                _ => None,
            },
            _ => None,
        }
    }

    /// Loop unrolling pass
    fn unroll_loops(&self, module: &Module) -> Result<Module> {
        // Loop unrolling in naga is complex and requires analyzing loop bounds
        // For now, we implement a marker that identifies candidates
        // Full implementation would:
        // 1. Analyze loop bounds to determine if they're constant
        // 2. Estimate unrolled code size
        // 3. Replicate loop body for each iteration
        // 4. Update variable references and phi nodes

        // Return module unchanged for safety
        // A production implementation would use naga's control flow analysis
        Ok(module.clone())
    }

    /// Common subexpression elimination pass
    fn eliminate_common_subexpressions(&self, module: &Module) -> Result<Module> {
        use std::collections::HashMap;

        let mut optimized = module.clone();

        // CSE implementation: Find duplicate expressions and reuse results
        // This is simplified - a full implementation would use value numbering

        for (_handle, function) in optimized.functions.iter_mut() {
            let mut expression_map: HashMap<u64, Vec<naga::Handle<naga::Expression>>> =
                HashMap::new();

            // Build map of expression hashes to handles
            for (handle, expr) in function.expressions.iter() {
                let hash = self.hash_expression(expr);
                expression_map.entry(hash).or_default().push(handle);
            }

            // Identify expressions that appear multiple times
            // In a full implementation, we would replace later occurrences
            // with references to the first occurrence
        }

        Ok(optimized)
    }

    /// Hash an expression for CSE
    fn hash_expression(&self, expr: &naga::Expression) -> u64 {
        use std::collections::hash_map::DefaultHasher;
        use std::hash::{Hash, Hasher};

        let mut hasher = DefaultHasher::new();
        // Hash the expression discriminant and key fields
        // This is simplified - a full implementation would hash all relevant fields
        std::mem::discriminant(expr).hash(&mut hasher);
        hasher.finish()
    }

    /// Instruction combining pass
    fn combine_instructions(&self, module: &Module) -> Result<Module> {
        use naga::{BinaryOperator, Expression, Literal};

        let mut optimized = module.clone();

        // Implement strength reduction and algebraic simplifications
        // Note: Full implementation would require rebuilding the expression arena
        // to properly replace expressions. For safety, we keep structure intact.
        for (_handle, function) in optimized.functions.iter_mut() {
            // Collect optimization opportunities
            let mut _optimization_candidates: Vec<(naga::Handle<naga::Expression>, &Expression)> =
                Vec::new();

            for (_expr_handle, expr) in function.expressions.iter() {
                // Pattern: x * 2 -> x + x (addition faster than multiplication)
                // Pattern: x * 1 -> x
                // Pattern: x + 0 -> x
                // Pattern: x * 0 -> 0

                if let Expression::Binary { op, left: _, right } = expr {
                    let right_val = function.expressions.try_get(*right);

                    // x * 1 = x
                    if matches!(op, BinaryOperator::Multiply) {
                        if let Ok(Expression::Literal(lit)) = right_val {
                            if matches!(lit, Literal::I32(1))
                                || matches!(lit, Literal::F32(v) if *v == 1.0)
                            {
                                // Identify candidate for replacement
                                // Full implementation would rebuild expression arena
                            }
                        }
                    }

                    // x + 0 = x
                    if matches!(op, BinaryOperator::Add) {
                        if let Ok(Expression::Literal(lit)) = right_val {
                            if matches!(lit, Literal::I32(0))
                                || matches!(lit, Literal::F32(v) if *v == 0.0)
                            {
                                // Identify candidate for replacement
                            }
                        }
                    }
                }
            }

            // In a full implementation, we would apply collected optimizations here
        }

        Ok(optimized)
    }

    /// Get optimization level preset
    pub fn get_level_preset(level: OptimizationLevel) -> Self {
        match level {
            OptimizationLevel::None => Self {
                enabled_passes: HashSet::new(),
            },
            OptimizationLevel::Basic => {
                let mut optimizer = Self::new();
                optimizer.enable_pass(OptimizationPass::DeadCodeElimination);
                optimizer.enable_pass(OptimizationPass::ConstantFolding);
                optimizer
            }
            OptimizationLevel::Aggressive => Self::new_aggressive(),
        }
    }
}

/// Optimization level presets
#[derive(Debug, Clone, Copy)]
pub enum OptimizationLevel {
    /// No optimizations
    None,
    /// Basic optimizations
    Basic,
    /// Aggressive optimizations
    Aggressive,
}

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

/// Optimization metrics
#[derive(Debug, Clone, Default)]
pub struct OptimizationMetrics {
    /// Instructions removed
    pub instructions_removed: usize,
    /// Constants folded
    pub constants_folded: usize,
    /// Loops unrolled
    pub loops_unrolled: usize,
    /// Common subexpressions eliminated
    pub cse_eliminated: usize,
    /// Register pressure reduced
    pub register_pressure_reduction: f32,
}

impl OptimizationMetrics {
    /// Create new metrics
    pub fn new() -> Self {
        Self::default()
    }

    /// Get total optimization count
    pub fn total_optimizations(&self) -> usize {
        self.instructions_removed
            + self.constants_folded
            + self.loops_unrolled
            + self.cse_eliminated
    }

    /// Print metrics
    pub fn print(&self) {
        println!("\nOptimization Metrics:");
        println!("  Instructions removed: {}", self.instructions_removed);
        println!("  Constants folded: {}", self.constants_folded);
        println!("  Loops unrolled: {}", self.loops_unrolled);
        println!("  CSE eliminated: {}", self.cse_eliminated);
        println!(
            "  Register pressure reduction: {:.1}%",
            self.register_pressure_reduction * 100.0
        );
        println!("  Total optimizations: {}", self.total_optimizations());
    }
}

/// Optimization configuration
#[derive(Debug, Clone)]
pub struct OptimizationConfig {
    /// Maximum loop unroll iterations
    pub max_unroll_iterations: usize,
    /// Enable aggressive inlining
    pub aggressive_inlining: bool,
    /// Target register count
    pub target_register_count: Option<usize>,
    /// Enable vectorization
    pub vectorization: bool,
}

impl Default for OptimizationConfig {
    fn default() -> Self {
        Self {
            max_unroll_iterations: 4,
            aggressive_inlining: false,
            target_register_count: None,
            vectorization: true,
        }
    }
}

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

    #[test]
    fn test_optimizer_creation() {
        let optimizer = ShaderOptimizer::new();
        assert!(optimizer.is_pass_enabled(OptimizationPass::DeadCodeElimination));
        assert!(optimizer.is_pass_enabled(OptimizationPass::ConstantFolding));
    }

    #[test]
    fn test_aggressive_optimizer() {
        let optimizer = ShaderOptimizer::new_aggressive();
        assert!(optimizer.is_pass_enabled(OptimizationPass::LoopUnrolling));
        assert!(optimizer.is_pass_enabled(OptimizationPass::CommonSubexpressionElimination));
    }

    #[test]
    fn test_pass_enable_disable() {
        let mut optimizer = ShaderOptimizer::new();
        optimizer.disable_pass(OptimizationPass::DeadCodeElimination);
        assert!(!optimizer.is_pass_enabled(OptimizationPass::DeadCodeElimination));

        optimizer.enable_pass(OptimizationPass::LoopUnrolling);
        assert!(optimizer.is_pass_enabled(OptimizationPass::LoopUnrolling));
    }

    #[test]
    fn test_optimization_metrics() {
        let metrics = OptimizationMetrics {
            instructions_removed: 10,
            constants_folded: 5,
            loops_unrolled: 2,
            cse_eliminated: 3,
            register_pressure_reduction: 0.15,
        };

        assert_eq!(metrics.total_optimizations(), 20);
    }
}