picolog/src/prover/constraints.rs

use std::fmt::Display;

use litemap::LiteMap;

use crate::ast::Body;
use crate::ast::Predicate;
use crate::ast::Variable;

#[derive(Clone, Debug)]
pub struct Constraints
{
    set: LiteMap<Variable, Predicate>,
}

impl Constraints
{
    pub fn none() -> Self
    {
        Constraints {
            set: LiteMap::new(),
        }
    }

    pub fn with(variable: Variable, predicate: Predicate) -> Self
    {
        let mut c = Constraints::none();
        c.set.insert(variable, predicate);
        c
    }

    pub fn try_append(&mut self, variable: &Variable, predicate: &Predicate) -> bool
    {
        if let Some(other_predicate) = self.set.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.set.insert(variable.clone(), predicate.clone());
                    true
                }
                else
                {
                    false
                }
            }
            else
            {
                // Unification is impossible, variable = predicates is contradictory under self
                false
            }
        }
        else
        {
            // No constraint
            self.set.insert(variable.clone(), predicate.clone());
            true
        }
    }

    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) in constraints.set.iter()
        {
            if !ok.try_append(var, pred)
            {
                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) in self.set.iter()
        {
            max_sub.set.insert(var.clone(), pred.substitute(self));
        }

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

        stripped
    }
}

impl Predicate
{
    pub fn substitute(&self, constraints: &Constraints) -> Predicate
    {
        match self
        {
            Predicate::Variable(name) =>
            {
                if let Some(pred) = constraints.set.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(),
            ),
        }
    }

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

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.set.len();
        for (i, (var, pred)) in self.set.iter().enumerate()
        {
            write!(f, "{} = {}", var, pred)?;
            if i != len - 1
            {
                write!(f, ", ")?;
            }
        }
        Ok(())
    }
}