E-class is in session

We'll implement makeset first. It needs to allocate a new ID, which we do by pushing onto the union-find's parent array. Extending the array "makes" a new valid index, which is the new ID:

class UnionFind:
    def makeset(self):
        new_id = len(self.parent)
        self.parent.append(new_id)
        return new_id

This can be a little confusing! Since we push onto the list, there will be one more valid index after makeset than before. That will be our new ID. We compute that new ID from the pre-pushing length of the list. For example, the first time we call this, the pre-pushing length is 0, and the new ID after we push is 0.

It isn't actually important that the IDs are allocated like this, it's just an implementation convenience. You could imagine using a hash table, or heap-allocated objects, to represent IDs, and allocate IDs randomly or use the object pointer as an ID or something like that.

The value we push should be the new ID's parent; by convention, we say that if an ID's parent is itself, then it has no parent, so it's the root of a tree. You could do something more intuitive, like having an array of Option[ID] elements, but this is a standard convention and it's a bit faster. Multiple IDs are allowed to have no parent; the full union-find is a collection of little trees.

Next, let's implement find. This method just finds the root of the tree that a given ID is in, and the implementation is simple and recursive:

class UnionFind:
    def find(self, id : ID) -> ID:
        if self.parent[id] == id:
            return id
        else:
            return self.find(self.parent[id])

Remember that self.parent[id] is morally id.parent, so this just walks up the tree until the root.

You can and should do path compression. It makes things faster, but it isn't critical; it is a micro-optimization. If you do an at-all-efficient union-find it will never be a bottleneck.

Philip likes the iterative version of find:

class UnionFind:
    def find(self, id : ID) -> ID:
        while id != self.parent[id]:
            id = self.parent[id]
        return id

One way to think about IDs is to say that equal IDs are represented by a canonical ID. That's what find returns: the canonical ID representing a given ID. Any two equal IDs, meaning they're in the same union-find tree, are represented by the same canonical ID, which is the root of that tree.

Now let's make two IDs equal. The first thing union does is it calls find on both IDs:

class UnionFind:
    def union(self, a : ID, b : ID) -> bool:
        a = self.find(a)
        b = self.find(b)
        # ...

We always want to operate on canonical IDs. Next, we check if the two are equal; if they are, union doesn't have to do anything, and we can return False:

class UnionFind:
    def union(self, a : ID, b : ID) -> bool:
        # ...
        if a == b: return False
        # ...

Otherwise, choose one to be the new parent. Philip points out that it's an interesting degree of freedom, which one you pick, and there's a literature here, but it won't matter much for us; we'll just pick a.

So now let's think it through. Currently a and b are both canonical, since we just found them. If a is going to be the canonical representative of both of them, we need b to not be canonical. So make its parent a:

class UnionFind:
    def union(self, a : ID, b : ID) -> bool:
        # ...
        self.parent[b] = a
        return True

We return True because we made a change to the union-find.

We're done with the union-find. Now is a good time to stop and do some tests. For example, you can make a union-find, make three nodes, union two pairs, and check that the third pair is equal:

test_uf = UnionFind()
a, b, c = test_uf.makeset(), test_uf.makeset(), test_uf.makeset()
test_uf.union(a, b)
test_uf.union(b, c)
assert test_uf.find(a) == test_uf.find(c)

Union-find is "solving" a theory, where the theory, basically, is transitivity. Union-find is a way to solve that cheaply. It works great.

You can also write a Graphviz generator for your union-find. That can be useful to get more familiar with the data structure.

You might want to add an are_equal(Id, Id) method, which takes two ids, calls find on both of them, and checks if they're equal. Philip notes that in more complicated e-graphs there's a degree of freedom here. Instead of returning a boolean, it could instead return something richer that tells you how the two IDs are related.

Anyway, our tests pass. Let's keep going. Add a union find field to the e-graph:

class EGraph:
    union_find : UnionFind
    def __init__(self):
        # ...
        self.union_find = UnionFind()

We're done with the first part of writing an e-graph, and are ready to move on to the second part.

Let's start the add method. To start, we're not even going to look at the children. Here's what we'll do: check if it's in the memo. If so, return its value in the memo. If it's not, put it in the memo, mapped to a new ID returned by makeset:

class EGraph:
    def add(self, node : Node) -> ID:
        if node in self.memo:
            return self.memo[node]
        else:
            new_id = self.union_find.makeset()
            self.memo[node] = new_id
            return new_id

Max points out that this is basically the only possible implementation, since makeset is our only "factory" for IDs. The only real tweak you can make is something with the node's children. We're not looking at the children at all, but we'll go back and fix that later.

union barely exists. Just call union on the union find:

class EGraph:
    def union(self, a : ID, b : ID) -> bool:
        return self.union_find.union(a, b)

Test that you're actually hash-consing. Make a node type called Leaf which takes a string and makes a Node with a given string and no children. That's useful for testing:

def Leaf(name : str) -> Node:
    return Node(name, ())

Insert two identical nodes and check that you get the same ID. Insert different nodes and make sure they're different:

test_eg = EGraph()
a, a_copy, b = Leaf("a"), Leaf("a"), Leaf("b")
a_id, a_copy_id, b_id = test_eg.add(a), test_eg.add(a_copy), test_eg.add(b)
assert a_id == a_copy_id
assert a_id != b_id

Philip points out that instead of exposing Node directly, you can also do a "smart constructor" approach, where the E-graph is a factory for nodes. Max points out that either way works; you can always munge the input Node.

How are we gonna do congruence closure? The stupidest way possible: we're going to reconstruct memo over and over again, until it stops changing. Specifically, we're going to walk over all the nodes in the current memo, and re-insert them into the new dictionary, which will become the new memo. As we do this, we're going to have old keys and old values. When we re-insert a node, it might already be there. It could be there with the current ID, or a different ID. If it's different, you found a congruence. In that case we'll make progress toward congruence closure.

So, for example, that's what's going to happen with the f(a) and f(b) example. When you insert f(b), you're going to find the ID of f(a). So you'll discover that f(a) and f(b) are equal, allowing you to union those two IDs. You just repeat this until fixed-point, until you stop finding new congruences.

To make this work well, we need the lookup in memo to find as many hits as possible. So it's time to look at the children in add. Add a new function, canonicalize_node, and call it in add:

class EGraph:
    def canonicalize_node(self, n : Node) -> Node:
        return Node(n.f, tuple(self.union_find.find(child) for child in n.l))

    def add(self, n : Node) -> ID:
        n = self.canonicalize_node(n)
        # ...

This line takes the input node's children and calls the union-find's find on them. This converts every child into its canonical representative, so if you call add twice with congruent but not equal nodes, you're more likely to get a "hit" in the hash cons. Most of the time, this won't do anything, but there are all sorts of reasons you might have stale IDs floating around, and this procedure helps ensure that the hash cons only stores canonical IDs.

We're trying to fix some kind of out-of-order-ness to our adds and unions. In our test, we add f(a) and f(b), then union a and b. But if we'd done it the other way around, if we'd unioned a and b first, then f(a) and f(b) would be immediately equal, because of the canonicalization in the add method. So what we're doing, basically, is doing a second pass of re-adding f(a) and f(b), but after the union on a and b, to fix that out-of-order-ness.

You can do this in various places. You can do this now or later, more or less often. It's not a bug, strictly speaking, to not call it. But you probably want to do it eventually. We're going to do it once per rewrite iteration.

So now we're ready to do rebuilding. Here's what we're going to do: take the hash-cons out of the e-graph, replacing it with an empty hash-cons; then walk the old e-graph and re-add every node to the new e-graph.

class EGraph:
    def rebuild(self) -> None:
        while True:
            old_memo = self.memo
            self.memo = {}
            # ...

Re-adding a node is very similar to adding a node, but with one difference: if the node isn't in the hash-cons, add calls makeset. But when we re-add a node, we can reuse the old ID:

class EGraph:
    def rebuild(self) -> None:
        while True:
            for old_node, old_id in old_memo.items():
                new_node = self.canonicalize_node(old_node)
                if new_node in self.memo:
                    new_id = self.memo[new_node]
                else:
                    new_id = self.union_find.find(old_id)
                    self.memo[new_node] = new_id
                # ...

Now we have an old ID and a new ID. Union them. This is the point where—if they're different—you've discovered a new congruence equality:

class EGraph:
    def rebuild(self) -> None:
        while True:
            for old_node, old_id in old_memo.items():
                # ...
                self.union(new_id, old_id)

Note that we're not invalidating the union-find (since it stores the unions). The union-find only stores opaque IDs. There is no way to go from an ID to a node, so the union-find isn't storing anything that could be invalid.

Now we just need to do this whole thing in a loop until the hash cons stops changing. Since each union returns a boolean, we just need to or that with a termination signal. If none of the union calls do anything, we can stop.

class EGraph:
    def rebuild(self) -> None:
        while True:
            keep_going = False
            for old_node, old_id in old_memo.items():
                # ...
                if self.union(new_id, old_id):
                    keep_going = True
            if not keep_going:
                break

The test should now pass. In fact, maybe add an assertion: when we're done, every ID in the hash cons, either in a node's children or the hash cons value, must be canonical.

Let me note that this was the buggiest part of the session so far. There were a couple of simultaneous implementations going and basically all of them ran into bugs. The sequence of steps here is very important, and if you get details wrong, it will infinite loop. You gotta break the right loop, for example. It's also tempting to call add, because most of re-adding a node is very similar, but you gotta reuse the old hashcons IDs, otherwise you will keep adding new IDs forever. Max and Philip made that mistake and it derailed us for a half hour. So you really do need to check that, at this point, everything works.

To start, we need patterns. A pattern is a syntax tree, with variables at the leaves. How you define this will depend on your language; Python doesn't have algebraic data types, so we'll make a Pattern class with two subclasses, but you can do it some other way if you're using some other language:

class Pattern: pass

@dataclass
class Var(Pattern):
    name : str

@dataclass
class App(Pattern):
    f : str
    l : list[Pattern]

Patterns have a related notion called a substitution, which is a map from variable names to IDs:

Substitution = dict[str, ID]

Let's add an instantiate method to the e-graph. It takes a pattern and a substitution and adds the pattern to the e-graph:

class EGraph:
    def instantiate(self, pat : Pattern, subst : Substitution) -> ID:
        match pat:
            case Var(name):
                return subst[name]
            case App(f, l):
                node = Node(f, tuple(self.instantiate(subpat, subst) for subpat in l))
                return self.add(node)

Again, there's basically only one way to write this function, which is a sign that it's a good function. This will be how the right-hand side of patterns works, the "apply" step. We're doing this first because it's easier, and it's also a nice API for adding nodes. It's not a bad idea here to add a test, where you use instantiate to add a node that you've already manually added, and make sure it's the same.

Now, the implementation of this method is complex enough, so let's put some helper methods on the board. First, nodes_in_class. It takes an ID and returns a list of nodes. Not a list of IDs—a list of nodes. This is important. Our implementation will be naive; a real implementation would add an index here to make it efficient, and you can do that later. But our implementation will be simple: it'll iterate over the hashcons and give us all the keys where the value is the input ID. We are reading the hash map backwards, basically:

class EGraph:
    def nodes_in_class(self, root : ID) -> list[Node]:
        out = []
        for node, id in self.memo.items():
            if id == root:
                out.append(node)
        return out

You don't need to do any find calls in here. We will only be doing this after calling rebuild.

There was a sidebar here between Philip and Max. The nodes_in_class is actually an e-analysis. In fact, Philip pointed out, it is the free analysis! A cool fact, though not relevant to our implementation.

Now let's get back to ematch_rec. It's recursive, so there are two cases.

class EGraph:
    def ematch_rec(self, pat : Pattern, root : ID, subst : Substitution) -> list[Substitution]:
        match pat:
            case Var(name):
                # ...
            case App(f, subpats):
                # ...

If the pattern is a variable, there are two cases: either it is mentioned in the input substitution, or it is not. If it's not mentioned, then we are unconstrained and can add a new binding to the substitution for it. Do not mutate the substitution, copy it, and return a singleton list with the new substitution. In this case we are always, successfully, extending the input substitution, and there's only one extension:

class EGraph:
    def ematch_rec(self, pat : Pattern, root : ID, subst : Substitution) -> list[Substitution]:
        match pat:
            case Var(name):
                if name not in subst:
                    return [ subst | { name: root } ]
                # ...

But what if it is in the substitution? Then we need to make sure that we agree with that constraint. If the substitution maps the variable to the current ID, then the input substitution is still valid, and we return it.

class EGraph:
    def ematch_rec(self, pat : Pattern, root : ID, subst : Substitution) -> list[Substitution]:
        match pat:
            case Var(name):
                # ...
                elif subst[name] == root:
                    return [subst]
                # ...

But if it maps the variable to a different ID, then the input substitution cannot match this pattern, so we return an empty list:

class EGraph:
    def ematch_rec(self, pat : Pattern, root : ID, subst : Substitution) -> list[Substitution]:
        match pat:
            case Var(name):
                # ...
                else:
                    return []

Note that there are three cases here: not in the substitution; in there and agree; in there and disagree. Each has a different behavior: extending the substitution, validating it, or dropping it.

Now, what if the pattern is not a variable? This is hairy code. We're gonna go line by line.

Write a loop. Get all the nodes in the current class, and loop through them. Filter out all the nodes where the function symbols don't match: check for equality and skip if not equal. Similarly do a length check: the node and the pattern need to have the same number of arguments, otherwise skip.

class EGraph:
    def ematch_rec(self, pat : Pattern, root : ID, subst : Substitution) -> list[Substitution]:
        match pat:
            case App(f, subpats):
                for node in self.nodes_in_class(root):
                    if node.f != f: continue
                    if len(node.l) != len(subpats): continue
                    # ...

Now what are we doing? We're looping over all the nodes in the class, and looking at only ones that could match. Next, we clearly we need to recurse. So what we're gonna do is, we're going to zip the pattern children and the node children. We've established that they have the same number of children:

class EGraph:
    def ematch_rec(self, pat : Pattern, root : ID, subst : Substitution) -> list[Substitution]:
        match pat:
            case App(f, subpats):
                for node in self.nodes_in_class(root):
                    # ...
                    for sub_id, sub_pat in zip(node.l, subpats):
                        # ...

Now, we have the child pattern, we have the child ID, we have a substitution. We could call ematch_rec. But this isn't quite right. What's the issue? Well, ematch_rec is going to return a list of substitutions. Which one is the right one? Or what if the list is empty? Also, we have a dependency here: when we recurse on the first child, and when we recurse on the second child, the substitutions they make have to agree, and we need to enforce that.

So what we're going to do is to maintain a list of substitutions that we thread through the children. (You could instead do an explicit join at the end. But we're not going to do that.) Create a singleton list with the initial substitution, called todo. It goes inside the nodes in class loop, but outside the zip loop:

class EGraph:
    def ematch_rec(self, pat : Pattern, root : ID, subst : Substitution) -> list[Substitution]:
        match pat:
            case App(f, subpats):
                for node in self.nodes_in_class(root):
                    # ...
                    todo = [subst]
                    for sub_id, sub_pat in zip(node.l, subpats):
                        # ...

When we look at each child/subpattern pair, we need to call ematch_rec for each substitution in the todo list. Combine the results into a single list, with flatmap or whatever your language's idiom is here. And then after each child/pattern iteration, you replace the todo list with the new list of matches:

class EGraph:
    def ematch_rec(self, pat : Pattern, root : ID, subst : Substitution) -> list[Substitution]:
        match pat:
            case App(f, subpats):
                for node in self.nodes_in_class(root):
                    # ...
                    for sub_id, sub_pat in zip(node.l, subpats):
                        results = []
                        for new_subst in todo:
                            results.extend(self.ematch_rec(sub_pat, sub_id, new_subst))
                        todo = results
                    # ...

This is tricky code! The extend here is doing the flatmap, and note the very common pattern where we are traversing todo to build a new version of it. You gotta get this right.

This is one of many different e-matching algorithms, and this one—top-down e-matching—is far from the best one. But it's a good starting-point for thinking about things; the better algorithms are more complex.

One we have left the zip loop, the todo list contains substitutions that agree with the initial substitution, and also all the subpatterns. So these are all good substitutions. So outside the node in class loop, make an empty list of outputs, and add this final todo list to the outputs, because each node in the class might potentially produce valid substitutions.

class EGraph:
    def ematch_rec(self, pat : Pattern, root : ID, subst : Substitution) -> list[Substitution]:
        match pat:
            case App(f, subpats):
                outputs = []
                for node in self.nodes_in_class(root):
                    # ...
                    outputs.extend(todo)
                return outputs

Oof! This is complicated, and it's another place where it's possible to screw up even with a type discipline. But just work through it line by line. But we are basically done now, at least with the core of pattern matching.

The difference from plain pattern matching is that at every stage there is a multiplicity of matches. This todo list is where that difference lives; it's what stores the set of possible substitutions.

Ok, so we are done with ematch_rec. We're out of the woods here. All we need is ematch, to start the recursion; all it does is call ematch_rec with an empty substitution.

class EGraph:
    def ematch(self, pat : Pattern, root : ID) -> list[Substitution]:
        return self.ematch_rec(pat, root, {})

We'll loop over IDs externally. Let's move on to that, now.