How to ensure coherent choices of representatives in an e-graph?

Question

E-graphs are a neat intermediate representation for program transformations and optimisations, which group program terms into equivalence classes (e-classes). An optimising compiler can then select the "best" representative from each e-class during code generation.

However, in some circumstances there may be two or more e-classes, where the correctness of the generated code depends on corresponding representatives being chosen from each class; i.e. the choice of representatives is not independent. For example, suppose our program has a set of small natural numbers, and we want the compiler to consider two possible representations: either a standard set data structure, or a bitmask. In that case, we have the following semantic equivalences:

• my_set = set() is equivalent to my_mask = 0
• my_set.add(x) is equivalent to my_mask |= 1 << x
• y in my_set is equivalent to (my_mask & (1 << y)) != 0

In order for the generated code to be correct, the compiler must consistently choose either the former or the latter representative from each e-class. If we select my_mask |= 1 << x from one e-class and y in my_set from another, then the generated code will be incorrect (even if the compiler knew to declare both variables). Put another way, the members of each e-class are only equivalent depending on which representatives are chosen from other e-classes.

How can an optimisation like this be implemented correctly using e-graphs?

Answer 1

You have to create one whole (sub) e-graph for the standard set version and another whole (sub) e-graph for the bitmask version.  These two e-graphs fragments are then equivalent.  Further optimization/transformation proceeds as normal, considering each of the two (sub) e-graphs, and finally the only option for later code generation is to choose one or the other.

Answer 2

An e-graph does not just group terms into equivalence classes, but congruence classes. This means that an equivalence between terms x and y also implies an equivalence between terms f(x) and f(y).

Congruence is important because e-graphs implicitly rely on it every time a rewrite is applied. All terms that use the rule's LHS are also immediately rewritten to use the RHS, and vice versa, because merging two e-classes preserves all uses of both.

Your example does not preserve congruence, and this will cause problems long before extraction time. For example, rewriting the term set() to 0 would also imply that x + 0 should be rewritten to x + set() (using x + _ for f).