egg教程（一）：e-graphs and equality saturation的概念

e-graph是一种支持equality saturation优化技术的数据结构。

e-graph（Gregory Nelson’s PhD Thesis, 1980）和equality saturation（Tate et. al., 2009）都是在此之前发明的。egg目的是让他们更加简单，易拓展和高效。

本教程将从高层次探讨这些概念。更多细节可参见论文。

为什么要这样做？

在构建几乎所有与编程语言相关的工具时，您通常会使用您语言中的term（我们将交替使用term、expression和program等词）。您可能会尝试:

1. 将term转换为 "更好的 "等价term（optimize）；

2. 根据规范从头构建一个term（synthesis）；

3. 证明两个term是等价的（verification）。

Term rewriting 是一种常用的处理技术，可以实现上述任何目标。你可以从一个term t 和一组rewriting开始，每个重写的形式都是 l → r。对于每个重写，我们都会尝试将模式 l 与term t 进行匹配，在某个subterm上生成一个替换 σ，然后将该替换应用到右侧的模式 r 上，并替换匹配的subterm。

例如，考虑term 42 × (7 + 7) 和重写 x + x → 2 × x：

1. 左侧将匹配子项 7 + 7，生成代词 σ = {x：7}。

2. 应用该替换后，r[σ] = 2 × 7（r[σ] 表示将替换词 σ 应用于 r）。

3. 用 r[σ] = 2 × 7 代替匹配子项 7 + 7，得到结果：42 × (2 × 7)。

Term rewriting（以及其他涉及term操作的编程语言技术）的一个问题是选择问题。重写是破坏性的，也就是说，一旦进行了重写，初始术语就会消失。如果你有多种重写方法可供选择，那么选择其中一种就像选择岔路口一样：返回并选择另一种方法是昂贵的，或者是不可能的，因此你必须坚持你所选择的道路。此外，在正确的时间选择正确的改写一般来说是非常困难的；有时，在局部看来不错的选择可能在整体上是错误的。

例如，C 语言优化器希望将 a * 2 重写为更便宜的 a << 1，但 (a * 2) / 2 又如何呢？贪婪的局部方法会 "转错弯"，结果是 (a << 1) / 2，它比输入表达式更好，但却掩盖了我们本可以取消乘除的事实。

有一些技术可以缓解这种情况（回溯和数值编号），但从根本上说，绕过它的唯一方法就是同时进行所有选择。但重写是破坏性的，因此对一个词条进行 n 次重写会产生 n 个term，然后对这 n 个term中的每个term进行 n 次重写又会产生 n2 个term，以此类推。如果你想探索这棵树的 m "层 "深处，你就必须存储 nm 个term，而这是非常昂贵的。

e-graph可以解决这个表示问题。当重写一堆term时，结果往往在结构上是相似的。e-graph可以非常紧凑地存储大量相似表达式的集合。使用e-graph，你可以同时应用许多重写，而不会造成空间的成倍增长。Equality saturation是一种为优化程序而进行此类重写的技术。

什么是e-graph?

e-graph 是一种数据结构，用于维护表达式的等价关系（实际上是全等关系，见下一节）。e-graph是一组等价类（e-classes），每个等价类都包含等价的 e-node。e-node是带有子节点的运算符，但子节点不是其他运算符或值，而是 e-class。在 egg 中，这些由 EGraph、EClass 和 Language（e-nodes）类型表示。

即使是小的e-graph也能表示大量的表达式，其数量与e-nodes的数量成指数关系。这种紧凑性使得e-graph成为一种引人注目的数据结构。我们可以将表示（或包含）一个term的含义定义如下：

1. 如果一个e-graph的任何一个e-class都代表一个term，那么这个e-graph就代表term。

2. 如果一个 e-class 的任何一个 e-node 都代表一个term，那么这个 e -class就代表term。

3. 如果 e-class ni 代表term ti，则 e -node f(n1, n2, ...) 代表term f(t1,t2,...)。

下面是一些e-graph。我们把 e-class 描绘成围绕着等价 e-node 的虚线框。请注意，边是从 e-node 到 e-class（而不是到其他 e-node），因为 e-node 的子节点是 e-class。先不用担心它们是如何构成的，我们很快就会重写。

可以通过称为 e-matching（等价匹配）的过程查询e-graph中的模式，该过程会在e-graph中搜索代表与给定模式匹配的term的 e-class。egg 的 e-matching（受 De Moura 和 Bjørner 的启发）由搜索器和模式 API 完成。修改e-graph的核心操作有两个：添加（add）和合并（union），前者用于向e-graph中添加e-node，后者用于合并两个e-class。

这些操作可以结合起来对e-graph进行重写。重要的是，这种重写过程只是加法，这意味着电子图永远不会忘记term的旧版本。此外，重写一次可以添加整类term。要在e-graph上执行 l → r 重写，首先要在电子图中搜索 l，得到 (c, σ) 对，其中 c 是匹配的 e-class，σ 是替换。然后，将术语 r[σ] 添加到 e 图中，并与匹配的 e-class c 结合。

下面我们就以上图中的四幅e-graph为例，来说明这一切。

1. 初始 e-graph 表示 term (a × 2) / 2 。由于每个 e-class 只有一个 e-node ，因此该 e-graph 基本上是一棵共享的抽象语法树（2 不重复）。

2. 应用重写 x × 2 → x << 1 已记录了 a × 2 = a << 1 的事实，而没有忘记 a × 2。请注意新添加的 a << 1 是如何指向现有的 "a " e-node的，而"<<" e-node已被联合到与模式 x × 2 匹配的"×"e-node相同的e-class中。

3. 应用重写 (x × y) / z → x × (y / z) 实现了除法与乘法的关联。这个重写对于发现我们正在寻找的 2 的抵消是至关重要的，尽管我们之前应用了 "错误的 "重写，但它仍然有效。

4. 应用 x / x → 1 和 x × 1 → x 重写并不会增加任何新的 e-node，因为所有的 e-node 都已经存在于 e-graph 中。结果只是联合了 e-class，也就是说，尽管现在 e-graph 代表了更多的项，但应用这些重写后，e-graph 实际上变小了。事实上，观察右上方"×" e-node 的左侧子节点就是它自己；这个循环意味着这个 e-class 代表了无限（！）的项集 a、a ×1、a × 1 ×1，以此类推。

不变式和重建

e-graph有两个修改e-graph的核心操作：add（向e-graph中添加e-node）和 union（合并两个e-class）。这些操作维护两个关键（相关）不变式：

1. 一致性

e-graphj不仅维护表达式的等价关系，还维护同构关系。同构关系的基本含义是，如果 x 等价于 y，那么 f(x) 必须等价于 f(y)。因此，当用户调用 union 时，除了给定的两个 e-class 外，其他许多 e-class 可能需要合并以保持一致性。

例如，假设术语 a + x 和 a + y 分别在 e-class 1 和 2 中表示。在以后的某个时刻，x 和 y 变得等价（也许用户在包含它们的 e-class 中调用了 union）。E-class 1 和 2 必须合并，因为现在两个 "+"运算符有了等价参数，使它们等价。

2. e-node的唯一性

在同一 e 类或不同 e 类中，不存在两个不同的 e 节点，它们具有相同的运算符和等同的子节点。这在一定程度上是通过 add 执行的散列整理以及 union 和 rebuild 执行的重复数据删除来实现的。

egg 采用延迟的方式来维护这些不变性。具体来说，调用 union（或应用调用 union 的重写）的效果可能不会立即反映出来。要恢复e-graph不变式并使这些效果可见，用户必须调用重建方法。

egg在这里的选择可以实现更高的性能。维持同构关系会使核心e-graph数据结构复杂化，并且在每次联合时都需要对e-graph进行昂贵的遍历。为了提高性能，egg 选择放宽这些不变式，只有在调用重建时才恢复不变式。关于 egg 的论文更详细地说明了为什么这种设计选择会让速度更快。有关 rebuild 的作用以及何时必须调用 rebuild 的更多信息，请参阅 rebuild 文档。还需注意的是，运行程序会在重写迭代之间调用 rebuild。

Equality Saturation

Equality saturation (Tate et. al., 2009) 是一种程序优化技术，egg通过类似下面伪代码的runner api中实现了equality saturation的变体。

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();
}

上文已经介绍了大部分内容，但我们需要定义两个新term：

1. 当e-graph检测到改写不再增加新信息时，就会出现饱和。考虑一下交换重写 x + y → y + x。在应用了一次之后，由于e-graph没有忘记最初的 x + y 项，所以第二次重写不会再增加新信息。如果所有的重写都处于这种状态，我们就说e-graph已经饱和，这意味着e-graph包含了从给定的重写中推导出的所有可能的等价关系。

2. Extraction是一种从e-graph中挑选单个代表term的过程，根据某些成本函数，该过程是最优的。egg的Extractors提供了这个功能。

总而言之，equality saturation会探索程序的所有可能变体，这些变体可以从一组重写中衍生出来，然后提取出最佳变体。这就解决了上文所述的选择问题，因为equality saturation基本上每次迭代都会应用每个重写，利用e-graph避免指数爆炸。

这听起来有点好得不像真的，虽然它肯定有其注意事项，但 egg 解决了这种方法的许多缺点，使其在实际问题中的应用变得可行。我们已经讨论了重建如何使 egg 的e-graph变得快速，稍后的教程将讨论分析如何使这种方法变得灵活，并能处理更多语法重写问题。