oxiz_solver/combination/

coordinator.rs

//! Theory Combination Coordinator
#![allow(missing_docs)] // Under development
//!
//! This module coordinates multiple theory solvers using the Nelson-Oppen method
//! with optimizations:
//! - Lazy vs eager theory combination
//! - Shared term management
//! - Equality sharing between theories
//! - Conflict minimization across theories

use rustc_hash::{FxHashMap, FxHashSet};
use std::collections::VecDeque;

/// Placeholder term identifier
pub type TermId = usize;

/// Theory identifier
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum TheoryId {
    Core,
    Arithmetic,
    BitVector,
    Array,
    Datatype,
    String,
    Uninterpreted,
}

/// Theory interface
pub trait TheorySolver {
    /// Get theory ID
    fn theory_id(&self) -> TheoryId;

    /// Assert a formula
    fn assert_formula(&mut self, formula: TermId) -> Result<(), String>;

    /// Check satisfiability
    fn check_sat(&mut self) -> Result<SatResult, String>;

    /// Get model (if SAT)
    fn get_model(&self) -> Option<FxHashMap<TermId, TermId>>;

    /// Get conflict explanation (if UNSAT)
    fn get_conflict(&self) -> Option<Vec<TermId>>;

    /// Backtrack to a level
    fn backtrack(&mut self, level: usize) -> Result<(), String>;

    /// Get implied equalities
    fn get_implied_equalities(&self) -> Vec<(TermId, TermId)>;

    /// Notify of external equality
    fn notify_equality(&mut self, lhs: TermId, rhs: TermId) -> Result<(), String>;
}

/// Satisfiability result
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SatResult {
    Sat,
    Unsat,
    Unknown,
}

/// Shared term between theories
#[derive(Debug, Clone)]
pub struct SharedTerm {
    /// The term
    pub term: TermId,
    /// Theories that use this term
    pub theories: FxHashSet<TheoryId>,
    /// Current equivalence class representative
    pub representative: TermId,
}

/// Equality propagation item
#[derive(Debug, Clone)]
pub struct EqualityProp {
    /// Left-hand side
    pub lhs: TermId,
    /// Right-hand side
    pub rhs: TermId,
    /// Source theory
    pub source: TheoryId,
    /// Explanation (justification)
    pub explanation: Vec<TermId>,
}

/// Statistics for theory combination
#[derive(Debug, Clone, Default)]
pub struct CoordinatorStats {
    pub check_sat_calls: u64,
    pub theory_conflicts: u64,
    pub equalities_propagated: u64,
    pub shared_terms_count: usize,
    pub theory_combination_rounds: u64,
}

/// Configuration for theory combination
#[derive(Debug, Clone)]
pub struct CoordinatorConfig {
    /// Use eager theory combination (propagate all equalities immediately)
    pub eager_combination: bool,
    /// Maximum theory combination rounds
    pub max_combination_rounds: usize,
    /// Enable conflict minimization across theories
    pub minimize_conflicts: bool,
}

impl Default for CoordinatorConfig {
    fn default() -> Self {
        Self {
            eager_combination: false,
            max_combination_rounds: 10,
            minimize_conflicts: true,
        }
    }
}

/// Theory combination coordinator
pub struct TheoryCoordinator {
    config: CoordinatorConfig,
    stats: CoordinatorStats,
    /// Registered theory solvers
    theories: FxHashMap<TheoryId, Box<dyn TheorySolver>>,
    /// Shared terms between theories
    shared_terms: FxHashMap<TermId, SharedTerm>,
    /// Pending equality propagations
    pending_equalities: VecDeque<EqualityProp>,
    /// Current decision level
    current_level: usize,
}

impl TheoryCoordinator {
    /// Create a new theory coordinator
    pub fn new(config: CoordinatorConfig) -> Self {
        Self {
            config,
            stats: CoordinatorStats::default(),
            theories: FxHashMap::default(),
            shared_terms: FxHashMap::default(),
            pending_equalities: VecDeque::new(),
            current_level: 0,
        }
    }

    /// Register a theory solver
    pub fn register_theory(&mut self, theory: Box<dyn TheorySolver>) {
        let theory_id = theory.theory_id();
        self.theories.insert(theory_id, theory);
    }

    /// Assert a formula to the appropriate theory
    pub fn assert_formula(&mut self, formula: TermId, theory: TheoryId) -> Result<(), String> {
        if let Some(solver) = self.theories.get_mut(&theory) {
            solver.assert_formula(formula)?;

            // Identify shared terms
            self.identify_shared_terms(formula)?;
        } else {
            return Err(format!("Theory {:?} not registered", theory));
        }

        Ok(())
    }

    /// Check satisfiability with theory combination
    pub fn check_sat(&mut self) -> Result<SatResult, String> {
        self.stats.check_sat_calls += 1;

        // Phase 1: Check individual theories
        for solver in self.theories.values_mut() {
            let result = solver.check_sat()?;

            match result {
                SatResult::Unsat => {
                    self.stats.theory_conflicts += 1;
                    return Ok(SatResult::Unsat);
                }
                SatResult::Unknown => {
                    return Ok(SatResult::Unknown);
                }
                SatResult::Sat => {
                    // Continue to next theory
                }
            }
        }

        // Phase 2: Theory combination via equality sharing
        if self.config.eager_combination {
            self.eager_theory_combination()
        } else {
            self.lazy_theory_combination()
        }
    }

    /// Eager theory combination: propagate all equalities immediately
    fn eager_theory_combination(&mut self) -> Result<SatResult, String> {
        let mut iteration = 0;

        loop {
            self.stats.theory_combination_rounds += 1;
            iteration += 1;

            if iteration > self.config.max_combination_rounds {
                return Ok(SatResult::Unknown);
            }

            // Collect implied equalities from all theories
            let mut new_equalities = Vec::new();

            for (theory_id, solver) in &self.theories {
                let equalities = solver.get_implied_equalities();

                for (lhs, rhs) in equalities {
                    // Only propagate equalities between shared terms
                    if self.is_shared_term(lhs) || self.is_shared_term(rhs) {
                        new_equalities.push(EqualityProp {
                            lhs,
                            rhs,
                            source: *theory_id,
                            explanation: vec![],
                        });
                    }
                }
            }

            // No new equalities: fixed point reached
            if new_equalities.is_empty() {
                return Ok(SatResult::Sat);
            }

            // Propagate equalities to all theories
            for eq in new_equalities {
                self.propagate_equality(eq)?;
            }

            // Re-check theories for conflicts
            for solver in self.theories.values_mut() {
                match solver.check_sat()? {
                    SatResult::Unsat => {
                        self.stats.theory_conflicts += 1;
                        return Ok(SatResult::Unsat);
                    }
                    SatResult::Unknown => {
                        return Ok(SatResult::Unknown);
                    }
                    SatResult::Sat => {}
                }
            }
        }
    }

    /// Lazy theory combination: propagate equalities on-demand
    fn lazy_theory_combination(&mut self) -> Result<SatResult, String> {
        // Process pending equalities
        while let Some(eq) = self.pending_equalities.pop_front() {
            self.propagate_equality(eq)?;

            // Check for conflicts after each propagation
            for solver in self.theories.values_mut() {
                match solver.check_sat()? {
                    SatResult::Unsat => {
                        self.stats.theory_conflicts += 1;
                        return Ok(SatResult::Unsat);
                    }
                    SatResult::Unknown => {
                        return Ok(SatResult::Unknown);
                    }
                    SatResult::Sat => {}
                }
            }
        }

        Ok(SatResult::Sat)
    }

    /// Propagate an equality to all relevant theories
    fn propagate_equality(&mut self, eq: EqualityProp) -> Result<(), String> {
        self.stats.equalities_propagated += 1;

        // Update equivalence classes
        self.merge_equivalence_classes(eq.lhs, eq.rhs)?;

        // Notify all theories that use these terms
        let theories_to_notify = self.get_theories_for_terms(eq.lhs, eq.rhs);

        for theory_id in theories_to_notify {
            if theory_id != eq.source
                && let Some(solver) = self.theories.get_mut(&theory_id)
            {
                solver.notify_equality(eq.lhs, eq.rhs)?;
            }
        }

        Ok(())
    }

    /// Identify shared terms in a formula
    fn identify_shared_terms(&mut self, _formula: TermId) -> Result<(), String> {
        // Placeholder: would traverse formula AST and identify terms used by multiple theories
        // For now, just update stats
        self.stats.shared_terms_count = self.shared_terms.len();
        Ok(())
    }

    /// Check if a term is shared between theories
    fn is_shared_term(&self, term: TermId) -> bool {
        self.shared_terms
            .get(&term)
            .is_some_and(|st| st.theories.len() > 1)
    }

    /// Get theories that use given terms
    fn get_theories_for_terms(&self, lhs: TermId, rhs: TermId) -> FxHashSet<TheoryId> {
        let mut theories = FxHashSet::default();

        if let Some(st) = self.shared_terms.get(&lhs) {
            theories.extend(&st.theories);
        }

        if let Some(st) = self.shared_terms.get(&rhs) {
            theories.extend(&st.theories);
        }

        theories
    }

    /// Merge equivalence classes for two terms
    fn merge_equivalence_classes(&mut self, lhs: TermId, rhs: TermId) -> Result<(), String> {
        // Get representatives
        let lhs_rep = self.find_representative(lhs);
        let rhs_rep = self.find_representative(rhs);

        if lhs_rep == rhs_rep {
            return Ok(());
        }

        // Union: make lhs_rep point to rhs_rep
        if let Some(st) = self.shared_terms.get_mut(&lhs_rep) {
            st.representative = rhs_rep;
        }

        Ok(())
    }

    /// Find equivalence class representative
    fn find_representative(&self, term: TermId) -> TermId {
        if let Some(st) = self.shared_terms.get(&term)
            && st.representative != term
        {
            // Path compression would be applied here
            return self.find_representative(st.representative);
        }
        term
    }

    /// Add a shared term
    pub fn add_shared_term(&mut self, term: TermId, theory: TheoryId) {
        self.shared_terms
            .entry(term)
            .or_insert_with(|| SharedTerm {
                term,
                theories: FxHashSet::default(),
                representative: term,
            })
            .theories
            .insert(theory);

        self.stats.shared_terms_count = self.shared_terms.len();
    }

    /// Enqueue an equality for propagation
    pub fn enqueue_equality(&mut self, lhs: TermId, rhs: TermId, source: TheoryId) {
        self.pending_equalities.push_back(EqualityProp {
            lhs,
            rhs,
            source,
            explanation: vec![],
        });
    }

    /// Backtrack all theories to a level
    pub fn backtrack(&mut self, level: usize) -> Result<(), String> {
        self.current_level = level;

        for solver in self.theories.values_mut() {
            solver.backtrack(level)?;
        }

        // Clear pending equalities
        self.pending_equalities.clear();

        Ok(())
    }

    /// Get combined model from all theories
    pub fn get_model(&self) -> Option<FxHashMap<TermId, TermId>> {
        let mut combined_model = FxHashMap::default();

        for solver in self.theories.values() {
            if let Some(model) = solver.get_model() {
                combined_model.extend(model);
            } else {
                return None;
            }
        }

        Some(combined_model)
    }

    /// Get combined conflict explanation
    pub fn get_conflict(&self) -> Option<Vec<TermId>> {
        // Collect conflicts from all theories
        let mut combined_conflict = Vec::new();

        for solver in self.theories.values() {
            if let Some(conflict) = solver.get_conflict() {
                combined_conflict.extend(conflict);
            }
        }

        if combined_conflict.is_empty() {
            None
        } else {
            // Minimize if enabled
            if self.config.minimize_conflicts {
                Some(self.minimize_conflict(combined_conflict))
            } else {
                Some(combined_conflict)
            }
        }
    }

    /// Minimize a conflict explanation
    fn minimize_conflict(&self, mut conflict: Vec<TermId>) -> Vec<TermId> {
        // Placeholder: would use resolution to minimize
        // For now, just remove duplicates
        conflict.sort();
        conflict.dedup();
        conflict
    }

    /// Get statistics
    pub fn stats(&self) -> &CoordinatorStats {
        &self.stats
    }

    /// Get current decision level
    pub fn current_level(&self) -> usize {
        self.current_level
    }

    /// Increment decision level
    pub fn increment_level(&mut self) {
        self.current_level += 1;
    }
}

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

    // Mock theory solver for testing
    struct MockTheory {
        id: TheoryId,
        sat_result: SatResult,
    }

    impl TheorySolver for MockTheory {
        fn theory_id(&self) -> TheoryId {
            self.id
        }

        fn assert_formula(&mut self, _formula: TermId) -> Result<(), String> {
            Ok(())
        }

        fn check_sat(&mut self) -> Result<SatResult, String> {
            Ok(self.sat_result)
        }

        fn get_model(&self) -> Option<FxHashMap<TermId, TermId>> {
            Some(FxHashMap::default())
        }

        fn get_conflict(&self) -> Option<Vec<TermId>> {
            None
        }

        fn backtrack(&mut self, _level: usize) -> Result<(), String> {
            Ok(())
        }

        fn get_implied_equalities(&self) -> Vec<(TermId, TermId)> {
            vec![]
        }

        fn notify_equality(&mut self, _lhs: TermId, _rhs: TermId) -> Result<(), String> {
            Ok(())
        }
    }

    #[test]
    fn test_coordinator_creation() {
        let config = CoordinatorConfig::default();
        let coordinator = TheoryCoordinator::new(config);
        assert_eq!(coordinator.stats.check_sat_calls, 0);
    }

    #[test]
    fn test_register_theory() {
        let config = CoordinatorConfig::default();
        let mut coordinator = TheoryCoordinator::new(config);

        let mock_theory = MockTheory {
            id: TheoryId::Arithmetic,
            sat_result: SatResult::Sat,
        };

        coordinator.register_theory(Box::new(mock_theory));
        assert!(coordinator.theories.contains_key(&TheoryId::Arithmetic));
    }

    #[test]
    fn test_check_sat_single_theory() {
        let config = CoordinatorConfig::default();
        let mut coordinator = TheoryCoordinator::new(config);

        let mock_theory = MockTheory {
            id: TheoryId::Arithmetic,
            sat_result: SatResult::Sat,
        };

        coordinator.register_theory(Box::new(mock_theory));

        let result = coordinator.check_sat();
        assert!(result.is_ok());
        assert_eq!(result.unwrap(), SatResult::Sat);
        assert_eq!(coordinator.stats.check_sat_calls, 1);
    }

    #[test]
    fn test_shared_term_management() {
        let config = CoordinatorConfig::default();
        let mut coordinator = TheoryCoordinator::new(config);

        coordinator.add_shared_term(1, TheoryId::Arithmetic);
        coordinator.add_shared_term(1, TheoryId::BitVector);

        assert!(coordinator.is_shared_term(1));
        assert_eq!(coordinator.stats.shared_terms_count, 1);
    }

    #[test]
    fn test_equivalence_classes() {
        let config = CoordinatorConfig::default();
        let mut coordinator = TheoryCoordinator::new(config);

        coordinator.add_shared_term(1, TheoryId::Arithmetic);
        coordinator.add_shared_term(2, TheoryId::Arithmetic);

        coordinator.merge_equivalence_classes(1, 2).unwrap();

        let rep1 = coordinator.find_representative(1);
        let rep2 = coordinator.find_representative(2);
        assert_eq!(rep1, rep2);
    }

    #[test]
    fn test_equality_propagation() {
        let config = CoordinatorConfig::default();
        let mut coordinator = TheoryCoordinator::new(config);

        coordinator.enqueue_equality(1, 2, TheoryId::Arithmetic);
        assert_eq!(coordinator.pending_equalities.len(), 1);
    }

    #[test]
    fn test_backtrack() {
        let config = CoordinatorConfig::default();
        let mut coordinator = TheoryCoordinator::new(config);

        let mock_theory = MockTheory {
            id: TheoryId::Arithmetic,
            sat_result: SatResult::Sat,
        };

        coordinator.register_theory(Box::new(mock_theory));
        coordinator.increment_level();
        coordinator.increment_level();

        assert_eq!(coordinator.current_level(), 2);

        coordinator.backtrack(0).unwrap();
        assert_eq!(coordinator.current_level(), 0);
    }

    #[test]
    fn test_get_model() {
        let config = CoordinatorConfig::default();
        let mut coordinator = TheoryCoordinator::new(config);

        let mock_theory = MockTheory {
            id: TheoryId::Arithmetic,
            sat_result: SatResult::Sat,
        };

        coordinator.register_theory(Box::new(mock_theory));

        let model = coordinator.get_model();
        assert!(model.is_some());
    }

    #[test]
    fn test_conflict_minimization() {
        let coordinator = TheoryCoordinator::new(CoordinatorConfig {
            minimize_conflicts: true,
            ..Default::default()
        });

        let conflict = vec![1, 2, 2, 3, 1, 4];
        let minimized = coordinator.minimize_conflict(conflict);

        assert_eq!(minimized, vec![1, 2, 3, 4]);
    }
}