// -*- mode: markdown;  markdown-fontify-code-block-default-mode: rustic-mode; evil-shift-width: 2; -*-
/*!

# Concepts: e-graphs and equality saturation

<style> em { display: inline-block } </style>

An _e-graph_ is a data structure that powers the _equality saturation_
optimization technique.

Both e-graphs
(Gregory Nelson's [PhD Thesis](https://courses.cs.washington.edu/courses/cse599f/06sp/papers/NelsonThesis.pdf), 1980)
and equality saturation
([Tate et. al., 2009](https://arxiv.org/abs/1012.1802))
were invented before.
`egg` aims to make them easy, extensible, and efficient.

This tutorial will approach these concepts at a high level.
More detail can be found in our [paper].

## Why any of this?

When building pretty much any tool relating to programming languages,
  you'll typically be working with _terms_ from your language
  (we will use the words _term_, _expression_, and _program_ interchangeably).
You may be trying to:
- Convert a term to a "better", equivalent term (optimization);
- Build a term from scratch according to a specification (synthesis);
- Or show that two terms are equivalent (verification).

_Term rewriting_ is common technique for working with terms
  toward any of these goals.
You start with a term _t_ and a set of rewrites,
  each of which take the form _l → r_.
For each rewrite,
you try to match the pattern _l_ against the term _t_,
  generating a substitution _σ_ at some subterm,
  and then you apply that substitution to the right-hand pattern _r_
  and replace the matching subterm.

For example, consider the term 42 × (7 + 7) and the rewrite _x + x → 2 × x_:
- The left-hand side would match at the subterm 7 + 7,
  generating the substitution _σ_ = {_x_: 7}.
- Applying the substitution gives _r\[σ\] =_ 2 × 7.
  (_r\[σ\]_ means apply the substitution _σ_ to _r_)
- Replacing the matched subterm 7 + 7 with _r\[σ\] =_ 2 × 7 gives the result: 42 × (2 × 7).

One issue with term rewriting
  (and other programming language techniques that involve manipulating terms)
  is the matter of choice.
Rewriting is _destructive_,
  meaning that once you do a rewrite, the initial term is gone.
If you have many rewrites to choose from,
  picking one is like choosing a fork in the road:
  going back and choosing another way is expensive or impossible,
  so you're committed to the path you've chosen.
Furthermore, picking the right rewrite at the right time is very difficult in general;
  sometimes, a choice that seems good locally may be wrong globally.

As an example, a C optimizer would like to rewrite `a * 2` to the cheaper `a << 1`.
But what about `(a * 2) / 2`?
A greedy, local approach will take a "wrong turn" resulting in `(a << 1) / 2`,
  which is better than the input expression
  but hides the fact that we could have canceled the multiplication and division.

There are some techniques to mitigate this
  (backtracking and [value-numbering](https://en.wikipedia.org/wiki/Value_numbering)),
  but fundamentally, the only way around it is to take _all choices at the same time_.
But rewriting is destructive,
  so applying _n_ rewrites to a term results in _n_ terms,
  and then applying _n_ rewrites to each of those _n_ terms yields _n_<sup>2</sup> more terms,
  and so on.
If you want to explore _m_ "layers" deep in this tree,
  you'll have to store _n_<sup><i>m</i></sup> terms,
  which is prohibitively expensive.

E-graphs can solve this representation problem.
When rewriting a bunch of terms,
  the results are often structurally similar.
E-graphs can store large sets of similar expressions very compactly.
Using e-graphs,
  you can apply many rewrites simultaneously without a
  multiplicative growth in space.
Equality saturation is a technique to do this kind of rewriting for program
  optimization.

## What's an e-graph?

An _e-graph_ ([/'igraf/][sound]) is a data structure to maintain
  an [equivalence relation]
  (really a [congruence relation], see the [next section](#invariants-and-rebuilding))
  over expressions.
An e-graph is a set of equivalence classes (_e-classes_),
each of which contains equivalent _e-nodes_.
An e-node is an operator with children, but instead of
children being other operators or values, the children are e-classes.
In `egg`, these are represented by the [`EGraph`], [`EClass`], and
[`Language`] (e-nodes) types.

Even small e-graphs can represent a large number of expressions, exponential in
the number of e-nodes.
This compactness is what makes e-graphs a compelling data structure.
We can define what it means to _represent_ (or _contain_) a term as follows:

- An e-graph represents a term if any of its e-classes do.
- An e-class represents a term if any of its e-nodes do.
- An e-node `f(n1, n2, ...)` represents a term `f(t1, t2, ...)` if e-class `ni` represents term `ti`.

Here are some e-graphs.
We picture e-classes as dotted boxes surrounding the equivalent e-nodes.
Note that edges go from e-nodes to e-classes (not to other e-nodes),
  since e-nodes have e-classes as children.
Don't worry about how they are made just yet, we'll get to rewriting soon.

<figure>
<img src="data:image/png;base64,[encoded data, 15612 characters]">
<figcaption style="text-align: center">
  An e-graph as it undergoes a series of rewrites. Gray indicates what didn't change from the previous image.
</figcaption>
</figure>

An e-graph can be queried for patterns through a procedure called _e-matching_
  (matching up to equivalence),
  which searches the e-graph for e-classes that represent terms that match a given pattern.
`egg`'s e-matching
  (inspired by
  [De Moura and Bjørner](http://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.800.2292&rep=rep1&type=pdf#page=193))
  is done by the [`Searcher`]  and [`Pattern`] APIs.
There are two core operations that modify the e-graph:
  [`add`] which adds e-nodes to the e-graph, and
  [`union`] which merges two e-classes.

These operations can be combined to do rewriting over the e-graph.
Critically, this rewriting process is only additive,
  meaning that the _e-graph never forgets_ old versions of terms.
Additionally, rewrites add entire _classes_ of terms at a time.
To perform a rewrite _l → r_ over an e-graph,
  you first search _l_ in the e-graph, yielding pairs of (_c_, _σ_),
  where _c_ is the matching e-class and _σ_ is the substitution.
Then, you `add` the term _r\[σ\]_ to the e-graph
  and `union` it with the matching e-class _c_.

Let's put it all together with an example referring to the four e-graphs in the
  image above.

1. The initial e-graph represents the term _(a × 2) / 2_.
   Since each e-class only has one e-node,
     the e-graph is basically an abstract syntax tree
    with sharing (the 2 is not duplicated).
2. Applying the rewrite _x × 2 → x << 1_
   has recorded the fact that _a × 2 = a << 1_
     without forgetting about _a × 2_.
   Note how the newly added _a << 1_ refers to the existing "_a_" e-node,
   and the "<<" e-node has been unioned into the same e-class
   as the equivalent "×" e-node where the pattern _x × 2_ matched.
3. Applying rewrite _(x × y) / z → x × (y / z)_ realizes that division
    associates with multiplication.
   This rewrite is critical to discovering the cancellation of 2s that we are looking for,
     and it still works despite the fact that we applied the "wrong" rewrite previously.
4. Applying rewrites _x / x → 1_ and _x × 1 → x_ doesn't add any new e-nodes,
     since all the e-nodes were already present in the e-graph.
   The result only unions e-classes,
     meaning that e-graph actually got _smaller_ from applying these rewrites,
     even though it now represents more terms.
   In fact, observe that the top-right "×" e-node's left child is _itself_;
     this cycle means the e-class represents the _infinite_ (!) set of terms
     _a_, _a × 1_, _a × 1 × 1_, and so on.

## Invariants and Rebuilding

An e-graph has two core operations that modify the e-graph:
[`add`] which adds e-nodes to the e-graph, and
[`union`] which merges two e-classes.
These operations maintains two key (related) invariants:

1. **Congruence**

   An e-graph maintains not just an [equivalence relation] over
   expressions, but a [congruence relation].
   Congruence basically states that if _x_ is equivalent to _y_,
     _f(x)_ must be equivalent to _f(y)_.
   So as the user calls [`union`], many e-classes other than the given
   two may need to merge to maintain congruence.

   For example, suppose terms _a + x_ and _a + y_ are represented in
   e-classes 1 and 2, respectively.
   At some later point, _x_ and _y_ become
   equivalent (perhaps the user called [`union`] on their containing
   e-classes).
   E-classes 1 and 2 must merge, because now the two "+"
   operators have equivalent arguments, making them equivalent.

2. **Uniqueness of e-nodes**

   There do not exist two distinct e-nodes with the same operators and equivalent
   children in the e-graph, either in the same e-class or different e-classes.
   This is maintained in part by the hashconsing performed by [`add`],
   and by deduplication performed by [`union`] and [`rebuild`].

`egg` takes a delayed approach to maintaining these invariants.
Specifically, the effects of calling [`union`] (or applying a rewrite,
which calls [`union`]) may not be reflected immediately.
To restore the e-graph invariants and make these effects visible, the
user *must* call the [`rebuild`] method.

`egg`'s choice here allows for a higher performance implementation.
Maintaining the congruence relation complicates the core e-graph data
structure and requires an expensive traversal through the e-graph on
every [`union`].
`egg` chooses to relax these invariants for better performance, only
restoring the invariants on a call to [`rebuild`].
The [paper] on `egg` goes into greater detail on why this design choice makes
  things faster.
See the [`rebuild`] documentation for more information on
  what it does and when you have to call it.
Note also that [`Runner`]s take care of this for you, calling
  [`rebuild`] between rewrite iterations.

## Equality Saturation


Equality saturation
  ([Tate et. al., 2009](https://arxiv.org/abs/1012.1802))
  is a technique for program optimization.
`egg` implements a variant of equality saturation in the [`Runner`] API that
  looks like the following pseudocode:

```ignore
fn equality_saturation(expr: Expression, rewrites: Vec<Rewrite>) -> Expression {
    let mut egraph = make_initial_egraph(expr);

    while !egraph.is_saturated_or_timeout() {
        let mut matches = vec![];

        // read-only phase, invariants are preserved
        for rw in rewrites {
            for (subst, eclass) in egraph.search(rw.lhs) {
                matches.push((rw, subst, eclass));
            }
        }

        // write-only phase, temporarily break invariants
        for (rw, subst, eclass) in matches {
            eclass2 = egraph.add(rw.rhs.subst(subst));
            egraph.union(eclass, eclass2);
        }

        // restore the invariants once per iteration
        egraph.rebuild();
    }

    return egraph.extract_best();
}
```

Most of this was covered above, but we need to define two new terms:

- _Saturation_ occurs when an e-graph detects that rewrites no longer add new information.
  Consider the commutative rewrite _x + y → y + x_.
  After applying it once, the second time adds no new information
    since the e-graph didn't forget about the initial _x + y_ terms.
  If all the rewrites are in this state, we say the e-graph is _saturated_,
    meaning that the e-graph encodes all possible equivalences derivable from
    the given rewrites.
- _Extraction_ is a procedure for picking a single represented term from an e-class
  that is optimal according to some cost function.
  `egg`'s [`Extractor`]s provide this functionality.

Put together,
  equality saturation explores all possible variants of a
  program that can be derived from a set of rewrites,
  and then it extracts the best one.
This solves the problem of choice described above,
  as equality saturation essentially applies every rewrite every iteration,
  using the e-graph to avoid exponential explosion.

It sounds a little too good to be true,
  and while it definitely has its caveats,
  `egg` addresses many of the shortcomings from this approach,
  making it feasible to apply on real-world problems.
We've already discussed how rebuilding makes `egg`'s e-graphs fast,
  and later tutorials will discuss how [`Analysis`] makes this approach flexible
  and able to handle more than just syntactic rewriting.

[`Searcher`]: super::super::Searcher
[`Pattern`]: super::super::Pattern
[`EGraph`]: super::super::EGraph
[`EClass`]: super::super::EClass
[`Rewrite`]: super::super::Rewrite
[`Runner`]: super::super::Runner
[`Extractor`]: super::super::Extractor
[`Language`]: super::super::Language
[`Analysis`]: super::super::Analysis
[`Id`]: super::super::Id
[`add`]: super::super::EGraph::add()
[`union`]: super::super::EGraph::union()
[`rebuild`]: super::super::EGraph::rebuild()
[equivalence relation]: https://en.wikipedia.org/wiki/Equivalence_relation
[congruence relation]: https://en.wikipedia.org/wiki/Congruence_relation
[dot]: super::super::Dot
[extract]: super::super::Extractor
[sound]: https://itinerarium.github.io/phoneme-synthesis/?w=/'igraf/
[paper]: https://arxiv.org/abs/2004.03082

*/
