Contextual Equality Saturation

• Contextual equality saturation is a technique that extends standard equality saturation by incorporating local context such as assumptions and control flow into rewrite rules.
• It leverages implementations like colored e-graphs and guarded rewriting to efficiently maintain contextual equivalences while sharing structure across contexts.
• This approach is applied in compiler, tensor, and query optimizations to adapt rewrites based on hardware constraints, resource bounds, and domain-specific analyses.

Searching arXiv for the cited CES-related papers to ground the article in current records. Contextual equality saturation is equality saturation performed under context rather than only under a single global congruence relation. In the broadest sense represented in recent work, it denotes rewrite saturation in which the applicability, scope, or profitability of an equality depends on assumptions, control-flow regions, binding structure, physical properties, side conditions, analyses, hardware models, or other contextual facts, while retaining the central e-graph idea of non-destructively accumulating equivalent forms and extracting a selected representative at the end (Singher et al., 2023). The term spans several distinct but related lines of work: conditional and assumption-indexed congruence layers in colored e-graphs, guarded equality analyses in the original equality-saturation framework, e-class-analysis-driven and relational formulations in egglog, persistent compiler embeddings in MLIR-like IRs, and domain-specific systems in tensor optimization, query optimization, symbolic rewriting, and higher-order or diagrammatic rewriting (Tate et al., 2010).

1. Definition and conceptual scope

In standard equality saturation, a set of universally quantified equalities and rewrite rules is applied to a common representation until no more useful equalities can be added, after which extraction selects a representative according to a cost model. The original compiler formulation replaces destructive pass sequencing with equality analyses that add equality information to a common intermediate representation, thereby avoiding phase-ordering concerns and deferring commitment to a single transformed program until global selection (Tate et al., 2010).

Contextual equality saturation extends that model by making equalities conditional on where, under which assumptions, or with which auxiliary facts they hold. One formulation is explicitly assumption-centric: contextual equality saturation is “the saturation of the term space under a family of contexts (assumption sets), maintaining which equalities hold under which assumptions” (Singher et al., 2023). Another formulation is guard-centric: rules carry side conditions dependent on analyses or external metadata, and a rewrite is applied only when those conditions hold (Cheli, 2021). A further formulation is relational: contextual equality saturation generalizes equality saturation to support rewrite rules conditioned on where a term appears in an expression, with contexts represented by relations over parents, siblings, analyses, and environment labels (Hou et al., 16 Jul 2025).

Across these formulations, the essential distinction from vanilla equality saturation is that equality is no longer uniformly global. A single e-graph congruence relation is insufficient when mutually inconsistent equalities, branch-local facts, type-sensitive rules, physical-property constraints, or binding environments must coexist. This suggests that contextual equality saturation is best viewed as a family of techniques for preserving the sharing benefits of e-graphs while admitting conditional, local, or layered equalities that traditional e-graphs do not directly encode (Singher et al., 2023).

2. Core semantics: contexts, congruence, and conditional equality

The most explicit semantic treatment in the supplied literature comes from colored e-graphs. There, a color $\kappa$ stands for a set of assumptions, the black layer represents a base congruence $\cong$, and each color defines an additional congruence $\cong_\kappa$ that coarsens the base relation, satisfying

$\cong \subseteq \cong_\kappa.$

If $\kappa \subseteq \kappa'$, then

$\cong_\kappa \subseteq \cong_{\kappa'},$

because additional assumptions can only add equalities, never split classes (Singher et al., 2023). In this semantics, contextual equality saturation is equivalent to maintaining a family of clone e-graphs, one per assumption context, but encoded inside one shared structure.

A second semantic line appears in the original equality-saturation framework through guarded equality analyses and a referentially transparent IR. In that setting, context is encoded semantically by constructs for branches, loops, effects, and semantic predicates. The operators $\phi$, $\theta_\ell$, $\mathrm{eval}_\ell$, and $\mathrm{pass}_\ell$ make control-flow and loop context explicit in the IR, while analyses such as MustEval and EvalCond identify where equalities hold safely (Tate et al., 2010). This does not introduce multiple congruence layers in the colored sense, but it does make equality reasoning context-aware at the level of rule triggers and semantic representation.

A third semantic line is relational and lattice-based. Egglog presents contextual equality saturation as equality saturation in which rewrite applicability and effect depend on context computed by analyses. Context can be encoded as relational predicates or lattice-valued functions, and side conditions are enforced by querying those facts during a single inflationary fixpoint that interleaves rewriting with analysis updates (Zhang et al., 2023). The relational CES perspective pushes this further by treating context as a monotone map

$\cong$0

from a lattice of contexts $\cong$1 to equivalence relations on a set $\cong$2, with monotonicity

$\cong$3

so that stronger context yields at least as many equalities (Hou et al., 16 Jul 2025).

In higher-order and diagrammatic settings, contextuality is tied to binding or monoidal context rather than assumptions or analyses alone. “E-Graphs With Bindings” treats contexts as morphisms into internal homs in a closed, semilattice-enriched symmetric monoidal category, making rewriting under binding context a first-class operation (Tiurin et al., 1 May 2025). “Equivalence Hypergraphs” generalizes contextual equality saturation from terms to string diagrams, where contexts are built from composition, tensor, symmetry, and join, and rewriting is realized by EDPOI rewriting on e-hypergraphs (Tiurin et al., 2024).

3. Principal implementation strategies

A central implementation strategy is the colored e-graph. The black layer stores the shared root e-graph, while each color maintains a lightweight union-find over black representatives together with color-local deltas. E-nodes are shared across colors; equalities specific to a context are represented by colored merges rather than by duplicating the whole structure. The paper characterizes this as a memory-efficient equivalent of multiple copies of an e-graph, attained by “sharing as much as possible between different cases, while carefully tracking which conclusion is true under which assumption” (Singher et al., 2023). This strategy is directly aimed at case splits, guards, and logical cuts.

A second strategy is guarded rewriting plus e-class analyses. Metatheory.jl adopts the egg pattern in a Julia setting: analyses compute attributes in a semi-lattice and attach them to e-classes; rules consult those attributes or external state through ordinary Julia guards; rebuilding restores congruence closure after batched additions and merges (Cheli, 2021). Egglog generalizes this by unifying Datalog-style inference and equality saturation in one database, so relations, lattices, function entries, and rewrites coexist in a single incremental fixpoint (Zhang et al., 2023). This makes contextual reasoning practical when the relevant context is already naturally expressed as relations such as type, purity, interval bounds, non-zero-ness, dominance, or shape.

A third strategy embeds the e-graph natively in compiler IR. The eqsat dialect for MLIR introduces eqsat.egraph, eqsat.eclass, and eqsat.yield, allowing the e-graph state to persist inside the compiler rather than being translated to an external library and discarded after extraction (Merckx et al., 14 May 2025). “E-Graphs as a Persistent Compiler Abstraction” extends this persistent view to xDSL/MLIR, reusing PDL-based pattern matching for e-matching and lifting SSA analyses as e-class analyses. In this setting, contextual equality saturation means that equalities survive across passes and abstraction levels, and that guards can be expressed in terms of types, dominance, purity, interval facts, and region scoping (Merckx et al., 18 Feb 2026).

A fourth strategy is to encode context by explicit auxiliary structure. Prior contextual approaches summarized in the relational CES paper include ASSUME nodes, top-down annotation with copying, and colored e-graphs (Hou et al., 16 Jul 2025). In concurrency and circuit settings, bond nodes or b-nodes play a related role: they store contextual relationships between otherwise disjoint computations, permitting non-local equality reasoning across concurrent subgraphs that are not syntactically connected (Lauko et al., 2022). In binding-sensitive and diagrammatic settings, hierarchy itself carries context: $\cong$4-boxes, e-boxes, and hierarchical hypergraph interfaces encode where equalities hold and how rewriting may cross those boundaries (Tiurin et al., 1 May 2025, Tiurin et al., 2024).

4. Context dimensions across major application domains

In compiler optimization, context is often semantic or operational. The original equality-saturation work already treated control-flow, loop, and effect context explicitly in PEGs and E-PEGs, and rule triggers could include guards such as loop invariance or must-evaluate predicates (Tate et al., 2010). Later compiler-integrated work broadens this with native IR persistence, typed SSA scoping, and interleaving with canonicalization, inlining, and common subexpression elimination (Merckx et al., 14 May 2025, Merckx et al., 18 Feb 2026). High-level Julia IR adds multiple dispatch, explicit CFG structure, dominance, and side effects as contextual dimensions; accordingly, its system propagates type information through e-classes, keeps a CFG skeleton outside the e-graph, and uses an ILP extraction that accounts for dominance-based reuse and acyclicity (Merckx et al., 24 Feb 2025).

In tensor graph optimization, contextual equality saturation is tied to tensor shapes, layouts, operator parameters, hardware cost, and node-budget constraints. Tensat applies rewrites only when shape, type, layout, axis, permutation, and parameter guards hold, and extracts with a device-aware cost model based on measured operator runtimes (Yang et al., 2021). The later MCTS-guided tensor work makes context explicitly resource-bounded and hardware-specific: the “context” includes GPU runtime estimates, tensor shapes and layouts, operator attributes, and memory budgets that limit e-graph growth (Hartmann et al., 2024). A plausible implication is that in this line of work, contextuality functions both as a soundness filter and as a search-control mechanism.

In relational query optimization, context includes relational algebra operator structure, estimated cardinalities and widths, physical operators, and the sequential history of rewrite application. Aurora treats the current e-graph topology and statistics as spatial context and the trajectory of prior rewrite and reset decisions as temporal context, using a GAT-plus-LSTM policy to guide equality saturation under memory and latency budgets (Bărbulescu et al., 2024). The relational CES agenda then shifts from heuristic guidance to semantically conditioned rewrites, such as replacing a merge join by a hash join only under an ancestor sort enforcer or recovering global equalities by meeting branch-local equalities (Hou et al., 16 Jul 2025).

In symbolic reasoning and theorem-proving-adjacent settings, context often means assumptions, guards, or semantic domain restrictions. Colored e-graphs address potentially inconsistent assumption sets directly (Singher et al., 2023). Ruler uses equality saturation to infer rewrite rules while respecting types, operator semantics, partiality, and side conditions; its cvec machinery and validation backends encode domain context so that inferred rules remain sound under the current semantics (Nandi et al., 2021). LGuess introduces another form of context: high-level intent supplied by an LLM as checkpoint preference. The e-graph still provides the certified low-level chain, but extraction is biased toward syntactic forms that the learned bigram model predicts are more promising under the LLM’s pairwise judgments (Peng et al., 1 Nov 2025).

In higher-order, concurrent, and diagrammatic domains, context is not merely metadata but part of the object language. Binding-sensitive equality saturation requires contexts under $\cong$5-abstraction, alpha-equivalence, and capture-avoiding substitution (Tiurin et al., 1 May 2025). Concurrent circuit optimization requires non-local contextual equivalence across disjoint subgraphs, handled by bond nodes and bond maps (Lauko et al., 2022). Monoidal and hypergraph formulations treat context as diagrammatic environment under composition and tensor, with EDPOI rewriting guaranteeing closure under such contexts (Tiurin et al., 2024).

5. Algorithms, extraction, and resource trade-offs

Contextual equality saturation preserves the usual saturation–rebuild–extract pattern but modifies all three phases. On the matching side, rules are no longer purely structural: they are filtered by guard predicates, relational joins, analysis facts, or context-layer membership. Colored e-graphs augment e-matching with on-demand colored traversal, so matching can fork when a black e-class has colored siblings under a context (Singher et al., 2023). Egglog uses relational e-matching and semi-naive evaluation to interleave context derivation with rewriting (Zhang et al., 2023). Persistent compiler approaches reinterpret PDL matching so that it backtracks over e-class members and checks constraints early, including constraints over newly created values (Merckx et al., 18 Feb 2026).

On the rebuilding side, contextual systems must maintain congruence per layer or per representation. Colored e-graphs rebuild the black structure first, then canonicalize color-local hash-cons structures with respect to color-specific union-find layers (Singher et al., 2023). Egglog uses a rebuilding operator $\cong$6 to repair equivalence and functional dependencies after inflationary rule application (Zhang et al., 2023). Eqsat-in-MLIR shifts much of rebuilding to existing compiler machinery, especially CSE, which can subsume explicit congruence restoration when equal operations with identical operands arise in the embedded e-graph IR (Merckx et al., 14 May 2025).

Extraction is also contextualized. In the colored setting, extraction is performed per color, so each assumption context gets its own optimal representative over the partition induced by that context’s union-find (Singher et al., 2023). In tensor optimization, extraction is device-aware and shape-aware, often via ILP or a greedy dynamic program corrected for common subexpressions (Yang et al., 2021, Hartmann et al., 2024). In Julia IR, extraction is formulated as an ILP over multiple roots with dominance reuse and cycle-breaking constraints (Merckx et al., 24 Feb 2025). In Aurora, extraction is latency-oriented and ILP-based, using relational query statistics propagated through e-classes (Bărbulescu et al., 2024). In LGuess, checkpoint extraction becomes an argmax over an additive score derived from a positional bigram model learned from LLM preferences (Peng et al., 1 Nov 2025).

The most persistent trade-off in the literature is between contextual expressivity and structural sharing. Clone-per-context approaches are straightforward but expensive. Colored e-graphs reduce memory by layering coarsened congruence relations over a shared base (Singher et al., 2023). ASSUME-node and top-down copying approaches make context explicit but can expand the e-graph considerably (Hou et al., 16 Jul 2025). Richer guidance, such as MCTS or RL, mitigates the reintroduced phase-ordering problem that appears when full saturation is infeasible under node limits or time budgets (Hartmann et al., 2024, Bărbulescu et al., 2024). This suggests that “contextual” frequently means not only “conditional” but also “resource-aware.”

6. Limitations, misconceptions, and active directions

A common misconception is that contextual equality saturation is a single data structure or algorithm. The supplied literature instead presents a spectrum: layered congruence relations indexed by assumptions, guard-driven e-class analyses, relational context predicates, persistent compiler embeddings, and category-theoretic representations for bindings and diagrams. These systems share a goal but not a single formalism (Singher et al., 2023, Hou et al., 16 Jul 2025).

Another misconception is that contextuality merely means adding more predicates to rewrite rules. In some domains that is accurate, but in others the challenge is representational. Traditional e-graphs struggle with mutually inconsistent assumptions, with variable binding, with non-local concurrent equivalences, or with region-scoped SSA values. Those cases require structural changes such as colored layers, hierarchical hypergraphs, bond nodes, or native IR regions rather than only additional rule guards (Singher et al., 2023, Lauko et al., 2022, Tiurin et al., 1 May 2025).

Several open issues recur across the literature. Large numbers of contexts increase per-context overhead, and colored e-graphs note that their layered union-find was simulated by copying the black union-find per color, leaving room for tighter representations (Singher et al., 2023). Relational CES identifies the difficulty of reconciling contextual layers with monotone fixpoint evaluation and of preventing programs from observing non-monotonicity (Hou et al., 16 Jul 2025). Persistent compiler approaches remain limited in handling side effects and region-bearing control-flow rewrites, although they identify effect-aware extraction and region-level matching as natural future directions (Merckx et al., 14 May 2025, Merckx et al., 18 Feb 2026). Tensor and query systems demonstrate that resource bounds reintroduce phase-ordering unless guidance is itself context-aware (Hartmann et al., 2024, Bărbulescu et al., 2024). Higher-order and categorical systems provide strong semantics, but the supplied descriptions do not report performance evaluations comparable to those in compiler or tensor settings (Tiurin et al., 1 May 2025, Tiurin et al., 2024).

A broader synthesis is that contextual equality saturation has become a unifying phrase for making equality saturation sensitive to information that ordinary congruence closure abstracts away: assumptions, branches, scopes, bindings, effects, dataflow properties, physical properties, hardware cost, and optimization intent. The literature suggests that the field is converging on two complementary directions. One direction internalizes context into the e-graph or its categorical generalizations through layers, hierarchy, or persistence. The other treats context as an external but incrementally maintained body of relational or analytical facts that rewrites query at match and extraction time. Both aim to retain the defining advantage of equality saturation—shared, non-destructive exploration of many equivalent forms—while making that exploration semantically local where necessary and globally reusable where possible (Singher et al., 2023, Zhang et al., 2023, Hou et al., 16 Jul 2025).