use haloumi_ir::Prime;
use crate::pcl::{
display::{TextRepresentable, TextRepresentation},
expr::{
self, Expr, Wrap,
impls::{BinaryExpr, ConstExpr, NegExpr},
traits::ConstantFolding,
},
};
use super::{OpFolder, OpLike, try_fold};
#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub enum BinaryOp {
Add,
#[allow(dead_code)]
Sub,
Mul,
#[allow(dead_code)]
Div,
}
impl OpLike for BinaryOp {
fn extraible(&self) -> bool {
true
}
}
impl BinaryOp {
fn fold_add(&self, lhs: &Expr, rhs: &Expr) -> Option<Expr> {
if lhs.is_zero() {
return Some(rhs.clone());
}
None
}
fn fold_mul(&self, lhs: &Expr, rhs: &Expr, prime: Prime) -> Option<Expr> {
if lhs.is_one() {
return Some(rhs.clone());
}
if lhs.is_zero() {
return Some(lhs.clone());
}
if lhs.is_minus_one() {
return try_fold(NegExpr(rhs.clone()), prime);
}
None
}
fn fold_sub(&self, lhs: &Expr, rhs: &Expr, prime: Prime) -> Option<Expr> {
if lhs.is_zero() && rhs.is_zero() {
return Some(Wrap::new(ConstExpr(lhs.as_const().unwrap())));
}
if lhs.is_zero() {
return try_fold(NegExpr(rhs.clone()), prime);
}
if rhs.is_zero() {
return Some(lhs.clone());
}
None
}
fn fold_impl(&self, lhs: &Expr, rhs: &Expr, prime: Prime) -> Option<Expr> {
match self {
BinaryOp::Add => self.fold_add(lhs, rhs),
BinaryOp::Sub => self.fold_sub(lhs, rhs, prime),
BinaryOp::Mul => self.fold_mul(lhs, rhs, prime),
BinaryOp::Div => None,
}
}
}
impl OpFolder for BinaryOp {
fn fold(&self, lhs: Expr, rhs: Expr, prime: Prime) -> Option<Expr> {
self.fold_impl(&lhs, &rhs, prime).or_else(|| {
self.flip(&lhs, &rhs)
.and_then(|e| e.op().fold_impl(&e.lhs(), &e.rhs(), prime))
})
}
fn commutative(&self) -> bool {
matches!(self, BinaryOp::Add | BinaryOp::Mul)
}
fn flip(&self, lhs: &Expr, rhs: &Expr) -> Option<BinaryExpr<Self>> {
match self {
BinaryOp::Add => Some(BinaryExpr::new(BinaryOp::Add, rhs.clone(), lhs.clone())),
BinaryOp::Sub => Some(BinaryExpr::new(BinaryOp::Add, expr::neg(rhs), lhs.clone())),
BinaryOp::Mul => Some(BinaryExpr::new(BinaryOp::Mul, rhs.clone(), lhs.clone())),
BinaryOp::Div => None,
}
}
}
impl TextRepresentable for BinaryOp {
fn to_repr(&self) -> TextRepresentation<'_> {
TextRepresentation::atom(match self {
BinaryOp::Add => "+",
BinaryOp::Sub => "-",
BinaryOp::Mul => "*",
BinaryOp::Div => "/",
})
}
fn width_hint(&self) -> usize {
1
}
}