Modelling mathematical expressions in Rust

Question

I've been trying so hard to come up with a data model that works for mathematical expressions (like x^2 + 2x - y * 4, no equals sign) in Rust. It's very different to other languages I'm most familiar with.

Eventually, I want to be able to

I think of mathematical expressions as recursive structures. As in the BNF

<expr> ::= <sum> | <product> | <reciprocal> | <number> | <variable>
<sum> ::= <expr> "+" <expr>
<product> ::= <expr> "*" <expr>
<reciprocal> ::= "1/" <expr>
<power> ::= <expr> "^" <expr>
<number> ::= ...|-4|-3|-2|-1|0|1|2|3|4|...
<variable> ::= "x" | "y" | "any string literal" | ...

Setting up a recursive data structure in the first place was hard for me. Maybe enums are the correct choice, maybe this concept is the wrong approach altogether for Rust.

I've also really struggled a lot with borrowing, dereferencing, ownership, etc.

Please let me know what you think of the code I've written, the approach I've taken and how amenable it will be to my ambitions with this program, and any tips you have for helping me to get into the Rust headspace.

use std::collections::BTreeSet;
use std::fmt;

#[derive(Debug, PartialEq, Eq, Hash, PartialOrd, Ord, Clone)]
enum Expression {
    Sum(BTreeSet<(usize, Expression)>),
    Product(BTreeSet<(usize, Expression)>),
    Negation(Box<Expression>),
    Reciprocal(Box<Expression>),
    Power(Box<Expression>, Box<Expression>),
    Variable(Box<str>),
    Number(isize),
}

use Expression::*;

impl fmt::Display for Expression {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Sum(summands) => {
                let mut flat_summands: Vec<String> = vec![];
                for (n, expr) in summands.iter() {
                    let stringified_expr = expr.to_string();
                    for _ in 0..*n {
                        flat_summands.push(format!("{}", stringified_expr));
                    }
                }
                fmt.write_str(&flat_summands.into_iter().map(|summand| summand.to_string()).collect::<Vec<_>>().join(" + "))
            },
            Product(factors) => {
                let mut flat_factors: Vec<String> = vec![];
                for (n, expr) in factors.iter() {
                    let stringified_expr = expr.to_string();
                    for _ in 0..*n {
                        flat_factors.push(format!("{}", stringified_expr));
                    }
                }
                fmt.write_str(&flat_factors.into_iter().map(|factor| factor.to_string()).collect::<Vec<_>>().join(" * "))
            },
            Reciprocal(ref expr) => fmt.write_fmt(format_args!("1 / {}", expr.to_string())),
            Power(ref base, ref exponent) => fmt.write_fmt(format_args!("{} ^ {}", base.to_string(), exponent.to_string())),
            Negation(ref expr) => fmt.write_fmt(format_args!("-{}", expr.to_string())),
            Variable(ref name) => fmt.write_fmt(format_args!("{}", name)),
            Number(n) => fmt.write_fmt(format_args!("{}", n)),
        }
    }
}

fn main() {
    let expr = Sum(BTreeSet::from([(1, Number(1)), (1, Product(BTreeSet::from([(1, Number(2)), (1, Variable(Box::from("x")))])))]));
    println!("{:?}", expr);
    println!("{}", expr);
    assert_eq!(format!("{}", expr), String::from("x * 2 + 1"));
}

I have made the arguments of Sum and Product sets of (count, expr) pairs, so that later mathematical manipulations like simplifying x + x + x + x + x to 5x, or x * y * z / y, will be easier (than having to search deeply-nested trees for equivalent expressions).

Answer 1

First things first: always run cargo clippy and cargo fmt! Clippy is usually very helpful and in this case it finds several errors in your code. Both commands will help you get into the habit of writing idiomatic and idiomatically formatted Rust code.

You derived PartialOrd and Ord for expressions. Are you sure that makes sense? What does it mean to compare two expressions for ordering?

I'm not a huge fan of the choice to make Sum and Product BTreeSets. You're not getting use out of the fact that they are sets, because you don't need to compare two (usize, Expression) values for equality. In fact, what you actually want is to match up terms that are the same expression, not the same pair of a value and expression, to combine those terms. So it would make sense to instead use this: HashMap<Expression, usize>.

An alternate design choice would be to use the concept of a normalized expression. If you have normalized expressions, you can enforce not only that expressions like x + x + x get collected into 3 * x, but also that multiplication is distributed over addition (e.g. 3 x + 3y instead of 3 ( x + y)).

There are multiple ways to do that, but the simplest way is to just stick with Sum and Product being binary, but then have a normalize function that rewrites the sum and product to the appropriate form. So:

enum Expression {
    Sum(Box<Expression>, Box<Expression>),
    Product(Box<Expression>, Box<Expression>),
    // ...
}

fn normalize(e: Expression) -> Expression

or

fn normalize(e: &mut Expression)

or even using a wrapper type:

struct Normalized(Expression);
fn normalize(e: Expression) -> NormalizedExpression

The point is that it will be easier to enforce your invariants this way, and likely more efficient. If you want to enforce that expressions have the form 5 x + 6 y + 7z, what you do is ensure that the expression is normalized to Sum(Prod(5, x), Sum(Prod(6, y), Prod(7, z))). To enforce this, you have to ensure: (1) that Sum is associated only left-to-right, so you rewrite Sum(Sum(a, b), c)) to Sum(a, Sum(b, c)); (2) that terms are written in some canonical order (I guess from x to y to z, so in alphabetical order); and (3) that each term is a constant multiplied by an expression.

You could also change the data model for Normalized expressions to be different, for example using HashMap<Expression, usize> to store each expression with its constant count.