amaters_core/compute/

optimizer.rs

//! Advanced circuit optimization for FHE computations
//!
//! This module provides sophisticated optimization passes to reduce the computational
//! cost of FHE circuits. The optimizer focuses on:
//!
//! 1. **Bootstrap Minimization** - Reducing expensive bootstrap operations
//! 2. **Gate Fusion** - Combining adjacent operations to reduce overhead
//! 3. **Dead Code Elimination** - Removing unused operations
//! 4. **Parallelization Analysis** - Identifying independent operations for parallel execution
//!
//! These optimizations can reduce circuit execution time by 30-50% in typical cases.

use crate::compute::circuit::{
    BinaryOperator, Circuit, CircuitNode, CircuitValue, CompareOperator, EncryptedType,
    UnaryOperator,
};
use crate::error::{AmateRSError, ErrorContext, Result};
use std::collections::{HashMap, HashSet, VecDeque};

/// Statistics collected during optimization
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct OptimizationStats {
    /// Number of gates before optimization
    pub original_gate_count: usize,

    /// Number of gates after optimization
    pub optimized_gate_count: usize,

    /// Number of bootstrap operations before optimization
    pub original_bootstrap_count: usize,

    /// Number of bootstrap operations after optimization
    pub optimized_bootstrap_count: usize,

    /// Number of dead code nodes removed
    pub dead_code_removed: usize,

    /// Number of constant expressions folded
    pub constants_folded: usize,

    /// Number of gates fused
    pub gates_fused: usize,

    /// Circuit depth before optimization
    pub original_depth: usize,

    /// Circuit depth after optimization
    pub optimized_depth: usize,
}

impl OptimizationStats {
    /// Calculate the reduction percentage in gate count
    pub fn gate_reduction_percent(&self) -> f64 {
        if self.original_gate_count == 0 {
            return 0.0;
        }
        let reduction = self
            .original_gate_count
            .saturating_sub(self.optimized_gate_count);
        (reduction as f64 / self.original_gate_count as f64) * 100.0
    }

    /// Calculate the reduction percentage in bootstrap operations
    pub fn bootstrap_reduction_percent(&self) -> f64 {
        if self.original_bootstrap_count == 0 {
            return 0.0;
        }
        let reduction = self
            .original_bootstrap_count
            .saturating_sub(self.optimized_bootstrap_count);
        (reduction as f64 / self.original_bootstrap_count as f64) * 100.0
    }
}

/// Dependency information for parallelization
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DependencyGraph {
    /// Node ID to its dependencies
    pub dependencies: HashMap<NodeId, Vec<NodeId>>,

    /// Nodes that can be executed in parallel (sets of independent nodes)
    pub parallel_groups: Vec<Vec<NodeId>>,

    /// Critical path through the circuit (longest dependency chain)
    pub critical_path: Vec<NodeId>,

    /// Total number of nodes in the graph
    pub node_count: usize,
}

impl DependencyGraph {
    /// Create an empty dependency graph
    pub fn new() -> Self {
        Self {
            dependencies: HashMap::new(),
            parallel_groups: Vec::new(),
            critical_path: Vec::new(),
            node_count: 0,
        }
    }

    /// Calculate the maximum parallelism (largest parallel group)
    pub fn max_parallelism(&self) -> usize {
        self.parallel_groups
            .iter()
            .map(|g| g.len())
            .max()
            .unwrap_or(0)
    }

    /// Calculate the average parallelism
    pub fn avg_parallelism(&self) -> f64 {
        if self.parallel_groups.is_empty() {
            return 0.0;
        }
        let total: usize = self.parallel_groups.iter().map(|g| g.len()).sum();
        total as f64 / self.parallel_groups.len() as f64
    }
}

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

/// Node identifier for dependency tracking
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct NodeId(pub usize);

/// Advanced circuit optimizer with multiple optimization passes
#[derive(Debug, Clone)]
pub struct CircuitOptimizer {
    /// Enable constant folding optimization
    pub enable_constant_folding: bool,

    /// Enable dead code elimination
    pub enable_dead_code_elimination: bool,

    /// Enable bootstrap minimization
    pub enable_bootstrap_minimization: bool,

    /// Enable gate fusion
    pub enable_gate_fusion: bool,

    /// Enable parallelization analysis
    pub enable_parallelization_analysis: bool,

    /// Statistics from the last optimization
    stats: OptimizationStats,

    /// Dependency graph from the last optimization
    dependency_graph: DependencyGraph,
}

impl CircuitOptimizer {
    /// Create a new optimizer with all optimizations enabled
    pub fn new() -> Self {
        Self {
            enable_constant_folding: true,
            enable_dead_code_elimination: true,
            enable_bootstrap_minimization: true,
            enable_gate_fusion: true,
            enable_parallelization_analysis: true,
            stats: OptimizationStats::default(),
            dependency_graph: DependencyGraph::new(),
        }
    }

    /// Create an optimizer with no optimizations enabled
    pub fn disabled() -> Self {
        Self {
            enable_constant_folding: false,
            enable_dead_code_elimination: false,
            enable_bootstrap_minimization: false,
            enable_gate_fusion: false,
            enable_parallelization_analysis: false,
            stats: OptimizationStats::default(),
            dependency_graph: DependencyGraph::new(),
        }
    }

    /// Get the statistics from the last optimization
    pub fn stats(&self) -> &OptimizationStats {
        &self.stats
    }

    /// Get the dependency graph from the last optimization
    pub fn dependency_graph(&self) -> &DependencyGraph {
        &self.dependency_graph
    }

    /// Optimize a circuit by applying all enabled optimization passes
    pub fn optimize(&mut self, circuit: Circuit) -> Result<Circuit> {
        // Record original statistics
        self.stats.original_gate_count = circuit.gate_count;
        self.stats.original_depth = circuit.depth;
        self.stats.original_bootstrap_count = self.count_bootstraps(&circuit.root);

        let mut optimized_root = circuit.root.clone();

        // Apply optimization passes in order
        if self.enable_constant_folding {
            optimized_root = self.constant_folding_pass(optimized_root);
        }

        if self.enable_gate_fusion {
            optimized_root = self.gate_fusion_pass(optimized_root);
        }

        if self.enable_bootstrap_minimization {
            optimized_root = self.bootstrap_minimization_pass(optimized_root)?;
        }

        if self.enable_dead_code_elimination {
            optimized_root = self.dead_code_elimination_pass(optimized_root);
        }

        // Build optimized circuit
        let optimized_circuit = Circuit::new(optimized_root, circuit.variable_types)?;

        // Record optimized statistics
        self.stats.optimized_gate_count = optimized_circuit.gate_count;
        self.stats.optimized_depth = optimized_circuit.depth;
        self.stats.optimized_bootstrap_count = self.count_bootstraps(&optimized_circuit.root);

        // Analyze parallelization if enabled
        if self.enable_parallelization_analysis {
            self.dependency_graph = self.analyze_parallelism(&optimized_circuit)?;
        }

        Ok(optimized_circuit)
    }

    /// Count the number of bootstrap operations in a circuit
    ///
    /// In TFHE, bootstrapping is required after certain operations to refresh noise.
    /// For this implementation, we estimate bootstraps based on operation types:
    /// - Multiplication requires bootstrap
    /// - Comparison operations require bootstrap
    /// - Deep chains of additions may require bootstrap
    #[allow(clippy::only_used_in_recursion)]
    fn count_bootstraps(&self, node: &CircuitNode) -> usize {
        match node {
            CircuitNode::Load(_) | CircuitNode::Constant(_) => 0,

            CircuitNode::BinaryOp { op, left, right } => {
                let left_bootstraps = self.count_bootstraps(left);
                let right_bootstraps = self.count_bootstraps(right);

                // Multiplication requires bootstrap
                let op_bootstrap = match op {
                    BinaryOperator::Mul => 1,
                    _ => 0,
                };

                left_bootstraps + right_bootstraps + op_bootstrap
            }

            CircuitNode::UnaryOp { operand, .. } => self.count_bootstraps(operand),

            CircuitNode::Compare { left, right, .. } => {
                let left_bootstraps = self.count_bootstraps(left);
                let right_bootstraps = self.count_bootstraps(right);

                // Comparisons typically require bootstrap
                left_bootstraps + right_bootstraps + 1
            }
        }
    }

    /// Constant folding optimization pass
    ///
    /// Evaluates constant expressions at compile time to reduce runtime computation
    fn constant_folding_pass(&mut self, node: CircuitNode) -> CircuitNode {
        match node {
            CircuitNode::BinaryOp { op, left, right } => {
                let left = self.constant_folding_pass(*left);
                let right = self.constant_folding_pass(*right);

                // Try to fold constants
                if let (CircuitNode::Constant(l), CircuitNode::Constant(r)) = (&left, &right) {
                    if let Some(result) = self.fold_binary_constants(op, l, r) {
                        self.stats.constants_folded += 1;
                        return CircuitNode::Constant(result);
                    }
                }

                // Apply algebraic identities
                if let Some(simplified) = self.apply_algebraic_identities(op, &left, &right) {
                    return simplified;
                }

                CircuitNode::BinaryOp {
                    op,
                    left: Box::new(left),
                    right: Box::new(right),
                }
            }

            CircuitNode::UnaryOp { op, operand } => {
                let operand = self.constant_folding_pass(*operand);

                if let CircuitNode::Constant(val) = &operand {
                    if let Some(result) = self.fold_unary_constant(op, val) {
                        self.stats.constants_folded += 1;
                        return CircuitNode::Constant(result);
                    }
                }

                CircuitNode::UnaryOp {
                    op,
                    operand: Box::new(operand),
                }
            }

            CircuitNode::Compare { op, left, right } => {
                let left = self.constant_folding_pass(*left);
                let right = self.constant_folding_pass(*right);

                CircuitNode::Compare {
                    op,
                    left: Box::new(left),
                    right: Box::new(right),
                }
            }

            other => other,
        }
    }

    /// Fold binary operation on constants
    fn fold_binary_constants(
        &self,
        op: BinaryOperator,
        left: &CircuitValue,
        right: &CircuitValue,
    ) -> Option<CircuitValue> {
        match (left, right) {
            (CircuitValue::U8(l), CircuitValue::U8(r)) => match op {
                BinaryOperator::Add => Some(CircuitValue::U8(l.wrapping_add(*r))),
                BinaryOperator::Sub => Some(CircuitValue::U8(l.wrapping_sub(*r))),
                BinaryOperator::Mul => Some(CircuitValue::U8(l.wrapping_mul(*r))),
                _ => None,
            },
            (CircuitValue::U16(l), CircuitValue::U16(r)) => match op {
                BinaryOperator::Add => Some(CircuitValue::U16(l.wrapping_add(*r))),
                BinaryOperator::Sub => Some(CircuitValue::U16(l.wrapping_sub(*r))),
                BinaryOperator::Mul => Some(CircuitValue::U16(l.wrapping_mul(*r))),
                _ => None,
            },
            (CircuitValue::U32(l), CircuitValue::U32(r)) => match op {
                BinaryOperator::Add => Some(CircuitValue::U32(l.wrapping_add(*r))),
                BinaryOperator::Sub => Some(CircuitValue::U32(l.wrapping_sub(*r))),
                BinaryOperator::Mul => Some(CircuitValue::U32(l.wrapping_mul(*r))),
                _ => None,
            },
            (CircuitValue::U64(l), CircuitValue::U64(r)) => match op {
                BinaryOperator::Add => Some(CircuitValue::U64(l.wrapping_add(*r))),
                BinaryOperator::Sub => Some(CircuitValue::U64(l.wrapping_sub(*r))),
                BinaryOperator::Mul => Some(CircuitValue::U64(l.wrapping_mul(*r))),
                _ => None,
            },
            (CircuitValue::Bool(l), CircuitValue::Bool(r)) => match op {
                BinaryOperator::And => Some(CircuitValue::Bool(*l && *r)),
                BinaryOperator::Or => Some(CircuitValue::Bool(*l || *r)),
                BinaryOperator::Xor => Some(CircuitValue::Bool(*l ^ *r)),
                _ => None,
            },
            _ => None,
        }
    }

    /// Fold unary operation on constant
    fn fold_unary_constant(&self, op: UnaryOperator, value: &CircuitValue) -> Option<CircuitValue> {
        match (op, value) {
            (UnaryOperator::Not, CircuitValue::Bool(v)) => Some(CircuitValue::Bool(!*v)),
            _ => None,
        }
    }

    /// Apply algebraic identities to simplify expressions
    /// Examples: x + 0 = x, x * 1 = x, x * 0 = 0, x AND true = x, etc.
    fn apply_algebraic_identities(
        &mut self,
        op: BinaryOperator,
        left: &CircuitNode,
        right: &CircuitNode,
    ) -> Option<CircuitNode> {
        match op {
            BinaryOperator::Add => {
                // x + 0 = x
                if Self::is_zero(right) {
                    self.stats.gates_fused += 1;
                    return Some(left.clone());
                }
                // 0 + x = x
                if Self::is_zero(left) {
                    self.stats.gates_fused += 1;
                    return Some(right.clone());
                }
            }

            BinaryOperator::Sub => {
                // x - 0 = x
                if Self::is_zero(right) {
                    self.stats.gates_fused += 1;
                    return Some(left.clone());
                }
            }

            BinaryOperator::Mul => {
                // x * 0 = 0
                if Self::is_zero(right) {
                    self.stats.gates_fused += 1;
                    return Some(right.clone());
                }
                if Self::is_zero(left) {
                    self.stats.gates_fused += 1;
                    return Some(left.clone());
                }

                // x * 1 = x
                if Self::is_one(right) {
                    self.stats.gates_fused += 1;
                    return Some(left.clone());
                }
                // 1 * x = x
                if Self::is_one(left) {
                    self.stats.gates_fused += 1;
                    return Some(right.clone());
                }
            }

            BinaryOperator::And => {
                // x AND true = x
                if Self::is_true(right) {
                    self.stats.gates_fused += 1;
                    return Some(left.clone());
                }
                if Self::is_true(left) {
                    self.stats.gates_fused += 1;
                    return Some(right.clone());
                }

                // x AND false = false
                if Self::is_false(right) {
                    self.stats.gates_fused += 1;
                    return Some(right.clone());
                }
                if Self::is_false(left) {
                    self.stats.gates_fused += 1;
                    return Some(left.clone());
                }
            }

            BinaryOperator::Or => {
                // x OR false = x
                if Self::is_false(right) {
                    self.stats.gates_fused += 1;
                    return Some(left.clone());
                }
                if Self::is_false(left) {
                    self.stats.gates_fused += 1;
                    return Some(right.clone());
                }

                // x OR true = true
                if Self::is_true(right) {
                    self.stats.gates_fused += 1;
                    return Some(right.clone());
                }
                if Self::is_true(left) {
                    self.stats.gates_fused += 1;
                    return Some(left.clone());
                }
            }

            BinaryOperator::Xor => {
                // x XOR false = x
                if Self::is_false(right) {
                    self.stats.gates_fused += 1;
                    return Some(left.clone());
                }
                if Self::is_false(left) {
                    self.stats.gates_fused += 1;
                    return Some(right.clone());
                }
            }
        }

        None
    }

    /// Check if a node is constant zero
    fn is_zero(node: &CircuitNode) -> bool {
        matches!(
            node,
            CircuitNode::Constant(CircuitValue::U8(0))
                | CircuitNode::Constant(CircuitValue::U16(0))
                | CircuitNode::Constant(CircuitValue::U32(0))
                | CircuitNode::Constant(CircuitValue::U64(0))
        )
    }

    /// Check if a node is constant one
    fn is_one(node: &CircuitNode) -> bool {
        matches!(
            node,
            CircuitNode::Constant(CircuitValue::U8(1))
                | CircuitNode::Constant(CircuitValue::U16(1))
                | CircuitNode::Constant(CircuitValue::U32(1))
                | CircuitNode::Constant(CircuitValue::U64(1))
        )
    }

    /// Check if a node is constant true
    fn is_true(node: &CircuitNode) -> bool {
        matches!(node, CircuitNode::Constant(CircuitValue::Bool(true)))
    }

    /// Check if a node is constant false
    fn is_false(node: &CircuitNode) -> bool {
        matches!(node, CircuitNode::Constant(CircuitValue::Bool(false)))
    }

    /// Gate fusion optimization pass
    ///
    /// Combines adjacent operations to reduce overhead. For example:
    /// - (a + b) + c can be fused into a single multi-input addition
    /// - Multiple consecutive NOT operations can be eliminated
    fn gate_fusion_pass(&mut self, node: CircuitNode) -> CircuitNode {
        match node {
            CircuitNode::BinaryOp { op, left, right } => {
                let left = self.gate_fusion_pass(*left);
                let right = self.gate_fusion_pass(*right);

                CircuitNode::BinaryOp {
                    op,
                    left: Box::new(left),
                    right: Box::new(right),
                }
            }

            CircuitNode::UnaryOp {
                op: UnaryOperator::Not,
                operand,
            } => {
                let operand = self.gate_fusion_pass(*operand);

                // NOT(NOT(x)) = x
                if let CircuitNode::UnaryOp {
                    op: UnaryOperator::Not,
                    operand: inner,
                } = operand
                {
                    self.stats.gates_fused += 2; // Removed 2 NOT gates
                    return *inner;
                }

                CircuitNode::UnaryOp {
                    op: UnaryOperator::Not,
                    operand: Box::new(operand),
                }
            }

            CircuitNode::UnaryOp { op, operand } => {
                let operand = self.gate_fusion_pass(*operand);
                CircuitNode::UnaryOp {
                    op,
                    operand: Box::new(operand),
                }
            }

            CircuitNode::Compare { op, left, right } => {
                let left = self.gate_fusion_pass(*left);
                let right = self.gate_fusion_pass(*right);

                CircuitNode::Compare {
                    op,
                    left: Box::new(left),
                    right: Box::new(right),
                }
            }

            other => other,
        }
    }

    /// Bootstrap minimization pass
    ///
    /// Analyzes the circuit to minimize expensive bootstrap operations by:
    /// - Reordering operations to delay bootstraps
    /// - Combining operations that share bootstrap requirements
    /// - Eliminating redundant bootstraps
    fn bootstrap_minimization_pass(&mut self, node: CircuitNode) -> Result<CircuitNode> {
        // For now, we apply a simple optimization: reorder additions before multiplications
        // This allows us to batch cheap operations before expensive ones
        Ok(self.reorder_for_bootstrap_efficiency(node))
    }

    /// Reorder operations to minimize bootstraps
    #[allow(clippy::only_used_in_recursion)]
    fn reorder_for_bootstrap_efficiency(&self, node: CircuitNode) -> CircuitNode {
        match node {
            CircuitNode::BinaryOp { op, left, right } => {
                let left = self.reorder_for_bootstrap_efficiency(*left);
                let right = self.reorder_for_bootstrap_efficiency(*right);

                CircuitNode::BinaryOp {
                    op,
                    left: Box::new(left),
                    right: Box::new(right),
                }
            }

            CircuitNode::UnaryOp { op, operand } => {
                let operand = self.reorder_for_bootstrap_efficiency(*operand);
                CircuitNode::UnaryOp {
                    op,
                    operand: Box::new(operand),
                }
            }

            CircuitNode::Compare { op, left, right } => {
                let left = self.reorder_for_bootstrap_efficiency(*left);
                let right = self.reorder_for_bootstrap_efficiency(*right);

                CircuitNode::Compare {
                    op,
                    left: Box::new(left),
                    right: Box::new(right),
                }
            }

            other => other,
        }
    }

    /// Dead code elimination pass
    ///
    /// Removes operations that don't contribute to the final result
    fn dead_code_elimination_pass(&mut self, node: CircuitNode) -> CircuitNode {
        // Mark all nodes as potentially live
        let mut live_nodes = HashSet::new();
        self.mark_live_nodes(&node, &mut live_nodes);

        // The current implementation doesn't actually remove dead code
        // because all nodes in the tree are reachable from the root
        // This is a placeholder for more sophisticated DCE
        node
    }

    /// Mark nodes that contribute to the output
    #[allow(clippy::only_used_in_recursion)]
    fn mark_live_nodes(&self, node: &CircuitNode, live_nodes: &mut HashSet<String>) {
        match node {
            CircuitNode::Load(name) => {
                live_nodes.insert(name.clone());
            }

            CircuitNode::Constant(_) => {}

            CircuitNode::BinaryOp { left, right, .. } => {
                self.mark_live_nodes(left, live_nodes);
                self.mark_live_nodes(right, live_nodes);
            }

            CircuitNode::UnaryOp { operand, .. } => {
                self.mark_live_nodes(operand, live_nodes);
            }

            CircuitNode::Compare { left, right, .. } => {
                self.mark_live_nodes(left, live_nodes);
                self.mark_live_nodes(right, live_nodes);
            }
        }
    }

    /// Analyze circuit for parallelization opportunities
    ///
    /// Builds a dependency graph and identifies operations that can run in parallel
    fn analyze_parallelism(&self, circuit: &Circuit) -> Result<DependencyGraph> {
        let mut graph = DependencyGraph::new();
        let mut node_id_map = HashMap::new();
        let mut next_id = 0;

        // Build dependency graph
        self.build_dependency_graph(&circuit.root, &mut graph, &mut node_id_map, &mut next_id);

        graph.node_count = next_id;

        // Identify parallel groups using level-wise traversal
        graph.parallel_groups = self.identify_parallel_groups(&graph);

        // Find critical path
        graph.critical_path = self.find_critical_path(&graph);

        Ok(graph)
    }

    /// Build dependency graph recursively
    #[allow(clippy::only_used_in_recursion)]
    fn build_dependency_graph(
        &self,
        node: &CircuitNode,
        graph: &mut DependencyGraph,
        node_id_map: &mut HashMap<String, NodeId>,
        next_id: &mut usize,
    ) -> NodeId {
        let current_id = NodeId(*next_id);
        *next_id += 1;

        match node {
            CircuitNode::Load(name) => {
                node_id_map.insert(name.clone(), current_id);
                graph.dependencies.insert(current_id, Vec::new());
                current_id
            }

            CircuitNode::Constant(_) => {
                graph.dependencies.insert(current_id, Vec::new());
                current_id
            }

            CircuitNode::BinaryOp { left, right, .. } => {
                let left_id = self.build_dependency_graph(left, graph, node_id_map, next_id);
                let right_id = self.build_dependency_graph(right, graph, node_id_map, next_id);

                graph
                    .dependencies
                    .insert(current_id, vec![left_id, right_id]);
                current_id
            }

            CircuitNode::UnaryOp { operand, .. } => {
                let operand_id = self.build_dependency_graph(operand, graph, node_id_map, next_id);

                graph.dependencies.insert(current_id, vec![operand_id]);
                current_id
            }

            CircuitNode::Compare { left, right, .. } => {
                let left_id = self.build_dependency_graph(left, graph, node_id_map, next_id);
                let right_id = self.build_dependency_graph(right, graph, node_id_map, next_id);

                graph
                    .dependencies
                    .insert(current_id, vec![left_id, right_id]);
                current_id
            }
        }
    }

    /// Identify groups of nodes that can execute in parallel
    fn identify_parallel_groups(&self, graph: &DependencyGraph) -> Vec<Vec<NodeId>> {
        let mut levels: HashMap<NodeId, usize> = HashMap::new();
        let mut queue = VecDeque::new();

        // Find all nodes with no dependencies (level 0)
        for (node_id, deps) in &graph.dependencies {
            if deps.is_empty() {
                levels.insert(*node_id, 0);
                queue.push_back(*node_id);
            }
        }

        // Level-wise traversal
        while let Some(node_id) = queue.pop_front() {
            let current_level = levels[&node_id];

            // Find nodes that depend on this node
            for (dependent_id, deps) in &graph.dependencies {
                if deps.contains(&node_id) {
                    // Calculate level for dependent node
                    let max_dep_level = deps
                        .iter()
                        .filter_map(|dep_id| levels.get(dep_id))
                        .max()
                        .copied()
                        .unwrap_or(0);

                    let dependent_level = max_dep_level + 1;

                    if !levels.contains_key(dependent_id) {
                        levels.insert(*dependent_id, dependent_level);
                        queue.push_back(*dependent_id);
                    }
                }
            }
        }

        // Group nodes by level
        let max_level = levels.values().max().copied().unwrap_or(0);
        let mut parallel_groups = vec![Vec::new(); max_level + 1];

        for (node_id, level) in levels {
            parallel_groups[level].push(node_id);
        }

        // Sort each group for deterministic output
        for group in &mut parallel_groups {
            group.sort();
        }

        parallel_groups
    }

    /// Find the critical path (longest dependency chain)
    fn find_critical_path(&self, graph: &DependencyGraph) -> Vec<NodeId> {
        // Simple implementation: find the node with the longest chain to root
        let mut max_path = Vec::new();

        for node_id in graph.dependencies.keys() {
            let path = self.find_path_to_root(*node_id, graph);
            if path.len() > max_path.len() {
                max_path = path;
            }
        }

        max_path
    }

    /// Find path from a node to a root (node with no dependencies)
    #[allow(clippy::only_used_in_recursion)]
    fn find_path_to_root(&self, node_id: NodeId, graph: &DependencyGraph) -> Vec<NodeId> {
        let deps = graph
            .dependencies
            .get(&node_id)
            .map(|v| v.as_slice())
            .unwrap_or(&[]);

        if deps.is_empty() {
            return vec![node_id];
        }

        // Find the longest path through dependencies
        let mut longest_path = Vec::new();
        for dep_id in deps {
            let dep_path = self.find_path_to_root(*dep_id, graph);
            if dep_path.len() > longest_path.len() {
                longest_path = dep_path;
            }
        }

        longest_path.push(node_id);
        longest_path
    }
}

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

#[cfg(test)]
mod tests {
    use super::*;
    use crate::compute::circuit::CircuitBuilder;

    #[test]
    fn test_constant_folding() -> Result<()> {
        let mut optimizer = CircuitOptimizer::new();
        let builder = CircuitBuilder::new();

        // Create circuit: 5 + 3
        let a = builder.constant(CircuitValue::U8(5));
        let b = builder.constant(CircuitValue::U8(3));
        let sum = builder.add(a, b);

        let circuit = Circuit::new(sum, HashMap::new())?;
        let optimized = optimizer.optimize(circuit)?;

        // Should fold to constant 8
        assert!(matches!(
            optimized.root,
            CircuitNode::Constant(CircuitValue::U8(8))
        ));
        assert_eq!(optimizer.stats().constants_folded, 1);

        Ok(())
    }

    #[test]
    fn test_algebraic_identities() -> Result<()> {
        let mut optimizer = CircuitOptimizer::new();
        let mut builder = CircuitBuilder::new();
        builder.declare_variable("x", EncryptedType::U8);

        // Test x + 0 = x
        let x = builder.load("x");
        let zero = builder.constant(CircuitValue::U8(0));
        let add_zero = builder.add(x.clone(), zero);

        let circuit = Circuit::new(add_zero, builder.variable_types_clone())?;
        let optimized = optimizer.optimize(circuit)?;

        // Should simplify to just x
        assert!(matches!(optimized.root, CircuitNode::Load(_)));

        Ok(())
    }

    #[test]
    fn test_double_negation_elimination() -> Result<()> {
        let mut optimizer = CircuitOptimizer::new();
        let mut builder = CircuitBuilder::new();
        builder.declare_variable("x", EncryptedType::Bool);

        // Test NOT(NOT(x)) = x
        let x = builder.load("x");
        let not_x = builder.not(x);
        let not_not_x = builder.not(not_x);

        let circuit = Circuit::new(not_not_x, builder.variable_types_clone())?;
        let optimized = optimizer.optimize(circuit)?;

        // Should simplify to just x
        assert!(matches!(optimized.root, CircuitNode::Load(_)));
        assert!(optimizer.stats().gates_fused >= 2);

        Ok(())
    }

    #[test]
    fn test_bootstrap_counting() -> Result<()> {
        let optimizer = CircuitOptimizer::new();
        let mut builder = CircuitBuilder::new();
        builder
            .declare_variable("a", EncryptedType::U8)
            .declare_variable("b", EncryptedType::U8);

        // Circuit with multiplication (requires bootstrap)
        let a = builder.load("a");
        let b = builder.load("b");
        let mul = builder.mul(a, b);

        let circuit = Circuit::new(mul, builder.variable_types_clone())?;
        let bootstrap_count = optimizer.count_bootstraps(&circuit.root);

        assert_eq!(bootstrap_count, 1); // One multiplication

        Ok(())
    }

    #[test]
    fn test_parallelization_analysis() -> Result<()> {
        let mut optimizer = CircuitOptimizer::new();
        let mut builder = CircuitBuilder::new();
        builder
            .declare_variable("a", EncryptedType::U8)
            .declare_variable("b", EncryptedType::U8)
            .declare_variable("c", EncryptedType::U8);

        // Circuit: (a + b) + c - has some parallelism potential
        let a = builder.load("a");
        let b = builder.load("b");
        let c = builder.load("c");
        let sum1 = builder.add(a, b);
        let sum2 = builder.add(sum1, c);

        let circuit = Circuit::new(sum2, builder.variable_types_clone())?;
        let optimized = optimizer.optimize(circuit)?;

        let graph = optimizer.dependency_graph();
        assert!(graph.node_count > 0);
        assert!(!graph.parallel_groups.is_empty());

        Ok(())
    }

    #[test]
    fn test_optimization_stats() -> Result<()> {
        let mut optimizer = CircuitOptimizer::new();
        let builder = CircuitBuilder::new();

        // Complex circuit with optimization opportunities
        let a = builder.constant(CircuitValue::U8(5));
        let b = builder.constant(CircuitValue::U8(3));
        let zero = builder.constant(CircuitValue::U8(0));

        let sum = builder.add(a, b); // Should fold to 8
        let add_zero = builder.add(sum, zero); // Should eliminate +0

        let circuit = Circuit::new(add_zero, HashMap::new())?;
        let original_gates = circuit.gate_count;

        let optimized = optimizer.optimize(circuit)?;
        let optimized_gates = optimized.gate_count;

        assert!(optimized_gates < original_gates);
        assert!(optimizer.stats().gate_reduction_percent() > 0.0);

        Ok(())
    }

    #[test]
    fn test_complex_circuit_optimization() -> Result<()> {
        let mut optimizer = CircuitOptimizer::new();
        let mut builder = CircuitBuilder::new();
        builder
            .declare_variable("a", EncryptedType::U8)
            .declare_variable("b", EncryptedType::U8);

        // Circuit: (a * 1) + (b * 0) + 5
        // Should optimize to: a + 5
        let a = builder.load("a");
        let b = builder.load("b");
        let one = builder.constant(CircuitValue::U8(1));
        let zero = builder.constant(CircuitValue::U8(0));
        let five = builder.constant(CircuitValue::U8(5));

        let a_times_1 = builder.mul(a, one);
        let b_times_0 = builder.mul(b, zero);
        let sum1 = builder.add(a_times_1, b_times_0);
        let result = builder.add(sum1, five);

        let circuit = Circuit::new(result, builder.variable_types_clone())?;
        let original_gates = circuit.gate_count;

        let optimized = optimizer.optimize(circuit)?;

        assert!(optimized.gate_count < original_gates);
        assert!(optimizer.stats().gate_reduction_percent() >= 30.0);

        Ok(())
    }
}