haloumi-picus 0.5.6

Haloumi backend to the Picus Constraint Language.
Documentation
use haloumi_ir::{Felt, Prime};

use crate::pcl::{
    display::{TextRepresentable, TextRepresentation},
    expr::{self, Expr, impls::BinaryExpr},
};

use super::{OpFolder, OpLike};

#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub enum ConstraintKind {
    Lt,
    Le,
    Gt,
    Ge,
    Eq,
    Ne,
}

impl OpLike for ConstraintKind {
    fn extraible(&self) -> bool {
        false
    }
}

fn zip_option<L, R>(lhs: Option<L>, rhs: Option<R>) -> Option<(L, R)> {
    lhs.and_then(|lhs| rhs.map(|rhs| (lhs, rhs)))
}

impl ConstraintKind {
    pub fn cmp_felts(&self, lhs: &Felt, rhs: &Felt) -> bool {
        match self {
            ConstraintKind::Lt => lhs < rhs,
            ConstraintKind::Le => lhs <= rhs,
            ConstraintKind::Gt => lhs > rhs,
            ConstraintKind::Ge => lhs >= rhs,
            ConstraintKind::Eq => lhs == rhs,
            ConstraintKind::Ne => lhs != rhs,
        }
    }

    fn fold_impl(&self, lhs: &Expr, rhs: &Expr) -> Option<bool> {
        zip_option(lhs.as_const(), rhs.as_const()).map(|(lhs, rhs)| self.cmp_felts(&lhs, &rhs))
    }
}

impl OpFolder for ConstraintKind {
    fn fold(&self, lhs: Expr, rhs: Expr, prime: Prime) -> Option<Expr> {
        self.fold_impl(&lhs, &rhs).map(|b| {
            if b {
                expr::r#true(prime)
            } else {
                expr::r#false(prime)
            }
        })
    }

    fn commutative(&self) -> bool {
        matches!(self, ConstraintKind::Eq)
    }

    fn flip(&self, lhs: &Expr, rhs: &Expr) -> Option<BinaryExpr<Self>> {
        match self {
            ConstraintKind::Lt => Some(BinaryExpr::new(Self::Ge, rhs.clone(), lhs.clone())),
            ConstraintKind::Le => Some(BinaryExpr::new(Self::Gt, rhs.clone(), lhs.clone())),
            ConstraintKind::Gt => Some(BinaryExpr::new(Self::Le, rhs.clone(), lhs.clone())),
            ConstraintKind::Ge => Some(BinaryExpr::new(Self::Lt, rhs.clone(), lhs.clone())),
            ConstraintKind::Eq => None,
            ConstraintKind::Ne => None,
        }
    }
}

impl TextRepresentable for ConstraintKind {
    fn to_repr(&self) -> TextRepresentation<'_> {
        TextRepresentation::atom(match self {
            ConstraintKind::Lt => "<",
            ConstraintKind::Le => "<=",
            ConstraintKind::Gt => ">",
            ConstraintKind::Ge => ">=",
            ConstraintKind::Eq => "=",
            ConstraintKind::Ne => "!=",
        })
    }

    fn width_hint(&self) -> usize {
        match self {
            ConstraintKind::Lt | ConstraintKind::Gt | ConstraintKind::Eq => 1,
            ConstraintKind::Le | ConstraintKind::Ge | ConstraintKind::Ne => 2,
        }
    }
}