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,
}
}
}