Symbolic Math dynamically at runtime with Rust

autodiff.rs

use std::ops::{Add, Sub, Mul, Div};
use std::fmt::{self, Debug};

#[derive(Clone)]
enum Expression {
    Zero,
    One,
    Constant(f64),
    Variable(&'static str),
    Add(Box<Expression>, Box<Expression>),
    Sub(Box<Expression>, Box<Expression>),
    Mul(Box<Expression>, Box<Expression>),
    Div(Box<Expression>, Box<Expression>),
}

impl fmt::Debug for Expression {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use Expression::*;
        match self {
            Zero => write!(f, "0"),
            One => write!(f, "1"),
            Constant(c) => write!(f, "{}", c),
            Variable(x) => write!(f, "{}", x),
            Add(a, b) => write!(f, "({:?} + {:?})", a, b),
            Sub(a, b) => write!(f, "({:?} - {:?})", a, b),
            Mul(a, b) => write!(f, "({:?} * {:?})", a, b),
            Div(a, b) => write!(f, "({:?} / {:?})", a, b),
        }
    }
}

impl Add for Expression {
    type Output = Expression;
    fn add(self, rhs: Expression) -> Self::Output {
        match (self, rhs) {
            (Expression::Zero, r) => r,
            (l, Expression::Zero) => l,
            (Expression::Constant(a), Expression::Constant(b)) => Expression::Constant(a + b),
            (a, b) => Expression::Add(Box::new(a), Box::new(b)),
        }
    }
}

impl Sub for Expression {
    type Output = Expression;
    fn sub(self, rhs: Expression) -> Self::Output {
        match (self, rhs) {
            (l, Expression::Zero) => l,
            (Expression::Constant(a), Expression::Constant(b)) => Expression::Constant(a + b),
            (a, b) => Expression::Sub(Box::new(a), Box::new(b)),
        }
    }
}

impl Mul for Expression {
    type Output = Expression;
    fn mul(self, rhs: Expression) -> Self::Output {
        match (self, rhs) {
            (Expression::Zero, _) => Expression::Zero,
            (_, Expression::Zero) => Expression::Zero,
            (Expression::One, r) => r,
            (l, Expression::One) => l,
            (Expression::Constant(a), Expression::Constant(b)) => Expression::Constant(a * b),
            (a, b) => Expression::Mul(Box::new(a), Box::new(b)),
        }
    }
}

impl Div for Expression {
    type Output = Expression;
    fn div(self, rhs: Expression) -> Self::Output {
        match (self, rhs) {
            (Expression::Zero, _) => Expression::Zero,
            (_, Expression::Zero) => panic!("Division by zero"),
            (l, Expression::One) => l,
            (Expression::Constant(a), Expression::Constant(b)) => Expression::Constant(a / b),
            (a, b) => Expression::Div(Box::new(a), Box::new(b)),
        }
    }
}

impl Expression {
    fn diff(&self, x: Expression) -> Expression {
        use Expression::*;
        let name = if let Variable(n) = x { n } else { panic!("need variable") };
        match self {
            Zero => Zero,
            One => Zero,
            Constant(_) => Zero,
            Variable(v) => if v == &name { One } else { Zero },
            Add(a, b) => a.diff(x.clone()) + b.diff(x),
            Sub(a, b) => a.diff(x.clone()) - b.diff(x),
            Mul(a, b) => a.diff(x.clone()) * *b.clone() + *a.clone() * b.diff(x),
            Div(a, b) => (a.diff(x.clone()) * *b.clone() - *a.clone() * b.diff(x)) / (*b.clone() * *b.clone()),
        }
    }
}

fn main() {
    use Expression::*;
    
    let f = (Variable("x1") * Variable("w1") + Variable("x2") * Variable("w2")) / (Variable("w1") + Variable("w2"));

    println!("{:?}", f);
    println!("{:?}", f.diff(Variable("w1")));
    println!("{:?}", f.diff(Variable("w2")));
}