use std::fmt;
use itertools::Itertools;
use log::debug;
use schemars::{JsonSchema, json_schema};
use serde::{Deserialize, Serialize};
use thiserror::Error;
use super::{ops::ParseExprError, *};
#[derive(Clone)]
pub struct Equation {
pub lhs: Box<dyn Expr>,
pub rhs: Box<dyn Expr>,
}
impl fmt::Display for Equation {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{} = {}", self.lhs.str(), self.rhs.str())
}
}
impl Serialize for Equation {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: serde::Serializer,
{
serializer.serialize_str(&self.str())
}
}
struct ExprDeserializeVisitor;
#[derive(Error, Debug)]
pub enum Error {
#[error("{0}")]
Message(String),
#[error("parsed expression is not an equation")]
NotAnEquation,
#[error("{0}")]
FailedParsing(#[from] ParseExprError),
}
impl<'de> serde::de::Visitor<'de> for ExprDeserializeVisitor {
type Value = Equation;
fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
formatter.write_str("a properly written equation")
}
fn visit_str<E>(self, v: &str) -> Result<Self::Value, E>
where
E: serde::de::Error,
{
Equation::from_str(v).map_err(|e| E::custom(e.to_string()))
}
}
impl<'de> Deserialize<'de> for Equation {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: serde::Deserializer<'de>,
{
deserializer.deserialize_str(ExprDeserializeVisitor)
}
}
impl fmt::Debug for Equation {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{}\n{}", self.str(), self.srepr())
}
}
impl Equation {
pub fn into_new(lhs: &Box<dyn Expr>, rhs: &Box<dyn Expr>) -> Equation {
Equation {
lhs: lhs.clone(),
rhs: rhs.clone(),
}
}
pub fn new(lhs: &dyn Expr, rhs: &dyn Expr) -> Equation {
Equation {
lhs: lhs.clone_box(),
rhs: rhs.clone_box(),
}
}
pub fn new_box(lhs: Box<dyn Expr>, rhs: Box<dyn Expr>) -> Box<dyn Expr> {
Box::new(Equation { lhs, rhs })
}
pub fn from_str(s: &str) -> Result<Equation, Error> {
ops::parse_expr(s)?.as_eq().ok_or(Error::NotAnEquation)
}
}
impl Expr for Equation {
fn name(&self) -> String {
"Eq".to_string()
}
fn get_ref<'a>(&'a self) -> &'a dyn Expr {
self as &dyn Expr
}
fn for_each_arg(&self, f: &mut dyn FnMut(&dyn Arg) -> ()) {
f(&*self.lhs);
f(&*self.rhs);
}
fn from_args(&self, args: Vec<Box<dyn Arg>>) -> Box<dyn Expr> {
Box::new(Equation {
lhs: args[0].clone().into(),
rhs: args[1].clone().into(),
})
}
fn clone_box(&self) -> Box<dyn Expr> {
Box::new(self.clone())
}
fn str(&self) -> String {
format!("{} = {}", self.lhs.str(), self.rhs.str())
}
}
impl std::ops::SubAssign<&dyn Expr> for Equation {
fn sub_assign(&mut self, rhs: &dyn Expr) {
self.lhs -= rhs;
self.rhs -= rhs;
}
}
impl std::ops::DivAssign<&dyn Expr> for Equation {
fn div_assign(&mut self, rhs: &dyn Expr) {
self.lhs /= rhs;
self.rhs /= rhs;
}
}
impl std::cmp::PartialEq for Equation {
fn eq(&self, other: &Self) -> bool {
&self.lhs == &other.lhs && &self.rhs == &other.rhs
}
}
#[derive(Error, Debug)]
#[error("failed to solve equation {equation} for unknowns {unknowns} : {reason}")]
pub struct SolvingError {
unknowns: String,
equation: String,
reason: String,
}
impl Equation {
pub fn solve<'a, S: IntoIterator<Item = &'a dyn Expr>>(
&self,
exprs: S,
) -> Result<Equation, SolvingError> {
let eq = (self.expand()).as_eq().expect("Should remain an eqation");
let symbols: Vec<_> = exprs.into_iter().collect();
debug!("solving equation {} for unknowns {:?}", self.str(), symbols);
debug!("expanded: {}", eq.str());
if symbols.is_empty() {
Err(SolvingError {
unknowns: "".to_string(),
equation: eq.str(),
reason: "no unknowns given".into(),
})?
}
let move_right: Vec<_> = eq
.lhs
.terms()
.filter(|e| symbols.iter().all(|s| !e.has(s.get_ref())))
.collect();
let move_left: Vec<_> = eq
.rhs
.terms()
.filter(|e| symbols.iter().any(|s| e.has(s.get_ref())))
.collect();
let mut res = eq.clone();
debug!("Equation: {}", res.str());
for t in move_right {
debug!("Moving {t} to the right");
res -= t;
debug!("Equation: {}", res.str());
}
for t in move_left {
debug!("Moving {t} to the left");
res -= t;
}
let (coeff, _) = (&res.lhs).get_coeff();
if coeff.is_zero() {
Err(SolvingError {
unknowns: symbols.iter().map(|e| e.str()).join(", "),
equation: res.str(),
reason: "failed to get unknown coefficient".into(),
})?
}
if !coeff.is_one() {
res /= coeff.get_ref();
}
let mut symbols_coeff = Vec::new();
if let KnownExpr::Mul(Mul { operands }) = KnownExpr::from_expr_box(&res.lhs) {
for op in operands {
if !symbols.iter().any(|s| op.has(s.get_ref())) {
symbols_coeff.push(op.clone_box());
}
}
}
let symbols_coeff = Mul {
operands: symbols_coeff,
};
res /= symbols_coeff.get_ref();
debug!("solved equation: {}", res.str());
Ok(res)
}
}
impl JsonSchema for Equation {
fn schema_name() -> std::borrow::Cow<'static, str> {
"Equation".into()
}
fn schema_id() -> std::borrow::Cow<'static, str> {
concat!(module_path!(), "::Equation").into()
}
fn json_schema(_generator: &mut schemars::SchemaGenerator) -> schemars::Schema {
json_schema!({
"type": "string",
"pattern": "^[^=]+=[^=]+$"
})
}
fn inline_schema() -> bool {
true
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{Integer, Symbol};
#[test]
fn test_solve_solved() {
let x = &Symbol::new("x");
let expr = Equation::new(x, &Integer::zero())
.solve([x.get_ref()])
.expect("solved equation");
let expected = "Eq(Symbol(x), Integer(0))";
assert_eq!(expr.srepr(), expected)
}
#[test]
fn test_solve_basic() {
let x = &Symbol::new("x");
let expr = Equation::new(&Integer::zero(), x)
.solve([x.get_ref()])
.expect("solved equation");
let expected = "Eq(Symbol(x), Integer(0))";
assert_eq!(expr.srepr(), expected)
}
#[test]
fn test_solve_normal() {
let x = &Symbol::new("x");
let y = &Symbol::new("y");
let two = &Integer::new(2);
let expr = Equation::new(y, &*(two * x));
let expected = Equation::new(x, &*(y / two));
assert_eq!(
expr.solve([x.get_ref()]).expect("solved equation"),
expected
)
}
#[test]
fn test_solve_non_number_coeff() {
let x = &Symbol::new("x");
let y = &Symbol::new("y");
let z = &Symbol::new("z");
let expr = Equation::new(y, &*(z * x));
let expected = Equation::new(x, &*(y / z));
assert_eq!(
expr.solve([x.get_ref()]).expect("solved equation"),
expected
)
}
}