Recursive Enum expression trees in Rust
I mused about an implementation of expression trees in Rust two days ago. This is to build a Genetic Programming library in Rust - an biologically-inspired approach to finding robust solutions to difficult problems.
As discussed elsewhere I’ve become a keen fan of Rust. This language’s limited generics pose interesting problems for expressing and generating equations as described above.
There are two common ways to represent a tree datastructure:
- The most obvious way is similar to a Linked List: allocate a separate portion of memory for each node in the tree and store pointers to the child/parent nodes. This makes it cheap to expand the tree and is a very natural representation.
- The other way is to put the tree into an array. Map the parent-child relationships into indices and you can store a tree as a single block of memory. Adding nodes can require copying all this data into a larger region of memory but this has a variety of benefits such as cache locality.
We can define the semantics of a heap-allocated expression tree at compile time. Here’s a cut-down example - you can get and play with the full code here.
#[derive(Clone, Debug, PartialEq, Eq)]
enum Equation {
Add(Node<Equation>, Node<Equation>), // `a + b`
Sub(Node<Equation>, Node<Equation>), // `a - b`
Mul(Node<Equation>, Node<Equation>), // `a * b`
Div(Node<Equation>, Node<Equation>), // `a / b`
Neg(Node<Equation>), // `-a`
Sin(Node<Equation>), // `sin(a)`
Float(f64), // evaluates to a constant 64-bit float
Input, // evaluates to an inputted 64-bit float
}
impl Equation {
fn eval(&self, env: &Self::Environment) -> Self::Action {
match *self {
Add(ref left, ref right) => left.eval(env) + right.eval(env),
Sub(ref left, ref right) => left.eval(env) - right.eval(env),
Mul(ref left, ref right) => left.eval(env) * right.eval(env),
Div(ref left, ref right) => protected_div(left.eval(env), right.eval(env)),
Neg(ref left) => -left.eval(env),
Sin(ref left) => left.eval(env).sin(),
Float(i) => i,
Input => *env,
}
}
}
fn main() {
// represent the equation `x` and evaluate for `x=5`
let tree1 = Equation::Input;
println!("{:?}", tree1.eval(5.0));
// represent the equation `x * sin(4.0)` and evaluate for `x=-7.3`
let tree2 =
Equation::Mul(
node(Equation::Input),
node(Equation::Sin(
node(Equation::Float(4.0))
))
);
println!("{:?}", tree2.eval(-7.3));
// represent the equation `x + sin(x)` and evaluate for `x=99.1` and `x=16.3`
let tree3 =
Equation::Add(
node(Equation::Input),
node(Equation::Sin(
node(Equation::Input)
))
);
println!("{:?}", tree3.eval(99.1), tree3.eval(16.3));
}
This works for Genetic Programming and I’ve devised nice ways to randomly generate these trees - see the Symbolic Regression example and (this might not be obvious) evco::gp::tree.
However this does limit flexibility. Using a generic tree datastructure would allow for easily introducing or removing sorts of nodes at compile time. I found that in DEAP this sort of flexibility helped for experimenting but tended not to stick around, so it might well be better.
That’s all for now, it’s time for bed. Sometime soon I’d like to write about diversity experiences in tech, and how EVCO generates trees.