picolog/src/prover/constraints.rs

use std::collections::HashMap;
use std::fmt::Display;

use crate::ast::Body;
use crate::ast::Predicate;
use crate::ast::Variable;
use crate::domain;
use crate::domain::Domain;
use crate::prover::predicate;

#[derive(Clone, Debug)]
pub struct Constraints
{
    pub(crate) predicates: HashMap<Variable, Predicate>,
    pub(crate) domains: HashMap<Variable, Domain>,
}

impl Constraints
{
    pub fn none() -> Self
    {
        Constraints {
            predicates: HashMap::new(),
            domains: HashMap::new(),
        }
    }

    pub fn with(variable: Variable, predicate: Option<Predicate>, domain: Option<Domain>) -> Self
    {
        let mut c = Constraints::none();
        if let Some(predicate) = predicate
        {
            c.predicates.insert(variable.clone(), predicate);
        }
        if let Some(domain) = domain
        {
            c.domains.insert(variable, domain);
        }
        c
    }

    pub fn try_append(
        &mut self,
        variable: &Variable,
        predicate: &Predicate,
        domain: &Domain,
    ) -> bool
    {
        let predicates = if let Some(other_predicate) = self.predicates.get(variable)
        {
            if predicate == other_predicate
            {
                return true;
            }
            // If variable is already contrained, we need to check if both contraints are
            // contradictory

            // Try to unify both predicates
            let unification = predicate.matches(other_predicate);
            if let Some(unification_contraints) = unification
            {
                // Instead of adding the variable = predicates contraint,
                // We can try adding the unification contraints which is implicitely the same
                if self.try_merge(&unification_contraints)
                {
                    self.predicates.insert(variable.clone(), predicate.clone());
                    true
                }
                else
                {
                    false
                }
            }
            else
            {
                // Unification is impossible, variable = predicates is contradictory under self
                false
            }
        }
        else
        {
            // No constraint
            self.predicates
                .insert(variable.clone(), (predicate.clone(), domain.clone()));
            true
        };

        // Check if domains are compatible
        let domains = if let Some((_, other_domain)) = self.predicates.get_mut(variable)
        {
            let intersection = domain.intersection(other_domain);
            if intersection.empty()
            {
                false
            }
            else
            {
                *other_domain = intersection;
                true
            }
        }
        else
        {
            true
        };

        domains && predicates
    }

    pub fn try_merge(&mut self, constraints: &Constraints) -> bool
    {
        // Trying to merge, is just trying to add all of the constraints into self
        let mut ok = self.clone();
        for (var, (pred, domain)) in constraints.predicates.iter()
        {
            if !ok.try_append(var, pred, domain)
            {
                return false;
            }
        }
        *self = ok;
        true
    }

    pub fn and(&self, constraints: &Constraints) -> Option<Constraints>
    {
        let mut new_contraints = self.clone();
        let valid = new_contraints.try_merge(constraints);
        if valid { Some(new_contraints) } else { None }
    }

    pub fn simplified(&self) -> Constraints
    {
        let mut max_sub = Constraints::none();
        for (var, (pred, domain)) in self.predicates.iter()
        {
            max_sub
                .predicates
                .insert(var.clone(), (pred.substitute(self), domain.clone()));
        }

        let mut stripped = max_sub.clone();
        'outer: for (var, _) in max_sub.predicates.iter()
        {
            if var.0.chars().next().is_some_and(|x| x == '_') || var.1.is_some()
            {
                for (_, (other_pred, _)) in max_sub.predicates.iter()
                {
                    if other_pred.contains_variable(var)
                    {
                        continue 'outer;
                    }
                }
                stripped.predicates.remove(var);
            }
        }

        stripped
    }
}

impl Predicate
{
    pub fn substitute(&self, constraints: &Constraints) -> Predicate
    {
        match self
        {
            Predicate::Variable(name) =>
            {
                if let Some((pred, _)) = constraints.predicates.get(name)
                {
                    pred.substitute(constraints)
                }
                else
                {
                    Predicate::Variable(name.clone())
                }
            }
            Predicate::Fixed(name, predicates) => Predicate::Fixed(
                name.clone(),
                predicates
                    .iter()
                    .map(|x| x.substitute(constraints))
                    .collect(),
            ),
            Predicate::Domain(domain) => Predicate::Domain(domain.clone()),
        }
    }

    pub fn contains_variable(&self, name: &Variable) -> bool
    {
        match self
        {
            Predicate::Variable(var_name) => name == var_name,
            Predicate::Fixed(_, predicates) => predicates.iter().any(|x| x.contains_variable(name)),
            Predicate::Domain(_) => false,
        }
    }
}

impl Body
{
    pub fn substitute(&self, constraints: &Constraints) -> Body
    {
        match self
        {
            Body::Term(predicate) => Body::Term(predicate.substitute(constraints)),
            Body::And(items) =>
            {
                Body::And(items.iter().map(|x| x.substitute(constraints)).collect())
            }
            Body::Or(items) => Body::Or(items.iter().map(|x| x.substitute(constraints)).collect()),
        }
    }
}

impl Default for Constraints
{
    fn default() -> Self
    {
        Self::none()
    }
}

impl Display for Constraints
{
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result
    {
        let len = self.predicates.len();
        for (i, (var, (pred, _))) in self.predicates.iter().enumerate()
        {
            write!(f, "{} = {}", var, pred)?;
            if i != len - 1
            {
                write!(f, ", ")?;
            }
        }
        write!(f, " ;; ")?;
        for (i, (var, (_, domain))) in self.predicates.iter().enumerate()
        {
            write!(f, "{} in {}", var, domain)?;
            if i != len - 1
            {
                write!(f, ", ")?;
            }
        }
        Ok(())
    }
}