Simplifying Expressions Bottom-Up
There's a technique for simplifying expression trees that I learned from an old manager of mine, have never seen elsewhere, and have never shown to anyone who had seen it before.
Despite that, it's very simple and, in my opinion, very obviously the best design for certain kinds of problems, so I'm certain it's known to many people.
The problem we're trying to solve is when you have to represent some kind of expression trees. Say you have a bunch of constructors like this:
const Number = (value) => ({
type: "number",
value,
});
const Variable = (name) => ({
type: "variable",
name,
});
const Plus = (left, right) => ({
type: "plus",
left,
right,
});
const Minus = (left, right) => ({
type: "minus",
left,
right,
});
const Times = (left, right) => ({
type: "times",
left,
right,
});
Then you construct expression trees with something like:
let expression = Plus(Number(0), Plus(Number(1), Variable("a")));
For this example we'll restrict ourselves to just like, very simple simplifications, like x+0=x
, and stuff, but this technique can handle more complicated transformations as well.
One way I've seen smart people implement this kind of thing before is something like the following (Copilot wrote most of this, which is I guess evidence of...something):
function simplify(node) {
switch (node.type) {
case "plus": {
const left = simplify(node.left);
const right = simplify(node.right);
// Commute variables to the left.
if (left.type !== "variable" && right.type === "variable") {
return Plus(right, left)
}
// Fold 0 + x => x.
if (left.type === "number" && left.value === 0) {
return right;
}
// Fold x + 0 => x.
if (right.type === "number" && right.value === 0) {
return left;
}
return Plus(left, right);
}
case "minus": {
const left = simplify(node.left);
const right = simplify(node.right);
return Minus(left, right);
}
case "times": {
const left = simplify(node.left);
const right = simplify(node.right);
return Times(left, right);
}
default:
return node;
}
}
And then you use it by running
expression = simplify(expression)
// => {
// type: 'plus',
// left: { type: 'variable', name: 'a' },
// right: { type: 'number', value: 2 }
//}
This code is actually broken, though, since we'd need to run it to a fixpoint in order to actually make sure no more transformations apply. We can see this with some inputs that don't actually output a fully simplified expression according to our rules:
let expression = Plus(Number(0), Variable("a"));
expression = simplify(expression);
console.log(expression);
// =>
// {
// type: 'plus',
// left: { type: 'variable', name: 'a' },
// right: { type: 'number', value: 0 }
// }
This is not an insurmountable problem—you can insert additional calls to simplify
in various places, or keep calling it until you hit a fixpoint, or something. But this kind of sucks: you have to remember to call it, and you have to re-walk the tree at each step, and in larger programs, in some contexts it's not clear whether you have a simplified or unsimplified tree. I just don't like it very much.
Here's the better way. Instead of having separate "construction" and "simplification" steps, do them at the same time by intercepting any attempts to construct an operator with simplification rules:
const Number = (value) => ({
type: "number",
value,
});
const Variable = (name) => ({
type: "variable",
name,
});
const Plus = (left, right) => {
// Commute variables to the left.
if (left.type !== "variable" && right.type === "variable") {
return Plus(right, left)
}
// Fold 0 + x => x.
if (left.type === "number" && left.value === 0) {
return right;
}
// Fold x + 0 => x.
if (right.type === "number" && right.value === 0) {
return left;
}
return {
type: "plus",
left,
right,
};
}
const Minus = (left, right) => ({
type: "minus",
left,
right,
});
const Times = (left, right) => ({
type: "times",
left,
right,
});
let expression = Plus(Number(0), Variable("a"));
console.log(expression);
// =>
// { type: 'variable', name: 'a' }
The key here is that the only place that is allowed to explicitly construct the actual concrete operator is the last line in each constructor. This guarantees that if any rules match in the initial construction of the expression, they'll re-call the constructor which will expose the new expression to simplification once again.
I want to highlight that even though we're intermingling construction and simplification temporally, this doesn't lead to any breach of separation of concerns, since the code that walks an AST or whatever and constructs the trees can be completely unaware of the optimizations going on as it's doing so.
This has some interesting implications:
- You will not necessarily get back the same kind of operator you attempted to construct. If I call
Plus(x, 0)
, I'll actually get back aVariable
operator. In CockroachDB when we used this technique, we had a "no peeking" rule, where you could not inspect the thing you got back from the constructors. - It's impossible to get your hands on a non-optimized tree, so there's no ambiguity about whether you need to simplify something or not (we did have global overrides to disable this, for testing purposes).
- This means that simplification rules don't need to be written to deal with "bad" trees: they can assume their inputs are always fully simplified.
- A lot of superfluous trees are simply never constructed, which can alleviate a lot of GC pressure: note that in the first implementation, we first construct the
1 + a
tree before replacing it with thea + 1
tree, but in the second implementation, we never create the1 + a
tree at all, when we try our simplification rule kicks in and we construct the correct tree instead. - For some non-local transformations (say, join reordering or column pruning, that require a bit more context) you still need some other mechanisms to move information around the tree, since this technique doesn't work super well for those.
- In some languages you do have to be a little careful about blowing out the stack for deeper queries. This is not a new problem to anyone who has ever implemented a production recursive-descent parser.
- Debugging can be a little more obnoxious sometimes, because of the interlacing of different parts of the planning process. It's important to have good tools to isolate various aspects of planning.
That's the technique. Like I said, I've never seen anyone talk about it but I'm sure it's known in at least some circles. I don't have a real name for it but I always have called it "the Kimballizer" after my manager who showed it to me.
NULL BITMAP is free. If you earn a software engineering salary and think you've gotten some value out of these posts, please consider tossing a couple dollars to the Palestine Children's Relief Fund.